Antigravity sidecar state dir from config

cmd/ax/internal/cliutil/cliutil.go

@@ -71,8 +71,9 @@ func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller.Cont
 		if address == "" {
 			address = "127.0.0.1:50053"
 		}
-		// Local mode: the harness owns the Python sidecar.
-		antigravityHarness, err = antigravity.New(ctx, address, true)
+		// Local mode: the harness owns the Python sidecar. Empty StateDir
+		// means the sidecar applies its own default
+		antigravityHarness, err = antigravity.New(ctx, address, cfg.Harnesses.Antigravity.StateDir, true)
 		if err != nil {
 			return nil, fmt.Errorf("antigravity harness: %w", err)
 		}

internal/cmd/e2e/main.go

@@ -75,7 +75,7 @@ func main() {
 		conn.Close()
 		fmt.Printf("Connected to Antigravity gRPC harness server at %s\n", address)
 		// TODO: add a companion demo that uses autoStart=true (mirrors ax exec local mode).
-		h, _ := antigravity.New(ctx, address, false)
+		h, _ := antigravity.New(ctx, address, "", false)
 		reg.RegisterHarness("antigravity", h)
 	})
 }

internal/config/config.go

@@ -97,7 +97,8 @@ type HarnessesConfig struct {
 // AntigravityHarnessConfig registers the built-in Antigravity harness.
 type AntigravityHarnessConfig struct {
 	Default  bool   `yaml:"default,omitempty"`
-	Endpoint string `yaml:"endpoint,omitempty"` // HarnessService address
+	Endpoint string `yaml:"endpoint,omitempty"`  // HarnessService address
+	StateDir string `yaml:"state_dir,omitempty"` // Trajectory storage dir (optional; sidecar defaults to ~/.ax/antigravity/conversations)
 }
 
 // AntigravityInteractionsHarnessConfig registers the built-in Antigravity

internal/config/config_test.go

@@ -140,3 +140,19 @@ eventlog:
 		t.Errorf("cfg.Version = %q, want %q", cfg.Version, "1.2.3")
 	}
 }
+
+// TestLoadFromFile_AntigravityStateDir: yaml -> AntigravityHarnessConfig.StateDir.
+func TestLoadFromFile_AntigravityStateDir(t *testing.T) {
+	data := `
+harnesses:
+  antigravity:
+    state_dir: /custom/path
+`
+	var cfg Config
+	if err := yaml.Unmarshal([]byte(data), &cfg); err != nil {
+		t.Fatalf("Unmarshal failed: %v", err)
+	}
+	if got, want := cfg.Harnesses.Antigravity.StateDir, "/custom/path"; got != want {
+		t.Errorf("StateDir = %q, want %q", got, want)
+	}
+}

internal/harness/antigravity/antigravity.go

@@ -48,8 +48,10 @@ type AntigravityHarness struct {
 }
 
 // New creates a new AntigravityHarness. When autoStart is true, New forks the
-// Antigravity Python sidecar and waits for it to become reachable.
-func New(ctx context.Context, address string, autoStart bool) (*AntigravityHarness, error) {
+// Antigravity Python sidecar and waits for it to become reachable. If
+// stateDir is non-empty, it is passed to the sidecar as --state-dir; empty
+// stateDir lets the sidecar apply its own default.
+func New(ctx context.Context, address, stateDir string, autoStart bool) (*AntigravityHarness, error) {
 	if address == "" {
 		address = "127.0.0.1:50053"
 	}
@@ -61,12 +63,13 @@ func New(ctx context.Context, address string, autoStart bool) (*AntigravityHarne
 	if err != nil {
 		return nil, fmt.Errorf("antigravity: invalid address %q: %w", address, err)
 	}
+	args := []string{"--host", host, "--port", port}
+	if stateDir != "" {
+		args = append(args, "--state-dir", stateDir)
+	}
 	cfg := pythonsidecar.Config{
-		Module: "python.antigravity.harness_server",
-		Args: []string{
-			"--host", host,
-			"--port", port,
-		},
+		Module:    "python.antigravity.harness_server",
+		Args:      args,
 		ReadyFunc: pythonsidecar.TCPReady(address),
 	}
 	sidecar := pythonsidecar.New(cfg)

internal/harness/antigravity/antigravity_test.go

@@ -30,7 +30,7 @@ func TestRun_AutoStartFalse_ServerOK_Succeeds(t *testing.T) {
 	srv := &harnesstest.MockHarnessServer{
 		Outputs: []*proto.Message{harnesstest.ThoughtText("Analyzing"), harnesstest.AssistantText("Hello world")},
 	}
-	harnessClient, err := New(context.Background(), harnesstest.StartHarnessServer(t, srv), false)
+	harnessClient, err := New(context.Background(), harnesstest.StartHarnessServer(t, srv), "", false)
 	if err != nil {
 		t.Fatalf("New: %v", err)
 	}
@@ -75,7 +75,7 @@ func TestRun_AutoStartFalse_ServerOK_Succeeds(t *testing.T) {
 
 func TestRun_AutoStartFalse_ServerErrorFrame_Fails(t *testing.T) {
 	srv := &harnesstest.MockHarnessServer{FailConnect: true, ErrMessage: "internal mock server crash"}
-	harnessClient, err := New(context.Background(), harnesstest.StartHarnessServer(t, srv), false)
+	harnessClient, err := New(context.Background(), harnesstest.StartHarnessServer(t, srv), "", false)
 	if err != nil {
 		t.Fatalf("New: %v", err)
 	}
@@ -98,7 +98,7 @@ func TestRun_AutoStartFalse_ServerErrorFrame_Fails(t *testing.T) {
 
 // TestNew_AutoStartFalse_NilSidecar: autoStart=false leaves sidecar nil.
 func TestNew_AutoStartFalse_NilSidecar(t *testing.T) {
-	h, err := New(context.Background(), "127.0.0.1:1", false)
+	h, err := New(context.Background(), "127.0.0.1:1", "", false)
 	if err != nil {
 		t.Fatalf("New: %v", err)
 	}

python/antigravity/harness_server.py

@@ -147,16 +147,15 @@ def _existing_sdk_conv_id(save_dir: str) -> str | None:
 class AntigravityHarnessServiceServicer(ax_pb2_grpc.HarnessServiceServicer):
     """Implements the ax.HarnessService protocol over gRPC."""
 
-    def __init__(self, default_config: AgentConfig):
+    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. Resume by SDK's own conv_id if a trajectory
-        # exists there. SDK auto-creates the directory.
-        # TODO(#269, #203): should be injected by the controller via
-        # harness_config; hardcoded here to mirror the Go-side
-        # antigravityinteractions.DefaultStateDir() convention.
-        save_dir = str(pathlib.Path.home() / ".ax" / "antigravity" / "conversations" / conversation_id)
+        # 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
@@ -337,9 +336,9 @@ def flush_thought():
             )
             return
 
-async def _serve(host: str, port: int, default_config: AgentConfig):
+async def _serve(host: str, port: int, default_config: AgentConfig, state_dir: pathlib.Path):
     server = grpc.aio.server()
-    servicer = AntigravityHarnessServiceServicer(default_config)
+    servicer = AntigravityHarnessServiceServicer(default_config, state_dir)
     ax_pb2_grpc.add_HarnessServiceServicer_to_server(servicer, server)
 
     # Serve the standard gRPC health protocol.
@@ -387,6 +386,7 @@ def main():
     parser = argparse.ArgumentParser(description="Antigravity gRPC Harness Server")
     parser.add_argument("--port", type=int, default=50053, help="Port to bind the server to")
     parser.add_argument("--host", default="localhost", help="Host to bind the server to")
+    parser.add_argument("--state-dir", default=str(pathlib.Path.home() / ".ax" / "antigravity" / "conversations"), help="Base directory for per-conversation trajectory storage")
     args = parser.parse_args()
 
     try:
@@ -407,7 +407,7 @@ def main():
     # having this entry even if it's the OCI image.
     _resolve_localhost()
 
-    asyncio.run(_serve(args.host, args.port, default_config))
+    asyncio.run(_serve(args.host, args.port, default_config, pathlib.Path(args.state_dir).expanduser()))
 
 if __name__ == "__main__":
     main()

python/antigravity/harness_server_test.py

@@ -24,11 +24,11 @@ def mock_config(monkeypatch):
     monkeypatch.setenv("GEMINI_API_KEY", "mock-api-key")
     return LocalAgentConfig(system_instructions="Test instructions")
 
-def test_grpc_connect_success(mock_config, monkeypatch):
+def test_grpc_connect_success(mock_config, monkeypatch, tmp_path):
     async def _run():
         # 1. Start temporary local gRPC server on random open port
         server = grpc.aio.server()
-        servicer = AntigravityHarnessServiceServicer(mock_config)
+        servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path)
         ax_pb2_grpc.add_HarnessServiceServicer_to_server(servicer, server)
         port = server.add_insecure_port("localhost:0")
         await server.start()
@@ -102,7 +102,7 @@ def test_grpc_connect_agent_per_turn_with_save_dir(mock_config, monkeypatch, tmp
 
     async def _run():
         server = grpc.aio.server()
-        servicer = AntigravityHarnessServiceServicer(mock_config)
+        servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path)
         ax_pb2_grpc.add_HarnessServiceServicer_to_server(servicer, server)
         port = server.add_insecure_port("localhost:0")
         await server.start()
@@ -166,13 +166,13 @@ async def req_iter():
     asyncio.run(_run())
 
 
-def test_health_check():
+def test_health_check(tmp_path):
     async def _run():
         from grpc_health.v1 import health, health_pb2, health_pb2_grpc
 
         cfg = LocalAgentConfig(system_instructions="health-check stub")
         server = grpc.aio.server()
-        ax_pb2_grpc.add_HarnessServiceServicer_to_server(AntigravityHarnessServiceServicer(cfg), server)
+        ax_pb2_grpc.add_HarnessServiceServicer_to_server(AntigravityHarnessServiceServicer(cfg, tmp_path), server)
         health_servicer = health.aio.HealthServicer()
         health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
         await health_servicer.set("", health_pb2.HealthCheckResponse.SERVING)
@@ -230,7 +230,7 @@ def test_has_credentials_vertex_express_mode(monkeypatch):
     assert _has_credentials(cfg) is True
 
 
-def test_grpc_connect_programmatic_credentials(monkeypatch):
+def test_grpc_connect_programmatic_credentials(monkeypatch, tmp_path):
     monkeypatch.delenv("GEMINI_API_KEY", raising=False)
     monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
     monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False)
@@ -241,7 +241,7 @@ def test_grpc_connect_programmatic_credentials(monkeypatch):
 
     async def _run():
         server = grpc.aio.server()
-        servicer = AntigravityHarnessServiceServicer(cfg)
+        servicer = AntigravityHarnessServiceServicer(cfg, tmp_path)
         ax_pb2_grpc.add_HarnessServiceServicer_to_server(servicer, server)
         port = server.add_insecure_port("localhost:0")
         await server.start()
@@ -316,10 +316,10 @@ def test_enhance_config_from_env(monkeypatch, tmp_path):
     assert str(skills_dir) in cfg.skills_paths
 
 
-def test_grpc_connect_buffering(mock_config, monkeypatch):
+def test_grpc_connect_buffering(mock_config, monkeypatch, tmp_path):
     async def _run():
         server = grpc.aio.server()
-        servicer = AntigravityHarnessServiceServicer(mock_config)
+        servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path)
         ax_pb2_grpc.add_HarnessServiceServicer_to_server(servicer, server)
         port = server.add_insecure_port("localhost:0")
         await server.start()
@@ -473,14 +473,14 @@ def test_servicer_requires_default_config():
         AntigravityHarnessServiceServicer()
 
 
-def test_run_turn_guards_against_missing_default_config(monkeypatch):
+def test_run_turn_guards_against_missing_default_config(monkeypatch, tmp_path):
     """If something sets _default_config to None at runtime (future bug in
     per-request layering, #194), _run_turn returns STATE_FAILED instead of
     crashing the server.
     """
     async def _run():
         cfg = LocalAgentConfig(system_instructions="will be set to None")
-        servicer = AntigravityHarnessServiceServicer(cfg)
+        servicer = AntigravityHarnessServiceServicer(cfg, tmp_path)
         servicer._default_config = None
         server = grpc.aio.server()
         ax_pb2_grpc.add_HarnessServiceServicer_to_server(servicer, server)