diff --git a/python/antigravity/harness_server.py b/python/antigravity/harness_server.py index 612f291..898b0d9 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. @@ -151,15 +162,47 @@ 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: - # 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) + def _build_config_for( + self, conversation_id: str, harness_config: bytes = b"" + ) -> LocalAgentConfig: + # 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 + + # 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 +248,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 +261,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 +366,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..cacce50 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 @@ -94,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() @@ -156,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] @@ -510,3 +510,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 == str(tmp_path / "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())