parse harness config from request

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(

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())