From 4169c6adda9819fb413d13e7231a44ba7dd7add5 Mon Sep 17 00:00:00 2001 From: joycel-github Date: Mon, 13 Jul 2026 21:04:58 +0000 Subject: [PATCH 1/2] antigravity: parse and validate request harness_config Parse HarnessStart.harness_config (JSON-in-bytes) and overlay it onto the server's default LocalAgentConfig before each turn. - Blocks a request from overriding AX-owned persistence fields (save_dir, conversation_id) via _AX_MANAGED_CONFIG_FIELDS; these are injected last so a request can't redirect trajectory storage. - Reconstructs the config (not model_copy) so the SDK re-validates the overlaid values and surfaces its own error; invalid config fails the turn with INVALID_ARGUMENT instead of crashing. - save_dir derives from the injected state_dir (#292) as state_dir / conversation_id. --- python/antigravity/harness_server.py | 78 +++++++++++++++++-- python/antigravity/harness_server_test.py | 94 +++++++++++++++++++++++ 2 files changed, 164 insertions(+), 8 deletions(-) diff --git a/python/antigravity/harness_server.py b/python/antigravity/harness_server.py index 612f291..fc152ed 100644 --- a/python/antigravity/harness_server.py +++ b/python/antigravity/harness_server.py @@ -19,6 +19,7 @@ import argparse import asyncio +import json import logging import os import pathlib @@ -40,6 +41,16 @@ from google.antigravity.types import Text, Thought, ToolCall +# AX owns these fields for resumption; a request must not set them. +# TODO: add validation for fields that are unsafe to set per execution +# (e.g. credentials, deployment routing) or that may only be set at +# conversation creation. +_AX_MANAGED_CONFIG_FIELDS = frozenset({"conversation_id", "save_dir"}) + + +class HarnessConfigError(ValueError): + """Raised when request harness_config is not a valid overlay.""" + class VertexKwargs(TypedDict, total=False): """Typed subset of LocalAgentConfig kwargs needed to enable Vertex AI. @@ -136,6 +147,30 @@ def _has_credentials(config: AgentConfig | None) -> bool: return False +def _parse_harness_config(raw_config: bytes) -> dict[str, object]: + if not raw_config: + return {} + + try: + config = json.loads(raw_config.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise HarnessConfigError(f"expected UTF-8 JSON: {exc}") from exc + + if not isinstance(config, dict): + raise HarnessConfigError("top-level JSON value must be an object") + + managed = sorted(set(config) & _AX_MANAGED_CONFIG_FIELDS) + if managed: + raise HarnessConfigError( + f"AX-managed field(s) cannot be set: {', '.join(managed)}" + ) + + # TODO: reject fields the SDK doesn't know. Pydantic defaults to + # extra="ignore", so a field-name typo is silently dropped instead of + # erroring; check against LocalAgentConfig.model_fields and raise. + return config + + def _existing_sdk_conv_id(save_dir: str) -> str | None: # SDK persists each conversation as {save_dir}/{sdk_conv_id}.db where sdk_conv_id # is SDK-picked (a hash), not our AX conversation_id. We give each AX conversation @@ -151,15 +186,29 @@ def __init__(self, default_config: AgentConfig, state_dir: pathlib.Path): self._default_config = default_config self._state_dir = state_dir - def _build_config_for(self, conversation_id: str) -> LocalAgentConfig: + def _build_config_for( + self, conversation_id: str, harness_config: bytes = b"" + ) -> LocalAgentConfig: + # Request overlay first, then AX-managed persistence values last so a + # request can never redirect trajectory storage. + overrides = _parse_harness_config(harness_config) + # Per-AX-conv save_dir under the configured state_dir base. Resume by # SDK's own conv_id if a trajectory exists there. SDK auto-creates the # directory. - save_dir = str(self._state_dir / conversation_id) - update = {"save_dir": save_dir} - if sdk_conv_id := _existing_sdk_conv_id(save_dir): - update["conversation_id"] = sdk_conv_id - return self._default_config.model_copy(update=update) + overrides["save_dir"] = str(self._state_dir / conversation_id) + if sdk_conv_id := _existing_sdk_conv_id(overrides["save_dir"]): + overrides["conversation_id"] = sdk_conv_id + + # Reconstruct (not model_copy) so the SDK re-validates overlaid values + # and surfaces its own error. + values = {name: getattr(self._default_config, name) + for name in self._default_config.model_fields_set} + values.update(overrides) + try: + return LocalAgentConfig(**values) + except (TypeError, ValueError) as exc: + raise HarnessConfigError(str(exc)) from exc async def Connect(self, request_iterator, context): # Each HarnessRequest{start} drives one stateless turn; the stream stays @@ -205,7 +254,6 @@ async def _run_turn(self, request): return latest_query_text = latest_message.content.text.text - # TODO(#194): parse and validate request.start.harness_config here. if not self._default_config: yield ax_pb2.HarnessResponse( conversation_id=request.conversation_id, @@ -219,7 +267,9 @@ async def _run_turn(self, request): ) return try: - per_conv_config = self._build_config_for(request.conversation_id) + per_conv_config = self._build_config_for( + request.conversation_id, request.start.harness_config + ) print(f"[gRPC] Starting Agent for conv_id={request.conversation_id}, save_dir={per_conv_config.save_dir}") async with Agent(per_conv_config) as agent: conversation = agent.conversation @@ -322,6 +372,18 @@ def flush_thought(): ) print("[gRPC] Turn completed successfully.") + except HarnessConfigError as exc: + yield ax_pb2.HarnessResponse( + conversation_id=request.conversation_id, + end=ax_pb2.HarnessEnd( + state=ax_pb2.STATE_FAILED, + error=ax_pb2.Error( + code=3, # INVALID_ARGUMENT + description=f"Invalid harness_config: {exc}", + ), + ), + ) + return except Exception as e: logging.exception("Error inside Connect servicer execution") yield ax_pb2.HarnessResponse( diff --git a/python/antigravity/harness_server_test.py b/python/antigravity/harness_server_test.py index e65f639..1223f38 100644 --- a/python/antigravity/harness_server_test.py +++ b/python/antigravity/harness_server_test.py @@ -13,10 +13,12 @@ # limitations under the License. import asyncio +import json import pytest import grpc from python.proto import ax_pb2, ax_pb2_grpc, content_pb2 from python.antigravity.harness_server import AntigravityHarnessServiceServicer +from python.antigravity.harness_server import HarnessConfigError from google.antigravity import LocalAgentConfig @pytest.fixture @@ -510,3 +512,95 @@ async def request_iter(): await server.stop(0) asyncio.run(_run()) + +def test_harness_config_empty_is_noop(mock_config, tmp_path): + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + assert servicer._build_config_for("conv-1", b"").system_instructions == ( + mock_config.system_instructions + ) + assert servicer._build_config_for("conv-1", b"{}").system_instructions == ( + mock_config.system_instructions + ) + + +def test_harness_config_overlay_applies(mock_config, tmp_path): + # Fields flow through to the SDK, which validates values. + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + raw = json.dumps({"system_instructions": "Answer in one sentence."}).encode() + + config = servicer._build_config_for("conv-1", raw) + + assert config.system_instructions == "Answer in one sentence." + + +def test_harness_config_overlay_keeps_ax_managed_save_dir(mock_config, tmp_path): + # A valid overlay must not disturb the AX-injected save_dir. + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + raw = json.dumps({"system_instructions": "x"}).encode() + + config = servicer._build_config_for("conv-1", raw) + + assert config.system_instructions == "x" + assert config.save_dir.endswith("/conv-1") + + +def test_harness_config_overlay_does_not_mutate_default(mock_config, tmp_path): + # Reconstruction must not mutate the shared server default. + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + servicer._build_config_for( + "conv-1", json.dumps({"system_instructions": "overridden"}).encode() + ) + assert mock_config.system_instructions == "Test instructions" + + +def test_harness_config_overlay_applies_multiple_fields(mock_config, tmp_path): + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + raw = json.dumps({ + "system_instructions": "x", + "model": "gemini-2.5-pro", + }).encode() + + config = servicer._build_config_for("conv-1", raw) + + assert config.system_instructions == "x" + assert config.model == "gemini-2.5-pro" + + +@pytest.mark.parametrize(("raw_config", "error"), [ + (b"{", "expected UTF-8 JSON"), + (b"\xff", "expected UTF-8 JSON"), + (json.dumps([]).encode(), "top-level JSON value must be an object"), + (json.dumps({"save_dir": "/tmp/other"}).encode(), "AX-managed field"), + (json.dumps({"conversation_id": "other"}).encode(), "AX-managed field"), + ( + json.dumps({"capabilities": {"enabled_tools": ["not-a-tool"]}}).encode(), + "validation error", + ), +]) +def test_harness_config_rejects(mock_config, tmp_path, raw_config, error): + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + with pytest.raises(HarnessConfigError, match=error): + servicer._build_config_for("conv-1", raw_config) + + +def test_run_turn_invalid_harness_config_maps_to_invalid_argument(mock_config, tmp_path): + async def _run(): + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + req = ax_pb2.HarnessRequest( + conversation_id="conv-1", + harness_id="antigravity", + start=ax_pb2.HarnessStart( + harness_config=b"{", + messages=[ax_pb2.Message( + role="user", + content=content_pb2.Content(text=content_pb2.TextContent(text="hi")), + )], + ), + ) + responses = [r async for r in servicer._run_turn(req)] + assert len(responses) == 1 + assert responses[0].end.state == ax_pb2.STATE_FAILED + assert responses[0].end.error.code == 3 + assert "Invalid harness_config" in responses[0].end.error.description + + asyncio.run(_run()) From 9a262f3aa1cc4c8d7e6dff3463e51262d433ef67 Mon Sep 17 00:00:00 2001 From: joycel-github Date: Mon, 13 Jul 2026 21:05:45 +0000 Subject: [PATCH 2/2] antigravity: address harness_config review comments - Inline _parse_harness_config into _build_config_for so the parsed dict stays a local intermediate; the method's only boundary type is the validated LocalAgentConfig (no dict[str, object] across a helper boundary). Addresses the "introduce a type for the config" comment. - Make per-conv save_dir test assertions portable: compare against str(tmp_path / conversation_id) instead of hardcoding "/conv-1", and drop the now-vestigial Path.home monkeypatch (save_dir derives from the injected state_dir since #292). Addresses the Windows path separator comment. _AX_MANAGED_CONFIG_FIELDS keeps guarding both conversation_id and save_dir on purpose: harness_config is a separate client-controlled surface from the request's conversation_id, so blocking it prevents a request from redirecting another conversation's trajectory storage. --- python/antigravity/harness_server.py | 56 ++++++++++------------- python/antigravity/harness_server_test.py | 12 ++--- 2 files changed, 30 insertions(+), 38 deletions(-) diff --git a/python/antigravity/harness_server.py b/python/antigravity/harness_server.py index fc152ed..898b0d9 100644 --- a/python/antigravity/harness_server.py +++ b/python/antigravity/harness_server.py @@ -147,30 +147,6 @@ def _has_credentials(config: AgentConfig | None) -> bool: return False -def _parse_harness_config(raw_config: bytes) -> dict[str, object]: - if not raw_config: - return {} - - try: - config = json.loads(raw_config.decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError) as exc: - raise HarnessConfigError(f"expected UTF-8 JSON: {exc}") from exc - - if not isinstance(config, dict): - raise HarnessConfigError("top-level JSON value must be an object") - - managed = sorted(set(config) & _AX_MANAGED_CONFIG_FIELDS) - if managed: - raise HarnessConfigError( - f"AX-managed field(s) cannot be set: {', '.join(managed)}" - ) - - # TODO: reject fields the SDK doesn't know. Pydantic defaults to - # extra="ignore", so a field-name typo is silently dropped instead of - # erroring; check against LocalAgentConfig.model_fields and raise. - return config - - def _existing_sdk_conv_id(save_dir: str) -> str | None: # SDK persists each conversation as {save_dir}/{sdk_conv_id}.db where sdk_conv_id # is SDK-picked (a hash), not our AX conversation_id. We give each AX conversation @@ -189,13 +165,31 @@ def __init__(self, default_config: AgentConfig, state_dir: pathlib.Path): def _build_config_for( self, conversation_id: str, harness_config: bytes = b"" ) -> LocalAgentConfig: - # Request overlay first, then AX-managed persistence values last so a - # request can never redirect trajectory storage. - overrides = _parse_harness_config(harness_config) - - # Per-AX-conv save_dir under the configured state_dir base. Resume by - # SDK's own conv_id if a trajectory exists there. SDK auto-creates the - # directory. + # Overlay the request's harness_config (JSON-in-bytes) onto the server + # default. The parsed dict is a local intermediate only; this method's + # boundary type is the validated LocalAgentConfig. + if harness_config: + try: + overrides = json.loads(harness_config.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise HarnessConfigError(f"expected UTF-8 JSON: {exc}") from exc + if not isinstance(overrides, dict): + raise HarnessConfigError("top-level JSON value must be an object") + managed = sorted(set(overrides) & _AX_MANAGED_CONFIG_FIELDS) + if managed: + raise HarnessConfigError( + f"AX-managed field(s) cannot be set: {', '.join(managed)}" + ) + # TODO: reject fields the SDK doesn't know. Pydantic defaults to + # extra="ignore", so a field-name typo is silently dropped instead + # of erroring; check against LocalAgentConfig.model_fields and raise. + else: + overrides = {} + + # AX-managed persistence values go on last so a request can never + # redirect trajectory storage. Per-AX-conv save_dir under the configured + # state_dir base; resume by the SDK's own conv_id if a trajectory + # already exists there. SDK auto-creates the directory. overrides["save_dir"] = str(self._state_dir / conversation_id) if sdk_conv_id := _existing_sdk_conv_id(overrides["save_dir"]): overrides["conversation_id"] = sdk_conv_id diff --git a/python/antigravity/harness_server_test.py b/python/antigravity/harness_server_test.py index 1223f38..cacce50 100644 --- a/python/antigravity/harness_server_test.py +++ b/python/antigravity/harness_server_test.py @@ -96,11 +96,9 @@ async def request_iter(): def test_grpc_connect_agent_per_turn_with_save_dir(mock_config, monkeypatch, tmp_path): - """Each turn spawns a fresh Agent with per-conv save_dir under - ~/.ax/antigravity/conversations/. Same AX conv_id -> same save_dir + """Each turn spawns a fresh Agent with per-conv save_dir under the + configured state_dir. Same AX conv_id -> same save_dir (SDK-native resume).""" - import pathlib as _pathlib - monkeypatch.setattr(_pathlib.Path, "home", lambda: tmp_path) async def _run(): server = grpc.aio.server() @@ -158,8 +156,8 @@ async def req_iter(): save_dirs = [a.config.save_dir for a in agent_instances] assert save_dirs[0] == save_dirs[1] assert save_dirs[0] != save_dirs[2] - assert save_dirs[0].endswith("/conv-1") - assert save_dirs[2].endswith("/conv-2") + assert save_dirs[0] == str(tmp_path / "conv-1") + assert save_dirs[2] == str(tmp_path / "conv-2") # conversation_id only passed when trajectory exists (not in these mocked runs). assert [a.config.conversation_id for a in agent_instances] == [None, None, None] @@ -541,7 +539,7 @@ def test_harness_config_overlay_keeps_ax_managed_save_dir(mock_config, tmp_path) config = servicer._build_config_for("conv-1", raw) assert config.system_instructions == "x" - assert config.save_dir.endswith("/conv-1") + assert config.save_dir == str(tmp_path / "conv-1") def test_harness_config_overlay_does_not_mutate_default(mock_config, tmp_path):