diff --git a/tests/runtime/codex/test_codex_runtime_smoke.py b/tests/runtime/codex/test_codex_runtime_smoke.py index 51ff3a2f1..965090f96 100644 --- a/tests/runtime/codex/test_codex_runtime_smoke.py +++ b/tests/runtime/codex/test_codex_runtime_smoke.py @@ -822,8 +822,11 @@ async def _drain(turn: Any) -> None: @pytest.mark.codex_smoke @pytest.mark.asyncio -async def test_real_codex_pauses_and_resumes_external_tool(monkeypatch) -> None: - """Exercise pending tool handoff over real Codex/shim processes.""" +@pytest.mark.parametrize("transport", ["shim", "direct"]) +async def test_real_codex_pauses_and_resumes_external_tool( + monkeypatch, transport: str +) -> None: + """Exercise pending tool handoff over both real Codex transports.""" if os.getenv("CODEX_RUN_SMOKE") != "1": pytest.skip("set CODEX_RUN_SMOKE=1 to spawn the real Codex binary") reason = _skip_reason() @@ -835,8 +838,11 @@ async def test_real_codex_pauses_and_resumes_external_tool(monkeypatch) -> None: from veadk import Agent, Runner from veadk.runtime.codex import runtime as runtime_module + from veadk.runtime.codex.mcp_bridge import shutdown_bridge from veadk.runtime.codex.proxy import ResponsesShim + from tests.runtime.codex.test_codex_mcp_bridge import _StubModel + executed = [] async def veadk_smoke_record_marker(marker: str) -> None: @@ -844,26 +850,34 @@ async def veadk_smoke_record_marker(marker: str) -> None: await asyncio.sleep(0.01) executed.append(marker) - backend = _StubResponsesBackend() - backend_url = await backend.start() - shim = ResponsesShim(api_base=f"{backend_url}/v1", api_key="veadk-smoke-key") - await shim.start() + shim = None + if transport == "direct": + backend = _StubModel(_ADK_TOOL_NAME, {"marker": _TOOL_MARKER}) + api_base = await backend.start() + final_answer = "done" + else: + backend = _StubResponsesBackend() + backend_url = await backend.start() + api_base = f"{backend_url}/v1" + final_answer = _FINAL_ANSWER + shim = ResponsesShim(api_base=api_base, api_key="veadk-smoke-key") + await shim.start() - async def lookup(*_args, **_kwargs): - return shim + async def lookup(*_args, **_kwargs): + return shim - monkeypatch.setattr(runtime_module, "get_shim", lookup) + monkeypatch.setattr(runtime_module, "get_shim", lookup) agent = Agent( name="external_tool_smoke", instruction="Call the tool, then answer.", runtime="codex", model_name="veadk-codex-smoke-model", - model_api_base=f"{backend_url}/v1", + model_api_base=api_base, model_api_key="veadk-smoke-key", model_api_key_name="", tools=[LongRunningFunctionTool(veadk_smoke_record_marker)], codex_runtime_config={ - "model_transport": "shim", + "model_transport": transport, "sandbox": "read_only", "network_access": False, "turn_timeout_seconds": 30, @@ -923,13 +937,172 @@ async def send(content): event.is_final_response() and event.content and any( - _FINAL_ANSWER in (part.text or "") for part in event.content.parts or [] + final_answer in (part.text or "") for part in event.content.parts or [] ) for event in second ) assert "external-smoke-completed" in json.dumps(backend.requests[-1]) assert executed == [_TOOL_MARKER] finally: - await shim.stop() + if shim is not None: + await shim.stop() + await runtime_module._PROCESS_POOL.close_all() + await shutdown_bridge() + await backend.stop() + if shim is not None: + assert shim._turns == {} + + +@pytest.mark.codex_smoke +@pytest.mark.asyncio +@pytest.mark.parametrize("response_order", [("first", "second"), ("second", "first")]) +@pytest.mark.parametrize("resumable", [False, True]) +@pytest.mark.parametrize("cold_resume", [False, True]) +async def test_real_codex_resumes_every_external_result( + response_order, resumable, cold_resume +): + """A persistent Codex thread must receive the whole externally answered batch.""" + if os.getenv("CODEX_RUN_SMOKE") != "1": + pytest.skip("set CODEX_RUN_SMOKE=1 to spawn the real Codex binary") + reason = _skip_reason() + if reason is not None: + pytest.skip(reason) + + from google.adk.apps.app import App, ResumabilityConfig + from google.adk.runners import Runner + from google.adk.sessions.in_memory_session_service import InMemorySessionService + from google.adk.tools.long_running_tool import LongRunningFunctionTool + from google.genai import types + + from veadk import Agent + from veadk.runtime.codex import runtime as runtime_module + from veadk.runtime.codex.mcp_bridge import shutdown_bridge + from tests.runtime.codex.test_codex_mcp_bridge import _StubModel + + class BatchModel(_StubModel): + def _script(self, body): + output = super()._script(body) + if len(self.requests) == 1: + return [ + { + **output[0], + "id": f"fc_{name}", + "call_id": f"call_{name}", + "arguments": json.dumps({"marker": name}), + } + for name in ("first", "second") + ] + payload = json.dumps(body) + complete = all(f"result-{name}" in payload for name in response_order) + output[0]["content"][0]["text"] = ( + "both-complete" if complete else "missing-result" + ) + return output + + executed = [] + both_started = asyncio.Event() + + async def external_job(marker: str) -> None: + """Start a job whose result will be submitted separately.""" + executed.append(marker) + if len(executed) == 2: + both_started.set() + await asyncio.wait_for(both_started.wait(), 5) + + def annotate_results(callback_context, llm_request): + for content in llm_request.contents: + if any(part.function_response for part in content.parts or []): + content.parts.append(types.Part(text="callback result annotation")) + + backend = BatchModel("external_job", {}) + api_base = await backend.start() + agent = Agent( + name="assistant", + runtime="codex", + model_name="partial-results-test", + model_api_base=api_base, + model_api_key="loopback-key", + model_api_key_name="", + tools=[LongRunningFunctionTool(external_job)], + before_model_callback=annotate_results, + codex_runtime_config={ + "model_transport": "direct", + "thread_mode": "resume", + "sandbox": "read_only", + "network_access": False, + "turn_timeout_seconds": 30, + }, + ) + sessions = InMemorySessionService() + if resumable: + runner = Runner( + app=App( + name="partial_results", + root_agent=agent, + resumability_config=ResumabilityConfig(is_resumable=True), + ), + session_service=sessions, + ) + else: + runner = Runner( + app_name="partial_results", agent=agent, session_service=sessions + ) + session_id = f"partial-results-{uuid.uuid4().hex[:8]}" + await sessions.create_session( + app_name="partial_results", user_id="test", session_id=session_id + ) + + async def send(content): + async with asyncio.timeout(_RUN_TIMEOUT_SECONDS): + return [ + event + async for event in runner.run_async( + user_id="test", session_id=session_id, new_message=content + ) + ] + + def result(name, will_continue=None): + return types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id=f"call_{name}", + name="external_job", + response={"value": f"result-{name}"}, + will_continue=will_continue, + ) + ) + ], + ) + + try: + initial = await send( + types.Content(role="user", parts=[types.Part(text="Run both jobs.")]) + ) + calls = [call for event in initial for call in event.get_function_calls()] + assert {call.id for call in calls} == {"call_first", "call_second"} + assert sorted(executed) == ["first", "second"] + assert len(backend.requests) == 1 + first, last = response_order + assert await send(result(first)) == [] + assert await send(result(last, True)) == [] + assert len(backend.requests) == 1 + if cold_resume: + await runtime_module._PROCESS_POOL.close_all() + events = await send(result(last, False)) + assert len(backend.requests) == 2 + assert "" not in json.dumps(backend.requests[-1]) + assert "callback result annotation" in json.dumps(backend.requests[-1]) + assert sorted(executed) == ["first", "second"] + assert not any(event.get_function_calls() for event in events) + assert any( + part.text == "both-complete" + for event in events + if not event.partial and event.content + for part in event.content.parts or [] + ) + finally: + await runtime_module._PROCESS_POOL.close_all() + await shutdown_bridge() await backend.stop() - assert shim._turns == {} diff --git a/tests/runtime/codex/test_codex_turn_contract.py b/tests/runtime/codex/test_codex_turn_contract.py index 2bf43c4a9..1660814ae 100644 --- a/tests/runtime/codex/test_codex_turn_contract.py +++ b/tests/runtime/codex/test_codex_turn_contract.py @@ -2294,6 +2294,35 @@ def after_model_callback(callback_context, llm_response): ) +@pytest.mark.asyncio +async def test_pending_direct_tool_interrupts_the_turn_only_once(monkeypatch): + from google.adk.tools.long_running_tool import LongRunningFunctionTool + + interrupts = [] + original_interrupt = fake_codex_sdk._DirectTurn.interrupt + + async def counted_interrupt(turn): + interrupts.append(turn.id) + await original_interrupt(turn) + + monkeypatch.setattr(fake_codex_sdk._DirectTurn, "interrupt", counted_interrupt) + + async def external_job() -> None: + """Wait for an external result.""" + await asyncio.sleep(0) + + events, _session, _backend, error = await _run_direct_turn( + monkeypatch, + plan=(Round(tool_calls=(("external_job", {}),)), Round(text="discard")), + agent_kwargs={"tools": [LongRunningFunctionTool(external_job)]}, + ) + assert error is None, error + assert any(event.long_running_tool_ids for event in events) + # Completion notifications still drain after pending interrupts the turn. + # They must not trigger another RPC against an already completed turn. + assert len(interrupts) == 1 + + @pytest.mark.asyncio async def test_warm_process_resumes_external_result_without_reexecuting_tool( monkeypatch, diff --git a/tests/runtime/conformance/test_pending_tool_conformance.py b/tests/runtime/conformance/test_pending_tool_conformance.py new file mode 100644 index 000000000..2762caca1 --- /dev/null +++ b/tests/runtime/conformance/test_pending_tool_conformance.py @@ -0,0 +1,274 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Every runtime waits for the complete external-tool batch before modeling.""" + +import pytest +from google.adk.apps.app import App, ResumabilityConfig +from google.adk.events.event import Event +from google.adk.runners import Runner +from google.adk.tools.long_running_tool import LongRunningFunctionTool +from google.genai import types + +from conformance_harness import Round + + +@pytest.mark.asyncio +@pytest.mark.parametrize("response_order", [("a", "b"), ("b", "a")]) +@pytest.mark.parametrize("packed_calls", [False, True]) +@pytest.mark.parametrize("resumable", [False, True]) +async def test_partial_tool_results_do_not_invoke_model( + harness, response_order, packed_calls, resumable +): + executed = [] + + async def external_job(value: str) -> None: + """Start an operation whose result is supplied later.""" + executed.append(value) + + scripted = harness.adapter.build( + [Round(text="Both results received.")], + tools=[LongRunningFunctionTool(external_job)], + ) + if resumable: + runner = Runner( + app=App( + name=harness.APP_NAME, + root_agent=scripted.agent, + resumability_config=ResumabilityConfig(is_resumable=True), + ), + session_service=harness.session_service, + ) + else: + runner = Runner( + app_name=harness.APP_NAME, + agent=scripted.agent, + session_service=harness.session_service, + ) + session_id = await harness.new_session() + session = await harness.session_service.get_session( + app_name=harness.APP_NAME, user_id=harness.USER_ID, session_id=session_id + ) + invocation = "inv-external-batch" + + async def append(event): + await harness.session_service.append_event(session, event) + + await append( + Event( + invocation_id=invocation, + author="user", + content=types.Content( + role="user", parts=[types.Part(text="Start both jobs.")] + ), + ) + ) + calls = [ + types.Part( + function_call=types.FunctionCall( + id=call_id, name="external_job", args={"value": call_id} + ) + ) + for call_id in ["a", "b"] + ] + for parts in [calls] if packed_calls else [[part] for part in calls]: + await append( + Event( + invocation_id=invocation, + author=scripted.agent.name, + long_running_tool_ids={part.function_call.id for part in parts}, + content=types.Content(role="model", parts=parts), + ) + ) + + async def send(call_id, *, will_continue=None): + message = types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id=call_id, + name="external_job", + response={"value": f"result-{call_id}", "isFinal": True}, + will_continue=will_continue, + ) + ) + ], + ) + return [ + event + async for event in runner.run_async( + user_id=harness.USER_ID, + session_id=session_id, + new_message=message, + ) + ] + + first, last = response_order + assert await send(first) == [] + assert harness.adapter.requests() == [] + assert await send(last, will_continue=True) == [] + assert harness.adapter.requests() == [] + events = await send(last, will_continue=False) + assert len(harness.adapter.requests()) == 1 + assert any( + event.content + and any( + part.text == "Both results received." for part in event.content.parts or [] + ) + for event in events + ) + assert not any(event.get_function_calls() for event in events) + assert executed == [] + saved = await harness.session_service.get_session( + app_name=harness.APP_NAME, user_id=harness.USER_ID, session_id=session_id + ) + assert [ + response.id + for event in saved.events + for response in event.get_function_responses() + ] == [first, last, last] + + # A new user turn after a cancelled/abandoned tool invocation stays runnable. + session = saved + await append( + Event( + invocation_id="abandoned-invocation", + author=scripted.agent.name, + long_running_tool_ids={"abandoned"}, + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id="abandoned", + name="external_job", + args={"value": "abandoned"}, + ) + ) + ], + ), + ) + ) + async for _ in runner.run_async( + user_id=harness.USER_ID, + session_id=session_id, + new_message=types.Content( + role="user", parts=[types.Part(text="Start a different task.")] + ), + ): + pass + assert len(harness.adapter.requests()) == 2 + + +@pytest.mark.asyncio +async def test_pending_agent_does_not_stop_parallel_sibling(harness): + from google.adk.agents.parallel_agent import ParallelAgent + + async def external_job(value: str) -> None: + """Wait for an external result.""" + raise AssertionError("Saved pending calls must not execute again") + + waiting = harness.adapter.build( + [Round(text="waiting finished")], + key="waiting", + tools=[LongRunningFunctionTool(external_job)], + ) + active = harness.adapter.build([Round(text="sibling finished")], key="active") + waiting.agent.name = "waiting" + active.agent.name = "active" + parallel = ParallelAgent(name="parallel", sub_agents=[waiting.agent, active.agent]) + runner = Runner( + app=App( + name=harness.APP_NAME, + root_agent=parallel, + resumability_config=ResumabilityConfig(is_resumable=True), + ), + session_service=harness.session_service, + ) + session_id = await harness.new_session() + session = await harness.session_service.get_session( + app_name=harness.APP_NAME, user_id=harness.USER_ID, session_id=session_id + ) + invocation = "inv-parallel" + await harness.session_service.append_event( + session, + Event( + author="user", + invocation_id=invocation, + content=types.Content( + role="user", parts=[types.Part(text="Run both branches.")] + ), + ), + ) + await harness.session_service.append_event( + session, + Event( + author="waiting", + invocation_id=invocation, + branch="parallel.waiting", + long_running_tool_ids={"pending"}, + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id="pending", name="external_job", args={"value": "pending"} + ) + ) + ], + ), + ), + ) + events = [ + event + async for event in runner.run_async( + user_id=harness.USER_ID, session_id=session_id, invocation_id=invocation + ) + ] + assert harness.adapter.requests("waiting") == [] + assert len(harness.adapter.requests("active")) == 1 + assert any( + event.author == "active" + and event.content + and any(part.text == "sibling finished" for part in event.content.parts or []) + for event in events + ) + resumed = [ + event + async for event in runner.run_async( + user_id=harness.USER_ID, + session_id=session_id, + new_message=types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id="pending", + name="external_job", + response={"value": "completed"}, + ) + ) + ], + ), + ) + ] + assert len(harness.adapter.requests("waiting")) == 1 + assert len(harness.adapter.requests("active")) == 1 + assert any( + event.author == "waiting" + and event.content + and any(part.text == "waiting finished" for part in event.content.parts or []) + for event in resumed + ) diff --git a/tests/runtime/conformance/test_runtime_conformance.py b/tests/runtime/conformance/test_runtime_conformance.py index c4d7491f6..9822304d9 100644 --- a/tests/runtime/conformance/test_runtime_conformance.py +++ b/tests/runtime/conformance/test_runtime_conformance.py @@ -244,6 +244,12 @@ async def test_streaming_partials_precede_one_final( assert all(i < finals[0] for i in partial_at), ( f"partial events at {partial_at} arrived after the final at {finals[0]}" ) + if harness.adapter.name != "adk": + # Database history sorts by timestamp, not the original yield order. + assert all( + turn.events[i].timestamp <= turn.events[finals[0]].timestamp + for i in partial_at + ) partial_text = "".join( p.text or "" for i in partial_at diff --git a/tests/runtime/piagent/README.md b/tests/runtime/piagent/README.md index af6461d82..2e9b7c358 100644 --- a/tests/runtime/piagent/README.md +++ b/tests/runtime/piagent/README.md @@ -11,8 +11,20 @@ PYTHONPYCACHEPREFIX=/private/tmp/veadk-pycache \ ## Real local smoke test -This test calls a real Pi binary and a real model. It is skipped unless -`PIAGENT_RUN_SMOKE=1` is set. +The native-history smoke runs a real Pi binary against a loopback model stub. +It verifies that pending tools resume through native assistant/tool-result +messages and that a new user request can call the same tool with the same +arguments again. It needs no model credentials: + +```bash +PIAGENT_RUN_SMOKE=1 \ + .venv/bin/python -m pytest \ + tests/runtime/piagent/test_piagent_runtime_smoke.py::test_real_piagent_resumes_native_tool_results \ + -p no:xdist -rs +``` + +The greeting smoke below calls a real model. Both smoke tests are skipped +unless `PIAGENT_RUN_SMOKE=1` is set; the greeting also requires credentials. ```bash export PIAGENT_BINARY=/private/tmp/veadk-piagent-binary/v0.80.6-darwin/extracted/pi/pi diff --git a/tests/runtime/piagent/test_piagent_context.py b/tests/runtime/piagent/test_piagent_context.py new file mode 100644 index 000000000..7815c7561 --- /dev/null +++ b/tests/runtime/piagent/test_piagent_context.py @@ -0,0 +1,262 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import json +import os +import stat +from datetime import datetime, timezone +from pathlib import Path + +import pytest +from google.adk.models.llm_request import LlmRequest +from google.genai import types + +from veadk.runtime.piagent.config import PiAgentModelConfig +from veadk.runtime.piagent.context_bridge import build_pi_messages, pi_context_extension + + +MODEL = PiAgentModelConfig( + provider_id="veadk", + model="test-model", + base_url="https://model.invalid", + api_key="private-model-credential", + api="openai-completions", + api_key_env="TEST_KEY", +) + + +def _request(*contents): + return LlmRequest( + contents=list(contents), + config=types.GenerateContentConfig(system_instruction="Wait for tool results."), + ) + + +def _call(call_id="call-1"): + return types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id=call_id, name="lookup", args={"key": "same"} + ) + ) + ], + ) + + +def _response(call_id="call-1"): + return types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id=call_id, name="lookup", response={"value": False} + ) + ) + ], + ) + + +def test_history_keeps_native_roles_and_tool_pairing(): + request = _request( + types.Content(role="user", parts=[types.Part(text="Look it up.")]), + _call(), + _response(), + types.Content( + role="model", + parts=[ + types.Part(text="private reasoning", thought=True), + types.Part(text="It is false."), + ], + ), + types.Content(role="user", parts=[types.Part(text="Look it up again.")]), + _call("call-2"), + _response("call-2"), + ) + messages = build_pi_messages(request, MODEL) + assert [m["role"] for m in messages] == [ + "user", + "assistant", + "toolResult", + "assistant", + "user", + "assistant", + "toolResult", + ] + assert messages[1]["content"][0]["id"] == messages[2]["toolCallId"] == "call-1" + assert messages[5]["content"][0]["id"] == messages[6]["toolCallId"] == "call-2" + assert json.loads(messages[6]["content"][0]["text"]) == {"value": False} + assert "private reasoning" not in json.dumps(messages) + assert messages[1]["provider"] == MODEL.provider_id + + +def test_missing_call_ids_are_paired_without_reusing_another_call(): + messages = build_pi_messages( + _request(_call(None), _response(None), _call(None), _response(None)), MODEL + ) + assert messages[0]["content"][0]["id"] == messages[1]["toolCallId"] + assert messages[2]["content"][0]["id"] == messages[3]["toolCallId"] + assert messages[1]["toolCallId"] != messages[3]["toolCallId"] + + +def test_generated_call_id_does_not_collide_with_explicit_parallel_id(): + messages = build_pi_messages( + _request( + _call("veadk-call-1"), + _call(None), + _response("veadk-call-1"), + _response(None), + ), + MODEL, + ) + ids = [call["id"] for call in messages[0]["content"]] + assert len(set(ids)) == 2 + assert [m["toolCallId"] for m in messages[1:]] == ids + + +def test_history_uses_the_current_pi_tool_alias(): + messages = build_pi_messages( + _request(_call(), _response()), MODEL, tool_names={"lookup": "veadk_lookup"} + ) + assert messages[0]["content"][0]["name"] == "veadk_lookup" + assert messages[1]["toolName"] == "veadk_lookup" + + +@pytest.mark.parametrize("packed_results", [False, True]) +def test_parallel_calls_share_one_assistant_and_keep_reversed_results(packed_results): + responses = [_response("call-b"), _response("call-a")] + if packed_results: + responses = [ + types.Content( + role="user", + parts=[part for result in responses for part in result.parts], + ) + ] + messages = build_pi_messages( + _request(_call("call-a"), _call("call-b"), *responses), MODEL + ) + assert [m["role"] for m in messages] == ["assistant", "toolResult", "toolResult"] + assert [call["id"] for call in messages[0]["content"]] == ["call-a", "call-b"] + assert [m["toolCallId"] for m in messages[1:]] == ["call-b", "call-a"] + + +@pytest.mark.parametrize("same_content", [False, True]) +def test_tool_result_precedes_new_user_text_without_merging_the_user_turn(same_content): + text = types.Part(text="Now ask a different question.") + result = _response() + following = [types.Content(role="user", parts=[text]), result] + if same_content: + following = [types.Content(role="user", parts=[text, *result.parts])] + messages = build_pi_messages(_request(_call(), *following), MODEL) + assert [m["role"] for m in messages] == ["assistant", "toolResult", "user"] + assert messages[-1]["content"] == [{"type": "text", "text": text.text}] + + +def test_history_preserves_pydantic_json_values(): + call, result = _call(), _response() + value = datetime(2026, 11, 3, tzinfo=timezone.utc) + call.parts[0].function_call.args = {"at": value} + result.parts[0].function_response.response = {"at": value} + messages = build_pi_messages(_request(call, result), MODEL) + assert messages[0]["content"][0]["arguments"] == {"at": "2026-11-03T00:00:00Z"} + assert json.loads(messages[1]["content"][0]["text"]) == { + "at": "2026-11-03T00:00:00Z" + } + + +def test_reused_call_id_in_a_new_turn_keeps_each_result_separate(): + second = _response() + second.parts[0].function_response.response = {"value": True} + messages = build_pi_messages( + _request( + _call(), + _response(), + types.Content(role="user", parts=[types.Part(text="Do it again.")]), + _call(), + second, + ), + MODEL, + ) + assert [m["role"] for m in messages] == [ + "assistant", + "toolResult", + "user", + "assistant", + "toolResult", + ] + assert json.loads(messages[1]["content"][0]["text"]) == {"value": False} + assert json.loads(messages[4]["content"][0]["text"]) == {"value": True} + + +def test_progress_and_final_result_share_one_protocol_result(): + final = _response() + final.parts[0].function_response.response = {"value": "completed"} + messages = build_pi_messages(_request(_call(), _response(), final), MODEL) + assert [m["role"] for m in messages] == ["assistant", "toolResult"] + assert [json.loads(p["text"]) for p in messages[1]["content"]] == [ + {"value": False}, + {"value": "completed"}, + ] + + +@pytest.mark.parametrize("raise_error", [False, True]) +def test_context_file_is_private_isolated_and_cleaned(raise_error): + request = _request(_call(), _response()) + paths = [] + try: + with pi_context_extension(request, MODEL, prompt="startup") as first: + with pi_context_extension(request, MODEL, prompt="startup") as second: + paths = [Path(first), Path(second)] + assert paths[0].parent != paths[1].parent + if os.name == "posix": + assert stat.S_IMODE(paths[0].stat().st_mode) == 0o600 + assert stat.S_IMODE(paths[0].parent.stat().st_mode) == 0o700 + payload = json.loads( + json.loads( + paths[0] + .read_text() + .splitlines()[0] + .removeprefix("const context = JSON.parse(") + .removesuffix(");") + ) + ) + assert payload["messages"][-1]["role"] == "toolResult" + assert payload["instruction"] == "Wait for tool results." + assert MODEL.api_key not in paths[0].read_text() + if raise_error: + raise ValueError("probe failure") + except ValueError: + assert raise_error + assert paths and all(not path.parent.exists() for path in paths) + + +@pytest.mark.asyncio +async def test_context_file_is_cleaned_on_cancellation(): + ready = asyncio.Event() + paths = [] + + async def run(): + with pi_context_extension(_request(), MODEL, prompt="startup") as filename: + paths.append(Path(filename)) + ready.set() + await asyncio.Event().wait() + + task = asyncio.create_task(run()) + await ready.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not paths[0].parent.exists() diff --git a/tests/runtime/piagent/test_piagent_runtime.py b/tests/runtime/piagent/test_piagent_runtime.py index df60071de..b774a27cf 100644 --- a/tests/runtime/piagent/test_piagent_runtime.py +++ b/tests/runtime/piagent/test_piagent_runtime.py @@ -1870,6 +1870,7 @@ def before_model_callback(callback_context, llm_request): agent_dir = tmp_path / "agent-home" monkeypatch.setenv("PIAGENT_BINARY", str(binary)) monkeypatch.setenv("PIAGENT_AGENT_DIR", str(agent_dir)) + monkeypatch.setenv("PIAGENT_DISABLE_TOOLS", "1") agent = Agent( name="assistant", @@ -1887,7 +1888,12 @@ def before_model_callback(callback_context, llm_request): _events = [event async for event in PiAgentRuntime().run_async(agent, ctx)] argv = json.loads(argv_path.read_text(encoding="utf-8")) - assert "--extension" not in argv + extensions = [ + argv[index + 1] for index, arg in enumerate(argv) if arg == "--extension" + ] + assert [Path(path).name for path in extensions] == ["veadk-context.ts"] + assert "--no-tools" in argv + assert not Path(extensions[0]).exists() @pytest.mark.asyncio @@ -1921,7 +1927,11 @@ def before_model_callback(callback_context, llm_request): _events = [event async for event in PiAgentRuntime().run_async(agent, ctx)] argv = json.loads(argv_path.read_text(encoding="utf-8")) - assert "--extension" not in argv + extensions = [ + argv[index + 1] for index, arg in enumerate(argv) if arg == "--extension" + ] + assert [Path(path).name for path in extensions] == ["veadk-context.ts"] + assert not Path(extensions[0]).exists() @pytest.mark.asyncio diff --git a/tests/runtime/piagent/test_piagent_runtime_smoke.py b/tests/runtime/piagent/test_piagent_runtime_smoke.py index a7e71b143..062916498 100644 --- a/tests/runtime/piagent/test_piagent_runtime_smoke.py +++ b/tests/runtime/piagent/test_piagent_runtime_smoke.py @@ -30,9 +30,13 @@ from __future__ import annotations +import asyncio +import json import os +from pathlib import Path import pytest +from fastapi import FastAPI, Request from veadk import Agent, Runner @@ -76,3 +80,366 @@ async def test_real_piagent_runtime_smoke(): assert isinstance(result, str) assert result.strip() + + +@pytest.mark.piagent_smoke +@pytest.mark.asyncio +async def test_real_piagent_resumes_native_tool_results(monkeypatch): + """A real Pi process must receive tool results as protocol messages.""" + if os.getenv("PIAGENT_RUN_SMOKE") != "1": + pytest.skip("set PIAGENT_RUN_SMOKE=1 to spawn the real Pi binary") + binary = Path( + os.getenv("PIAGENT_BINARY", "~/.cache/veadk/piagent/pi/pi") + ).expanduser() + if not binary.is_file(): + pytest.skip("set PIAGENT_BINARY to a locally installed Pi binary") + + import uvicorn + from fastapi.responses import StreamingResponse + from google.adk.tools.long_running_tool import LongRunningFunctionTool + from google.genai import types + from veadk.runtime.piagent.client import PiAgentRpcClient + + monkeypatch.setenv("PIAGENT_BINARY", str(binary)) + monkeypatch.setenv("PIAGENT_DISABLE_BUILTIN_TOOLS", "1") + requests = [] + tool_metadata = {"__proto__": {"value": "must remain"}, "ordinary": "ok"} + pi_events = [] + original_prompt = PiAgentRpcClient.prompt + + async def capture_prompt(client, message): + async for event in original_prompt(client, message): + pi_events.append(event) + yield event + + monkeypatch.setattr(PiAgentRpcClient, "prompt", capture_prompt) + executions = [] + app = FastAPI() + + @app.post("/v1/chat/completions") + async def complete(request: Request): + body = await request.json() + requests.append(body) + latest = body["messages"][-1] + if latest["role"] == "tool": + delta = {"role": "assistant", "content": latest["content"]} + finish = "stop" + else: + delta = { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "id": f"call-{len(requests)}", + "type": "function", + "function": { + "name": "external_job", + "arguments": json.dumps( + {"city": "Beijing", "metadata": tool_metadata} + ), + }, + } + ], + } + finish = "tool_calls" + + def stream(): + for update, reason in ((delta, None), ({}, finish)): + payload = { + "id": "stub-response", + "object": "chat.completion.chunk", + "created": 1, + "model": body["model"], + "choices": [{"index": 0, "delta": update, "finish_reason": reason}], + } + yield f"data: {json.dumps(payload)}\n\n" + yield "data: [DONE]\n\n" + + return StreamingResponse(stream(), media_type="text/event-stream") + + server = uvicorn.Server( + uvicorn.Config( + app, host="127.0.0.1", port=0, log_level="warning", lifespan="off" + ) + ) + server.install_signal_handlers = lambda: None + serving = asyncio.create_task(server.serve()) + while not server.started: + if serving.done(): + serving.result() + await asyncio.sleep(0.01) + port = server.servers[0].sockets[0].getsockname()[1] + + async def external_job(city: str, metadata: dict) -> None: + """Start a job completed by the caller.""" + await asyncio.sleep(0) + assert metadata == tool_metadata + executions.append(city) + + agent = Agent( + name="assistant", + instruction="Wait for the external result before answering.", + runtime="piagent", + model_name="pi-native-history-smoke", + model_api_base=f"http://127.0.0.1:{port}/v1", + model_api_key="local-smoke-key", + model_api_key_name="", + tools=[LongRunningFunctionTool(external_job)], + ) + runner = Runner(agent=agent, app_name="pi_native_history_smoke") + await runner.short_term_memory.create_session( + app_name="pi_native_history_smoke", user_id="user", session_id="session" + ) + + async def send(content): + async with asyncio.timeout(40): + return [ + event + async for event in runner.run_async( + user_id="user", session_id="session", new_message=content + ) + ] + + try: + for index in range(2): + pending = await send( + types.Content( + role="user", parts=[types.Part(text="Start the Beijing job.")] + ) + ) + calls = [call for event in pending for call in event.get_function_calls()] + assert len(calls) == 1, [ + event.get("message", {}).get("errorMessage") + for event in pi_events + if isinstance(event.get("message"), dict) + and event["message"].get("errorMessage") + ] + result = f"external-result-{index}" + resumed = await send( + types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id=calls[0].id, + name="external_job", + response={"value": result}, + ) + ) + ], + ) + ) + assert not any(e.get_function_calls() for e in resumed) + assert any( + e.is_final_response() + and e.content + and any(result in (p.text or "") for p in e.content.parts or []) + for e in resumed + ) + latest = requests[-1]["messages"][-1] + assert latest["role"] == "tool" + assert latest["tool_call_id"] == calls[0].id + restored_calls = [ + call + for message in requests[-1]["messages"] + for call in message.get("tool_calls", []) + if call["id"] == calls[0].id + ] + assert len(restored_calls) == 1 + assert ( + json.loads(restored_calls[0]["function"]["arguments"])["metadata"] + == tool_metadata + ) + # Identical arguments are legitimate in the next user turn. + assert executions == ["Beijing"] * (index + 1) + resumed_requests = [ + body + for body in requests + if body["messages"][-1]["role"] == "tool" + and "external-result-" in body["messages"][-1]["content"] + ] + assert len(resumed_requests) == 2 + assert ( + "Wait for the external result before answering." + in requests[0]["messages"][0]["content"] + ) + finally: + server.should_exit = True + await serving + + +@pytest.mark.piagent_smoke +@pytest.mark.asyncio +@pytest.mark.parametrize("response_order", [("first", "second"), ("second", "first")]) +@pytest.mark.parametrize("resumable", [False, True]) +async def test_real_piagent_waits_for_all_external_results( + response_order, resumable, monkeypatch, tmp_path +): + """Partial results must never become Pi's synthetic missing-result message.""" + if os.getenv("PIAGENT_RUN_SMOKE") != "1": + pytest.skip("set PIAGENT_RUN_SMOKE=1 to spawn the real Pi binary") + binary = Path( + os.getenv("PIAGENT_BINARY", "~/.cache/veadk/piagent/pi/pi") + ).expanduser() + if not binary.is_file(): + pytest.skip("set PIAGENT_BINARY to a locally installed Pi binary") + + import uvicorn + from fastapi.responses import StreamingResponse + from google.adk.apps.app import App, ResumabilityConfig + from google.adk.events.event import Event + from google.adk.runners import Runner as ResumableRunner + from google.adk.sessions.in_memory_session_service import InMemorySessionService + from google.adk.tools.long_running_tool import LongRunningFunctionTool + from google.genai import types + + monkeypatch.setenv("PIAGENT_BINARY", str(binary)) + monkeypatch.setenv("PIAGENT_DISABLE_BUILTIN_TOOLS", "1") + monkeypatch.setenv("PIAGENT_AGENT_DIR", str(tmp_path / "pi-home")) + requests = [] + app = FastAPI() + + @app.post("/v1/chat/completions") + async def complete(request: Request): + body = await request.json() + requests.append(body) + + def stream(): + for delta, reason in [ + ({"role": "assistant", "content": "both complete"}, None), + ({}, "stop"), + ]: + payload = { + "id": "stub-response", + "object": "chat.completion.chunk", + "created": 1, + "model": body["model"], + "choices": [{"index": 0, "delta": delta, "finish_reason": reason}], + } + yield f"data: {json.dumps(payload)}\n\n" + yield "data: [DONE]\n\n" + + return StreamingResponse(stream(), media_type="text/event-stream") + + server = uvicorn.Server( + uvicorn.Config( + app, host="127.0.0.1", port=0, log_level="warning", lifespan="off" + ) + ) + server.install_signal_handlers = lambda: None + serving = asyncio.create_task(server.serve()) + while not server.started: + if serving.done(): + serving.result() + await asyncio.sleep(0.01) + port = server.servers[0].sockets[0].getsockname()[1] + + async def external_job(value: str) -> None: + """Wait for a result supplied by the caller.""" + raise AssertionError("Persisted calls must not execute again") + + agent = Agent( + name="assistant", + runtime="piagent", + model_name="partial-results-test", + model_api_base=f"http://127.0.0.1:{port}/v1", + model_api_key="loopback-key", + model_api_key_name="", + tools=[LongRunningFunctionTool(external_job)], + ) + sessions = InMemorySessionService() + if resumable: + runner = ResumableRunner( + app=App( + name="partial_results", + root_agent=agent, + resumability_config=ResumabilityConfig(is_resumable=True), + ), + session_service=sessions, + ) + else: + runner = ResumableRunner( + app_name="partial_results", agent=agent, session_service=sessions + ) + session = await sessions.create_session( + app_name="partial_results", user_id="test", session_id="test" + ) + invocation = "inv-partial-results" + await sessions.append_event( + session, + Event( + invocation_id=invocation, + author="user", + content=types.Content( + role="user", parts=[types.Part(text="Run both jobs.")] + ), + ), + ) + for call_id in ["first", "second"]: + await sessions.append_event( + session, + Event( + invocation_id=invocation, + author=agent.name, + long_running_tool_ids={call_id}, + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id=call_id, name="external_job", args={"value": call_id} + ) + ) + ], + ), + ), + ) + + async def send(call_id, will_continue=None): + message = types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id=call_id, + name="external_job", + response={"value": f"result-{call_id}"}, + will_continue=will_continue, + ) + ) + ], + ) + async with asyncio.timeout(40): + return [ + event + async for event in runner.run_async( + user_id="test", + session_id="test", + new_message=message, + ) + ] + + try: + first, last = response_order + assert await send(first) == [] + assert requests == [] + assert await send(last, True) == [] + assert requests == [] + events = await send(last, False) + assert len(requests) == 1 + assert not any(event.get_function_calls() for event in events) + messages = requests[0]["messages"] + tool_results = [message for message in messages if message["role"] == "tool"] + assert {message["tool_call_id"] for message in tool_results} == { + "first", + "second", + } + assert len(tool_results) == 2 + assert "No result provided" not in json.dumps(messages) + assert all( + f"result-{message['tool_call_id']}" in message["content"] + for message in tool_results + ) + finally: + server.should_exit = True + await serving diff --git a/tests/runtime/test_external_tool_continuation.py b/tests/runtime/test_external_tool_continuation.py index f34b94dbc..88c583326 100644 --- a/tests/runtime/test_external_tool_continuation.py +++ b/tests/runtime/test_external_tool_continuation.py @@ -16,6 +16,7 @@ import asyncio import importlib +import json import pytest from google.adk.agents.invocation_context import InvocationContext @@ -123,6 +124,89 @@ def test_tool_result_survives_prompt_translation(runtime, include_history): assert "finished-value" in prompt +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize( + "callback_positions", [(), (0,), (1,), (2,), (0, 2), (0, 1, 2)] +) +def test_codex_resume_includes_every_result_from_the_current_batch( + reverse, callback_positions +): + from veadk.runtime.codex.translate import build_prompt_from_llm_request + + first, second, third = _response(), _response(), _response() + first.parts[0].function_response.id = "call-first" + first.parts[0].function_response.response = {"value": "callback-updated"} + second.parts[0].function_response.id = "call-second" + third.parts[0].function_response.id = "call-third" + results = [first, second, third] + if reverse: + results.reverse() + for index in callback_positions: + results[index].parts.append(types.Part(text=f"callback note {index}")) + request = LlmRequest(contents=[_content("start the job"), *results]) + before = request.model_dump() + + prompt = build_prompt_from_llm_request(request, include_history=False) + + records = json.loads(prompt.split("\n")[1]) + assert [ + record["id"] for record in records if record["type"] == "function_response" + ] == [result.parts[0].function_response.id for result in results] + for index in range(len(results)): + assert (f"callback note {index}" in prompt) == (index in callback_positions) + assert "callback-updated" in prompt + assert "start the job" not in prompt + assert request.model_dump() == before + + +@pytest.mark.parametrize( + "boundary", ["text", "model_text", "reasoning", "call", "empty"] +) +def test_codex_resume_does_not_merge_results_across_a_turn_boundary(boundary): + from veadk.runtime.codex.translate import build_prompt_from_llm_request + + earlier, current = _response(), _response() + earlier.parts[0].function_response.response = {"value": "previous-batch"} + if boundary == "text": + separator = _content("do it again") + elif boundary == "empty": + separator = types.Content(role="user", parts=[]) + elif boundary == "model_text": + separator = types.Content(role="model", parts=[types.Part(text="done")]) + elif boundary == "reasoning": + separator = types.Content( + role="model", parts=[types.Part(text="next step", thought=True)] + ) + else: + separator = types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id="call-external", name="external_job" + ) + ) + ], + ) + prompt = build_prompt_from_llm_request( + LlmRequest(contents=[earlier, separator, current]), include_history=False + ) + records = json.loads(prompt.split("\n")[1]) + assert len(records) == 1 + assert "finished-value" in prompt + assert "previous-batch" not in prompt + + +def test_codex_new_user_turn_does_not_repeat_prior_tool_results(): + from veadk.runtime.codex.translate import build_prompt_from_llm_request + + prompt = build_prompt_from_llm_request( + LlmRequest(contents=[_response(), _response(), _content("next request")]), + include_history=False, + ) + assert prompt == "next request" + + @pytest.mark.asyncio @pytest.mark.parametrize("runtime", ["codex", "piagent"]) @pytest.mark.parametrize("kind", ["long_running", "deferred", "normal"]) diff --git a/tests/runtime/test_model_callbacks.py b/tests/runtime/test_model_callbacks.py index 43e770dcf..ec92576ed 100644 --- a/tests/runtime/test_model_callbacks.py +++ b/tests/runtime/test_model_callbacks.py @@ -173,6 +173,72 @@ def before_model_callback(callback_context, llm_request): assert event.actions.state_delta == {"guard": "hit"} +@pytest.mark.asyncio +async def test_finalized_event_keeps_emission_order_after_sqlite_reload(tmp_path): + from google.adk.sessions.database_session_service import DatabaseSessionService + + service = DatabaseSessionService( + db_url=f"sqlite+aiosqlite:///{tmp_path / 'events.sqlite'}" + ) + try: + session = await service.create_session( + app_name="order", user_id="user", session_id="session" + ) + agent = LlmAgent(name="assistant", model="test-model") + ctx = InvocationContext( + agent=agent, + invocation_id="inv-order", + session_service=service, + session=session, + ) + runtime_call = await build_runtime_llm_request(agent, ctx, model="test-model") + shell = runtime_call.model_response_event + shell.timestamp -= 1 + shell.actions.state_delta["callback_value"] = "kept" + original_timestamp = shell.timestamp + reasoning = Event( + author=agent.name, + invocation_id=ctx.invocation_id, + content=types.Content( + role="model", parts=[types.Part(text="reasoning", thought=True)] + ), + ) + final = llm_response_to_event( + runtime_call.llm_request, + LlmResponse(content=_content("final", role="model")), + shell, + ) + for event in [reasoning, final]: + await service.append_event(session, event) + restored = await service.get_session( + app_name="order", user_id="user", session_id="session" + ) + assert [event.id for event in restored.events] == [reasoning.id, final.id] + assert final.timestamp >= reasoning.timestamp + assert final.id == shell.id + assert final.actions.state_delta == {"callback_value": "kept"} + assert shell.timestamp == original_timestamp + finally: + await service.close() + + +@pytest.mark.parametrize("partial", [False, True]) +def test_finalizing_response_updates_time_without_changing_flags_or_shell(partial): + shell = Event(author="assistant", invocation_id="inv", branch="branch", timestamp=1) + shell.actions.state_delta["value"] = "kept" + finalized = llm_response_to_event( + SimpleNamespace(tools_dict={}), + LlmResponse(content=_content("response", role="model"), partial=partial), + shell, + ) + assert finalized.timestamp > shell.timestamp + assert finalized.id == shell.id + assert finalized.branch == shell.branch + assert finalized.actions == shell.actions + assert finalized.partial is partial + assert shell.timestamp == 1 + + @pytest.mark.asyncio async def test_plugin_before_model_short_circuits_agent_callbacks() -> None: calls: list[str] = [] diff --git a/tests/runtime/test_output_state.py b/tests/runtime/test_output_state.py index c359f815a..a23e2289a 100644 --- a/tests/runtime/test_output_state.py +++ b/tests/runtime/test_output_state.py @@ -17,9 +17,11 @@ from types import SimpleNamespace import pytest +from google.adk.agents.invocation_context import InvocationContext from google.adk.events.event import Event from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.sessions.session import Session from google.genai import types from veadk import Agent @@ -27,6 +29,15 @@ from veadk.runtime.output_state import maybe_save_output_to_state +def _empty_context(agent: Agent) -> InvocationContext: + return InvocationContext( + agent=agent, + invocation_id="inv-1", + session_service=InMemorySessionService(), + session=Session(id="session", app_name="test", user_id="test"), + ) + + def _text_event( text: str, *, @@ -189,7 +200,7 @@ async def run_async(self, agent, ctx): output_key="result", ) - events = [event async for event in agent._run_async_impl(SimpleNamespace())] + events = [event async for event in agent._run_async_impl(_empty_context(agent))] assert events == turn written = [ @@ -228,7 +239,7 @@ async def run_async(self, agent, ctx): output_key="result", ) - events = [event async for event in agent._run_async_impl(SimpleNamespace())] + events = [event async for event in agent._run_async_impl(_empty_context(agent))] assert [e.actions.state_delta["result"] for e in events] == [ "Let me check that.", diff --git a/tests/runtime/test_pending_tools.py b/tests/runtime/test_pending_tools.py new file mode 100644 index 000000000..fe1cad74c --- /dev/null +++ b/tests/runtime/test_pending_tools.py @@ -0,0 +1,212 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pending tool barriers belong to one agent in one resumable invocation.""" + +import pytest +from google.adk.agents.invocation_context import InvocationContext +from google.adk.apps.app import ResumabilityConfig +from google.adk.events.event import Event +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.sessions.session import Session +from google.genai import types + +from veadk import Agent +from veadk.runtime.pending_tools import has_pending_tool_calls + + +def _call(call_id="call", **kwargs): + return Event( + author=kwargs.pop("author", "assistant"), + invocation_id=kwargs.pop("invocation_id", "inv"), + long_running_tool_ids={call_id}, + content=types.Content( + role="model", + parts=[ + types.Part(function_call=types.FunctionCall(id=call_id, name="job")) + ], + ), + **kwargs, + ) + + +def _response(call_id="call", *, will_continue=None, **kwargs): + return Event( + author=kwargs.pop("author", "user"), + invocation_id=kwargs.pop("invocation_id", "inv"), + content=types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id=call_id, + name="job", + response={"value": "result", "isFinal": True}, + will_continue=will_continue, + ) + ) + ], + ), + **kwargs, + ) + + +@pytest.fixture +def context(): + agent = Agent(name="assistant", model_api_key="unit-test-key") + + def build(*events, branch=None, resumable=True): + return InvocationContext( + session_service=InMemorySessionService(), + invocation_id="inv", + agent=agent, + branch=branch, + resumability_config=ResumabilityConfig(is_resumable=resumable), + session=Session( + id="session", + app_name="test", + user_id="test", + state={}, + events=list(events), + ), + ) + + return build + + +@pytest.mark.parametrize("first_result", ["first", "second"]) +def test_partial_batch_stays_pending_in_either_response_order(context, first_result): + assert has_pending_tool_calls( + context(_call("first"), _call("second"), _response(first_result)) + ) + + +@pytest.mark.parametrize("will_continue", [None, False]) +def test_only_final_result_resolves_its_call(context, will_continue): + assert not has_pending_tool_calls( + context(_call(), _response(will_continue=will_continue)) + ) + assert has_pending_tool_calls(context(_call(), _response(will_continue=True))) + + +def test_progress_after_final_does_not_reopen_completed_call(context): + assert not has_pending_tool_calls( + context(_call(), _response(), _response(will_continue=True)) + ) + + +def test_multiple_questions_wait_until_every_answer_is_final(context): + events = [_call("a"), _call("b"), _call("c"), _response("b"), _response("a")] + assert has_pending_tool_calls(context(*events)) + assert has_pending_tool_calls(context(*events, _response("c", will_continue=True))) + assert not has_pending_tool_calls(context(*events, _response("c"))) + + +def test_partial_events_do_not_complete_or_create_pending_calls(context): + assert has_pending_tool_calls(context(_call(), _response(partial=True))) + assert not has_pending_tool_calls(context(_call(partial=True))) + + +def test_old_response_cannot_resolve_reused_call_id(context): + assert has_pending_tool_calls(context(_call(), _response(), _call())) + assert not has_pending_tool_calls( + context(_call(), _response(), _call(), _response()) + ) + + +def test_only_current_agent_can_own_a_pending_call(context): + assert not has_pending_tool_calls(context(_call(author="other"))) + assert has_pending_tool_calls(context(_call(), _response(author="other"))) + + +def test_parallel_branch_does_not_block_or_resolve_current_branch(context): + assert not has_pending_tool_calls(context(_call(branch="other"), branch="current")) + assert has_pending_tool_calls( + context(_call(branch="current"), _response(branch="other"), branch="current") + ) + assert not has_pending_tool_calls( + context(_call(branch="current"), _response(), branch="current") + ) + + +def test_old_invocation_after_cancel_cannot_block_new_turn(context): + assert not has_pending_tool_calls(context(_call(invocation_id="cancelled-inv"))) + assert has_pending_tool_calls(context(_call(), _response(invocation_id="old-inv"))) + + +def test_unmarked_regular_tools_do_not_block_but_default_runner_tools_do(context): + call = _call() + call.long_running_tool_ids = None + assert not has_pending_tool_calls(context(call)) + assert has_pending_tool_calls(context(_call(), resumable=False)) + + +def test_final_text_or_business_flags_do_not_resolve_another_call(context): + text = Event( + author="assistant", + invocation_id="inv", + content=types.Content(role="model", parts=[types.Part(text="Still waiting.")]), + ) + assert text.is_final_response() + assert has_pending_tool_calls(context(_call(), text, _response("unrelated"))) + + +def test_nested_tool_input_can_resume_parent_without_unblocking_siblings(context): + call = _call("parent", branch="current") + sibling = _call("sibling", branch="current") + call.content.parts.extend(sibling.content.parts) + call.long_running_tool_ids.add("sibling") + nested_input = Event( + author="user", + invocation_id="inv", + branch="current.job@parent.child", + content=types.Content(role="user", parts=[types.Part(text="Nested answer.")]), + ) + assert has_pending_tool_calls(context(call, nested_input, branch="current")) + assert not has_pending_tool_calls( + context(call, nested_input, _response("sibling"), branch="current") + ) + nested_input.invocation_id = "other-invocation" + assert has_pending_tool_calls( + context(call, nested_input, _response("sibling"), branch="current") + ) + + +def test_explicit_invocation_wins_over_reused_ids_in_a_later_invocation(context): + response = _response("a") + ctx = context( + _call("a"), + _call("b"), + _call("a", invocation_id="later-inv"), + _response("a", invocation_id="later-inv"), + response, + resumable=False, + ) + ctx.user_content = response.content + assert has_pending_tool_calls(ctx) + + +@pytest.mark.parametrize("branch", [None, "root"]) +def test_explicit_child_response_cannot_resolve_root_call_with_reused_id( + context, branch +): + child_branch = f"{branch or 'other'}.child" + assert has_pending_tool_calls( + context( + _call(branch=branch), + _call(author="child", branch=child_branch), + _response(branch=child_branch), + branch=branch, + ) + ) diff --git a/veadk/agent.py b/veadk/agent.py index c8fbf8893..c39202443 100644 --- a/veadk/agent.py +++ b/veadk/agent.py @@ -880,6 +880,11 @@ async def _run_async_impl( external agent harness (e.g. the Claude Code SDK) back into the ADK event stream, so the surrounding ``Runner`` is unaffected. """ + from veadk.runtime.pending_tools import has_pending_tool_calls + + if has_pending_tool_calls(ctx): + return + if self.runtime == "adk": # A transfer can close this wrapper before the LLM stream ends. # Close it in the same task so tracing contexts do not leak into @@ -891,6 +896,7 @@ async def _run_async_impl( from veadk.runtime import get_runtime from veadk.runtime.compat import check_agent_runtime_support + from veadk.runtime.model_callbacks import is_final_model_text_event from veadk.runtime.output_state import maybe_save_output_to_state # Authoritative support-matrix gate. It has to live here rather than in @@ -904,10 +910,27 @@ async def _run_async_impl( run_config=getattr(ctx, "run_config", None), ) + completed = False async for event in get_runtime(self.runtime).run_async(self, ctx): maybe_save_output_to_state(self, event) + if ( + event.author == self.name + and not event.partial + and not event.long_running_tool_ids + and ( + is_final_model_text_event(event, self.name) + or event.actions.skip_summarization + ) + ): + completed = True yield event + if ctx.is_resumable and completed and not has_pending_tool_calls(ctx): + # Parallel/workflow parents use this persisted ADK state to avoid + # rerunning finished siblings when another branch's tools resume. + ctx.set_agent_state(self.name, end_of_agent=True) + yield self._create_agent_state_event(ctx) + async def _run_live_impl( self, ctx: "InvocationContext" ) -> AsyncGenerator["Event", None]: diff --git a/veadk/runtime/codex/runtime.py b/veadk/runtime/codex/runtime.py index f8bbe0208..fb1ab0462 100644 --- a/veadk/runtime/codex/runtime.py +++ b/veadk/runtime/codex/runtime.py @@ -1242,6 +1242,7 @@ async def _pump_codex() -> None: if queued is _TOOL_PENDING: if not tool_pending: tool_pending = True + bridge_interrupted = True final_text_events.clear() await interrupt_turn( turn, diff --git a/veadk/runtime/codex/translate.py b/veadk/runtime/codex/translate.py index 4fd41a0ce..21f9e2250 100644 --- a/veadk/runtime/codex/translate.py +++ b/veadk/runtime/codex/translate.py @@ -122,22 +122,37 @@ def build_prompt_from_llm_request( ) -> str: """Render callback-mutated LlmRequest contents into Codex turn input. - ``include_history=False`` renders the current message only, for a resumed - Codex thread that already holds the earlier turns itself. + ``include_history=False`` renders the current message and its adjacent + external results; a resumed thread has not seen a partially answered batch. """ records = [_content_event_record(content) for content in llm_request.contents] - records = [record for record in records if record["parts"]] - if not records: + visible = [index for index, record in enumerate(records) if record["parts"]] + if not visible: return NO_TEXT_PROMPT - current_record = records[-1]["parts"] + current_record = records[visible[-1]]["parts"] + if not include_history and any( + part.get("type") == "function_response" for part in current_record + ): + # Keep model/reasoning and ordinary user messages as batch boundaries, + # including records whose text is hidden from the serialized prompt. + for record in reversed(records[: visible[-1]]): + if ( + record["role"] != "user" + or not record["parts"] + or not any( + part["type"] == "function_response" for part in record["parts"] + ) + ): + break + current_record = record["parts"] + current_record current_text = "\n".join( str(part["text"]) for part in current_record if part.get("type") == "text" and part.get("text") ).strip() - history = records[:-1] if include_history else [] + history = [records[index] for index in visible[:-1]] if include_history else [] continuation = ( "\nContinue the paused request using these tool results; " "this is not a new user request." diff --git a/veadk/runtime/model_callbacks.py b/veadk/runtime/model_callbacks.py index da4a3142b..fe51beb5e 100644 --- a/veadk/runtime/model_callbacks.py +++ b/veadk/runtime/model_callbacks.py @@ -18,6 +18,7 @@ import copy import inspect +import time from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -237,6 +238,9 @@ def llm_response_to_event( { **model_response_event.model_dump(exclude_none=True), **llm_response.model_dump(exclude_none=True), + # The callback shell predates streaming; persisted order must + # reflect when this response is finalized, not when the turn began. + "timestamp": time.time(), } ) diff --git a/veadk/runtime/pending_tools.py b/veadk/runtime/pending_tools.py new file mode 100644 index 000000000..b8f0f3ebd --- /dev/null +++ b/veadk/runtime/pending_tools.py @@ -0,0 +1,94 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Keep partial external-tool results from restarting an agent's model.""" + +from google.adk.agents.invocation_context import InvocationContext +from google.adk.events.event import Event + + +def has_pending_tool_calls(ctx: InvocationContext) -> bool: + """Wait for this agent's unresolved calls, preserving nested branch resume.""" + branch_events = [ + event for event in ctx._get_events(current_branch=True) if not event.partial + ] + invocation_id = ctx.invocation_id + if not ctx.is_resumable and ctx.user_content: + response_ids = { + part.function_response.id + for part in ctx.user_content.parts or [] + if part.function_response and part.function_response.id + } + # The default Runner gives each submitted result a new invocation ID. + # An explicit result still resumes its original call's external batch. + matching_calls = [ + event + for event in reversed(branch_events) + if event.author == ctx.agent.name + and any(call.id in response_ids for call in event.get_function_calls()) + ] + if matching_calls: + owner = next( + ( + event + for event in matching_calls + if event.invocation_id == invocation_id + ), + matching_calls[0], + ) + invocation_id = owner.invocation_id + + external_resume = invocation_id != ctx.invocation_id + pending: dict[str, Event] = {} + owned_ids: set[str] = set() + events = [] + for event in branch_events: + owns_invocation = event.invocation_id == invocation_id + responses = event.get_function_responses() + if not owns_invocation and not ( + external_resume + and event.author in {"user", ctx.agent.name} + and any(response.id in owned_ids for response in responses) + ): + continue + events.append(event) + if owns_invocation and event.author == ctx.agent.name: + for call in event.get_function_calls(): + if call.id and call.id in (event.long_running_tool_ids or ()): + pending[call.id] = event + owned_ids.add(call.id) + if event.author in {"user", ctx.agent.name}: + for response in responses: + owner = pending.get(response.id) + if ( + owner is not None + and response.will_continue is not True + and (event.branch is None or event.branch == owner.branch) + ): + pending.pop(response.id, None) + + if not pending: + return False + + # ADK exempts a parent tool when input is resuming its nested branch. Scope + # that check to this invocation, and to only the still-unresolved call IDs. + scoped_ctx = ctx.model_copy( + update={"session": ctx.session.model_copy(update={"events": events})} + ) + return any( + scoped_ctx.should_pause_invocation( + event.model_copy(update={"long_running_tool_ids": {call_id}}) + ) + for call_id, event in pending.items() + ) diff --git a/veadk/runtime/piagent/context_bridge.py b/veadk/runtime/piagent/context_bridge.py new file mode 100644 index 000000000..abc029706 --- /dev/null +++ b/veadk/runtime/piagent/context_bridge.py @@ -0,0 +1,260 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Preserve ADK conversation roles across stateless Pi invocations.""" + +from __future__ import annotations + +import json +import os +import tempfile +from collections import defaultdict, deque +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator + +from google.adk.models.llm_request import LlmRequest + +from veadk.runtime.piagent.config import PiAgentModelConfig +from veadk.runtime.piagent.translate import _system_instruction_text + + +def build_pi_messages( + request: LlmRequest, + model: PiAgentModelConfig, + *, + tool_names: dict[str, str] | None = None, +) -> list[dict[str, Any]]: + """Keep tool call IDs and results in Pi's native message protocol.""" + messages: list[dict[str, Any]] = [] + tool_names = tool_names or {} + pending_ids: dict[str, list[str]] = {} + call_number = 0 + used_ids = { + part.function_call.id + for content in request.contents + for part in content.parts or [] + if part.function_call and part.function_call.id + } + for content in request.contents: + blocks: list[dict[str, Any]] = [] + + def flush() -> None: + if not blocks: + return + has_calls = any(block["type"] == "toolCall" for block in blocks) + message: dict[str, Any] = { + "role": "user" + if content.role == "user" and not has_calls + else "assistant", + "content": list(blocks), + "timestamp": 0, + } + if message["role"] == "assistant": + message.update( + api=model.api, + provider=model.provider_id, + model=model.model, + stopReason="toolUse" if has_calls else "stop", + usage={ + "input": 0, + "output": 0, + "cacheRead": 0, + "cacheWrite": 0, + "totalTokens": 0, + "cost": { + "input": 0, + "output": 0, + "cacheRead": 0, + "cacheWrite": 0, + "total": 0, + }, + }, + ) + if ( + message["role"] == "assistant" + and messages + and messages[-1]["role"] == "assistant" + ): + previous = messages[-1] + previous["content"].extend(message["content"]) + if has_calls: + previous["stopReason"] = "toolUse" + else: + messages.append(message) + blocks.clear() + + for part in content.parts or []: + if part.text and not part.thought: + blocks.append({"type": "text", "text": part.text}) + elif call := part.function_call: + while not call.id and f"veadk-call-{call_number}" in used_ids: + call_number += 1 + call_id = call.id or f"veadk-call-{call_number}" + call_number += 1 + used_ids.add(call_id) + pending_ids.setdefault(call.name or "", []).append(call_id) + blocks.append( + { + "type": "toolCall", + "id": call_id, + "name": tool_names.get(call.name or "", call.name), + "arguments": call.model_dump(mode="json").get("args") or {}, + } + ) + elif response := part.function_response: + flush() + pending = pending_ids.get(response.name or "", []) + call_id = response.id or (pending[0] if pending else None) + if call_id in pending: + pending.remove(call_id) + messages.append( + { + "role": "toolResult", + "toolCallId": call_id, + "toolName": tool_names.get(response.name or "", response.name), + "content": [ + { + "type": "text", + "text": json.dumps( + response.model_dump(mode="json")["response"], + ensure_ascii=False, + ), + } + ], + "isError": False, + "timestamp": 0, + } + ) + flush() + return _pair_tool_results(messages) + + +def _pair_tool_results(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Providers require a call's results before the next user/assistant message. + + ADK may emit each parallel call separately, mix user text with results, or + deliver a result after intervening user input. Keep user order intact while + moving each result to its owning call, rather than letting Pi fabricate a + missing-result placeholder and later repeat the real result. + """ + results: dict[str, deque[tuple[int, dict[str, Any]]]] = defaultdict(deque) + next_call: dict[str, int] = {} + boundaries: dict[tuple[int, str], int] = {} + for index in range(len(messages) - 1, -1, -1): + message = messages[index] + if message["role"] == "assistant": + for part in message["content"]: + if part["type"] == "toolCall": + call_id = part["id"] + boundaries[index, call_id] = next_call.get(call_id, len(messages)) + next_call[call_id] = index + for index, message in enumerate(messages): + if message["role"] == "toolResult": + results[message["toolCallId"]].append((index, message)) + + consumed: set[int] = set() + ordered: list[dict[str, Any]] = [] + for index, message in enumerate(messages): + if index in consumed: + continue + ordered.append(message) + if message["role"] != "assistant": + continue + following = [] + for part in message["content"]: + if part["type"] != "toolCall": + continue + call_id = part["id"] + candidates = results[call_id] + while candidates and candidates[0][0] < index: + candidates.popleft() + updates = [] + while candidates and candidates[0][0] < boundaries[index, call_id]: + position, result = candidates.popleft() + consumed.add(position) + updates.append((position, result)) + if updates: + position, result = updates[0] + # Long-running tools can send progress and a final result for + # one call ID. Preserve every update within one protocol result. + following.append( + ( + position, + { + **result, + "content": [ + block + for _, update in updates + for block in update["content"] + ], + }, + ) + ) + ordered.extend( + result for _, result in sorted(following, key=lambda item: item[0]) + ) + return ordered + + +@contextmanager +def pi_context_extension( + request: LlmRequest, + model: PiAgentModelConfig, + *, + prompt: str, + tool_names: dict[str, str] | None = None, +) -> Iterator[str]: + """Replace only the RPC's startup message, preserving later Pi messages.""" + payload = { + "messages": build_pi_messages(request, model, tool_names=tool_names), + "prompt": prompt, + "instruction": _system_instruction_text(request.config.system_instruction), + } + # The extension contains conversation and tool-result data, so neither its + # directory nor its file may be accessible to another local user. + with tempfile.TemporaryDirectory(prefix="veadk-piagent-context-") as directory: + path = Path(directory) / "veadk-context.ts" + source = ( + "const context = JSON.parse(" + + json.dumps(json.dumps(payload, ensure_ascii=True)) + + ");\n" + ) + source += """ +export default function(pi) { + if (context.instruction) { + pi.on("before_agent_start", event => ({ + systemPrompt: event.systemPrompt + "\\n\\n" + context.instruction, + })); + } + pi.on("context", event => { + const index = event.messages.findIndex(message => { + if (message.role !== "user") return false; + const text = typeof message.content === "string" ? message.content : + message.content.filter(part => part.type === "text").map(part => part.text).join(""); + return text === context.prompt; + }); + if (index < 0) return; + return {messages: [ + ...event.messages.slice(0, index), + ...structuredClone(context.messages), + ...event.messages.slice(index + 1), + ]}; + }); +} +""" + descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(descriptor, "w", encoding="utf-8") as stream: + stream.write(source) + yield str(path) diff --git a/veadk/runtime/piagent/runtime.py b/veadk/runtime/piagent/runtime.py index e5ab6ddbb..5dde9f221 100644 --- a/veadk/runtime/piagent/runtime.py +++ b/veadk/runtime/piagent/runtime.py @@ -17,6 +17,8 @@ from __future__ import annotations import asyncio +import contextlib +from dataclasses import replace from typing import TYPE_CHECKING, AsyncGenerator from google.adk.agents.invocation_context import LlmCallsLimitExceededError @@ -42,6 +44,7 @@ ) from veadk.runtime.piagent.client import PiAgentRpcClient from veadk.runtime.piagent.config import PiAgentConfig, prepare_piagent_home +from veadk.runtime.piagent.context_bridge import pi_context_extension from veadk.runtime.piagent.installer import resolve_or_install_piagent_binary from veadk.runtime.piagent.skills import materialize_skills_for_pi from veadk.runtime.piagent.tool_runtime import PiToolRuntime @@ -176,7 +179,18 @@ async def _on_tool_pending() -> None: # that bookkeeping somewhere real to land; see # `_merge_turn_bookkeeping`. merge_target: "Event | None" = None - async with PiToolRuntime(tool_bundle) as tools: + async with contextlib.AsyncExitStack() as resources: + context_extension = resources.enter_context( + pi_context_extension( + runtime_call.llm_request, + config.model, + prompt=prompt, + tool_names={ + spec.original_name: spec.name for spec in tool_bundle.specs + }, + ) + ) + tools = await resources.enter_async_context(PiToolRuntime(tool_bundle)) run_config = ( config.with_skills(skill_paths=list(skill_bundle.paths)) if skill_bundle.paths @@ -187,6 +201,10 @@ async def _on_tool_pending() -> None: if tools.enabled else run_config ) + run_config = replace( + run_config, + extensions=(*run_config.extensions, context_extension), + ) async with PiAgentRpcClient(run_config) as client: pump: asyncio.Task[None] | None = None