fix(runtime): resume pending tools in Codex and PiAgent

tests/runtime/codex/test_codex_runtime_smoke.py

@@ -818,3 +818,118 @@ def do_POST(self) -> None:  # noqa: N802 - http.server API
 async def _drain(turn: Any) -> None:
     async for _ in turn.stream():
         pass
+
+
+@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."""
+    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.tools.long_running_tool import LongRunningFunctionTool
+    from google.genai import types
+
+    from veadk import Agent, Runner
+    from veadk.runtime.codex import runtime as runtime_module
+    from veadk.runtime.codex.proxy import ResponsesShim
+
+    executed = []
+
+    async def veadk_smoke_record_marker(marker: str) -> None:
+        """Start a job whose result arrives from outside this process."""
+        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()
+
+    async def lookup(*_args, **_kwargs):
+        return shim
+
+    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_key="veadk-smoke-key",
+        model_api_key_name="",
+        tools=[LongRunningFunctionTool(veadk_smoke_record_marker)],
+        codex_runtime_config={
+            "model_transport": "shim",
+            "sandbox": "read_only",
+            "network_access": False,
+            "turn_timeout_seconds": 30,
+        },
+    )
+    runner = Runner(agent=agent, app_name="external_tool_smoke")
+    session_id = f"external-tool-{uuid.uuid4().hex[:8]}"
+    await runner.short_term_memory.create_session(
+        app_name="external_tool_smoke", user_id="user", session_id=session_id
+    )
+
+    async def send(content):
+        events = []
+        async with contextlib.aclosing(
+            runner.run_async(user_id="user", session_id=session_id, new_message=content)
+        ) as stream:
+            async for event in stream:
+                events.append(event)
+        return events
+
+    try:
+        first = await asyncio.wait_for(
+            send(
+                types.Content(
+                    role="user", parts=[types.Part(text="Start the external job.")]
+                )
+            ),
+            _RUN_TIMEOUT_SECONDS,
+        )
+        calls = [
+            call
+            for event in first
+            for call in event.get_function_calls()
+            if call.name == _ADK_TOOL_NAME
+        ]
+        assert len(calls) == 1 and executed == [_TOOL_MARKER]
+        assert not any(
+            event.is_final_response()
+            and event.content
+            and any(part.text for part in event.content.parts or [])
+            for event in first
+        )
+        result = types.Content(
+            role="user",
+            parts=[
+                types.Part(
+                    function_response=types.FunctionResponse(
+                        id=calls[0].id,
+                        name=_ADK_TOOL_NAME,
+                        response={"value": "external-smoke-completed"},
+                    )
+                )
+            ],
+        )
+        second = await asyncio.wait_for(send(result), _RUN_TIMEOUT_SECONDS)
+        assert any(
+            event.is_final_response()
+            and event.content
+            and any(
+                _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()
+        await backend.stop()
+    assert shim._turns == {}

tests/runtime/codex/test_codex_turn_contract.py

@@ -1156,12 +1156,16 @@ async def _direct_session(
     )
     runner = Runner(app_name="contract", agent=agent, session_service=session_service)
 
-    async def send(text: str) -> BaseException | None:
+    async def send(text: str | types.Content) -> BaseException | None:
         try:
             async for _ in runner.run_async(
                 user_id="user",
                 session_id=session_id,
-                new_message=types.Content(role="user", parts=[types.Part(text=text)]),
+                new_message=(
+                    text
+                    if isinstance(text, types.Content)
+                    else types.Content(role="user", parts=[types.Part(text=text)])
+                ),
             ):
                 pass
         except BaseException as e:  # noqa: BLE001 - the error is the observable
@@ -2133,20 +2137,17 @@ async def fail_tool_setup(*args: Any, **kwargs: Any) -> Any:
         raise RuntimeError("simulated tool setup failure")
 
     monkeypatch.setattr(runtime_module, "build_executable_tools", fail_tool_setup)
+    # This exercises permit cleanup, not setup speed on a loaded CI worker.
+    # Admission and setup deadlines have their own tests below.
     _events, _session, _backend, error = await _run_turn(
         monkeypatch,
         plan=(Round(text="never reached"),),
-        agent_kwargs={
-            "codex_runtime_config": {
-                "turn_timeout_seconds": 0.05,
-            }
-        },
     )
     assert isinstance(error, RuntimeError), error
     assert str(error) == "simulated tool setup failure"
 
     async def acquire_from_another_task() -> None:
-        async with slots.hold(timeout=0.05):
+        async with slots.hold(timeout=5):
             pass
 
     # Child tasks inherit ContextVars, so this also proves reentrancy is
@@ -2253,3 +2254,119 @@ async def acquire_from_another_task() -> None:
             pass
 
     await asyncio.create_task(acquire_from_another_task())
+
+
+@pytest.mark.asyncio
+async def test_long_running_tool_pauses_after_async_initialization(monkeypatch):
+    from google.adk.tools.long_running_tool import LongRunningFunctionTool
+
+    initialized = []
+
+    async def external_job() -> None:
+        """Start a job completed outside the agent process."""
+        await asyncio.sleep(0.01)
+        initialized.append("started")
+
+    def after_model_callback(callback_context, llm_response):
+        callback_context.state["after_model"] = True
+
+    events, session, _backend, error = await _run_direct_turn(
+        monkeypatch,
+        plan=(
+            Round(tool_calls=(("external_job", {}),), usage=(2, 1)),
+            Round(text="premature answer", usage=(2, 1)),
+        ),
+        agent_kwargs={
+            "tools": [LongRunningFunctionTool(external_job)],
+            "after_model_callback": after_model_callback,
+        },
+    )
+    assert error is None, error
+    assert initialized == ["started"]
+    assert session.state["after_model"] is True
+    assert any(e.long_running_tool_ids for e in events)
+    assert not any(e.get_function_responses() for e in events)
+    assert not any(
+        e.is_final_response()
+        and e.content
+        and any(p.text for p in e.content.parts or [])
+        for e in events
+    )
+
+
+@pytest.mark.asyncio
+async def test_warm_process_resumes_external_result_without_reexecuting_tool(
+    monkeypatch,
+):
+    from google.adk.tools.long_running_tool import LongRunningFunctionTool
+    from veadk.runtime.codex.mcp_bridge import shutdown_bridge
+
+    executions = []
+
+    async def external_job() -> None:
+        """Start a job completed externally."""
+        await asyncio.sleep(0)
+        executions.append("started")
+
+    codex_class = _warm_capable_codex()
+    send, backend, _, _ = await _direct_session(
+        monkeypatch,
+        (
+            Round(tool_calls=(("external_job", {}),), usage=(2, 1)),
+            Round(text="received result", usage=(2, 1)),
+            Round(text="remembered result", usage=(2, 1)),
+        ),
+        tools=[LongRunningFunctionTool(external_job)],
+        codex_class=codex_class,
+    )
+    try:
+        assert await send("start the job") is None
+        assert len(backend.calls) == 1
+        response = types.Content(
+            role="user",
+            parts=[
+                types.Part(
+                    function_response=types.FunctionResponse(
+                        id="call-0-0",
+                        name="external_job",
+                        response={"value": "external-result-value"},
+                    )
+                )
+            ],
+        )
+        assert await send(response) is None
+        assert "external-result-value" in json.dumps(
+            backend.raw_requests[-1], default=str
+        )
+        assert await send("what was the result?") is None
+        assert "external-result-value" in json.dumps(
+            backend.raw_requests[-1], default=str
+        )
+        assert executions == ["started"]
+        # Pending bridge calls retire the first process; the restored thread
+        # can then reuse its replacement for ordinary turns.
+        assert _starts_and_resumes(codex_class) == [(1, 0), (0, 1)]
+        assert codex_class.instances[0].closed
+        assert not codex_class.instances[1].closed
+    finally:
+        await _close_process_pool()
+        await shutdown_bridge()
+
+
+@pytest.mark.asyncio
+async def test_pending_tool_does_not_hide_a_broken_codex_stream(monkeypatch):
+    from google.adk.tools.long_running_tool import LongRunningFunctionTool
+
+    async def external_job() -> None:
+        """Wait for an external result."""
+        return None
+
+    events, _session, _backend, error = await _run_turn(
+        monkeypatch,
+        plan=(Round(tool_calls=(("external_job", {}),)), Round(text="discard")),
+        agent_kwargs={"tools": [LongRunningFunctionTool(external_job)]},
+        codex_class=_BrokenStreamCodex,
+    )
+    assert isinstance(error, RuntimeError)
+    assert str(error) == "codex stream dropped"
+    assert any(event.long_running_tool_ids for event in events)

tests/runtime/piagent/test_piagent_runtime.py

@@ -2039,3 +2039,92 @@ async def test_piagent_runtime_loads_and_cleans_materialized_skills(
     assert "--no-skills" in argv
     assert skill_path.name == "demo-skill"
     assert not skill_path.exists()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("returns_initial_result", [False, True])
+async def test_long_running_bridge_pauses_then_resumes_with_history(
+    tmp_path, monkeypatch, returns_initial_result
+):
+    from google.adk.tools.long_running_tool import LongRunningFunctionTool
+
+    initialized = []
+
+    async def external_job(city: str):
+        """Start an external job for a city."""
+        await asyncio.sleep(0.01)
+        initialized.append(city)
+        return {"job_id": "job-1"} if returns_initial_result else None
+
+    def after_model_callback(callback_context, llm_response):
+        callback_context.state["after_model"] = True
+
+    binary = _make_fake_pi_that_calls_bridge(
+        tmp_path, tool_name="external_job", text_after_tool="premature answer"
+    )
+    monkeypatch.setenv("PIAGENT_BINARY", str(binary))
+    monkeypatch.setenv("PIAGENT_AGENT_DIR", str(tmp_path / "agent-home"))
+    agent = Agent(
+        name="assistant",
+        instruction="Answer.",
+        model_name="model-a",
+        model_api_base="https://ark.example.com/api/v3/",
+        model_api_key="test-key",
+        model_api_key_name="",
+        runtime="piagent",
+        tools=[LongRunningFunctionTool(external_job)],
+        after_model_callback=after_model_callback,
+    )
+    ctx = _ctx(agent, _user_event("start the job"))
+    events = [e async for e in PiAgentRuntime().run_async(agent, ctx)]
+    assert initialized == ["Beijing"]
+    assert any(e.long_running_tool_ids for e in events)
+    assert any(e.actions.state_delta.get("after_model") for e in events)
+    assert (
+        bool(
+            [
+                e
+                for e in events
+                if e.is_final_response()
+                and e.content
+                and any(p.text for p in e.content.parts or [])
+            ]
+        )
+        is returns_initial_result
+    )
+    assert not any(
+        not returns_initial_result
+        and e.is_final_response()
+        and e.content
+        and any(p.text for p in e.content.parts or [])
+        for e in events
+    )
+    assert (
+        bool([e for e in events if e.get_function_responses()])
+        is returns_initial_result
+    )
+
+    binary, prompt_path = _make_fake_pi_with_prompt_capture(tmp_path)
+    response = Event(
+        author="user",
+        invocation_id=ctx.invocation_id,
+        content=types.Content(
+            role="user",
+            parts=[
+                types.Part(
+                    function_response=types.FunctionResponse(
+                        id="call-weather",
+                        name="external_job",
+                        response={"value": "finished-value"},
+                    )
+                )
+            ],
+        ),
+    )
+    ctx.session.events.extend(events + [response])
+    resumed = [e async for e in PiAgentRuntime().run_async(agent, ctx)]
+    assert any(e.is_final_response() for e in resumed)
+    prompt = prompt_path.read_text()
+    assert "finished-value" in prompt and "call-weather" in prompt
+    assert prompt.count("start the job") == 1
+    assert initialized == ["Beijing"]