diff --git a/tests/runtime/codex/test_codex_runtime.py b/tests/runtime/codex/test_codex_runtime.py index bd721c4ea..3c13eab44 100644 --- a/tests/runtime/codex/test_codex_runtime.py +++ b/tests/runtime/codex/test_codex_runtime.py @@ -23,7 +23,9 @@ import httpx import pytest +from google.adk.agents import RunConfig from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.invocation_context import LlmCallsLimitExceededError from google.adk.agents.llm_agent import LlmAgent from google.adk.auth.auth_credential import AuthCredential from google.adk.auth.auth_credential import AuthCredentialTypes @@ -867,6 +869,128 @@ async def test_shim_rejects_unknown_invocation_token() -> None: assert response.status_code == 401 +@pytest.mark.asyncio +async def test_shim_counts_each_backend_call_before_tool_loop_retry( + monkeypatch, +) -> None: + shim = ResponsesShim("https://backend.invalid/v1", "backend-key") + agent = LlmAgent(name="assistant", model="model") + ctx = InvocationContext( + session_service=InMemorySessionService(), + invocation_id="inv-1", + agent=agent, + session=Session( + id="session-1", + appName="app", + userId="user", + state={}, + events=[], + ), + run_config=RunConfig(max_llm_calls=1), + ) + backend_calls = 0 + + async def executor(args, call_id): + return "{}" + + token = shim.register_turn( + [{"type": "function", "name": "loop", "parameters": {}}], + {"loop": executor}, + before_model_call=ctx.increment_llm_call_count, + ) + + async def fake_aresponses(**kwargs): + nonlocal backend_calls + backend_calls += 1 + if backend_calls > 1: + raise AssertionError("second backend call should be blocked") + return { + "id": "tool", + "model": "model", + "output": [ + { + "id": "fc", + "call_id": "call-loop", + "type": "function_call", + "name": "loop", + "arguments": "{}", + "status": "completed", + } + ], + } + + monkeypatch.setattr("veadk.runtime.codex.proxy.litellm.aresponses", fake_aresponses) + transport = httpx.ASGITransport(app=shim._app) + try: + async with httpx.AsyncClient( + transport=transport, base_url="http://shim" + ) as client: + with pytest.raises(LlmCallsLimitExceededError): + await client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {token}"}, + json={ + "model": "model", + "input": [{"type": "message", "role": "user", "content": "go"}], + }, + ) + finally: + shim.unregister_turn(token) + + assert backend_calls == 1 + + +@pytest.mark.asyncio +async def test_shim_counts_one_plain_backend_call(monkeypatch) -> None: + shim = ResponsesShim("https://backend.invalid/v1", "backend-key") + backend_calls = 0 + counted_calls = 0 + + def before_model_call() -> None: + nonlocal counted_calls + counted_calls += 1 + + token = shim.register_turn([], {}, before_model_call=before_model_call) + + async def fake_aresponses(**kwargs): + nonlocal backend_calls + backend_calls += 1 + return { + "id": "final", + "model": "model", + "output": [ + { + "id": "msg", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done"}], + } + ], + } + + monkeypatch.setattr("veadk.runtime.codex.proxy.litellm.aresponses", fake_aresponses) + transport = httpx.ASGITransport(app=shim._app) + try: + async with httpx.AsyncClient( + transport=transport, base_url="http://shim" + ) as client: + response = await client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {token}"}, + json={ + "model": "model", + "input": [{"type": "message", "role": "user", "content": "go"}], + }, + ) + finally: + shim.unregister_turn(token) + + assert response.status_code == 200 + assert backend_calls == 1 + assert counted_calls == 1 + + @pytest.mark.asyncio async def test_shim_completes_turn_after_transfer_without_second_model_call( monkeypatch, diff --git a/tests/runtime/codex/test_codex_runtime_sdk.py b/tests/runtime/codex/test_codex_runtime_sdk.py index 83a3f057f..c8ebdacf3 100644 --- a/tests/runtime/codex/test_codex_runtime_sdk.py +++ b/tests/runtime/codex/test_codex_runtime_sdk.py @@ -46,6 +46,7 @@ def register_turn( *, max_tool_iterations, invocation_id, + before_model_call=None, ): self.registered.append( { @@ -53,6 +54,7 @@ def register_turn( "executors": executors, "max_tool_iterations": max_tool_iterations, "invocation_id": invocation_id, + "before_model_call": before_model_call, } ) return "opaque-turn-token" @@ -60,6 +62,9 @@ def register_turn( def unregister_turn(self, token): self.unregistered.append(token) + def pop_turn_error(self, token): + return None + class _EmptyStream: def __aiter__(self): @@ -129,6 +134,9 @@ class _Context(SimpleNamespace): def _get_events(self, **kwargs): return list(self.session.events) + def increment_llm_call_count(self): + return None + @pytest.mark.asyncio async def test_runtime_passes_isolated_config_and_safe_sdk_controls( diff --git a/tests/runtime/piagent/test_piagent_runtime.py b/tests/runtime/piagent/test_piagent_runtime.py index 11f9a3a19..f6c1de4db 100644 --- a/tests/runtime/piagent/test_piagent_runtime.py +++ b/tests/runtime/piagent/test_piagent_runtime.py @@ -24,8 +24,10 @@ from types import SimpleNamespace import pytest +from google.adk.agents import RunConfig from google.adk.agents.base_agent import BaseAgent from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.invocation_context import LlmCallsLimitExceededError from google.adk.events.event import Event from google.adk.models.llm_response import LlmResponse from google.adk.sessions.in_memory_session_service import InMemorySessionService @@ -83,10 +85,13 @@ def _fake_ctx(*events: Event): session=SimpleNamespace(events=list(events), state={}), branch=None, plugin_manager=None, + increment_llm_call_count=lambda: None, ) -def _ctx(agent, *events: Event, user_content=None) -> InvocationContext: +def _ctx( + agent, *events: Event, user_content=None, run_config=None +) -> InvocationContext: return InvocationContext( session_service=InMemorySessionService(), invocation_id="inv-1", @@ -99,6 +104,7 @@ def _ctx(agent, *events: Event, user_content=None) -> InvocationContext: state={}, events=list(events), ), + run_config=run_config, ) @@ -1462,6 +1468,94 @@ async def test_piagent_runtime_text_only_end_to_end(tmp_path, monkeypatch): assert models["providers"]["veadk"]["models"][0]["id"] == "model-a" +@pytest.mark.asyncio +async def test_piagent_runtime_counts_one_visible_prompt(tmp_path, monkeypatch): + binary = _make_fake_pi(tmp_path) + agent_dir = tmp_path / "agent-home" + monkeypatch.setenv("PIAGENT_BINARY", str(binary)) + monkeypatch.setenv("PIAGENT_AGENT_DIR", str(agent_dir)) + + agent = Agent( + name="assistant", + instruction="Answer briefly.", + model_name="model-a", + model_api_base="https://ark.example.com/api/v3/", + model_api_key="test-key", + model_api_key_name="", + runtime="piagent", + ) + ctx = _fake_ctx(_user_event("ping")) + counted_calls = 0 + + def increment_llm_call_count() -> None: + nonlocal counted_calls + counted_calls += 1 + + ctx.increment_llm_call_count = increment_llm_call_count + + events = [event async for event in PiAgentRuntime().run_async(agent, ctx)] + + assert len(events) == 3 + assert counted_calls == 1 + + +@pytest.mark.asyncio +async def test_piagent_runtime_blocks_prompt_when_llm_call_limit_exceeded( + tmp_path, + monkeypatch, +): + binary = _make_fake_pi(tmp_path) + agent_dir = tmp_path / "agent-home" + monkeypatch.setenv("PIAGENT_BINARY", str(binary)) + monkeypatch.setenv("PIAGENT_AGENT_DIR", str(agent_dir)) + prompt_called = False + + class FakePiAgentRpcClient: + def __init__(self, config): + self.config = config + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return None + + def prompt(self, prompt): + nonlocal prompt_called + prompt_called = True + + async def _events(): + yield {"type": "agent_settled"} + + return _events() + + monkeypatch.setattr( + "veadk.runtime.piagent.runtime.PiAgentRpcClient", + FakePiAgentRpcClient, + ) + + agent = Agent( + name="assistant", + instruction="Answer briefly.", + model_name="model-a", + model_api_base="https://ark.example.com/api/v3/", + model_api_key="test-key", + model_api_key_name="", + runtime="piagent", + ) + ctx = _ctx( + agent, + _user_event("ping"), + run_config=RunConfig(max_llm_calls=1), + ) + ctx.increment_llm_call_count() + + with pytest.raises(LlmCallsLimitExceededError): + [event async for event in PiAgentRuntime().run_async(agent, ctx)] + + assert prompt_called is False + + @pytest.mark.asyncio async def test_piagent_runtime_emits_canonical_bridge_tool_events( tmp_path, diff --git a/veadk/runtime/codex/proxy.py b/veadk/runtime/codex/proxy.py index aa9e7f95f..72ea17df2 100644 --- a/veadk/runtime/codex/proxy.py +++ b/veadk/runtime/codex/proxy.py @@ -31,7 +31,7 @@ import secrets import time from dataclasses import dataclass -from typing import Any, AsyncIterator +from typing import Any, AsyncIterator, Callable import litellm import uvicorn @@ -119,6 +119,7 @@ class ShimTurnContext: executors: dict[str, Any] max_tool_iterations: int invocation_id: str = "" + before_model_call: Callable[[], None] | None = None class ResponsesShim: @@ -144,6 +145,7 @@ def __init__(self, api_base: str, api_key: str) -> None: # subprocess as its provider API key and arrives as a Bearer token, so # concurrent turns can never overwrite one another's tools/context. self._turns: dict[str, ShimTurnContext] = {} + self._turn_errors: dict[str, BaseException] = {} self._app = self._build_app() def register_turn( @@ -153,6 +155,7 @@ def register_turn( *, max_tool_iterations: int = _AGENT_TOOL_MAX_ITERS, invocation_id: str = "", + before_model_call: Callable[[], None] | None = None, ) -> str: """Register immutable routing state and return its opaque bearer token.""" token = secrets.token_urlsafe(32) @@ -161,6 +164,7 @@ def register_turn( executors=dict(executors or {}), max_tool_iterations=max(1, max_tool_iterations), invocation_id=invocation_id, + before_model_call=before_model_call, ) logger.debug( "codex_shim_turn_registered invocation_id=%s tool_count=%d", @@ -172,12 +176,27 @@ def register_turn( def unregister_turn(self, token: str) -> None: """Remove one invocation's routing state.""" context = self._turns.pop(token, None) + self._turn_errors.pop(token, None) if context is not None: logger.debug( "codex_shim_turn_unregistered invocation_id=%s", context.invocation_id, ) + def pop_turn_error(self, token: str) -> BaseException | None: + """Return and clear an invocation-scoped shim error, if one exists.""" + return self._turn_errors.pop(token, None) + + def _before_model_call(self, token: str, context: ShimTurnContext) -> None: + callback = context.before_model_call + if callback is None: + return + try: + callback() + except BaseException as e: + self._turn_errors[token] = e + raise + def _build_app(self) -> FastAPI: app = FastAPI() @@ -258,6 +277,7 @@ async def responses(request: Request) -> Any: max_iters = turn_context.max_tool_iterations if agent_executors else 0 iters = 0 while True: + self._before_model_call(token, turn_context) result = await litellm.aresponses(**call_kwargs) resp = _to_dict(result) if max_iters <= 0: diff --git a/veadk/runtime/codex/runtime.py b/veadk/runtime/codex/runtime.py index c077fd59b..868279fb3 100644 --- a/veadk/runtime/codex/runtime.py +++ b/veadk/runtime/codex/runtime.py @@ -159,6 +159,13 @@ async def run_async( async def _emit_tool_event(event: "Event") -> None: await event_queue.put(event) + def _raise_shim_turn_error() -> None: + if turn_token is None: + return + turn_error = shim.pop_turn_error(turn_token) + if turn_error is not None: + raise turn_error + try: tool_bundle = await build_executable_tools( agent, @@ -269,6 +276,7 @@ async def _emit_tool_event(event: "Event") -> None: tool_bundle.executors, max_tool_iterations=runtime_config.max_tool_iterations, invocation_id=ctx.invocation_id, + before_model_call=ctx.increment_llm_call_count, ) # Keep privileged instructions out of the user transcript. The SDK @@ -408,10 +416,12 @@ async def _pump_codex() -> None: break if transfer_requested: await pump + _raise_shim_turn_error() if deferred_transfer_event is not None: yield deferred_transfer_event return await pump + _raise_shim_turn_error() if final_text_events: llm_response = final_events_to_llm_response(final_text_events) llm_response = await run_after_model_callbacks( @@ -446,6 +456,7 @@ async def _pump_codex() -> None: ctx.invocation_id, type(e).__name__, ) + _raise_shim_turn_error() if isinstance(e, Exception) and "runtime_call" in locals(): fallback = await run_on_model_error_callbacks( agent, diff --git a/veadk/runtime/piagent/runtime.py b/veadk/runtime/piagent/runtime.py index b65a41f5c..707a5a53b 100644 --- a/veadk/runtime/piagent/runtime.py +++ b/veadk/runtime/piagent/runtime.py @@ -158,6 +158,10 @@ async def _emit_tool_event(event: Event) -> None: async def _pump_pi() -> None: try: + # Pi may perform multiple internal model calls in its + # subprocess; veADK can only enforce one budget unit + # for each visible prompt boundary. + ctx.increment_llm_call_count() async for pi_event in client.prompt(prompt): for event in translator.event_to_adk_events(pi_event): await event_queue.put(event)