From 770d4f00740b017f636ef1104a4e843272203a86 Mon Sep 17 00:00:00 2001 From: Jaideep <67646710+jdchawla29@users.noreply.github.com> Date: Thu, 15 Oct 2026 03:03:06 -0700 Subject: [PATCH 1/4] fix(agents): leave Gemini sampling to the model and report an uncontinued cap as length GeminiConfig no longer sends temperature 1.0, top_p 0.95, and top_k 40 on every request; unset, each model uses its own sampling. Gemini 4 models ignore sampling parameters, and Gemini 3's recommended temperature is its default. A Gemini 4 turn that reaches the per-request output cap ends with finish reason CONTINUATION. When the client cannot continue it (a google-genai older than 2.30.1), the turn is a truncation, so the run now reports stop_reason=length instead of an empty "done". --- docs/v6/reference/agents.mdx | 4 +++- hud/agents/tests/test_gemini_agent.py | 18 ++++++++++++++---- hud/agents/tests/test_tool_agent.py | 8 +++++++- hud/agents/tool_agent.py | 7 +++++-- hud/agents/types.py | 6 +++--- 5 files changed, 32 insertions(+), 11 deletions(-) diff --git a/docs/v6/reference/agents.mdx b/docs/v6/reference/agents.mdx index 62a93f03b..44370dde8 100644 --- a/docs/v6/reference/agents.mdx +++ b/docs/v6/reference/agents.mdx @@ -57,7 +57,7 @@ agent = ClaudeAgent(ClaudeConfig(model="claude-sonnet-4-5", max_steps=30)) |-------|--------|---------------| | `ClaudeAgent` | `ClaudeConfig` | `claude-sonnet-4-6` | | `OpenAIAgent` | `OpenAIConfig` | `gpt-5.6` | -| `GeminiAgent` | `GeminiConfig` | `gemini-3-pro-preview` | +| `GeminiAgent` | `GeminiConfig` | `gemini-3.1-pro-preview` | | `OpenAIChatAgent` | `OpenAIChatConfig` | `gpt-5.4-mini` | | `ClaudeCLIAgent` | `ClaudeCLIConfig` | `claude-sonnet-5` | | `CodexCLIAgent` | `CodexCLIConfig` | `gpt-5.6-sol` | @@ -72,6 +72,8 @@ config; `__call__(run)` takes only the run. `OpenAIConfig.max_output_tokens` and `GeminiConfig.max_output_tokens` are unset by default, so each request gets the model's own output limit. Set them to cap output per turn. Gemini counts thinking tokens toward the cap, so a low cap can truncate a thinking model before it answers. +`GeminiConfig.temperature`, `top_p`, and `top_k` are unset by default too, so each request uses the model's own sampling; Gemini 4 models ignore sampling parameters. `GeminiConfig.thinking_level` accepts `minimal`, `low`, `medium`, or `high`, but each model supports its own subset (Gemini 4 Argon rejects `minimal`); unset leaves the model's default. Gemini 4 long decode, which lets a turn run past the model's per-request output cap, needs `max_output_tokens` above 65,536 and `google-genai` 2.30.1 or later; the SDK then continues the turn across requests, and its hops are recorded as one turn. A turn that stops at the per-request cap, as every long turn does with `max_output_tokens` unset, ends the run with `stop_reason="length"`. + `OpenAIConfig` generates a stable `prompt_cache_key` for the agent and sends it with every Responses request. Pass your own key when multiple agent instances should share the same cached prompt prefix. diff --git a/hud/agents/tests/test_gemini_agent.py b/hud/agents/tests/test_gemini_agent.py index 30b672dda..06f80ca82 100644 --- a/hud/agents/tests/test_gemini_agent.py +++ b/hud/agents/tests/test_gemini_agent.py @@ -177,7 +177,7 @@ async def generate_content(self, **kwargs: Any) -> Any: return self._response -async def _sent_output_cap(**config: Any) -> int | None: +async def _sent_config(**config: Any) -> Any: resp_content = SimpleNamespace( role="model", parts=[SimpleNamespace(function_call=None, text="answer", thought=None)], @@ -196,15 +196,25 @@ async def _sent_output_cap(**config: Any) -> int | None: ) ) await agent.get_response(_state(agent)) - return models.configs[0].max_output_tokens + return models.configs[0] async def test_output_cap_defaults_to_the_model_limit() -> None: - assert await _sent_output_cap() is None + assert (await _sent_config()).max_output_tokens is None async def test_configured_output_cap_is_sent() -> None: - assert await _sent_output_cap(max_output_tokens=4096) == 4096 + assert (await _sent_config(max_output_tokens=4096)).max_output_tokens == 4096 + + +async def test_sampling_defaults_to_the_model() -> None: + sent = await _sent_config() + assert (sent.temperature, sent.top_p, sent.top_k) == (None, None, None) + + +async def test_configured_sampling_is_sent() -> None: + sent = await _sent_config(temperature=0.5, top_p=0.9, top_k=20) + assert (sent.temperature, sent.top_p, sent.top_k) == (0.5, 0.9, 20) class RecordingModels: diff --git a/hud/agents/tests/test_tool_agent.py b/hud/agents/tests/test_tool_agent.py index 19ff303b8..1ca164622 100644 --- a/hud/agents/tests/test_tool_agent.py +++ b/hud/agents/tests/test_tool_agent.py @@ -725,7 +725,13 @@ async def test_loop_marks_length_finish_as_truncated() -> None: # A final turn cut off at the provider token cap (e.g. mid-tool-call) ends the # rollout normally but is a truncation, not a natural finish — across every # provider's finish-reason vocabulary. - for finish_reason in ("length", "max_output_tokens", "max_tokens", "MAX_TOKENS"): + for finish_reason in ( + "length", + "max_output_tokens", + "max_tokens", + "MAX_TOKENS", + "CONTINUATION", + ): agent = DictAgent([AgentStep(content="partial", done=True, finish_reason=finish_reason)]) run = cast("Run", _FakeRun()) diff --git a/hud/agents/tool_agent.py b/hud/agents/tool_agent.py index 198776e0a..d909fe424 100644 --- a/hud/agents/tool_agent.py +++ b/hud/agents/tool_agent.py @@ -31,8 +31,11 @@ logger = logging.getLogger(__name__) # Output-token-cap finish reasons: OpenAI chat "length", Responses "max_output_tokens", -# Claude "max_tokens", Gemini "MAX_TOKENS". -TRUNCATION_FINISH_REASONS = frozenset({"length", "max_output_tokens", "max_tokens", "MAX_TOKENS"}) +# Claude "max_tokens", Gemini "MAX_TOKENS", and Gemini "CONTINUATION" (a per-request +# cap hit by a turn that was not continued). +TRUNCATION_FINISH_REASONS = frozenset( + {"length", "max_output_tokens", "max_tokens", "MAX_TOKENS", "CONTINUATION"} +) MAX_CONSECUTIVE_SSH_FAILURES = 3 MessageT = TypeVar("MessageT") diff --git a/hud/agents/types.py b/hud/agents/types.py index 5be380bf8..a40b06494 100644 --- a/hud/agents/types.py +++ b/hud/agents/types.py @@ -146,9 +146,9 @@ class GeminiConfig(ToolAgentConfig): model_name: str = "Gemini" model: str = Field(default="gemini-3.1-pro-preview", validation_alias=_model_alias) - temperature: float = 1.0 - top_p: float = 0.95 - top_k: int = 40 + temperature: float | None = None + top_p: float | None = None + top_k: int | None = None max_output_tokens: int | None = None excluded_predefined_functions: list[str] = Field(default_factory=list) thinking_level: Literal["minimal", "low", "medium", "high"] | None = None From 1e20e601f6c618ca321190406cb24a85ea99a01b Mon Sep 17 00:00:00 2001 From: Jaideep <67646710+jdchawla29@users.noreply.github.com> Date: Thu, 15 Oct 2026 07:05:09 -0700 Subject: [PATCH 2/4] feat(agents): stream GeminiAgent turns GeminiAgent made one non-streaming generateContent call per turn. A Gemini 4 turn of 32,768 output tokens took about 645 seconds in measurement, longer than non-streaming request timeouts in front of the model (the HUD gateway allows 600 seconds), and Google recommends streaming for long generations. Each turn now streams and joins the streamed parts into one model turn, kept as received so thought signatures stay on the parts they arrived with. Text and thought fragments concatenate without separators, since streamed fragments split mid-word. Usage, grounding, and finish reason come from the last chunk that carries them. --- docs/v6/reference/agents.mdx | 4 +- hud/agents/gemini/agent.py | 45 +++++--- hud/agents/tests/test_gemini_agent.py | 155 ++++++++++++++++++++++++-- 3 files changed, 181 insertions(+), 23 deletions(-) diff --git a/docs/v6/reference/agents.mdx b/docs/v6/reference/agents.mdx index 44370dde8..764d0b71f 100644 --- a/docs/v6/reference/agents.mdx +++ b/docs/v6/reference/agents.mdx @@ -72,7 +72,9 @@ config; `__call__(run)` takes only the run. `OpenAIConfig.max_output_tokens` and `GeminiConfig.max_output_tokens` are unset by default, so each request gets the model's own output limit. Set them to cap output per turn. Gemini counts thinking tokens toward the cap, so a low cap can truncate a thinking model before it answers. -`GeminiConfig.temperature`, `top_p`, and `top_k` are unset by default too, so each request uses the model's own sampling; Gemini 4 models ignore sampling parameters. `GeminiConfig.thinking_level` accepts `minimal`, `low`, `medium`, or `high`, but each model supports its own subset (Gemini 4 Argon rejects `minimal`); unset leaves the model's default. Gemini 4 long decode, which lets a turn run past the model's per-request output cap, needs `max_output_tokens` above 65,536 and `google-genai` 2.30.1 or later; the SDK then continues the turn across requests, and its hops are recorded as one turn. A turn that stops at the per-request cap, as every long turn does with `max_output_tokens` unset, ends the run with `stop_reason="length"`. +`GeminiConfig.temperature`, `top_p`, and `top_k` are unset by default, so each request uses the model's own sampling; Gemini 4 models ignore them. `GeminiConfig.thinking_level` accepts `minimal`, `low`, `medium`, or `high`, subject to the model (Gemini 4 Argon rejects `minimal`); unset uses the model's default. + +`GeminiAgent` streams each turn. For Gemini 4 long decode, set `max_output_tokens` above 65,536 and install `google-genai` 2.30.1 or later; the turn then runs past the per-request output cap and is recorded as one turn. Otherwise a turn that reaches the cap ends the run with `stop_reason="length"`. Keep `include_thoughts` on for long-thinking runs, since without it the model sends nothing while it thinks and a long think can hit the stream's first-chunk deadline. `OpenAIConfig` generates a stable `prompt_cache_key` for the agent and sends it with every Responses request. Pass your own key when multiple agent instances should share the same cached prompt prefix. diff --git a/hud/agents/gemini/agent.py b/hud/agents/gemini/agent.py index 649ce27a8..b508329d1 100644 --- a/hud/agents/gemini/agent.py +++ b/hud/agents/gemini/agent.py @@ -201,22 +201,40 @@ async def get_response( thinking_config=thinking_config, ) - api_response = await self.gemini_client.aio.models.generate_content( + stream = await self.gemini_client.aio.models.generate_content_stream( model=self.config.model, contents=cast("Any", messages), config=generate_config, ) - if not api_response.candidates: + # The turn keeps every streamed part as received: thought signatures belong to the + # parts they arrive on, including empty text parts. + parts: list[genai_types.Part] = [] + candidate_seen = False + model_version: str | None = None + usage_meta: genai_types.GenerateContentResponseUsageMetadata | None = None + grounding_meta: genai_types.GroundingMetadata | None = None + finish_reason: genai_types.FinishReason | None = None + async for chunk in stream: + model_version = chunk.model_version or model_version + if chunk.usage_metadata is not None: + usage_meta = chunk.usage_metadata + if not chunk.candidates: + continue + candidate_seen = True + candidate = chunk.candidates[0] + if candidate.content is not None and candidate.content.parts: + parts.extend(candidate.content.parts) + if candidate.grounding_metadata is not None: + grounding_meta = candidate.grounding_metadata + if candidate.finish_reason is not None: + finish_reason = candidate.finish_reason + if not candidate_seen: raise RuntimeError(f"Gemini returned no candidates for model {self.config.model}") - - candidate = api_response.candidates[0] - content = candidate.content - if content is not None: - messages.append(content) + if parts: + messages.append(genai_types.Content(role="model", parts=parts)) result = AgentStep(content="", done=True) - result.model = api_response.model_version or self.config.model - usage_meta = api_response.usage_metadata + result.model = model_version or self.config.model if usage_meta is not None: result.usage = Usage( prompt_tokens=usage_meta.prompt_token_count, @@ -226,7 +244,7 @@ async def get_response( text_parts: list[str] = [] thought_parts: list[str] = [] - for part in (content.parts or []) if content else []: + for part in parts: function_call = part.function_call if function_call is not None: tc = self._make_tool_call(function_call, computer_tool) @@ -241,14 +259,13 @@ async def get_response( result.content = "".join(text_parts) if thought_parts: - result.reasoning = "\n".join(thought_parts) + result.reasoning = "".join(thought_parts) - grounding_meta = candidate.grounding_metadata if grounding_meta is not None: result.citations = _grounding_citations(grounding_meta) - if candidate.finish_reason is not None: - result.finish_reason = candidate.finish_reason.name + if finish_reason is not None: + result.finish_reason = finish_reason.name if result.finish_reason == "MALFORMED_FUNCTION_CALL": result.error = "Provider returned a malformed function call" result.stop_reason = "malformed_tool_call" diff --git a/hud/agents/tests/test_gemini_agent.py b/hud/agents/tests/test_gemini_agent.py index 06f80ca82..00ce428d3 100644 --- a/hud/agents/tests/test_gemini_agent.py +++ b/hud/agents/tests/test_gemini_agent.py @@ -1,13 +1,17 @@ -"""``GeminiAgent`` — ``get_response`` parsing over a fake Generate Content client, +"""``GeminiAgent`` — ``get_response`` over a fake streaming Generate Content client, plus ``_make_tool_call`` mapping and ``_grounding_citations``. """ from __future__ import annotations +import json from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast +import httpx import mcp.types as mcp_types +from google import genai +from google.genai import types as genai_types from hud.agents.gemini.agent import GeminiAgent, _grounding_citations from hud.agents.tool_agent import RunState @@ -15,15 +19,22 @@ from hud.types import MCPToolResult, Step, Trace if TYPE_CHECKING: + from collections.abc import AsyncIterator + from hud.eval.run import Run +async def _chunks(*chunks: Any) -> AsyncIterator[Any]: + for chunk in chunks: + yield chunk + + class FakeModels: def __init__(self, response: Any) -> None: self._response = response - async def generate_content(self, **_kwargs: Any) -> Any: - return self._response + async def generate_content_stream(self, **_kwargs: Any) -> AsyncIterator[Any]: + return _chunks(self._response) class FakeGenai: @@ -109,7 +120,7 @@ async def test_get_response_done_text_only() -> None: async def test_get_response_no_candidates_raises() -> None: - agent = _agent(SimpleNamespace(candidates=[])) + agent = _agent(SimpleNamespace(candidates=[], model_version=None, usage_metadata=None)) try: await agent.get_response(_state(agent)) except RuntimeError: @@ -151,6 +162,134 @@ def test_grounding_citations() -> None: assert cites[0].type == "grounding" +def _chunk(*parts: genai_types.Part, finish_reason: str | None = None) -> Any: + return genai_types.GenerateContentResponse( + candidates=[ + genai_types.Candidate( + content=genai_types.Content(role="model", parts=list(parts)), + finish_reason=genai_types.FinishReason(finish_reason) if finish_reason else None, + ) + ], + model_version="gemini-test-v2", + ) + + +async def test_streamed_chunks_form_one_model_turn() -> None: + """Every streamed part joins one model turn as received; text and thoughts concatenate.""" + parts = [ + genai_types.Part(text="Checking ", thought=True), + genai_types.Part(text="divisors.", thought=True), + genai_types.Part(text="It is 17 *"), + genai_types.Part(text=" 23."), + genai_types.Part( + function_call=genai_types.FunctionCall(name="bash", args={"command": "ls"}), + thought_signature=b"sig-1", + ), + genai_types.Part(text="", thought_signature=b"sig-2"), + ] + stream = _chunks( + _chunk(*parts[:2]), + _chunk(*parts[2:4]), + _chunk(*parts[4:], finish_reason="STOP"), + ) + + async def generate_content_stream(**_kwargs: Any) -> AsyncIterator[Any]: + return stream + + agent = GeminiAgent( + GeminiConfig( + model="gemini-test", + model_client=SimpleNamespace( + aio=SimpleNamespace( + models=SimpleNamespace(generate_content_stream=generate_content_stream) + ) + ), + ) + ) + state = _state(agent) + + result = await agent.get_response(state) + + assert result.content == "It is 17 * 23." + assert result.reasoning == "Checking divisors." + assert [tc.name for tc in result.tool_calls] == ["bash"] + assert result.finish_reason == "STOP" + assert result.model == "gemini-test-v2" + turn = state.messages[-1] + assert turn.role == "model" + assert turn.parts == parts + + +async def test_turn_streams_through_the_sdk() -> None: + """A turn arrives over the SDK's SSE stream and is recorded as one model turn.""" + events: list[dict[str, Any]] = [ + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "Thinking.", "thought": True}]}} + ] + }, + {"candidates": [{"content": {"role": "model", "parts": [{"text": "39"}]}}]}, + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "1"}, {"text": "", "thoughtSignature": "c2ln"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 7, + "candidatesTokenCount": 2, + "thoughtsTokenCount": 5, + "totalTokenCount": 14, + }, + "modelVersion": "gemini-4-argon", + }, + ] + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + content=b"".join( + b"data: " + json.dumps(event).encode() + b"\r\n\r\n" for event in events + ), + headers={"content-type": "text/event-stream"}, + ) + + client = genai.Client( + api_key="hud-key", + http_options=genai_types.HttpOptions( + base_url="https://gateway.test", + api_version="v1beta", + async_client_args={"transport": httpx.MockTransport(handler)}, + ), + ) + agent = GeminiAgent(GeminiConfig(model="gemini-4-argon", model_client=client)) + state = _state(agent) + + result = await agent.get_response(state) + + [request] = requests + assert request.url.path == "/v1beta/models/gemini-4-argon:streamGenerateContent" + assert request.url.params["alt"] == "sse" + assert result.content == "391" + assert result.reasoning == "Thinking." + assert result.finish_reason == "STOP" + assert result.usage is not None + assert result.usage.prompt_tokens == 7 + turn = state.messages[-1] + assert [(part.text, part.thought, part.thought_signature) for part in turn.parts or []] == [ + ("Thinking.", True, None), + ("39", None, None), + ("1", None, None), + ("", None, b"sig"), + ] + + async def test_malformed_function_call_is_a_normalized_error() -> None: response = _api_response( SimpleNamespace( @@ -172,9 +311,9 @@ def __init__(self, response: Any) -> None: self._response = response self.configs: list[Any] = [] - async def generate_content(self, **kwargs: Any) -> Any: + async def generate_content_stream(self, **kwargs: Any) -> AsyncIterator[Any]: self.configs.append(kwargs["config"]) - return self._response + return _chunks(self._response) async def _sent_config(**config: Any) -> Any: @@ -224,10 +363,10 @@ def __init__(self, responses: list[Any]) -> None: self._responses = list(responses) self.requests: list[Any] = [] - async def generate_content(self, **kwargs: Any) -> Any: + async def generate_content_stream(self, **kwargs: Any) -> AsyncIterator[Any]: # Snapshot: the agent goes on mutating the list it passed. self.requests.append(list(kwargs["contents"])) - return self._responses.pop(0) + return _chunks(self._responses.pop(0)) class RecordingGenai: From 29f3eefdda07462993b36a45e93cc4caffbd79c9 Mon Sep 17 00:00:00 2001 From: Jaideep <67646710+jdchawla29@users.noreply.github.com> Date: Thu, 15 Oct 2026 07:58:36 -0700 Subject: [PATCH 3/4] fix(agents): build Claude web-fetch citations as a plain dict anthropic 1.x renamed the request TypedDict for citations to BetaCitationsConfigParamParam; BetaCitationsConfigParam is now the response model. A plain {"enabled": True} is the citations parameter on both 0.x and 1.x. --- hud/agents/claude/tools/hosted.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/hud/agents/claude/tools/hosted.py b/hud/agents/claude/tools/hosted.py index d7076acb3..9f4b33e23 100644 --- a/hud/agents/claude/tools/hosted.py +++ b/hud/agents/claude/tools/hosted.py @@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any from anthropic.types.beta import ( - BetaCitationsConfigParam, BetaToolSearchToolBm25_20251119Param, BetaToolUnionParam, BetaWebFetchTool20250910Param, @@ -75,7 +74,7 @@ def to_params(self) -> BetaWebFetchTool20250910Param: if self.max_content_tokens is not None: params["max_content_tokens"] = self.max_content_tokens if self.citations_enabled: - params["citations"] = BetaCitationsConfigParam(enabled=True) + params["citations"] = {"enabled": True} return params From 72554d1129b7c513b82a9721cb47f7c57bedf1d0 Mon Sep 17 00:00:00 2001 From: Jaideep <67646710+jdchawla29@users.noreply.github.com> Date: Thu, 15 Oct 2026 07:58:36 -0700 Subject: [PATCH 4/4] build: test hud apart from the browseruse extra Every browser-use release pins google-genai below 2.0 (0.11.13) or to 1.65.0 (0.12.1 and later), and the dev environment resolved together with the browseruse extra. That held SDK tests at google-genai 1.75, openai 2.x, and anthropic 0.x, while installing hud without that extra gets the latest of each. The browseruse extra now resolves apart from dev, and CI runs every other extra. dev requires google-genai 2.30.1 or later, so GeminiAgent's tests cover a stream that stops at the per-request cap and continues with its continuation token as one turn. hud's own google-genai floor is unchanged, so hud[browseruse] still installs. On openai 3.x and anthropic 1.x the SDK clients run on httpx2, so the tests that inject transports or build error responses use httpx2. --- .github/workflows/ci.yml | 2 +- hud/agents/tests/test_claude_agent.py | 4 +- hud/agents/tests/test_gemini_agent.py | 115 +++++++++++++++++++++----- hud/utils/tests/test_gateway.py | 19 +++-- pyproject.toml | 5 ++ 5 files changed, 113 insertions(+), 32 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ee2e8dd62..d709bf006 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -27,7 +27,7 @@ jobs: run: uv python install ${{ matrix.python-version }} - name: Run tests - run: uv run --python ${{ matrix.python-version }} --all-extras pytest --cov --cov-report='' + run: uv run --python ${{ matrix.python-version }} --all-extras --no-extra browseruse pytest --cov --cov-report='' lint-ruff: runs-on: ubuntu-latest diff --git a/hud/agents/tests/test_claude_agent.py b/hud/agents/tests/test_claude_agent.py index 38571b214..c9a07b4c7 100644 --- a/hud/agents/tests/test_claude_agent.py +++ b/hud/agents/tests/test_claude_agent.py @@ -104,10 +104,10 @@ def _state(agent: ClaudeAgent) -> Any: def _inline_error(type_: str, *, headers: dict[str, str] | None = None) -> APIStatusError: body = {"type": "error", "error": {"type": type_, "message": "stream failed"}} - response = httpx.Response( + response = httpx2.Response( 200, headers=headers, - request=httpx.Request("POST", "https://api.anthropic.com/v1/messages"), + request=httpx2.Request("POST", "https://api.anthropic.com/v1/messages"), ) return APIStatusError(str(body), response=response, body=body) diff --git a/hud/agents/tests/test_gemini_agent.py b/hud/agents/tests/test_gemini_agent.py index 00ce428d3..45d5e0c2c 100644 --- a/hud/agents/tests/test_gemini_agent.py +++ b/hud/agents/tests/test_gemini_agent.py @@ -220,6 +220,32 @@ async def generate_content_stream(**_kwargs: Any) -> AsyncIterator[Any]: assert turn.parts == parts +def _sse_client(*hops: list[dict[str, Any]]) -> tuple[genai.Client, list[httpx.Request]]: + """A ``genai.Client`` whose n-th request streams the n-th hop's events over SSE.""" + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + events = hops[len(requests) - 1] + return httpx.Response( + 200, + content=b"".join( + b"data: " + json.dumps(event).encode() + b"\r\n\r\n" for event in events + ), + headers={"content-type": "text/event-stream"}, + ) + + client = genai.Client( + api_key="hud-key", + http_options=genai_types.HttpOptions( + base_url="https://gateway.test", + api_version="v1beta", + async_client_args={"transport": httpx.MockTransport(handler)}, + ), + ) + return client, requests + + async def test_turn_streams_through_the_sdk() -> None: """A turn arrives over the SDK's SSE stream and is recorded as one model turn.""" events: list[dict[str, Any]] = [ @@ -248,26 +274,7 @@ async def test_turn_streams_through_the_sdk() -> None: "modelVersion": "gemini-4-argon", }, ] - requests: list[httpx.Request] = [] - - def handler(request: httpx.Request) -> httpx.Response: - requests.append(request) - return httpx.Response( - 200, - content=b"".join( - b"data: " + json.dumps(event).encode() + b"\r\n\r\n" for event in events - ), - headers={"content-type": "text/event-stream"}, - ) - - client = genai.Client( - api_key="hud-key", - http_options=genai_types.HttpOptions( - base_url="https://gateway.test", - api_version="v1beta", - async_client_args={"transport": httpx.MockTransport(handler)}, - ), - ) + client, requests = _sse_client(events) agent = GeminiAgent(GeminiConfig(model="gemini-4-argon", model_client=client)) state = _state(agent) @@ -290,6 +297,74 @@ def handler(request: httpx.Request) -> httpx.Response: ] +async def test_turn_past_the_per_request_cap_continues_as_one_turn() -> None: + """A stream that stops at the per-request cap resumes with its continuation token, and + both hops form one model turn.""" + first_hop: list[dict[str, Any]] = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Working it out.", "thought": True}], + } + } + ] + }, + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "", "thoughtSignature": "c2lnLTE="}], + }, + "finishReason": "CONTINUATION", + "continuationToken": "dG9rZW4tMQ==", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "thoughtsTokenCount": 65526, + "totalTokenCount": 65536, + }, + }, + ] + second_hop: list[dict[str, Any]] = [ + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "42"}]}, "finishReason": "STOP"} + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 2, + "totalTokenCount": 12, + }, + "modelVersion": "gemini-4-argon", + } + ] + client, requests = _sse_client(first_hop, second_hop) + agent = GeminiAgent(GeminiConfig(model="gemini-4-argon", model_client=client)) + state = _state(agent) + + result = await agent.get_response(state) + + first, second = (json.loads(request.content) for request in requests) + assert second == {**first, "continuationToken": "dG9rZW4tMQ=="} + assert result.content == "42" + assert result.reasoning == "Working it out." + assert result.done is True + assert result.finish_reason == "STOP" + assert result.usage is not None + assert result.usage.prompt_tokens == 20 + turns = [content for content in state.messages if content.role == "model"] + assert len(turns) == 1 + assert [(part.text, part.thought_signature) for part in turns[0].parts or []] == [ + ("Working it out.", None), + ("", b"sig-1"), + ("42", None), + ] + + async def test_malformed_function_call_is_a_normalized_error() -> None: response = _api_response( SimpleNamespace( diff --git a/hud/utils/tests/test_gateway.py b/hud/utils/tests/test_gateway.py index f452994f6..4ba61f9cf 100644 --- a/hud/utils/tests/test_gateway.py +++ b/hud/utils/tests/test_gateway.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import httpx +import httpx2 import pytest from hud.settings import settings @@ -30,13 +31,13 @@ async def test_openai_client_resolves_trace_id_per_request( ) -> None: seen_headers: dict[str, str] = {} - async def handler(request: httpx.Request) -> httpx.Response: + async def handler(request: httpx2.Request) -> httpx2.Response: trace_id = request.headers["Trace-Id"] await asyncio.sleep(0) seen_headers[request.url.path] = trace_id - return httpx.Response(200, json={"object": "list", "data": []}) + return httpx2.Response(200, json={"object": "list", "data": []}) - transport = httpx.MockTransport(handler) + transport = httpx2.MockTransport(handler) real_client_factory = gateway.DefaultAsyncHttpxClient monkeypatch.setattr( gateway, @@ -69,12 +70,12 @@ async def test_openai_client_trace_context_overrides_empty_header( ) -> None: seen_trace_id: str | None = None - async def handler(request: httpx.Request) -> httpx.Response: + async def handler(request: httpx2.Request) -> httpx2.Response: nonlocal seen_trace_id seen_trace_id = request.headers["Trace-Id"] - return httpx.Response(200, json={"object": "list", "data": []}) + return httpx2.Response(200, json={"object": "list", "data": []}) - transport = httpx.MockTransport(handler) + transport = httpx2.MockTransport(handler) real_client_factory = gateway.DefaultAsyncHttpxClient monkeypatch.setattr( gateway, @@ -103,12 +104,12 @@ async def test_openai_client_sends_child_and_parent_trace_ids( ) -> None: seen_headers: dict[str, str] = {} - async def handler(request: httpx.Request) -> httpx.Response: + async def handler(request: httpx2.Request) -> httpx2.Response: seen_headers["Trace-Id"] = request.headers["Trace-Id"] seen_headers["X-HUD-Parent-Trace-Id"] = request.headers["X-HUD-Parent-Trace-Id"] - return httpx.Response(200, json={"object": "list", "data": []}) + return httpx2.Response(200, json={"object": "list", "data": []}) - transport = httpx.MockTransport(handler) + transport = httpx2.MockTransport(handler) real_client_factory = gateway.DefaultAsyncHttpxClient monkeypatch.setattr( gateway, diff --git a/pyproject.toml b/pyproject.toml index b4ccdbfe5..25caa69d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,6 +129,8 @@ bedrock = [ dev = [ "hud[agents]", # Include agents for dev "hud[robot]", # the robot end-to-end gates (hud/tests/test_robot.py) run the real wire + # GeminiAgent's long-decode tests need the SDK's automatic continuation. + "google-genai>=2.30.1", "dotenv>=0.9.9", # Testing and linting "ruff==0.16.1", @@ -173,6 +175,9 @@ train = [ [tool.uv] required-version = "==0.12.2" +# Every browser-use release pins google-genai below 2.0, so the browseruse extra +# resolves apart from dev, which needs google-genai 2.30.1 or later. +conflicts = [[{ extra = "browseruse" }, { extra = "dev" }]] [tool.ruff] target-version = "py311"