Add max LLM calls to server commands

src/google/adk/cli/api_server.py

@@ -839,6 +839,7 @@ def __init__(
       auto_create_session: bool = False,
       trigger_sources: Optional[list[str]] = None,
       default_llm_model: Optional[str] = None,
+      max_llm_calls: Optional[int] = None,
   ):
     self.agent_loader = agent_loader
     self.session_service = session_service
@@ -859,6 +860,7 @@ def __init__(
     self.auto_create_session = auto_create_session
     self.trigger_sources = trigger_sources
     self.default_llm_model = default_llm_model
+    self.max_llm_calls = max_llm_calls
     self.default_app_name = os.getenv("ADK_DEFAULT_APP_NAME")
 
   async def get_runner_async(self, app_name: str) -> Runner:
@@ -1807,11 +1809,12 @@ async def run_agent(req: RunAgentRequest, request: Request) -> list[Event]:
       self.current_app_name_ref.value = req.app_name
       runner = await self.get_runner_async(req.app_name)
       _set_telemetry_context_if_needed(runner)
-      run_config = (
-          RunConfig(custom_metadata=req.custom_metadata)
-          if req.custom_metadata
-          else None
-      )
+      run_config_kwargs: dict[str, Any] = {}
+      if req.custom_metadata:
+        run_config_kwargs["custom_metadata"] = req.custom_metadata
+      if self.max_llm_calls is not None:
+        run_config_kwargs["max_llm_calls"] = self.max_llm_calls
+      run_config = RunConfig(**run_config_kwargs) if run_config_kwargs else None
 
       async def worker():
         try:
@@ -1907,6 +1910,11 @@ async def event_generator():
                   run_config=RunConfig(
                       streaming_mode=stream_mode,
                       custom_metadata=req.custom_metadata,
+                      **(
+                          {"max_llm_calls": self.max_llm_calls}
+                          if self.max_llm_calls is not None
+                          else {}
+                      ),
                   ),
                   invocation_id=req.invocation_id,
               )
@@ -2061,6 +2069,11 @@ async def forward_events():
             ),
             save_live_blob=save_live_blob,
             explicit_vad_signal=explicit_vad_signal,
+            **(
+                {"max_llm_calls": self.max_llm_calls}
+                if self.max_llm_calls is not None
+                else {}
+            ),
         )
         async with Aclosing(
             runner.run_live(

src/google/adk/cli/cli_tools_click.py

@@ -1927,6 +1927,15 @@ def decorator(func):
         ),
         default=None,
     )
+    @click.option(
+        "--max_llm_calls",
+        type=int,
+        help=(
+            "Optional. Maximum number of LLM calls allowed for each agent"
+            " run. Values less than or equal to zero disable the limit."
+        ),
+        default=None,
+    )
     # Parsed into list[str] by the wrapper below (server commands need a list).
     @click.option(
         "--trigger_sources",
@@ -2012,6 +2021,7 @@ def cli_web(
     logo_text: str | None = None,
     logo_image_url: str | None = None,
     trigger_sources: list[str] | None = None,
+    max_llm_calls: int | None = None,
 ):
   """Starts a FastAPI server with Web UI for agents.
 
@@ -2082,6 +2092,7 @@ async def _lifespan(app: FastAPI):
       logo_image_url=logo_image_url,
       trigger_sources=trigger_sources,
       default_llm_model=default_llm_model,
+      max_llm_calls=max_llm_calls,
   )
   config = uvicorn.Config(
       app,
@@ -2163,6 +2174,7 @@ def cli_api_server(
     with_ui: bool = False,
     gemini_enterprise_app_name: str | None = None,
     express_mode: bool = False,
+    max_llm_calls: int | None = None,
 ):
   """Starts a FastAPI server for agents.
 
@@ -2223,6 +2235,7 @@ async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
           trigger_sources=trigger_sources,
           gemini_enterprise_app_name=gemini_enterprise_app_name,
           express_mode=express_mode,
+          max_llm_calls=max_llm_calls,
           lifespan=_lifespan,
       ),
       host=host,

src/google/adk/cli/fast_api.py

@@ -122,6 +122,7 @@ def get_fast_api_app(
     default_llm_model: str | None = None,
     gemini_enterprise_app_name: str | None = None,
     express_mode: bool = False,
+    max_llm_calls: int | None = None,
 ) -> FastAPI:
   """Constructs and returns a FastAPI application for serving ADK agents.
 
@@ -178,6 +179,8 @@ def get_fast_api_app(
     gemini_enterprise_app_name: The Gemini Enterprise app name to use for the
       agent.
     express_mode: Whether to enable express mode.
+    max_llm_calls: Maximum number of LLM calls allowed for each agent run.
+      When None, ``RunConfig`` resolves its normal default.
 
   Returns:
     The configured FastAPI application instance.
@@ -300,6 +303,7 @@ def get_fast_api_app(
       auto_create_session=auto_create_session,
       trigger_sources=trigger_sources,
       default_llm_model=default_llm_model,
+      max_llm_calls=max_llm_calls,
   )
 
   # In single agent mode, use that agent as the default app.

tests/unittests/cli/test_adk_web_server_run_live.py

@@ -61,7 +61,7 @@ async def run_live(
     yield Event(author="runner")
 
 
-def test_run_live_applies_run_config_query_options():
+def test_run_live_applies_server_and_query_run_config_options():
   session_service = InMemorySessionService()
   asyncio.run(
       session_service.create_session(
@@ -82,6 +82,7 @@ def test_run_live_applies_run_config_query_options():
       eval_sets_manager=types.SimpleNamespace(),
       eval_set_results_manager=types.SimpleNamespace(),
       agents_dir=".",
+      max_llm_calls=37,
   )
 
   async def _get_runner_async(_self, _app_name: str):
@@ -122,6 +123,7 @@ async def _get_runner_async(_self, _app_name: str):
   assert run_config.session_resumption.transparent is True
   assert run_config.save_live_blob is True
   assert run_config.explicit_vad_signal is True
+  assert run_config.max_llm_calls == 37
 
 
 @pytest.mark.parametrize(

tests/unittests/cli/utils/test_cli_tools_click.py

@@ -1589,6 +1589,28 @@ def test_cli_api_server_invokes_uvicorn(
   assert _patch_uvicorn.calls, "uvicorn.Server.run must be called"
 
 
+@pytest.mark.parametrize("command", ["web", "api_server"])
+def test_cli_server_passes_max_llm_calls(
+    tmp_path: Path,
+    _patch_uvicorn: _Recorder,
+    monkeypatch: pytest.MonkeyPatch,
+    command: str,
+) -> None:
+  """Both server commands pass the requested LLM call limit to the app."""
+  agents_dir = tmp_path / "agents"
+  agents_dir.mkdir()
+  mock_get_app = _Recorder()
+  monkeypatch.setattr("google.adk.cli.fast_api.get_fast_api_app", mock_get_app)
+
+  result = CliRunner().invoke(
+      cli_tools_click.main,
+      [command, "--max_llm_calls", "37", str(agents_dir)],
+  )
+
+  assert result.exit_code == 0, (result.output, repr(result.exception))
+  assert mock_get_app.calls[0][1]["max_llm_calls"] == 37
+
+
 def test_cli_web_passes_service_uris(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch, _patch_uvicorn: _Recorder
 ) -> None:
@@ -2632,10 +2654,21 @@ def test_fast_api_common_options_documented_defaults() -> None:
   assert captured["a2a"] is False
   assert captured["allow_origins"] == ()
   assert captured["log_level"] == "INFO"
+  assert captured["max_llm_calls"] is None
   # --verbose is consumed while folding it into log_level.
   assert "verbose" not in captured
 
 
+def test_fast_api_common_options_parses_max_llm_calls() -> None:
+  """The common server option parses an explicit per-run LLM call limit."""
+  command, captured = _fast_api_command()
+
+  result = CliRunner().invoke(command, ["--max_llm_calls", "37"])
+
+  assert result.exit_code == 0, (result.output, repr(result.exception))
+  assert captured["max_llm_calls"] == 37
+
+
 # adk test
 @pytest.fixture
 def fake_pytest_run(monkeypatch: pytest.MonkeyPatch):