fix: enforce llm call limits for non-adk runtimes

Author: FirstayZheng

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,

tests/runtime/codex/test_codex_runtime_sdk.py

@@ -46,20 +46,25 @@ def register_turn(
         *,
         max_tool_iterations,
         invocation_id,
+        before_model_call=None,
     ):
         self.registered.append(
             {
                 "specs": specs,
                 "executors": executors,
                 "max_tool_iterations": max_tool_iterations,
                 "invocation_id": invocation_id,
+                "before_model_call": before_model_call,
             }
         )
         return "opaque-turn-token"
 
     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(

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,

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: