diff --git a/cmd/ax/internal/cliutil/cliutil.go b/cmd/ax/internal/cliutil/cliutil.go index 9e63e2e..3d2e9d1 100644 --- a/cmd/ax/internal/cliutil/cliutil.go +++ b/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) } diff --git a/internal/cmd/e2e/main.go b/internal/cmd/e2e/main.go index 4c0563d..0e48dc6 100644 --- a/internal/cmd/e2e/main.go +++ b/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) }) } diff --git a/internal/config/config.go b/internal/config/config.go index c8d82ed..7190d6b 100644 --- a/internal/config/config.go +++ b/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 diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 0c88010..dec9db3 100644 --- a/internal/config/config_test.go +++ b/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) + } +} diff --git a/internal/harness/antigravity/antigravity.go b/internal/harness/antigravity/antigravity.go index 4f9829f..3b1e8b2 100644 --- a/internal/harness/antigravity/antigravity.go +++ b/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) diff --git a/internal/harness/antigravity/antigravity_test.go b/internal/harness/antigravity/antigravity_test.go index 39cfae1..a96b47f 100644 --- a/internal/harness/antigravity/antigravity_test.go +++ b/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) } diff --git a/python/antigravity/harness_server.py b/python/antigravity/harness_server.py index c9daebc..612f291 100644 --- a/python/antigravity/harness_server.py +++ b/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() diff --git a/python/antigravity/harness_server_test.py b/python/antigravity/harness_server_test.py index fedefc7..43bab59 100644 --- a/python/antigravity/harness_server_test.py +++ b/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)