diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..14a14253 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,8 @@ +.git +.idea +.vscode +*.log +eventlog/ +tmp/ +node_modules/ +bin/ diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml new file mode 100644 index 00000000..3fd354fc --- /dev/null +++ b/.github/workflows/go.yml @@ -0,0 +1,73 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This workflow will build a golang project +# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-go + +name: Go + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + +jobs: + + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Cache addlicense + id: cache-addlicense + uses: actions/cache@v4 + with: + path: ~/go/bin/addlicense + key: ${{ runner.os }}-addlicense-v1.2.0 + + - name: Install addlicense + if: steps.cache-addlicense.outputs.cache-hit != 'true' + run: go install github.com/google/addlicense@v1.2.0 + + - name: Check license headers + run: | + $(go env GOPATH)/bin/addlicense -check . + + - name: Set up Go + uses: actions/setup-go@v4 + with: + go-version-file: go.mod + + - name: Verify go mod tidy + run: | + go mod tidy + git diff --exit-code go.mod go.sum + + - name: Build + run: go build -v ./... + + - name: Test + run: go test -v ./... + + - name: Check for binary files in PR + run: | + # Get the list of files where Git shows no line count (binary) + BINARIES=$(git diff --numstat origin/${{ github.base_ref }}...HEAD | grep "^-" | cut -f3) + + if [ -n "$BINARIES" ]; then + echo "Binary files detected in diff:" + echo "$BINARIES" + exit 1 # Fail the check if binaries are found + fi diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml new file mode 100644 index 00000000..135f8911 --- /dev/null +++ b/.github/workflows/python.yml @@ -0,0 +1,51 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Runs pytest against the antigravity harness sidecar +# (python/antigravity/harness_server_test.py). Split from go.yml so the +# Go workflow stays Go-only and this can evolve independently as more +# Python surface area lands. + +name: Python + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + +jobs: + + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install Python deps + # Runtime deps from python/antigravity/requirements.txt plus test-only + # tooling. Version floors on pytest match the upstream google-antigravity + # SDK's own dev extras (see google-antigravity/antigravity-sdk-python + # pyproject.toml) so this doesn't drift. pytest-timeout is added here + # only -- our tests spin up real gRPC servers on random ports and can hang. + run: pip install -r python/antigravity/requirements.txt 'pytest>=7.0' 'pytest-timeout>=2.0' + + - name: Run Python tests + # --timeout guards against hung gRPC servers so a stuck test can't + # eat the whole workflow budget. + run: python -m pytest python/antigravity/ --timeout=30 --timeout-method=thread -v diff --git a/.gitignore b/.gitignore index 7fd6df11..727d1a2e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,8 @@ bin/* -eventlog/* \ No newline at end of file +eventlog/* +/ax +*.log +ax-bin +python/__pycache__ +python/proto/__pycache__ +__pycache__/ diff --git a/.ko.yaml b/.ko.yaml new file mode 100644 index 00000000..0041fed6 --- /dev/null +++ b/.ko.yaml @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +defaultBaseImage: gcr.io/distroless/static diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index bb80dabe..6b31efbf 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,6 +1,17 @@ # How to contribute -We'd love to accept your patches and contributions to this project. +We'd love to accept your patches and contributions to this project. However, please note our current temporary policy below. + +> [!IMPORTANT] +> **Temporary Pause on External Contributions** +> +> AX is currently undergoing a significant architectural redesign, and the codebase is highly unstable. To ensure we can focus on stabilizing the core framework, we are temporarily **pausing the acceptance of external Pull Requests**. +> +> Please do not open Pull Requests at this time, as they are likely to conflict with ongoing changes and may be closed without review. Instead, we encourage you to: +> - Open an **Issue** to propose feature requests or report bugs. +> - Join discussions on existing issues to share your feedback. +> +> We will update this guide once the core architecture stabilizes and we are ready to welcome contributions again. Thank you for your understanding and patience! ## Before you begin @@ -49,8 +60,81 @@ make proto make test ``` -### Run local agent +### Install the ax CLI ```bash -make run-local +make install +``` + +Ensure that $GOPATH/bin is in your $PATH. + +```bash +export PATH="$PATH:$(go env GOPATH)/bin" +``` + +If you have the `GOBIN` environment variable set, it will be installed there instead. +Make sure the installation directory is in your `$PATH`. + +To add the default location to your `PATH` for the current session, run: + +```bash +export PATH="$PATH:$(go env GOPATH)/bin" +``` + +### Creating a pull request + +First, clone the repo: + +``` +git clone git@github.com:google/ax.git +``` + +If you already have cloned the repo locally, make sure that +your main branch is up to date: + +``` +git checkout main +git pull -q -r origin main +``` + +Check a new feature branch: + +``` +git checkout -b my-feature +``` + +Make edits to files, and test them locally. Add the changes (e.g. git add .) to stage. +Commit the changes once you staged the changes: + +``` +git commit -m "Describe he changes made" +``` + +Push the branch to the origin and open a pull request: + +``` +git push origin my-feature +``` + +Visit https://github.com/google/ax to open a pull request. + + +## Troubleshooting + +### Outdated table schema + +AX is still under heavy development and the database schema is not yet stable. If you encounter errors related to outdated table schemas, you can reset the database by deleting the `eventlog` directory. + +An example: + +```bash +ax exec --input "hello" + +Error: error creating controller: failed to create event log: sqlite_eventlog: create index exec_checkpoint_id: SQL logic error: no such column: checkpoint_id (1) +``` + +Delete the `eventlog` directory and try again. + +```bash +rm -rf ./eventlog ``` diff --git a/Makefile b/Makefile index 45d1e015..ce04e638 100644 --- a/Makefile +++ b/Makefile @@ -1,33 +1,40 @@ -.PHONY: all build proto test clean install +.PHONY: all build proto test test-python clean install deps clean-logs # Build all binaries all: proto build # Build binaries build: - @echo "Building gar..." + @echo "Building ax..." @mkdir -p bin - @go build -o bin/gar ./cmd/gar - @echo "Building local agent example..." - @go build -o bin/local_agent ./examples/local_agent - @echo "Building remote agent example..." - @go build -o bin/remote_agent ./examples/remote_agent + @go build -o bin/ax ./cmd/ax @echo "Build complete!" + # Generate protobuf code proto: @echo "Generating protobuf code..." @export PATH=$$PATH:$$(go env GOPATH)/bin && \ protoc --go_out=. --go_opt=paths=source_relative \ --go-grpc_out=. --go-grpc_opt=paths=source_relative \ - proto/gar.proto + proto/ax.proto proto/content.proto + @python3 -m grpc_tools.protoc -I. --python_out=python --grpc_python_out=python proto/ax.proto proto/content.proto @echo "Protobuf generation complete!" -# Run tests +# Run Go tests test: - @echo "Running tests..." + @echo "Running Go tests..." @go test -v ./... +# Run Python tests for the antigravity harness sidecar. +# Assumes deps are installed for the same interpreter as `python3`, e.g.: +# python3 -m pip install -r python/antigravity/requirements.txt \ +# 'pytest>=7.0' 'pytest-timeout>=2.0' +# --timeout guards against hung gRPC servers. +test-python: + @echo "Running Python tests..." + @python3 -m pytest python/antigravity/ --timeout=30 --timeout-method=thread + # Clean build artifacts clean: @echo "Cleaning..." @@ -35,20 +42,12 @@ clean: @rm -rf eventlog/ @echo "Clean complete!" -# Install gar to GOPATH/bin +# Install ax to GOPATH/bin install: - @echo "Installing gar..." - @go install ./cmd/gar + @echo "Installing ax..." + @go install ./cmd/ax @echo "Install complete!" -# Run local agent example -run-local: - @go run ./examples/local_agent - -# Run remote agent example -run-remote: - @go run ./examples/remote_agent - # Install dependencies deps: @echo "Installing dependencies..." @@ -56,3 +55,8 @@ deps: @go install google.golang.org/protobuf/cmd/protoc-gen-go@latest @go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@latest @echo "Dependencies installed!" + +clean-logs: + @echo "Cleaning the event logs..." + rm -rf ./eventlog + mkdir ./eventlog diff --git a/README.md b/README.md index bcbfbf64..f142d3a9 100644 --- a/README.md +++ b/README.md @@ -1,406 +1,385 @@ -# gar - -GAR, a short for Google Agent Runtime, is a single-writer agent orchestrator system built in Go. It provides a minimal runtime that coordinates agentic loops, manages sessions with event logging, and communicates with both local and remote agents via streaming protocols. +# Agent Executor (AX) + +> [!WARNING] +> 🚧 **AX is in active early development.** +> +> We are actively refining our core, resumption protocols, +> and runtime specifications, which will introduce major breaking +> changes prior to a stable release. +> +> **Temporary Policy:** We are temporarily pausing the acceptance of external Pull Requests while we stabilize the core architecture. We warmly encourage you to open Issues for feedback and feature requests instead. +> +> We will announce this project +> widely soon. If you are interested in collaborating with us, +> please reach out to **ax-dev@google.com**! + +AX, short for Agent Executor, is a distributed harness runtime. +It dynamically provisions isolated environments from suspendable/resumable +images to execute harnesses and agents. +AX is designed for reliability, with native support for recovery +and execution resumption, even in distributed setups. ## Features -- **Streaming**: gRPC bidirectional streaming for agent communication -- **Local & Remote Agents**: Support for both in-process and remote agent deployment -- **Session Management**: Start, pause, resume, and inspect agentic loop sessions -- **Agent Registry**: Automatic health monitoring and agent discovery +- **Distributed Runtime**: Harnesses, skills, tools, and agents can execute in isolation +- **Resumption**: Automatic recovery from failures or interruptions +- **Built-in Harnesses**: Support for frontier harnesses and custom implementations +- **Auditing & Policy**: All user and agentic calls are coordinated by a common controller, easy to control and audit the overall execution and skill/tool/agent calls +- **Portability**: Runs anywhere, scales to small and large deployments +- **Customizability**: Agnostic of harness and model Built-in consistency and resumability features: -- **Single-Writer Architecture**: Centralized controller ensures consistent state management -- **Event Log**: Durable session state with automatic recovery +- **Single-Writer Architecture**: Single controller ensures consistent state management +- **Event Log**: Durable execution state with automatic recovery +- **Advanced Resumption**: Support for compute-layer actor resumption on compatible platforms + +## Demo + +[![Demo](https://i.imgur.com/ADiU1OP.png)](https://www.youtube.com/watch?v=L5Iw1IrZ6Nc) + +Watch our demo to see AX works when deployed on [Agent Substrate](https://github.com/agent-substrate/substrate). ## Overview +```mermaid +%%{init: {"flowchart": {"diagramPadding": 80}}}%% +graph LR + Client + + subgraph Cluster[" "] + Server["AX Server
(multi-tenant)"] + DB[("Event Log"
Storage)] + ControlService["Control API"] + Actor["AX Harness Server
(stateful session-tenant)"] + end + + SnapshotService["Snapshots"] + HarnessService["Harness or model
service"] + MCPServer["MCP server"] + + Client <-->|resumable stream| Server + Server <-->|scan/append| DB + Server --> ControlService + ControlService -->|resume/suspend| Actor + Server <-->|resumable stream| Actor + ControlService <-->|read/write| SnapshotService + Actor -.-> HarnessService + Actor --> MCPServer + Actor -.-> Environment ``` -┌────────────────────────┐ -│ [Controller] │ -│ - Session Manager │ -│ - Event Log │ -│ - Loop Executor │ -│ - Agent Registry │ -└──────┬──────────┬──────┘ - │ │ - (in-process) (gRPC stream) - │ │ - ┌───────┐ ┌───────┐ - │ Local │ │Remote │ - │ Agent │ │ Agent │ - └───────┘ └───────┘ -``` + +As agents evolve from simple assistants to autonomous long running workers, +developers need a robust runtime to manage state, ensure reliability, +and audit execution. As we are moving away from monolithic agents towards +distributed harnesses where tools, skills and agents are deployed as +isolated actors, a distributed runtime with dynamically spawned isolated +workers becomes a necessity. AX provides the foundational layer to fill these gaps. + +While compute-agnostic, AX is aiming to provide the best +experience on Kubernetes. + +We expect every sophisticated agentic application will need the +capabilities provided by AX. +We are building this layer as a widely available foundation, +enabling developers to focus on building their applications rather +than infrastructure. We decided to build this project in public to +validate every design decision before a stable release is cut. +We highly encourage you to give us feedback. ## Installation -Install the gar CLI directly from the repository: +Install the ax CLI directly from the repository: ```bash -go install github.com/google/gar/cmd/gar@latest +go install github.com/google/ax/cmd/ax@latest ``` ### Verify Installation -Check that gar is installed correctly: +Check that ax is installed correctly: ```bash -gar --help +ax --help ``` -You should see the gar CLI usage information. +You should see the ax CLI usage information. -## Quick Start - -### 1. Run Local Agent Example +### Kubernetes -```bash -# Run the local agent example by linking the local agent with controller. -go run examples/local_agent/main.go -``` +AX is natively supported on +[Agent Substrate](https://github.com/agent-substrate/substrate) +on Kubernetes and it's the recommended deployment option for production +use. For more details on setup and configuration, see the +[deployment guide](./manifests/README.md). +Read more about [this new layer](https://cloud.google.com/blog/products/containers-kubernetes/bringing-you-agent-sandbox-on-gke-and-agent-substrate) +that provides higher density to agentic workloads on Kubernetes. -### 2. Run Remote Agent with GAR Server +### Built-in Antigravity harness -This example demonstrates how the gar server triggers remote agents through the -AgentService.Process RPC. - -**Terminal 1** - Start the remote agent server: -```bash -go run examples/remote_agent/main.go -``` +Antigravity SDK is a reference harness implementation. Local execution needs +`python3` and `pip` available on your `PATH`. AX handles the rest: on first +`ax exec` it starts the harness server as a Python sidecar and installs the +pinned Antigravity SDK dependencies for you. -The remote agent runs as a gRPC server implementing AgentService on port :50051. +## Authentication -**Terminal 2** - Start the gar controller server: -```bash -gar serve -``` +The built-in Antigravity harness supports Google AI Studio and Vertex AI. -The gar server exposes the GARService on port :8494. +For Google AI Studio, set a Gemini API key: -**Terminal 3** - Register the remote agent and trigger a session: ```bash -# Register the remote agent with gar -gar register \ - --agent-id remote-echo-agent \ - --agent-name "Echo Agent" \ - --agent-description "Echoes input in uppercase" \ - --agent-addr localhost:50051 - -# Trigger a session - gar will trigger the remote agent via Process RPC -gar trigger \ - --session-id session123 \ - --input "Hello remote agent" - -# Inspect session details -gar inspect --session-id session123 +export GEMINI_API_KEY="your-api-key" ``` -## Usage - -### GAR CLI - -The `gar` command provides several subcommands: - -#### Trigger a Session +For Vertex AI, configure Application Default Credentials and the target project +and location: ```bash -gar trigger \ - --input \ - [--session-id ] \ - [--checkpoint ] \ - [--server
] +gcloud auth application-default login +export GOOGLE_CLOUD_PROJECT="your-project-id" +export GOOGLE_CLOUD_LOCATION="us-central1" +export GOOGLE_GENAI_USE_VERTEXAI=true ``` -Triggers a new agentic loop session or automatically resumes an existing one. If the session ID already exists, the session will be resumed from its last checkpoint (or a specific checkpoint if provided) with the new input. +## Quick Start -Options: -- `--input`: Input message to send to agents (required) -- `--session-id`: Unique session identifier (optional, generates UUID if not provided, or resumes if exists) -- `--checkpoint`: Resume from specific checkpoint (empty for latest) -- `--server`: gRPC controller server address (default: "localhost:8494") +### Execute -**Examples:** +The CLI starts the built-in Antigravity harness automatically. No separate harness server setup is required. ```bash -# Trigger a new session -gar trigger --input "Hello agent" - -# Resume an existing session with new input -gar trigger --session-id abc123 --input "Continue processing" +# Using the checked-in ax.yaml, which sets Antigravity as the default harness. +ax exec --input "Can you list this directory?" -# Resume from a specific checkpoint (useful for undoing mistakes or exploring alternatives) -gar trigger --session-id abc123 \ - --checkpoint "550e8400-e29b-41d4-a716-446655440000" \ - --input "Try a different approach" +# Using exec with an AX server +ax exec --input "Can you list this directory?" --server localhost:8494 ``` -#### Inspect a Session +Conversations can be continued any time: ```bash -gar inspect --session-id [--server
] +ax exec \ + --conversation d85a4b4e-c53b-4c84-b879-f10d905bce40 \ + --input "Show me the contents of README.md" ``` -Options: -- `--session-id`: Session identifier to inspect (required) -- `--server`: gRPC controller server address (default: "localhost:8494") +If the client gets disconnected, pass the last sequence it saw to +replay the events it missed. This catches the client up; it does not +rewind the conversation. -#### Register a Remote Agent +In this example, we catch up a client from sequence number 12: ```bash -gar register \ - --agent-id \ - --agent-addr
\ - --agent-name \ - --agent-description \ - [--server
] +ax exec \ + --conversation d85a4b4e-c53b-4c84-b879-f10d905bce40 \ + --last-seq 12 \ + --resume ``` -Options: -- `--agent-id`: Unique agent identifier (required) -- `--agent-addr`: gRPC agent server address (e.g., "localhost:50051") (required) -- `--agent-name`: Human-readable name for the agent (required) -- `--agent-description`: Description of agent capabilities (required) -- `--server`: gRPC controller server address (default: "localhost:8494") +Instead of running the default harness, you can start executing +any registered harness: -#### Run Server +```bash +ax exec \ + --harness antigravity \ + --input "Can you write me a simple HTTP server in Python?" +``` +If anything goes wrong during the execution of a harness, +you can resume an incomplete execution in a conversation: ```bash -gar serve [--config ] +ax exec \ + --conversation edf98ef5-4bb1-4a9e-a091-3a77e03727e6 \ + --harness antigravity \ + --resume ``` -Starts the controller as a gRPC server using a YAML configuration file. -Options: -- `--config`: Path to YAML configuration file (default: "gar.yaml") +## Usage -Example configuration file (`gar.yaml`): -```yaml -server: - address: ":8494" +The `ax` command provides several subcommands: -eventlog: - dir: "eventlog" - -# Maximum steps per trigger -max_steps: 50 - -# Health check interval for agents -health_check_interval: 30s - -# Remote agents to register on startup -remote_agents: - - id: "text-processing-agent" - name: "Text Processing Agent" - description: "An agent for text processing" - address: "localhost:50051" - metadata: - version: "1.0" -``` +### Execute -Example: ```bash -# Start server with default config (gar.yaml) -gar serve - -# Start server with custom config -gar serve --config my-config.yaml +ax exec \ + [--input ] \ + [--conversation ] \ + [--harness ] \ + [--harness-config ] \ + [--harness-config-json ] \ + [--server
] \ + [--config ] \ + [--resume] \ + [--last-seq ] ``` -### Checkpoints +Executes a new harness execution or automatically resumes an existing one. If the conversation ID already exists, the execution will be resumed from its last state. -Checkpoints provide a mechanism to save and resume session state at specific points. Every content event (both `CONTENT_IN` and `CONTENT_OUT`) automatically creates a checkpoint with a unique UUID. +Options: +- `--input`: Input message to send to agents (optional if `--resume` is set, otherwise required) +- `--conversation`: Conversation ID (optional, generates UUID if not provided, or resumes if exists) +- `--harness`: Harness ID to use (optional, defaults to the default harness) +- `--harness-config`: Path to a JSON file with per-request harness configuration (optional) +- `--harness-config-json`: Per-request harness configuration as an inline JSON string (optional, mutually exclusive with `--harness-config`) +- `--server`: gRPC controller server address (optional. If not provided, runs with a local built-in AX server) +- `--config`: Path to YAML configuration file (only used with a local built-in AX server, default: "ax.yaml") +- `--resume`: Resume a conversation without inputs (optional, mutually exclusive with `--input`) +- `--last-seq`: Last sequence number seen by the client to resume from (optional). The server replays any later events so the client can catch up after a disconnect. -**Usage Examples:** +**Examples:** ```bash -# Inspect a session to see available checkpoints -gar inspect --session-id session123 +# Execute a new execution +ax exec --input "Hello agents!" -# Resume from a specific checkpoint -gar trigger --session-id session123 \ - --checkpoint "550e8400-e29b-41d4-a716-446655440000" \ - --input "Try different approach" -``` +# Resume an existing execution with new input +ax exec --conversation a53d4db3-1165-4925-87da-be6c72bbdeb1 --input "Ok, now let's do something else..." +# Execute using server mode +ax exec --server localhost:8494 --input "Hello agents!" -### Event Log Format +# Execute using a specific harness +ax exec --harness antigravity --input "Write me a cool Go program!" -Event logs use JSON Lines format (one JSON object per line). Each entry includes the session ID, a monotonic sequence number, and checkpoint ID (for content events) for traceability: +# Execute with per-request harness config +ax exec \ + --harness-config-json '{"system_instructions":"Answer in one sentence.","model":"gemini-3.5-flash"}' \ + --input "Explain durable execution." -```json -{"session_id": "session123", "timestamp": "2026-01-02T10:30:00Z", "seq": 0, "type": "CONTENT_IN", "checkpoint_id": "550e8400-e29b-41d4-a716-446655440000", "data": {...}} -{"session_id": "session123", "timestamp": "2026-01-02T10:30:01Z", "seq": 1, "type": "CONTENT_OUT", "checkpoint_id": "6ba7b810-9dad-11d1-80b4-00c04fd430c8", "data": {...}} -{"session_id": "session123", "timestamp": "2026-01-02T10:30:02Z", "seq": 2, "type": "STATE", "data": {...}} +# To keep the same JSON in a file, use `--harness-config` instead: +ax exec --harness-config antigravity.json --input "Explain durable execution." ``` -**Event Types:** -- `CONTENT_IN`: Incoming content from user or external source (includes `checkpoint_id`) -- `CONTENT_OUT`: Outgoing content from agents (includes `checkpoint_id`) -- `STATE`: Session state changes, e.g. running, completed, failed. - -## Building Custom Agents - -### Local Agent - -```go -import ( - "context" - "github.com/google/gar/agent" - "github.com/google/gar/proto" -) - -// Define your process function using callback handler -processFunc := func(ctx context.Context, sessionID string, inputs []*proto.Content, handler agent.OutputHandler) error { - for _, content := range inputs { - output := &proto.Content{ - Role: "assistant", - Type: "text", - Mimetype: "text/plain", - Data: "Your response: " + content.Data, - } - if err := handler(output); err != nil { - return err - } - } - return nil -} - -// Define health check function (optional) -healthCheckFunc := func(ctx context.Context) error { - // Return nil if healthy, error otherwise - return nil -} - -// Create the agent -myAgent, err := agent.NewLocalAgent(agent.LocalAgentConfig{ - ID: "my-agent", - ProcessFunc: processFunc, - HealthCheckFunc: healthCheckFunc, // optional -}) -``` +### Serve -### Remote Agent - -Remote agents run as gRPC servers implementing the `AgentService` interface defined in `proto/gar.proto`. The gar controller triggers remote agents by calling their `Process` RPC with bidirectional streaming. - -```go -type server struct { - proto.UnimplementedAgentServiceServer -} - -// Process handles bidirectional streaming - gar controller calls this RPC -func (s *server) Process(stream proto.AgentService_ProcessServer) error { - for { - // Receive input content from gar controller - content, err := stream.Recv() - if err == io.EOF { - return nil - } - if err != nil { - return err - } - - // Process the content - output := &proto.Content{ - Role: "assistant", - Type: "text", - Mimetype: "text/plain", - Data: "Processed: " + content.Data, - } - - // Send response back to gar controller - if err := stream.Send(output); err != nil { - return err - } - } -} - -func (s *server) HealthCheck(ctx context.Context, req *proto.HealthCheckRequest) (*proto.HealthCheckResponse, error) { - // Return health status for gar controller health monitoring - return &proto.HealthCheckResponse{ - Healthy: true, - Message: "Agent is healthy", - }, nil -} +```bash +ax serve [--config ] ``` -**Workflow:** -1. Remote agent starts as gRPC server on a port (e.g., :50051) -2. Start gar controller: `gar serve` -3. Register with gar: `gar register --agent-id my-agent --agent-name "My Agent" --agent-description "Agent description" --agent-addr localhost:50051` -4. When gar triggers a session, it calls the agent's `Process` RPC -5. GAR streams input content → Agent processes → Agent streams output back - -See `examples/remote_agent/main.go` for a complete implementation. - -### Remote Python Agent +Starts the controller as a gRPC server using a YAML configuration file. -Python agents can be built using the GAR agent framework. First, install dependencies and generate Python gRPC code: +Options: +- `--config`: Path to YAML configuration file (default: "ax.yaml") -```bash -# Install dependencies -pip install grpcio grpcio-tools +Example configuration file (`ax.yaml`): +```yaml +version: v1alpha -# Generate Python code from proto file -python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. proto/gar.proto -``` +server: + address: ":8494" -Then implement your agent using the framework: - -```python -from gar import Agent -import proto.gar_pb2 as pb2 - -def process(session_id, inputs): - """Process incoming content list and yield responses""" - for content in inputs: - yield pb2.Content( - role="assistant", - type="text", - mimetype="text/plain", - data=f"Python processed: {content.data.upper()}" - ) - -def health_check(): - """Health check function that always returns healthy""" - return True, "OK", {} - -# Create and start the agent -agent = Agent( - process_func=process, - health_check_func=health_check -) -agent.serve(port=50051) +eventlog: + sqlite: + filename: "eventlog/log.sqlite" ``` -**Register and use:** +Example: ```bash -# Start the Python agent -python agent.py - -# Register with gar (in another terminal) -gar register \ - --agent-id "text-processing-agent" \ - --agent-name "Text Processing Agent" \ - --agent-description "An agent that processes text" \ - --agent-addr localhost:50051 - -# Trigger a session -gar trigger \ - --session-id session123 \ - --input "Hello, I heard that there is an agent that can help with processing this text!" -``` +# Start server with default config (ax.yaml) +ax serve -## Future Enhancements +# Start server with custom config +ax serve --config my-config.yaml +``` -- [ ] Replace Content with Interactions Content -- [ ] Observability and trajectory collection -- [ ] TLS support for remote agents -- [ ] gar deploy from container -- [ ] Fork session when resuming from a checkpoint that isn't the latest -- [ ] Create a package of local agents (listing files, reading files, grep'ing the directory) -- [ ] Web UI +## Extensions + +### Harnesses + +AX provides built-in harnesses (e.g. Antigravity) but you can bring your +own harness implementation by implementing `HarnessService`. On supported +compute services (e.g. Agent Substrate), AX automatically runs the +harness in isolation with automatic resumption and suspension. + +Traditional agents (e.g. tool use or workflow agents), or +language models can be implemented as harnesses. + +### Skills + +Built-in harnesses like Antigravity includes built-in support for +Agent Skills. See [Skills](examples/skills) for more. + +### MCP Tools + +Built-in harnesses like Antigravity provides support for discovering +and making calls to MCP tools when they are configured. + +## What AX is NOT? +* A managed service. AX is self-hosted and not a managed service. + We aim to make it easy for users to deploy and operate it on + their Kubernetes clusters. +* An agentic framework. AX is agnostic of the framework used + to build agents. +* A specific harness like a specific coding agent, e.g. Antigravity. + AX provides the serving layer around harnesses and is agnostic of the + harness implementation. Soon, we will allow users to bring their own + harnesses. +* A model specific controller. AX is agnostic of the models used. + +## Roadmap + +Below is an overview of our upcoming features and planned changes: + +1. Support for more frontier harnesses besides Antigravity +1. Support for BYOH (Bring Your Own Harness) +1. Support for tool call approvals from harnesses +1. Improvements to resumption protocols +1. Forking from event log and snapshots +1. Trajectory exposition +1. Better telemetry exposition +1. Integrations for policy, auditing, and more + +## Contributing + +Please refer to the [CONTRIBUTING.md](CONTRIBUTING.md) guide for instructions +on how to contribute to this project. + +We are currently undergoing a significant architectural redesign, and external contributions are temporarily paused. +However, in the meantime, we warmly encourage you to file bugs and +send feature requests. + +## History + +Over the years, teams across Google built and operated several +distributed execution engines. As these systems evolved, certain +architectural patterns consistently stood the test of time. +The teams realized they were repeatedly solving similar +orchestration problems, prompting the push to extract these +lessons into common runtime layer. + +While this common layer was taking shape, the AI landscape underwent +a massive shift. Applications were transitioning from +statless tool use agents to autonomous, long-running, self improving +agents that often need isolated resumable execution environments. +Also, from an efficiency standpoint, agentic workloads are inherently bursty. +An agent might compute intensively for a minute, then sit idle for +hours or days awaiting human approval. Keeping a stateful actor active +during these long idle periods is highly inefficient and cost-prohibitive +at scale. +Over the last 10 years, Kubernetes has become the standard for +large scale job orchestration, but it was fundamentally designed for +stateless microservices or predictable batch jobs -- +not for suspending and resuming stateful, sandboxed agent actors. + +Driven by these dual challenges, we decided to build a robust, common +agentic orchestrator designed specifically for the new compute +layers we are developing on Kubernetes. Our goal is to ease +the productionization of agents, allowing developers +and researchers to focus on building and evaluating their +applications rather than dealing with underlying infrastructure. + +AX is developed and maintained by the team actively working on +Google's internal runtime. Although the two projects operate +at different layers today, we are applying our knowledge +and insights to AX in the public domain every day. + +## Acknowledgements + +We thank Google DeepMind for their earlier work in distributed harnesses which +heavily influenced AX. +We thank the Google Kubernetes Engine team for their deep contributions +regarding isolation, resumption and job scheduling. ## License diff --git a/agent/agent.go b/agent/agent.go deleted file mode 100644 index 97c6aeda..00000000 --- a/agent/agent.go +++ /dev/null @@ -1,43 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Package agent provides interfaces and implementations for local and remote agents. -// Agents process content through callback handlers and can be registered with the controller. -package agent - -import ( - "context" - - "github.com/google/gar/proto" -) - -// OutputHandler is a callback function that handles output content from an agent. -// It is called for each piece of content the agent generates. -type OutputHandler func(content *proto.Content) error - -// Agent defines the common interface for both local and remote agents. -// Agents process content using callback handlers. -type Agent interface { - // Process handles processing of input content. - // It calls the output handler for each piece of content generated. - // The handler may be called multiple times during processing. - Process(ctx context.Context, sessionID string, inputs []*proto.Content, handler OutputHandler) error - - // HealthCheck checks if the agent is healthy and responsive. - // Returns an error if the agent is unhealthy or unreachable. - HealthCheck(ctx context.Context) error - - // Close gracefully shuts down the agent and releases resources. - Close() error -} diff --git a/agent/local.go b/agent/local.go deleted file mode 100644 index c9af3432..00000000 --- a/agent/local.go +++ /dev/null @@ -1,68 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package agent - -import ( - "context" - "fmt" - - "github.com/google/gar/proto" -) - -// LocalAgent wraps a local (in-process) agent implementation. -// It implements the Agent interface for agents running in the same process as the controller. -type LocalAgent struct { - processFunc func(ctx context.Context, sessionID string, inputs []*proto.Content, handler OutputHandler) error - healthCheckFunc func(ctx context.Context) error -} - -// LocalAgentConfig configures a local agent. -type LocalAgentConfig struct { - ProcessFunc func(ctx context.Context, sessionID string, inputs []*proto.Content, handler OutputHandler) error - HealthCheckFunc func(ctx context.Context) error -} - -// NewLocalAgent creates a new local agent with the provided configuration. -func NewLocalAgent(config LocalAgentConfig) (*LocalAgent, error) { - if config.ProcessFunc == nil { - return nil, fmt.Errorf("ProcessFunc cannot be nil") - } - - // Provide defaults for optional functions - if config.HealthCheckFunc == nil { - config.HealthCheckFunc = func(ctx context.Context) error { return nil } - } - - return &LocalAgent{ - processFunc: config.ProcessFunc, - healthCheckFunc: config.HealthCheckFunc, - }, nil -} - -// Process handles processing of input content with callback handler. -func (a *LocalAgent) Process(ctx context.Context, sessionID string, inputs []*proto.Content, handler OutputHandler) error { - return a.processFunc(ctx, sessionID, inputs, handler) -} - -// HealthCheck checks if the agent is healthy. -func (a *LocalAgent) HealthCheck(ctx context.Context) error { - return a.healthCheckFunc(ctx) -} - -// Close gracefully shuts down the agent. -func (a *LocalAgent) Close() error { - // Local agents don't typically need cleanup, but this can be extended - return nil -} diff --git a/agent/remote.go b/agent/remote.go deleted file mode 100644 index d375dffe..00000000 --- a/agent/remote.go +++ /dev/null @@ -1,170 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package agent - -import ( - "context" - "fmt" - "io" - "sync" - "time" - - "google.golang.org/grpc" - "google.golang.org/grpc/credentials/insecure" - "google.golang.org/grpc/metadata" - - "github.com/google/gar/proto" -) - -// RemoteAgent is a gRPC client that implements the Agent interface. -// It communicates with remote agent services over gRPC. -type RemoteAgent struct { - address string - conn *grpc.ClientConn - client proto.AgentServiceClient - mu sync.Mutex - reconnect bool - maxRetries int - dialOpts []grpc.DialOption -} - -// RemoteAgentConfig configures a remote agent client. -type RemoteAgentConfig struct { - Address string // gRPC server address (e.g., "localhost:50051") - Reconnect bool // Whether to automatically reconnect on failures - MaxRetries int // Maximum number of retry attempts (0 = infinite) - DialOpts []grpc.DialOption // gRPC dial options for customizing the connection -} - -// NewRemoteAgent creates a new remote agent client. -func NewRemoteAgent(config RemoteAgentConfig) (*RemoteAgent, error) { - if config.Address == "" { - return nil, fmt.Errorf("agent address cannot be empty") - } - - agent := &RemoteAgent{ - address: config.Address, - reconnect: config.Reconnect, - maxRetries: config.MaxRetries, - dialOpts: config.DialOpts, - } - - if err := agent.connect(); err != nil { - return nil, fmt.Errorf("failed to connect to remote agent: %w", err) - } - - return agent, nil -} - -// connect establishes a gRPC connection to the remote agent. -func (a *RemoteAgent) connect() error { - a.mu.Lock() - defer a.mu.Unlock() - - if a.conn != nil { - a.conn.Close() - } - - // Use provided dial options, or default to insecure credentials - opts := a.dialOpts - if len(opts) == 0 { - opts = []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())} - } - - conn, err := grpc.NewClient(a.address, opts...) - if err != nil { - return fmt.Errorf("failed to dial: %w", err) - } - - a.conn = conn - a.client = proto.NewAgentServiceClient(conn) - return nil -} - -// Process handles processing of input content with the remote agent. -func (a *RemoteAgent) Process(ctx context.Context, sessionID string, inputs []*proto.Content, handler OutputHandler) error { - // Add session_id to gRPC metadata - md := metadata.Pairs("session-id", sessionID) - ctx = metadata.NewOutgoingContext(ctx, md) - - stream, err := a.client.Process(ctx) - if err != nil { - return fmt.Errorf("failed to create stream: %w", err) - } - - // Send all inputs to the remote agent - for _, content := range inputs { - if err := stream.Send(content); err != nil { - return fmt.Errorf("failed to send content: %w", err) - } - } - - // Close the send direction to signal we're done sending - if err := stream.CloseSend(); err != nil { - return fmt.Errorf("failed to close send: %w", err) - } - - // Receive outputs and call handler for each - for { - content, err := stream.Recv() - if err == io.EOF { - // Stream completed successfully - return nil - } - if err != nil { - return fmt.Errorf("failed to receive content: %w", err) - } - - // Call the handler with the received content - if err := handler(content); err != nil { - return fmt.Errorf("handler error: %w", err) - } - - // Check for context cancellation - select { - case <-ctx.Done(): - return ctx.Err() - default: - } - } -} - -// HealthCheck checks if the remote agent is healthy. -func (a *RemoteAgent) HealthCheck(ctx context.Context) error { - ctx, cancel := context.WithTimeout(ctx, 5*time.Second) - defer cancel() - - resp, err := a.client.HealthCheck(ctx, &proto.HealthCheckRequest{}) - if err != nil { - return fmt.Errorf("health check failed: %w", err) - } - - if !resp.Healthy { - return fmt.Errorf("agent unhealthy: %s", resp.Message) - } - - return nil -} - -// Close gracefully shuts down the remote agent connection. -func (a *RemoteAgent) Close() error { - a.mu.Lock() - defer a.mu.Unlock() - - if a.conn != nil { - return a.conn.Close() - } - return nil -} diff --git a/ax.yaml b/ax.yaml new file mode 100644 index 00000000..56a980a1 --- /dev/null +++ b/ax.yaml @@ -0,0 +1,34 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Sample LOCAL configuration for the AX *harness* build (compiled with -tags=harness). +# Usage: ax serve --config ax.yaml +# ax exec --config ax.yaml --input "hello" # uses harnesses.default +version: v1alpha + +server: + address: ":8494" + +eventlog: + sqlite: + filename: "eventlog/log.sqlite" + +harnesses: + antigravity: + default: true + antigravity-interactions: {} + +telemetry: + otlp: + enabled: true diff --git a/cmd/ax/Dockerfile b/cmd/ax/Dockerfile new file mode 100644 index 00000000..16a1b901 --- /dev/null +++ b/cmd/ax/Dockerfile @@ -0,0 +1,101 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Multi-target AX image build: +# --target antigravity : python:3.13-slim + the Antigravity Python sidecar +# (SDK, localharness, agent). Used by the antigravity +# harness actor (`ax harness antigravity`). +# --target ax : debian:stable-slim + the Go `ax` binary. Used by +# the ax-server (`ax serve`) and the Go interactions +# harness actor (`ax harness antigravity-interactions`). +# +# Build context: repository root. Images target the cluster's linux/amd64 nodes, +# so build with `--platform linux/amd64`. +# For a standalone build: +# docker build --platform linux/amd64 --target antigravity -f cmd/ax/Dockerfile -t . && docker push +# docker build --platform linux/amd64 --target ax -f cmd/ax/Dockerfile -t . && docker push + +# --- Stage 1: build the ax Go binary (with the harness build tag) ------------- +FROM --platform=$BUILDPLATFORM golang:1.26 AS build +ARG TARGETOS=linux +ARG TARGETARCH=amd64 +WORKDIR /src +COPY go.mod go.sum ./ +RUN go mod download +COPY . . +RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} \ + go build -o /out/ax ./cmd/ax + +# --- antigravity image ------- +FROM python:3.13-slim AS antigravity +WORKDIR /workspace + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + ca-certificates \ + curl \ + less \ + procps \ + wget \ + && rm -rf /var/lib/apt/lists/* + +# Antigravity runtime deps, installed from PyPI at build time. +COPY python/antigravity/requirements.txt /tmp/requirements.txt +RUN pip install --no-cache-dir -r /tmp/requirements.txt + +# Python sidecar (with generated proto stubs), the agent it serves, and the ax +# binary built in stage 1. +COPY python/ /ax-app/python + +# You can copy skills or use skills managers (e.g. npx skills) +# to download them into the image. +# Example: COPY examples/skills/ /ax/skills/ + +ENV SKILLS_DIR=/ax/skills +COPY --from=build /out/ax /ax-app/ax + +# `from python.proto import ...` needs /ax-app on the path; the generated stubs do +# `from proto import ...`, which needs /ax-app/python. +ENV PYTHONPATH=/ax-app:/ax-app/python + +# Flush stdout/stderr immediately so server logs surface in container logs. +ENV PYTHONUNBUFFERED=1 + +EXPOSE 80 + +# Default command for local docker runs. Substrate ignores the image CMD and runs +# the ActorTemplate `command` instead. `ax harness` forks the Python sidecar, +# which serves the HarnessService on :80. +CMD ["/ax-app/ax", "harness", \ + "--host", "0.0.0.0", "--port", "80"] + +# --- ax image (Go-only) for the interactions harness and ax-server ----------- +FROM debian:stable-slim AS ax +WORKDIR /workspace + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + ca-certificates \ + curl \ + less \ + procps \ + wget \ + && rm -rf /var/lib/apt/lists/* + +COPY --from=build /out/ax /ax-app/ax + +EXPOSE 80 + +CMD ["/ax-app/ax", "harness", "antigravity-interactions", \ + "--host", "0.0.0.0", "--port", "80"] diff --git a/cmd/ax/dashboard.go b/cmd/ax/dashboard.go new file mode 100644 index 00000000..e518f74b --- /dev/null +++ b/cmd/ax/dashboard.go @@ -0,0 +1,475 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "database/sql" + _ "embed" + "encoding/json" + "fmt" + "log/slog" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "runtime" + "time" + + "github.com/google/ax/cmd/ax/internal/cliutil" + "github.com/google/ax/internal/controller/eventlog" + "github.com/google/ax/proto" + "github.com/spf13/cobra" + _ "modernc.org/sqlite" +) + +//go:embed web/index.html +var dashboardHTML string + +var ( + dashboardAddr string + dashboardConfigFile string +) + +var dashboardCmd = &cobra.Command{ + Use: "dashboard", + Short: "Start the AX Dashboard", + Long: `Start a local HTTP server to display AX conversations and executions dashboard.`, + RunE: runDashboard, +} + +func init() { + dashboardCmd.Flags().StringVar(&dashboardAddr, "addr", "localhost:8080", "Server address to listen on") + dashboardCmd.Flags().StringVar(&dashboardConfigFile, "config", "ax.yaml", "Path to YAML configuration file") +} + +type ConversationResponse struct { + ID string `json:"id"` + Agent string `json:"agent"` + Status string `json:"status"` + LastSeq int32 `json:"last_seq"` + Duration string `json:"duration"` +} + +func runDashboard(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, nil))) + + cfg, err := newConfig(cmd, dashboardConfigFile) + if err != nil { + return err + } + + dbPath := cfg.EventLog.SQLiteConfig.Filename + slog.InfoContext(ctx, "Opening event log database", slog.String("path", dbPath)) + + if err := os.MkdirAll(filepath.Dir(dbPath), 0755); err != nil { + return fmt.Errorf("failed to create database directory: %w", err) + } + + db, err := sql.Open("sqlite", dbPath) + if err != nil { + return fmt.Errorf("failed to open sqlite database: %w", err) + } + defer db.Close() + + // Verify database connection + if err := db.PingContext(ctx); err != nil { + return fmt.Errorf("failed to ping database: %w", err) + } + + // Create tables if they don't exist (to avoid crashes on fresh setup) + if _, err := db.Exec(` + CREATE TABLE IF NOT EXISTS conversation_log ( + conversation_id TEXT NOT NULL, + seq INTEGER NOT NULL, + payload TEXT NOT NULL, + PRIMARY KEY (conversation_id, seq) + )`); err != nil { + return fmt.Errorf("failed to initialize conversation_log table: %w", err) + } + + if _, err := db.Exec(` + CREATE TABLE IF NOT EXISTS execution_log ( + exec_id TEXT NOT NULL, + payload TEXT NOT NULL, + timestamp DATETIME NOT NULL + )`); err != nil { + return fmt.Errorf("failed to initialize execution_log table: %w", err) + } + + if _, err := db.Exec(`CREATE INDEX IF NOT EXISTS idx_execution_log_exec_id ON execution_log(exec_id)`); err != nil { + return fmt.Errorf("failed to create index on execution_log: %w", err) + } + + // Setup API handlers + mux := http.NewServeMux() + mux.HandleFunc("/api/conversations", func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + convs, err := fetchConversations(r.Context(), db) + if err != nil { + slog.ErrorContext(r.Context(), "Failed to fetch conversations", slog.Any("error", err)) + http.Error(w, "Internal server error", http.StatusInternalServerError) + return + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(convs); err != nil { + slog.ErrorContext(r.Context(), "Failed to encode conversations response", slog.Any("error", err)) + } + }) + + mux.HandleFunc("/api/trace", func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + convID := r.URL.Query().Get("conversation") + if convID == "" { + http.Error(w, "Missing conversation ID", http.StatusBadRequest) + return + } + + data, err := loadTraceData(r.Context(), cfg, convID) + if err != nil { + slog.ErrorContext(r.Context(), "Failed to load trace data", slog.String("conversation_id", convID), slog.Any("error", err)) + http.Error(w, "Internal server error", http.StatusInternalServerError) + return + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(data); err != nil { + slog.ErrorContext(r.Context(), "Failed to encode trace data response", slog.Any("error", err)) + } + }) + + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprint(w, dashboardHTML) + }) + + listener, err := net.Listen("tcp", dashboardAddr) + if err != nil { + return fmt.Errorf("failed to bind server to %s: %w", dashboardAddr, err) + } + defer listener.Close() + + addr := listener.Addr().String() + host, port, err := net.SplitHostPort(addr) + if err == nil && (host == "::" || host == "0.0.0.0" || host == "" || host == "[::]") { + addr = fmt.Sprintf("localhost:%s", port) + } + url := fmt.Sprintf("http://%s", addr) + + slog.InfoContext(ctx, "AX Dashboard started", slog.String("url", url)) + + go openBrowser(url) + + server := &http.Server{ + Handler: mux, + } + + return server.Serve(listener) +} + +func fetchConversations(ctx context.Context, db *sql.DB) ([]ConversationResponse, error) { + query := ` +SELECT + c.conversation_id, + c.last_seq, + c.state, + e.agent_id, + e.start_time, + e.end_time +FROM ( + SELECT conversation_id, seq AS last_seq, + json_extract(payload, '$.exec_id') AS exec_id, + json_extract(payload, '$.state') AS state + FROM conversation_log + WHERE (conversation_id, seq) IN ( + SELECT conversation_id, MAX(seq) + FROM conversation_log + GROUP BY conversation_id + ) +) c +LEFT JOIN ( + SELECT + exec_id, + json_extract(payload, '$.agent_id') AS agent_id, + MIN(timestamp) AS start_time, + MAX(timestamp) AS end_time + FROM execution_log + GROUP BY exec_id +) e ON c.exec_id = e.exec_id; +` + rows, err := db.QueryContext(ctx, query) + if err != nil { + return nil, err + } + defer rows.Close() + + convs := []ConversationResponse{} + for rows.Next() { + var id string + var lastSeq int32 + var state string + var agentID sql.NullString + var startTimeStr, endTimeStr sql.NullString + + err := rows.Scan(&id, &lastSeq, &state, &agentID, &startTimeStr, &endTimeStr) + if err != nil { + return nil, err + } + + agent := "unknown" + if agentID.Valid && agentID.String != "" { + agent = agentID.String + // Strip special prefix if it starts with "__" + if len(agent) > 2 && agent[:2] == "__" { + agent = agent[2:] + } + } + + durationStr := "N/A" + if startTimeStr.Valid && endTimeStr.Valid { + startTime, err1 := parseSQLiteTime(startTimeStr.String) + endTime, err2 := parseSQLiteTime(endTimeStr.String) + if err1 == nil && err2 == nil { + duration := endTime.Sub(startTime) + durationStr = fmt.Sprintf("%.1fs", duration.Seconds()) + } else { + slog.WarnContext(ctx, "Failed to parse sqlite timestamps", slog.String("start", startTimeStr.String), slog.String("end", endTimeStr.String), slog.Any("err1", err1), slog.Any("err2", err2)) + } + } + + status := state + if len(status) > 6 && status[:6] == "STATE_" { + status = status[6:] + } + if status == "PENDING" { + status = "RUNNING" + } + + convs = append(convs, ConversationResponse{ + ID: id, + Agent: agent, + Status: status, + LastSeq: lastSeq, + Duration: durationStr, + }) + } + + return convs, nil +} + +func parseSQLiteTime(s string) (time.Time, error) { + layouts := []string{ + time.RFC3339Nano, + time.RFC3339, + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02 15:04:05.999999999 -0700 MST", + "2006-01-02 15:04:05", + } + var err error + var t time.Time + for _, layout := range layouts { + t, err = time.Parse(layout, s) + if err == nil { + return t, nil + } + } + return time.Time{}, err +} + +func openBrowser(url string) { + var err error + switch runtime.GOOS { + case "linux": + err = exec.Command("xdg-open", url).Start() + case "windows": + err = exec.Command("rundll32", "url.dll,FileProtocolHandler", url).Start() + case "darwin": + err = exec.Command("open", url).Start() + } + if err != nil { + fmt.Printf("Failed to open browser: %v\n", err) + } +} + +type Text struct { + Text string `json:"text"` +} + +type Approval struct { + Approved bool `json:"approved"` +} + +type Confirmation struct { + ID string `json:"id"` + Question string `json:"question,omitempty"` + Approval *Approval `json:"approval,omitempty"` +} + +type Content struct { + Role string `json:"role"` + Text *Text `json:"text,omitempty"` + Confirmation *Confirmation `json:"confirmation,omitempty"` +} + +type ExecutionEvent struct { + ExecID string `json:"exec_id"` + AgentID string `json:"agent_id"` + Inputs []Content `json:"inputs"` + Outputs []Content `json:"outputs"` + State string `json:"state"` + Timestamp time.Time `json:"timestamp"` +} + +type ExecTrace struct { + ExecID string `json:"exec_id"` + AgentID string `json:"agent_id"` + Events []ExecutionEvent `json:"events"` +} + +type TraceData struct { + ConversationID string `json:"conversation_id"` + RootExecID string `json:"root_exec_id"` + Execs []ExecTrace `json:"execs"` +} + +func loadTraceData(ctx context.Context, cfg *cliutil.Config, convID string) (*TraceData, error) { + events, rootExecID, execIDs, err := fetch(ctx, cfg, convID) + if err != nil { + return nil, err + } + + // TODO(jbd): Trace view incorrectly displays graph executions. We are not + // changing the EventLog interface to fix this because the executor is being + // removed soon in favor of a linear execution model. We will adopt a different + // style of visualization once that's done. + return &TraceData{ + ConversationID: convID, + RootExecID: rootExecID, + Execs: buildExecTraces(execIDs, events), + }, nil +} + +func fetch(ctx context.Context, cfg *cliutil.Config, convID string) ([]*proto.ConversationEvent, string, []string, error) { + evLog, err := eventlog.OpenSQLiteEventLog(cfg.EventLog.SQLiteConfig.Filename) + if err != nil { + return nil, "", nil, fmt.Errorf("could not open sqlite eventlog: %w", err) + } + defer evLog.Close() + + convEvents, err := evLog.Events(ctx, convID) + if err != nil { + return nil, "", nil, fmt.Errorf("failed to query conversation events: %w", err) + } + + var execIDs []string + seen := make(map[string]bool) + for _, ev := range convEvents { + if ev.ExecId != "" && !seen[ev.ExecId] { + execIDs = append(execIDs, ev.ExecId) + seen[ev.ExecId] = true + } + } + + if len(execIDs) == 0 { + return nil, "", nil, fmt.Errorf("no executions found for conversation: %s", convID) + } + + // Use the first execID as the rootExecID as requested by user + rootExecID := execIDs[0] + + return convEvents, rootExecID, execIDs, nil +} + +func buildExecTraces(execIDs []string, events []*proto.ConversationEvent) []ExecTrace { + execsMap := make(map[string][]ExecutionEvent) + harnessIDs := make(map[string]string) + + for _, protoEv := range events { + exID := protoEv.ExecId + if protoEv.HarnessId != "" { + harnessIDs[exID] = protoEv.HarnessId + } + ev := extractExecutionEvent(exID, protoEv) + execsMap[exID] = append(execsMap[exID], ev) + } + + var execs []ExecTrace + for _, execID := range execIDs { + if evs, ok := execsMap[execID]; ok { + agentID := harnessIDs[execID] + execs = append(execs, ExecTrace{ + ExecID: execID, + AgentID: agentID, + Events: evs, + }) + } + } + + return execs +} + +func extractMsgs(protoContents []*proto.Message) []Content { + var results []Content + for _, c := range protoContents { + content := Content{Role: c.Role} + msgContent := c.GetContent() + if msgContent == nil { + continue + } + if textC := msgContent.GetText(); textC != nil { + content.Text = &Text{Text: textC.Text} + } else if conf := msgContent.GetConfirmation(); conf != nil { + content.Confirmation = &Confirmation{ + ID: conf.Id, + Question: conf.Question, + } + if app := conf.GetApproval(); app != nil { + content.Confirmation.Approval = &Approval{Approved: app.Approved} + } else if dec := conf.GetDecline(); dec != nil { + content.Confirmation.Approval = &Approval{Approved: !dec.Declined} + } + } + results = append(results, content) + } + return results +} + +func extractExecutionEvent(execID string, protoEv *proto.ConversationEvent) ExecutionEvent { + ev := ExecutionEvent{ + ExecID: execID, + AgentID: protoEv.HarnessId, + } + + ev.State = fmt.Sprint(protoEv.State) + if protoEv.HarnessId != "" { + ev.Inputs = extractMsgs(protoEv.Messages) + } else { + ev.Outputs = extractMsgs(protoEv.Messages) + } + + return ev +} diff --git a/cmd/ax/dashboard_test.go b/cmd/ax/dashboard_test.go new file mode 100644 index 00000000..df412296 --- /dev/null +++ b/cmd/ax/dashboard_test.go @@ -0,0 +1,139 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "database/sql" + "testing" + "time" + + _ "modernc.org/sqlite" +) + +func TestFetchConversations(t *testing.T) { + // Create an in-memory SQLite database for testing + db, err := sql.Open("sqlite", ":memory:") + if err != nil { + t.Fatalf("failed to open in-memory db: %v", err) + } + defer db.Close() + + // Initialize the schema + schema := ` + CREATE TABLE conversation_log ( + conversation_id TEXT, + seq INTEGER, + timestamp DATETIME, + payload TEXT + ); + + CREATE TABLE execution_log ( + exec_id TEXT, + timestamp DATETIME, + payload TEXT + ); + ` + if _, err := db.Exec(schema); err != nil { + t.Fatalf("failed to create schema: %v", err) + } + + // Insert test data + // Conversation 1: Only conversation_log (V2 execution without execution_log) + _, err = db.Exec(` + INSERT INTO conversation_log (conversation_id, seq, timestamp, payload) + VALUES ('conv-1', 1, ?, '{"exec_id": "exec-1", "state": "STATE_PENDING"}') + `, time.Now().Format(time.RFC3339)) + if err != nil { + t.Fatalf("failed to insert conv-1: %v", err) + } + + // Conversation 2: Both conversation_log and execution_log (V1 execution) + now := time.Now() + start := now.Add(-5 * time.Second) + end := now + + _, err = db.Exec(` + INSERT INTO conversation_log (conversation_id, seq, timestamp, payload) + VALUES ('conv-2', 5, ?, '{"exec_id": "exec-2", "state": "STATE_COMPLETED"}') + `, end.Format(time.RFC3339)) + if err != nil { + t.Fatalf("failed to insert conv-2: %v", err) + } + + _, err = db.Exec(` + INSERT INTO execution_log (exec_id, timestamp, payload) + VALUES + ('exec-2', ?, '{"agent_id": "my-agent"}'), + ('exec-2', ?, '{"agent_id": "my-agent"}') + `, start.Format(time.RFC3339), end.Format(time.RFC3339)) + if err != nil { + t.Fatalf("failed to insert exec-2: %v", err) + } + + // Fetch the conversations + ctx := context.Background() + convs, err := fetchConversations(ctx, db) + if err != nil { + t.Fatalf("fetchConversations failed: %v", err) + } + + if len(convs) != 2 { + t.Fatalf("expected 2 conversations, got %d", len(convs)) + } + + // Create a map for easy lookup + convMap := make(map[string]ConversationResponse) + for _, c := range convs { + convMap[c.ID] = c + } + + // Check conv-1 + c1, ok := convMap["conv-1"] + if !ok { + t.Fatalf("conv-1 not found") + } + if c1.Status != "RUNNING" { // STATE_PENDING -> RUNNING + t.Errorf("conv-1 expected status RUNNING, got %q", c1.Status) + } + if c1.Agent != "unknown" { + t.Errorf("conv-1 expected agent unknown, got %q", c1.Agent) + } + if c1.Duration != "N/A" { + t.Errorf("conv-1 expected duration N/A, got %q", c1.Duration) + } + if c1.LastSeq != 1 { + t.Errorf("conv-1 expected last_seq 1, got %d", c1.LastSeq) + } + + // Check conv-2 + c2, ok := convMap["conv-2"] + if !ok { + t.Fatalf("conv-2 not found") + } + if c2.Status != "COMPLETED" { // STATE_COMPLETED -> COMPLETED + t.Errorf("conv-2 expected status COMPLETED, got %q", c2.Status) + } + if c2.Agent != "my-agent" { + t.Errorf("conv-2 expected agent my-agent, got %q", c2.Agent) + } + // Duration should be roughly 5.0s + if c2.Duration != "5.0s" { + t.Errorf("conv-2 expected duration 5.0s, got %q", c2.Duration) + } + if c2.LastSeq != 5 { + t.Errorf("conv-2 expected last_seq 5, got %d", c2.LastSeq) + } +} diff --git a/cmd/ax/exec.go b/cmd/ax/exec.go new file mode 100644 index 00000000..947b6b18 --- /dev/null +++ b/cmd/ax/exec.go @@ -0,0 +1,476 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "os/signal" + "strings" + "sync" + "syscall" + "time" + + "github.com/google/ax/cmd/ax/internal" + "github.com/google/ax/cmd/ax/internal/cliutil" + "github.com/google/ax/proto" + "github.com/google/uuid" + "github.com/spf13/cobra" +) + +var ( + execConversationID string + execHarnessID string + execHarnessConfig string + execHarnessConfigJSON string + execInput string + execServerAddr string + execConfigFile string + execResume bool // allow resuming an execution without inputs + execLastSeq int32 +) + +var execCmd = &cobra.Command{ + Use: "exec", + Short: "Execute a conversation or resume an existing one", + Long: `Execute a new conversation or resume an existing one. +If no conversation ID is provided, a new UUID will be generated.`, + SilenceUsage: true, + RunE: runExec, +} + +func init() { + execCmd.Flags().StringVar(&execConversationID, "conversation", "", "Conversation ID (optional, generates UUID if not provided)") + execCmd.Flags().StringVar(&execHarnessID, "harness", "", "Harness ID (optional, default harness is used if not specified)") + execCmd.Flags().StringVar(&execHarnessConfig, "harness-config", "", "Path to a JSON file with per-request harness configuration") + execCmd.Flags().StringVar(&execHarnessConfigJSON, "harness-config-json", "", "Per-request harness configuration as an inline JSON string (mutually exclusive with --harness-config)") + execCmd.Flags().StringVar(&execInput, "input", "", "Input message to send (optional)") + execCmd.Flags().StringVar(&execServerAddr, "server", "", "gRPC controller server address (if specified, connects to remote server; otherwise runs with a local built-in AX server)") + execCmd.Flags().StringVar(&execConfigFile, "config", "ax.yaml", "Path to YAML configuration file (only used with a local built-in AX server)") + execCmd.Flags().BoolVar(&execResume, "resume", false, "Resume a conversation without inputs") + execCmd.Flags().Int32Var(&execLastSeq, "last-seq", 0, "Last sequence number seen by the client") + execCmd.MarkFlagsMutuallyExclusive("input", "resume") + execCmd.MarkFlagsMutuallyExclusive("harness-config", "harness-config-json") +} + +// TODO(jbd): Add multimodal input flags, e.g. --input-image. + +var ( + // The concrete type is *controller.Controller + execController cliutil.Controller + interruptHandler = NewInterruptHandler() +) + +func runExec(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + + if execConversationID == "" { + execConversationID = uuid.NewString() + } + + // Setup signal handling for graceful shutdown + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + sigChan := make(chan os.Signal, 2) // Buffer to not miss signals + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + + go func() { + for { + sig := <-sigChan + if sig == syscall.SIGTERM { + fmt.Println("\nReceived SIGTERM, exiting...") + interruptHandler.exit() + } + + if !interruptHandler.TriggerCancel() { + if interruptHandler.HandleInterrupt() { + fmt.Println("\nExiting...") + interruptHandler.exit() + } + } + } + }() + + if execServerAddr == "" { + cfg, err := newConfig(cmd, execConfigFile) + if err != nil { + return err + } + // Validate configuration (matches `ax serve`). + if err := cfg.Validate(); err != nil { + return fmt.Errorf("invalid configuration: %w", err) + } + c, err := cliutil.NewControllerFromConfig(ctx, cfg) + if err != nil { + return fmt.Errorf("error creating controller: %w", err) + } + execController = c + } + + var harnessConfig []byte + if execHarnessConfig != "" { + b, err := os.ReadFile(execHarnessConfig) + if err != nil { + return fmt.Errorf("failed to read harness config %q: %w", execHarnessConfig, err) + } + harnessConfig = b + } else if execHarnessConfigJSON != "" { + harnessConfig = []byte(execHarnessConfigJSON) + } + + return execLoop(ctx, execConversationID, execHarnessID, harnessConfig, execInput, execLastSeq) +} + +func execLoop(ctx context.Context, id string, harnessID string, harnessConfig []byte, input string, lastSeq int32) error { + d := internal.NewDisplay(id, os.Stdout) + d.DisplayHeader() + + var inputs []*proto.Message + if !execResume { + var quit bool + var err error + input, harnessConfig, quit, err = promptUser(d, input, harnessConfig) + if err != nil { + return err + } + if quit { + return nil + } + inputs = []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{ + Text: input, + }, + }, + }, + }, + } + } + + for { + reqCtx, cancel := context.WithCancel(ctx) + interruptHandler.SetActiveCancel(cancel) + + conf, err := runAutoExec(reqCtx, d, &proto.ExecRequest{ + ConversationId: id, + HarnessId: harnessID, + HarnessConfig: harnessConfig, + Inputs: inputs, + LastSeq: lastSeq, + }) + lastSeq = 0 // disable resuming from sequence, user sees the seq on the screen + + interruptHandler.ClearActiveCancel() + cancel() + + if err != nil { + if errors.Is(err, context.Canceled) { + fmt.Println("Request canceled.") + inputs = nil + continue + } + return err + } + + if conf != nil { + for { + approved, err := d.PromptForApproval(conf.Question) + if err != nil { + if errors.Is(err, internal.ErrUserAborted) { + if interruptHandler.HandleInterrupt() { + return nil + } + continue + } + return err + } + var decision []*proto.Message + if approved { + decision = []*proto.Message{{ + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Confirmation{ + Confirmation: &proto.ConfirmationContent{ + Id: conf.Id, + Decision: &proto.ConfirmationContent_Approval{ + Approval: &proto.ApprovalDecision{Approved: true}, + }, + }, + }, + }, + }} + } else { + decision = []*proto.Message{{ + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Confirmation{ + Confirmation: &proto.ConfirmationContent{ + Id: conf.Id, + Decision: &proto.ConfirmationContent_Decline{ + Decline: &proto.DeclineDecision{Declined: true}, + }, + }, + }, + }, + }} + } + + reqCtx, cancel := context.WithCancel(ctx) + interruptHandler.SetActiveCancel(cancel) + + conf, err = runAutoExec(reqCtx, d, &proto.ExecRequest{ + ConversationId: id, + HarnessId: harnessID, + HarnessConfig: harnessConfig, + Inputs: decision, + }) + + interruptHandler.ClearActiveCancel() + cancel() + + if err != nil { + if errors.Is(err, context.Canceled) { + fmt.Println("Request canceled.") + break + } + return err + } + if conf == nil { + break + } + } + } + + // Per-request config: clear the config after each turn. + harnessConfig = nil + + var quit bool + input, harnessConfig, quit, err = promptUser(d, "", harnessConfig) + if err != nil { + return err + } + if quit { + return nil + } + + inputs = []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{ + Text: input, + }, + }, + }, + }, + } + } +} + +func runAutoExec(ctx context.Context, d *internal.Display, req *proto.ExecRequest) (*proto.ConfirmationContent, error) { + fn := runExecHeadless + if execServerAddr != "" { + fn = runExecServer + } + return fn(ctx, d, req) +} + +func runExecHeadless(ctx context.Context, d *internal.Display, req *proto.ExecRequest) (*proto.ConfirmationContent, error) { + var confirmation *proto.ConfirmationContent + var lastSeq int32 + outputHandler := cliutil.ExecHandler(func(resp *proto.ExecResponse) error { + for _, m := range resp.Outputs { + if conf := m.GetContent().GetConfirmation(); conf != nil { + confirmation = conf + } + } + lastSeq = resp.Seq + displayContents(d, resp.Outputs) + return nil + }) + if err := execController.Exec(ctx, req, outputHandler); err != nil { + return nil, fmt.Errorf("error executing with local server: %w", err) + } + + if confirmation == nil { + d.FinishOutput(fmt.Sprintf("seq=%d", lastSeq)) + } + return confirmation, nil +} + +func runExecServer(ctx context.Context, d *internal.Display, req *proto.ExecRequest) (*proto.ConfirmationContent, error) { + conn, err := connect(execServerAddr) + if err != nil { + return nil, err + } + defer conn.Close() + + client := proto.NewExecutionServiceClient(conn) + stream, err := client.Exec(ctx, req) + if err != nil { + return nil, fmt.Errorf("error executing: %w", err) + } + + var confirmation *proto.ConfirmationContent + var lastSeq int32 + for { + resp, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + return nil, fmt.Errorf("error receiving response: %w", err) + } + lastSeq = resp.Seq + if resp.Outputs != nil { + for _, m := range resp.Outputs { + if conf := m.GetContent().GetConfirmation(); conf != nil { + confirmation = conf + } + } + displayContents(d, resp.Outputs) + } + } + if confirmation == nil { + d.FinishOutput(fmt.Sprintf("seq=%d", lastSeq)) + } + return confirmation, nil +} + +func displayContents(d *internal.Display, contents []*proto.Message) { + for _, output := range contents { + if content := output.GetContent(); content != nil { + d.Display(content) + } + } +} + +// promptUser loops until the user provides a non-empty input string. +// The "/config" command opens the harness config menu. +// It returns: +// - string: the valid user input +// - []byte: the (possibly updated) harness config +// - bool: true if the user entered a quit command +// - error: any error that occurred during prompting +func promptUser(d *internal.Display, input string, harnessConfig []byte) (string, []byte, bool, error) { + for { + for strings.TrimSpace(input) == "" { + var err error + input, err = d.PromptForInput() + if err != nil { + if errors.Is(err, internal.ErrUserAborted) { + if interruptHandler.HandleInterrupt() { + return "", harnessConfig, true, nil + } + input = "" // Continue loop to prompt again + continue + } + return "", harnessConfig, false, err + } + } + + trimmed := strings.TrimSpace(input) + if trimmed == "/config" { + cfg, err := runConfigMenu(d, harnessConfig) + if err != nil { + return "", harnessConfig, false, err + } + harnessConfig = cfg + input = "" // Re-prompt after handling the config. + continue + } + + d.DisplayInput(input) + if strings.ToLower(trimmed) == "q" { + d.ShowResumption(execConversationID, execServerAddr) + return "", harnessConfig, true, nil + } + return input, harnessConfig, false, nil + } +} + +// InterruptHandler encapsulates the cancellation and signal handling state. +type InterruptHandler struct { + mu sync.Mutex + activeCancel context.CancelFunc + interruptCount int32 +} + +// NewInterruptHandler creates a new InterruptHandler. +func NewInterruptHandler() *InterruptHandler { + return &InterruptHandler{} +} + +// SetActiveCancel sets the active cancel function. +func (h *InterruptHandler) SetActiveCancel(cancel context.CancelFunc) { + h.mu.Lock() + defer h.mu.Unlock() + h.activeCancel = cancel +} + +// ClearActiveCancel clears the active cancel function. +func (h *InterruptHandler) ClearActiveCancel() { + h.mu.Lock() + defer h.mu.Unlock() + h.activeCancel = nil +} + +// HandleInterrupt increments the interrupt count and returns true if the process should exit. +// It also starts a timer to reset the count. +func (h *InterruptHandler) HandleInterrupt() bool { + h.mu.Lock() + defer h.mu.Unlock() + h.interruptCount++ + if h.interruptCount == 1 { + fmt.Println("\nPress Ctrl+C again to exit.") + go func() { + time.Sleep(2 * time.Second) + h.mu.Lock() + h.interruptCount = 0 + h.mu.Unlock() + }() + return false + } + return true +} + +// TriggerCancel triggers cancellation if there is an active cancel function. +// It returns true if it triggered cancellation, or false if there was no active cancellation. +func (h *InterruptHandler) TriggerCancel() bool { + h.mu.Lock() + defer h.mu.Unlock() + if h.activeCancel != nil { + fmt.Println("\nCanceling current request...") + h.activeCancel() + return true + } + return false +} + +// exit gracefully shuts down the controller (if any) and terminates the process. +func (h *InterruptHandler) exit() { + if execController != nil { + execController.Close() + } + os.Exit(1) +} diff --git a/cmd/ax/harness.go b/cmd/ax/harness.go new file mode 100644 index 00000000..c028007b --- /dev/null +++ b/cmd/ax/harness.go @@ -0,0 +1,215 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package main: the `ax harness` command. It runs one harness in-process, +// selected by the optional [harness-id] argument: "antigravity" (default) +// or "antigravity-interactions" +package main + +import ( + "context" + "encoding/base64" + "fmt" + "log" + "net" + "net/http" + "os" + "os/signal" + "strconv" + "syscall" + "time" + + "github.com/google/ax/internal/config" + "github.com/google/ax/internal/harness/antigravity" + "github.com/google/ax/internal/harness/antigravityinteractions" + "github.com/google/ax/internal/pythonsidecar" + "github.com/spf13/cobra" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + healthpb "google.golang.org/grpc/health/grpc_health_v1" +) + +var ( + harnessPort int + harnessHost string +) + +// harnessReadyzPort is the port for the HTTP /readyz readiness endpoint. +const harnessReadyzPort = 8081 + +var harnessCmd = &cobra.Command{ + Use: "harness [harness-id]", + Short: "Run the harness gRPC server", + Args: cobra.MaximumNArgs(1), + Hidden: true, + RunE: runHarness, +} + +func init() { + harnessCmd.Flags().IntVar(&harnessPort, "port", 50053, "Port for the HarnessService to listen on") + harnessCmd.Flags().StringVar(&harnessHost, "host", "127.0.0.1", "Host interface for the HarnessService to bind") + rootCmd.AddCommand(harnessCmd) +} + +// setHarnessWorkDir changes the process working directory to AX_HARNESS_WORKDIR +// when it is set. The forked Python sidecar inherits it, which scopes the agent's +// default workspace (os.getcwd()) away from its own source tree. +func setHarnessWorkDir() error { + dir := os.Getenv("AX_HARNESS_WORKDIR") + if dir == "" { + return nil + } + if err := os.Chdir(dir); err != nil { + return fmt.Errorf("set harness working directory %q: %w", dir, err) + } + log.Printf("harness working directory set to %s", dir) + return nil +} + +// runHarness runs the harness selected by the optional [harness-id] argument: +// "antigravity" (default) or "antigravity-interactions". +func runHarness(cmd *cobra.Command, args []string) error { + harnessID := config.AntigravityHarnessID + if len(args) > 0 { + harnessID = args[0] + } + switch harnessID { + case config.AntigravityInteractionsHarnessID: + return runAntigravityInteractionsHarness(cmd.Context()) + case config.AntigravityHarnessID: + return runAntigravityHarness(cmd) + default: + return fmt.Errorf("unknown harness %q (want %q or %q)", + harnessID, config.AntigravityHarnessID, config.AntigravityInteractionsHarnessID) + } +} + +// runAntigravityHarness forks the Antigravity Python sidecar server, which serves +// the HarnessService (and gRPC health) on the configured port. ax harness +// supervises the child: it forwards termination signals and exits with its status. +func runAntigravityHarness(cmd *cobra.Command) error { + if err := setHarnessWorkDir(); err != nil { + return err + } + + stateDir, err := antigravity.DefaultStateDir() + if err != nil { + return fmt.Errorf("failed to resolve antigravity state dir: %w", err) + } + + cfg := pythonsidecar.Config{ + Module: "python.antigravity.harness_server", + Args: []string{ + "--host", harnessHost, + "--port", strconv.Itoa(harnessPort), + "--state-dir", stateDir, + }, + Stdout: os.Stdout, + Stderr: os.Stderr, + ReadyFunc: pythonsidecar.TCPReady(net.JoinHostPort("127.0.0.1", strconv.Itoa(harnessPort))), + } + + sidecar := pythonsidecar.New(cfg) + if err := sidecar.Start(cmd.Context(), ""); err != nil { + return fmt.Errorf("failed to start antigravity harness server: %w", err) + } + log.Printf("forked antigravity harness server (pid %d) on %s:%d", sidecar.Pid(), harnessHost, harnessPort) + + // Serve the /readyz endpoint that substrate's readiness probe polls (during + // golden snapshotting and per-actor Run/Restore). + go serveReadyz(harnessReadyzPort, harnessPort) + + // Forward termination signals to the child so substrate can stop the actor. + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + go func() { + <-sigChan + _ = sidecar.Stop() + }() + + if err := sidecar.Wait(); err != nil { + return fmt.Errorf("antigravity harness server exited: %w", err) + } + return nil +} + +// runAntigravityInteractionsHarness runs the Go Antigravity Interactions harness +// server (HarnessService + gRPC health + HTTP /readyz) on the configured ports. +func runAntigravityInteractionsHarness(ctx context.Context) error { + if err := setHarnessWorkDir(); err != nil { + return err + } + + // Read harness config from ax.yaml handed to the actor via AX_CONFIG_CONTENT. + // Fall back to built-in defaults. + var hc config.AntigravityInteractionsHarnessConfig + if raw := os.Getenv("AX_CONFIG_CONTENT"); raw != "" { + if data, err := base64.StdEncoding.DecodeString(raw); err != nil { + log.Printf("AX_CONFIG_CONTENT: base64 decode failed, using defaults: %v", err) + } else if cfg, err := config.LoadFromBytes(data); err != nil { + log.Printf("AX_CONFIG_CONTENT: parse failed, using defaults: %v", err) + } else { + hc = cfg.Harnesses.AntigravityInteractions + } + } + agent := hc.Agent + if agent == "" { + agent = antigravityinteractions.DefaultAgent + } + // AX owns the resume-cursor path (a harness implementation detail), so it's + // derived internally rather than read from ax.yaml. + stateDir, err := antigravityinteractions.DefaultStateDir() + if err != nil { + return err + } + + cfg := antigravityinteractions.AntigravityInteractionsConfig{ + Agent: agent, + StateDir: stateDir, + } + return antigravityinteractions.Serve(ctx, cfg, harnessHost, harnessPort, harnessReadyzPort) +} + +// serveReadyz serves the HTTP /readyz endpoint on readyzPort that substrate's +// readiness probe polls (during golden snapshotting and per-actor Run/Restore). +// Each request forwards to the forked harness's gRPC health Check. +func serveReadyz(readyzPort, grpcPort int) { + conn, err := grpc.NewClient( + net.JoinHostPort("127.0.0.1", strconv.Itoa(grpcPort)), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + log.Printf("readyz: failed to create gRPC health client: %v", err) + return + } + healthClient := healthpb.NewHealthClient(conn) + + mux := http.NewServeMux() + mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) { + ctx, cancel := context.WithTimeout(r.Context(), time.Second) + defer cancel() + resp, err := healthClient.Check(ctx, &healthpb.HealthCheckRequest{}) + if err != nil || resp.GetStatus() != healthpb.HealthCheckResponse_SERVING { + http.Error(w, "not ready", http.StatusServiceUnavailable) + return + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok\n")) + }) + + addr := net.JoinHostPort(harnessHost, strconv.Itoa(readyzPort)) + if err := http.ListenAndServe(addr, mux); err != nil { + log.Printf("readyz server on %s exited: %v", addr, err) + } +} diff --git a/cmd/ax/harness_test.go b/cmd/ax/harness_test.go new file mode 100644 index 00000000..599f5abb --- /dev/null +++ b/cmd/ax/harness_test.go @@ -0,0 +1,128 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "fmt" + "net" + "net/http" + "os" + "path/filepath" + "testing" + "time" + + "google.golang.org/grpc" + "google.golang.org/grpc/health" + healthpb "google.golang.org/grpc/health/grpc_health_v1" +) + +func TestSetHarnessWorkDir(t *testing.T) { + // os.Chdir is global process state; save and restore it around the test. + orig, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + t.Cleanup(func() { _ = os.Chdir(orig) }) + + t.Run("unset leaves the working directory unchanged", func(t *testing.T) { + t.Setenv("AX_HARNESS_WORKDIR", "") + before, _ := os.Getwd() + if err := setHarnessWorkDir(); err != nil { + t.Fatalf("setHarnessWorkDir: %v", err) + } + after, _ := os.Getwd() + if before != after { + t.Errorf("working directory changed: %q -> %q", before, after) + } + }) + + t.Run("set changes the working directory", func(t *testing.T) { + dir := t.TempDir() + t.Setenv("AX_HARNESS_WORKDIR", dir) + if err := setHarnessWorkDir(); err != nil { + t.Fatalf("setHarnessWorkDir: %v", err) + } + got, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + want, _ := filepath.EvalSymlinks(dir) + gotResolved, _ := filepath.EvalSymlinks(got) + if gotResolved != want { + t.Errorf("working directory = %q, want %q", gotResolved, want) + } + }) + + t.Run("missing directory returns an error", func(t *testing.T) { + t.Setenv("AX_HARNESS_WORKDIR", filepath.Join(t.TempDir(), "does-not-exist")) + if err := setHarnessWorkDir(); err == nil { + t.Error("expected an error for a missing directory, got nil") + } + }) +} + +func TestServeReadyz(t *testing.T) { + // Start a real gRPC server exposing the standard health service; this stands + // in for the Antigravity Python harness that serveReadyz probes. + grpcListener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen for gRPC: %v", err) + } + healthServer := health.NewServer() + grpcServer := grpc.NewServer() + healthpb.RegisterHealthServer(grpcServer, healthServer) + go func() { _ = grpcServer.Serve(grpcListener) }() + t.Cleanup(grpcServer.Stop) + grpcPort := grpcListener.Addr().(*net.TCPAddr).Port + + // Reserve an ephemeral port for /readyz, then hand it to serveReadyz to bind. + readyzListener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen for readyz: %v", err) + } + readyzPort := readyzListener.Addr().(*net.TCPAddr).Port + _ = readyzListener.Close() + + // Start with the harness not serving, then run serveReadyz against it. + healthServer.SetServingStatus("", healthpb.HealthCheckResponse_NOT_SERVING) + go serveReadyz(readyzPort, grpcPort) + url := fmt.Sprintf("http://127.0.0.1:%d/readyz", readyzPort) + + // /readyz re-derives readiness from gRPC health on every request: 503 until + // the harness is SERVING, 200 once it is, and back to 503 if it stops serving. + waitReadyz(t, url, http.StatusServiceUnavailable) + healthServer.SetServingStatus("", healthpb.HealthCheckResponse_SERVING) + waitReadyz(t, url, http.StatusOK) + healthServer.SetServingStatus("", healthpb.HealthCheckResponse_NOT_SERVING) + waitReadyz(t, url, http.StatusServiceUnavailable) +} + +func waitReadyz(t *testing.T, url string, want int) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + got := 0 + for time.Now().Before(deadline) { + resp, err := http.Get(url) + if err == nil { + got = resp.StatusCode + _ = resp.Body.Close() + if got == want { + return + } + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("GET %s = %d, want %d", url, got, want) +} diff --git a/cmd/ax/harnessclient.go b/cmd/ax/harnessclient.go new file mode 100644 index 00000000..dc0f778e --- /dev/null +++ b/cmd/ax/harnessclient.go @@ -0,0 +1,130 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package main implements a simple client for the fake HarnessService. +// It is intended for testing purposes only and should be replaced with +// the actual ax client implementation. +// TODO(wjjclaud): Update or replace this file with ax client implementation. +package main + +import ( + "bufio" + "context" + "fmt" + "io" + "log" + "os" + + "github.com/google/ax/proto" + "github.com/google/uuid" + "github.com/spf13/cobra" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" +) + +var ( + harnessServerAddr string + harnessClientID string +) + +var harnessClientCmd = &cobra.Command{ + Use: "harnessclient", + Short: "Run the harness client to connect to the server", + Hidden: true, + RunE: runHarnessClient, +} + +func init() { + harnessClientCmd.Flags().StringVar(&harnessServerAddr, "server", "localhost:50053", "The server address for the gRPC HarnessService.") + harnessClientCmd.Flags().StringVar(&harnessClientID, "harness", "testharness", "The harness id to send on the request envelope.") + rootCmd.AddCommand(harnessClientCmd) +} + +func runHarnessClient(cmd *cobra.Command, args []string) error { + ctx := context.Background() + + log.Printf("Connecting to HarnessService at %s...", harnessServerAddr) + conn, err := grpc.NewClient(harnessServerAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + return fmt.Errorf("failed to connect to server: %v", err) + } + defer conn.Close() + + client := proto.NewHarnessServiceClient(conn) + + fmt.Print("Client > ") + scanner := bufio.NewScanner(os.Stdin) + scanner.Scan() + input := scanner.Text() + + stream, err := client.Connect(ctx) + if err != nil { + return fmt.Errorf("failed to open connection stream: %v", err) + } + + // A single HarnessRequest{start} initiates the turn. + start := &proto.HarnessRequest{ + ConversationId: uuid.NewString(), + HarnessId: harnessClientID, + Type: &proto.HarnessRequest_Start{ + Start: &proto.HarnessStart{ + Messages: []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{Text: &proto.TextContent{Text: input}}, + }, + }, + }, + }, + }, + } + if err := stream.Send(start); err != nil { + return fmt.Errorf("failed to send start: %v", err) + } + if err := stream.CloseSend(); err != nil { + return fmt.Errorf("failed to close send side: %v", err) + } + + // Drain HarnessResponse frames until HarnessEnd / EOF. + for { + resp, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + return fmt.Errorf("failed to receive response: %v", err) + } + switch payload := resp.Type.(type) { + case *proto.HarnessResponse_Outputs: + for i, m := range payload.Outputs.Messages { + var text string + if tb, ok := m.Content.Type.(*proto.Content_Text); ok { + text = tb.Text.Text + } + fmt.Printf("Server > message[%d] (%s): %s\n", i, m.Role, text) + } + case *proto.HarnessResponse_End: + if errDetail := payload.End.GetError(); errDetail != nil { + fmt.Printf("Server > [end] state=%s error=(code=%d description=%q)\n", + payload.End.GetState(), errDetail.GetCode(), errDetail.GetDescription()) + } else { + fmt.Printf("Server > [end] state=%s\n", payload.End.GetState()) + } + } + } + + log.Println("Stream closed successfully by server.") + return nil +} diff --git a/cmd/ax/harnessconfig.go b/cmd/ax/harnessconfig.go new file mode 100644 index 00000000..24492ab6 --- /dev/null +++ b/cmd/ax/harnessconfig.go @@ -0,0 +1,138 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "os" + "strings" + + "github.com/google/ax/cmd/ax/internal" +) + +// runConfigMenu shows the /config menu and returns the (possibly updated) config. +// An updated config is sent on subsequent requests. +func runConfigMenu(d *internal.Display, harnessConfig []byte) ([]byte, error) { + for { + action, err := d.PromptForConfigAction() + if err != nil { + if errors.Is(err, internal.ErrUserAborted) { + return harnessConfig, nil // Esc/Ctrl+C on the menu cancels /config. + } + return harnessConfig, err + } + + switch action { + case "edit": + cfg, done, err := editHarnessConfig(d, harnessConfig) + if err != nil { + return harnessConfig, err + } + if done { + return cfg, nil + } + case "load": + cfg, done, err := loadHarnessConfig(d) + if err != nil { + return harnessConfig, err + } + if done { + return cfg, nil + } + default: // "cancel" or anything else + return harnessConfig, nil + } + } +} + +// editHarnessConfig opens the JSON editor pre-filled with the current config. It +// returns the updated config with done=true if the config was updated, or +// done=false (config ignored) if the user cancelled back to the menu. Invalid +// JSON is reported and the editor re-opens with the user's draft so they can fix +// it. +func editHarnessConfig(d *internal.Display, harnessConfig []byte) ([]byte, bool, error) { + draft := prettyHarnessConfig(harnessConfig) + for { + edited, err := d.PromptForConfigEdit(draft) + if err != nil { + if errors.Is(err, internal.ErrUserAborted) { + return nil, false, nil // Back to the menu. + } + return nil, false, err + } + normalized, err := normalizeHarnessConfigJSON(edited) + if err != nil { + d.ShowNotice(fmt.Sprintf("Invalid config: %v", err)) + draft = edited // Preserve the user's input so they can fix it. + continue + } + return normalized, true, nil + } +} + +// loadHarnessConfig lets the user pick a JSON file and loads it. It returns the +// loaded config with done=true, or done=false (config ignored) if the user +// cancelled back to the menu or the file could not be used. +func loadHarnessConfig(d *internal.Display) ([]byte, bool, error) { + path, err := d.PromptForConfigFile() + if err != nil { + if errors.Is(err, internal.ErrUserAborted) { + return nil, false, nil // Back to the menu. + } + return nil, false, err + } + b, err := os.ReadFile(strings.TrimSpace(path)) + if err != nil { + d.ShowNotice(fmt.Sprintf("Failed to read file: %v", err)) + return nil, false, nil + } + normalized, err := normalizeHarnessConfigJSON(string(b)) + if err != nil { + d.ShowNotice(fmt.Sprintf("Invalid config: %v", err)) + return nil, false, nil + } + return normalized, true, nil +} + +// normalizeHarnessConfigJSON trims and validates the given JSON config, returning +// the bytes to send on the wire. Empty input clears the config (returns nil). The +// config must be a JSON object. +func normalizeHarnessConfigJSON(s string) ([]byte, error) { + s = strings.TrimSpace(s) + if s == "" { + return nil, nil + } + var obj map[string]any + if err := json.Unmarshal([]byte(s), &obj); err != nil { + return nil, err + } + return []byte(s), nil +} + +// prettyHarnessConfig returns an indented JSON rendering of the config bytes for +// display, falling back to the raw bytes if they cannot be parsed. +func prettyHarnessConfig(b []byte) string { + if len(b) == 0 { + return "" + } + var buf bytes.Buffer + if err := json.Indent(&buf, b, "", " "); err != nil { + return string(b) + } + return buf.String() +} diff --git a/cmd/ax/harnessconfig_test.go b/cmd/ax/harnessconfig_test.go new file mode 100644 index 00000000..869759bf --- /dev/null +++ b/cmd/ax/harnessconfig_test.go @@ -0,0 +1,81 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "strings" + "testing" +) + +func TestNormalizeHarnessConfigJSON(t *testing.T) { + tests := []struct { + name string + in string + want string + wantErr bool + }{ + {name: "empty", in: "", want: ""}, + {name: "whitespace clears", in: " \n\t ", want: ""}, + {name: "valid object", in: `{"model":"gemini"}`, want: `{"model":"gemini"}`}, + {name: "trims surrounding whitespace", in: " {\"model\":\"gemini\"}\n", want: `{"model":"gemini"}`}, + {name: "nested object", in: `{"a":{"b":1},"c":[1,2]}`, want: `{"a":{"b":1},"c":[1,2]}`}, + {name: "invalid json", in: `{bad}`, wantErr: true}, + {name: "non-object array", in: `[1,2,3]`, wantErr: true}, + {name: "non-object scalar", in: `42`, wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := normalizeHarnessConfigJSON(tt.in) + if tt.wantErr { + if err == nil { + t.Fatalf("expected error, got nil (result %q)", got) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if string(got) != tt.want { + t.Errorf("got %q, want %q", got, tt.want) + } + if tt.want == "" && got != nil { + t.Errorf("expected nil for cleared config, got %q", got) + } + }) + } +} + +func TestPrettyHarnessConfig(t *testing.T) { + if got := prettyHarnessConfig(nil); got != "" { + t.Errorf("nil: got %q, want empty string", got) + } + if got := prettyHarnessConfig([]byte{}); got != "" { + t.Errorf("empty: got %q, want empty string", got) + } + + // Invalid JSON falls back to the raw bytes. + if got := prettyHarnessConfig([]byte("not json")); got != "not json" { + t.Errorf("invalid: got %q, want raw passthrough", got) + } + + // Valid JSON is rendered multi-line and indented. + got := prettyHarnessConfig([]byte(`{"model":"gemini"}`)) + if !strings.Contains(got, "model") || !strings.Contains(got, "gemini") { + t.Errorf("valid: got %q, want it to contain the key and value", got) + } + if !strings.Contains(got, "\n") { + t.Errorf("valid: got %q, want indented multi-line output", got) + } +} diff --git a/cmd/ax/internal/cliutil/cliutil.go b/cmd/ax/internal/cliutil/cliutil.go new file mode 100644 index 00000000..a5637d34 --- /dev/null +++ b/cmd/ax/internal/cliutil/cliutil.go @@ -0,0 +1,182 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cliutil + +import ( + "context" + "fmt" + "os" + + "github.com/google/ax/internal/config" + "github.com/google/ax/internal/controller" + "github.com/google/ax/internal/controller/eventlog" + "github.com/google/ax/internal/harness" + "github.com/google/ax/internal/harness/antigravity" + "github.com/google/ax/internal/harness/antigravityinteractions" + "github.com/google/ax/internal/harness/substrate" + "github.com/google/ax/internal/skills/geminienterprise" +) + +// Controller is the active controller type for this build. +type Controller = *controller.Controller + +// ExecHandler is the handler type accepted by Controller.Exec. +type ExecHandler = controller.ExecHandler + +// Config is the configuration type for this build. +type Config = config.Config + +// LoadFromFile loads configuration from a YAML file. +func LoadFromFile(path string) (*Config, error) { + return config.LoadFromFile(path) +} + +// DefaultConfig returns a configuration with default values set. +func DefaultConfig() *Config { + return config.DefaultConfig() +} + +// NewControllerFromConfig creates a controller.Controller instance based on the provided configuration. +func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller.Controller, error) { + reg := controller.NewRegistry() + + // AX_SUBSTRATE selects how built-in harnesses run: locally (unset) or as + // substrate actors ("1"). + substrateMode := os.Getenv("AX_SUBSTRATE") == "1" + + // Validate config before local-mode antigravity.New forks the Python + // sidecar, so a config error surfaces without spawning a subprocess. + if len(cfg.Harnesses.Substrate) > 0 && !substrateMode { + return nil, fmt.Errorf("custom substrate harnesses require AX_SUBSTRATE=1") + } + + var defaultHarnessID string + var err error + + // Materialize registry skills once, up front (skills config is top-level and + // harness-agnostic; each actor runs a single harness that consumes the + // materialized folder). Unconditional when configured. Fail-safe: a registry + // error degrades capability but never blocks harness creation. The + // interactions harness (no SKILLS_DIR concept) is told where the materialized + // skills are via a pointer appended to its system instruction. Only the local + // path materializes; substrate/pod does not yet read ax.yaml. + // + // TODO(joycel): wire the Antigravity SDK harness too. It discovers skills via + // SKILLS_DIR, so its SKILLS_DIR needs to be pointed at the materialized + // target_dir; currently only the interactions harness is fully wired. + var skillsPointer string + if !substrateMode { + skillsPointer = antigravityinteractions.SkillsSystemInstruction(geminienterprise.Materialize(ctx, cfg.Skills)) + } + + // Built-in Antigravity harness. + var antigravityHarness harness.Harness + if !substrateMode { + address := cfg.Harnesses.Antigravity.Endpoint + if address == "" { + address = "127.0.0.1:50053" + } + // Local mode: the harness owns the Python sidecar. AX owns the + // trajectory storage path (a harness implementation detail), so it's + // derived internally rather than exposed in ax.yaml. + stateDir, sErr := antigravity.DefaultStateDir() + if sErr != nil { + return nil, fmt.Errorf("antigravity harness: %w", sErr) + } + antigravityHarness, err = antigravity.New(ctx, address, stateDir, true) + if err != nil { + return nil, fmt.Errorf("antigravity harness: %w", err) + } + } else { + antigravityHarness, err = substrate.New(config.AntigravityHarnessID, "", "", config.AntigravityHarnessTemplate, 80) + if err != nil { + return nil, fmt.Errorf("antigravity harness: %w", err) + } + } + if err := reg.RegisterHarness(config.AntigravityHarnessID, antigravityHarness); err != nil { + return nil, fmt.Errorf("register antigravity harness: %w", err) + } + if cfg.Harnesses.Antigravity.Default { + defaultHarnessID = config.AntigravityHarnessID + } + + // Built-in Antigravity Interactions harness. + var antigravityInteractionsHarness harness.Harness + if !substrateMode { + aiCfg := cfg.Harnesses.AntigravityInteractions + agent := aiCfg.Agent + if agent == "" { + agent = antigravityinteractions.DefaultAgent + } + // AX owns the resume-cursor path (a harness implementation detail), so + // it's derived internally rather than exposed in ax.yaml. + stateDir, sErr := antigravityinteractions.DefaultStateDir() + if sErr != nil { + return nil, fmt.Errorf("antigravity-interactions harness: %w", sErr) + } + // skillsPointer was built once, up front, from the top-level skills + // config (see above). Append it to any configured system instruction. + antigravityInteractionsHarness, err = antigravityinteractions.New(antigravityinteractions.AntigravityInteractionsConfig{ + Agent: agent, + SystemInstruction: antigravityinteractions.JoinSystemInstruction(aiCfg.SystemInstruction, skillsPointer), + StateDir: stateDir, + }) + } else { + antigravityInteractionsHarness, err = substrate.New(config.AntigravityInteractionsHarnessID, "", "", config.AntigravityInteractionsTemplate, 80) + } + if err != nil { + return nil, fmt.Errorf("antigravity-interactions harness: %w", err) + } + if err := reg.RegisterHarness(config.AntigravityInteractionsHarnessID, antigravityInteractionsHarness); err != nil { + return nil, fmt.Errorf("register antigravity-interactions harness: %w", err) + } + if cfg.Harnesses.AntigravityInteractions.Default { + defaultHarnessID = config.AntigravityInteractionsHarnessID + } + + for _, sc := range cfg.Harnesses.Substrate { + h, err := sc.NewHarness("") + if err != nil { + return nil, fmt.Errorf("substrate harness %q: %w", sc.ID, err) + } + if err := reg.RegisterHarness(sc.ID, h); err != nil { + return nil, fmt.Errorf("register substrate harness %q: %w", sc.ID, err) + } + if sc.Default { + defaultHarnessID = sc.ID + } + } + + // Set the default harness. + if defaultHarnessID != "" { + if err := reg.SetDefaultHarness(defaultHarnessID); err != nil { + return nil, fmt.Errorf("set default harness %q: %w", defaultHarnessID, err) + } + } + + return controller.New(ctx, controller.Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { + if cfg.EventLog.PostgresConfig.DSN != "" { + dsn := os.ExpandEnv(cfg.EventLog.PostgresConfig.DSN) + if dsn == "" { + return nil, fmt.Errorf("eventlog: postgres dsn %q expanded to empty", cfg.EventLog.PostgresConfig.DSN) + } + return eventlog.OpenPostgresEventLog(dsn) + } + return eventlog.OpenSQLiteEventLog(cfg.EventLog.SQLiteConfig.Filename) + }, + }) +} diff --git a/cmd/ax/internal/cliutil/cliutil_test.go b/cmd/ax/internal/cliutil/cliutil_test.go new file mode 100644 index 00000000..31fe2e72 --- /dev/null +++ b/cmd/ax/internal/cliutil/cliutil_test.go @@ -0,0 +1,129 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cliutil + +import ( + "context" + "path/filepath" + "strings" + "testing" + + "github.com/google/ax/internal/config" +) + +func TestNewControllerFromConfig_BuiltinSubstrate(t *testing.T) { + t.Setenv("AX_SUBSTRATE", "1") + + cfg := &config.Config{ + EventLog: config.EventLogConfig{ + SQLiteConfig: config.SQLiteConfig{ + Filename: filepath.Join(t.TempDir(), "log.sqlite"), + }, + }, + Harnesses: config.HarnessesConfig{ + Antigravity: config.AntigravityHarnessConfig{ + Default: true, + }, + }, + } + + c, err := NewControllerFromConfig(context.Background(), cfg) + if err != nil { + t.Fatalf("NewControllerFromConfig: %v", err) + } + if c == nil { + t.Fatal("expected non-nil controller") + } + c.Close() +} + +func TestNewControllerFromConfig_InteractionsSubstrate(t *testing.T) { + t.Setenv("AX_SUBSTRATE", "1") + + cfg := &config.Config{ + EventLog: config.EventLogConfig{ + SQLiteConfig: config.SQLiteConfig{ + Filename: filepath.Join(t.TempDir(), "log.sqlite"), + }, + }, + Harnesses: config.HarnessesConfig{ + Antigravity: config.AntigravityHarnessConfig{Default: true}, + }, + } + + c, err := NewControllerFromConfig(context.Background(), cfg) + if err != nil { + t.Fatalf("NewControllerFromConfig: %v", err) + } + if c == nil { + t.Fatal("expected non-nil controller") + } + defer c.Close() + + if _, err := c.Registry().Harness(config.AntigravityInteractionsHarnessID); err != nil { + t.Errorf("interactions harness not registered in substrate mode: %v", err) + } +} + +func TestNewControllerFromConfig_CustomHarnessRequiresSubstrateMode(t *testing.T) { + t.Setenv("AX_SUBSTRATE", "") + + cfg := &config.Config{ + EventLog: config.EventLogConfig{ + SQLiteConfig: config.SQLiteConfig{ + Filename: filepath.Join(t.TempDir(), "log.sqlite"), + }, + }, + Harnesses: config.HarnessesConfig{ + Substrate: []config.SubstrateHarnessConfig{ + {ID: "custom", Namespace: "team-ns", Template: "custom-template"}, + }, + }, + } + + _, err := NewControllerFromConfig(context.Background(), cfg) + if err == nil { + t.Fatal("expected error for custom substrate harness without AX_SUBSTRATE=1, got nil") + } + if !strings.Contains(err.Error(), "AX_SUBSTRATE=1") { + t.Errorf("expected error to mention AX_SUBSTRATE=1, got: %v", err) + } +} + +func TestNewControllerFromConfig_CustomHarnessInSubstrateMode(t *testing.T) { + t.Setenv("AX_SUBSTRATE", "1") + + cfg := &config.Config{ + EventLog: config.EventLogConfig{ + SQLiteConfig: config.SQLiteConfig{ + Filename: filepath.Join(t.TempDir(), "log.sqlite"), + }, + }, + Harnesses: config.HarnessesConfig{ + Substrate: []config.SubstrateHarnessConfig{ + {ID: "custom", Namespace: "team-ns", Template: "custom-template"}, + }, + }, + } + + c, err := NewControllerFromConfig(context.Background(), cfg) + if err != nil { + t.Fatalf("NewControllerFromConfig: %v", err) + } + if c == nil { + t.Fatal("expected non-nil controller") + } + c.Close() +} diff --git a/cmd/ax/internal/display.go b/cmd/ax/internal/display.go new file mode 100644 index 00000000..376e25bf --- /dev/null +++ b/cmd/ax/internal/display.go @@ -0,0 +1,329 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package internal + +import ( + "fmt" + "io" + "os" + + "charm.land/bubbles/v2/key" + tea "charm.land/bubbletea/v2" + "charm.land/huh/v2" + "charm.land/lipgloss/v2" + "github.com/google/ax/proto" +) + +const ( + boxWidth = 100 +) + +var ( + isDark = lipgloss.HasDarkBackground(os.Stdin, os.Stdout) + lightDark = lipgloss.LightDark(isDark) + + purple = lightDark(lipgloss.Color("#5A56E0"), lipgloss.Color("#7571F9")) + comment = lightDark(lipgloss.Color("#a0a0a0"), lipgloss.Color("#6d6d6d")) +) + +// ErrUserAborted is returned when the user aborts a prompt. +var ErrUserAborted = huh.ErrUserAborted + +type displayState int + +const ( + stateNone displayState = iota + stateText + stateThought +) + +type Display struct { + id string + w io.Writer // Target output writer, e.g., os.Stdout or a test buffer + + userStyle lipgloss.Style + checkpointStyle lipgloss.Style + idStyle lipgloss.Style + resumeStyle lipgloss.Style + + state displayState // Tracks the last printed chunk type to correctly format transition newlines +} + +func NewDisplay(id string, w io.Writer) *Display { + if w == nil { + w = os.Stdout + } + return &Display{ + id: id, + w: w, + userStyle: lipgloss.NewStyle().Foreground(purple), + checkpointStyle: lipgloss.NewStyle().Foreground(comment), + idStyle: lipgloss.NewStyle().Foreground(comment), + resumeStyle: lipgloss.NewStyle().Foreground(comment), + state: stateNone, + } +} + +// DisplayInput displays the user input. +func (d *Display) DisplayInput(text string) { + if d.state != stateNone { + fmt.Fprintln(d.w) + } + d.state = stateNone + fmt.Fprintf(d.w, "%s %s\n", + d.userStyle.Render("⏺"), + text, + ) + fmt.Fprintln(d.w) +} + +// Display prints a content block according to its type. +func (d *Display) Display(content *proto.Content) { + if content == nil { + return + } + switch o := content.Type.(type) { + case *proto.Content_Text: + if d.state == stateThought { + fmt.Fprintln(d.w) // end the thinking line + } + d.state = stateText + fmt.Fprint(d.w, o.Text.Text) + + case *proto.Content_Confirmation: + // Let the confirmation prompt handle displaying the question. + + case *proto.Content_ToolCall: + // Tool calls aren't rendered, but they mark a boundary between + // contiguous text/thought blocks. Terminate the current line so the + // next response starts fresh instead of running into the previous one + // (e.g. "...configured.I will list..."). + if d.state != stateNone { + fmt.Fprintln(d.w) + d.state = stateNone + } + + case *proto.Content_ToolResult: + // Only print if the tool returned an error, otherwise skip + tr := o.ToolResult + if fr := tr.GetFunctionResult(); fr != nil { + if fr.GetResponse() != nil { + respMap := fr.GetResponse().AsMap() + if errStr, ok := respMap["error"]; ok { + d.displaySystem(fmt.Sprintf("[TOOL ERROR for %s]\n%v", fr.Name, errStr)) + } + } + } + + case *proto.Content_Thought: + for _, summary := range o.Thought.GetSummary() { + if textContent := summary.GetText(); textContent != nil { + if d.state != stateThought { + if d.state == stateText { + fmt.Fprintln(d.w) + } + fmt.Fprint(d.w, "Thinking: ") + } + d.state = stateThought + fmt.Fprint(d.w, textContent.Text) + } + } + + case *proto.Content_Image, *proto.Content_Audio, *proto.Content_Video, *proto.Content_Document: + d.displaySystem(fmt.Sprintf("unsupported output type for display: %T", o)) + + default: + d.displaySystem(fmt.Sprintf("unknown output type: %v", o)) + } +} + +// displaySystem prints a system/error message on a new line. +func (d *Display) displaySystem(text string) { + if d.state != stateNone { + fmt.Fprintln(d.w) + } + d.state = stateNone + fmt.Fprintln(d.w, text) +} + +// ShowNotice prints an informational message (e.g. the current harness config +// or a validation error) on its own line. +func (d *Display) ShowNotice(text string) { + d.displaySystem(text) +} + +// FinishOutput completes the streaming output and shows info if provided +func (d *Display) FinishOutput(info string) { + if d.state != stateNone { + fmt.Fprintln(d.w) + } + d.state = stateNone + if info != "" { + fmt.Fprintln(d.w, d.checkpointStyle.Render(info)) + } + fmt.Fprintln(d.w) +} + +func (d *Display) DisplayHeader() { + fmt.Fprintln(d.w, d.idStyle.Render("Conversation: "+d.id)) + fmt.Fprintln(d.w) +} + +// PromptForApproval shows an accept/reject dialog +// Returns true if accepted, false if rejected, and an error if cancelled (Ctrl+C) +func (d *Display) PromptForApproval(question string) (bool, error) { + var accepted bool + form := huh.NewForm( + huh.NewGroup( + huh.NewConfirm(). + Title(question). + Affirmative("Accept"). + Negative("Reject"). + Value(&accepted), + ), + ).WithWidth(boxWidth) + + if err := form.Run(); err != nil { + return false, err + } + return accepted, nil +} + +// PromptForInput shows the input box and returns the user input +// Returns the input string and an error if the user cancelled (Ctrl+C) +func (d *Display) PromptForInput() (string, error) { + var userInput string + + // Rebind tab to complete the suggestion and enter to submit. + keymap := huh.NewDefaultKeyMap() + keymap.Input.AcceptSuggestion = key.NewBinding(key.WithKeys("tab"), key.WithHelp("tab", "complete")) + keymap.Input.Next = key.NewBinding(key.WithKeys("enter"), key.WithHelp("enter", "submit")) + + form := huh.NewForm( + huh.NewGroup( + huh.NewInput(). + Placeholder("Enter prompt... (type `q` to quit)"). + Suggestions([]string{"/config"}). + Value(&userInput), + ), + ).WithWidth(boxWidth).WithShowHelp(false).WithKeyMap(keymap) + + if err := form.Run(); err != nil { + return "", err + } + return userInput, nil +} + +// configKeyMap returns a huh keymap for the /config menu that cancels on esc in +// addition to the default ctrl+c, and shows a consistent "esc close" hint in +// every field's help footer. +func configKeyMap() *huh.KeyMap { + km := huh.NewDefaultKeyMap() + km.Quit = key.NewBinding(key.WithKeys("ctrl+c", "esc"), key.WithHelp("esc", "close")) + + // Surface the esc hint through a slot that stays enabled. + closeHint := key.NewBinding(key.WithKeys("esc"), key.WithHelp("esc", "close")) + km.Select.Filter = closeHint + km.Text.Editor = closeHint + return km +} + +// clearMenuOnCancel converts the interrupt huh emits on cancel (esc/ctrl+c) into +// a graceful quit. +var clearMenuOnCancel = tea.WithFilter(func(_ tea.Model, msg tea.Msg) tea.Msg { + if _, ok := msg.(tea.InterruptMsg); ok { + return tea.QuitMsg{} + } + return msg +}) + +// PromptForConfigAction shows the /config menu and returns the chosen action, +// one of "edit", "load", or "cancel". It returns an error if the user cancelled +// (Ctrl+C or Esc). +func (d *Display) PromptForConfigAction() (string, error) { + var action string + form := huh.NewForm( + huh.NewGroup( + huh.NewSelect[string](). + Title("Harness config"). + Options( + huh.NewOption("View/edit config", "edit"), + huh.NewOption("Load from file", "load"), + huh.NewOption("Cancel", "cancel"), + ). + Value(&action), + ), + ).WithWidth(boxWidth).WithKeyMap(configKeyMap()).WithProgramOptions(clearMenuOnCancel) + + if err := form.Run(); err != nil { + return "", err + } + return action, nil +} + +// PromptForConfigEdit shows a multi-line editor pre-filled with current and +// returns the edited text. It returns an error if the user cancelled (Ctrl+C or +// Esc). +func (d *Display) PromptForConfigEdit(current string) (string, error) { + text := current + form := huh.NewForm( + huh.NewGroup( + huh.NewText(). + Title("Harness config (JSON, leave empty to clear)"). + Lines(10). + ExternalEditor(true). + Value(&text), + ), + ).WithWidth(boxWidth).WithKeyMap(configKeyMap()).WithProgramOptions(clearMenuOnCancel) + + if err := form.Run(); err != nil { + return "", err + } + return text, nil +} + +// PromptForConfigFile shows a file browser for selecting a JSON config file and +// returns the chosen path. It returns an error if the user cancelled (Ctrl+C or +// Esc). +func (d *Display) PromptForConfigFile() (string, error) { + var path string + form := huh.NewForm( + huh.NewGroup( + huh.NewFilePicker(). + Title("Load harness config"). + CurrentDirectory("."). + AllowedTypes([]string{".json"}). + FileAllowed(true). + DirAllowed(false). + Picking(true). + Height(10). + Value(&path), + ), + ).WithWidth(boxWidth).WithKeyMap(configKeyMap()).WithProgramOptions(clearMenuOnCancel) + + if err := form.Run(); err != nil { + return "", err + } + return path, nil +} + +func (d *Display) ShowResumption(id string, server string) { + fmt.Fprintln(d.w, d.resumeStyle.Render("To resume the conversation,")) + if server != "" { + fmt.Fprintln(d.w, d.resumeStyle.Render(fmt.Sprintf("ax exec --conversation %s --server %s", id, server))) + } else { + fmt.Fprintln(d.w, d.resumeStyle.Render(fmt.Sprintf("ax exec --conversation %s", id))) + } +} diff --git a/cmd/ax/internal/display_test.go b/cmd/ax/internal/display_test.go new file mode 100644 index 00000000..edd4b462 --- /dev/null +++ b/cmd/ax/internal/display_test.go @@ -0,0 +1,196 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package internal + +import ( + "bytes" + "testing" + + "github.com/google/ax/proto" +) + +func TestDisplay_Streaming(t *testing.T) { + textContent := func(txt string) *proto.Content { + return &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: txt}}} + } + thoughtContent := func(txt string) *proto.Content { + return &proto.Content{Type: &proto.Content_Thought{Thought: &proto.ThoughtContent{ + Summary: []*proto.ThoughtSummaryContent{ + {Type: &proto.ThoughtSummaryContent_Text{Text: &proto.TextContent{Text: txt}}}, + }, + }}} + } + toolCallContent := func() *proto.Content { + return &proto.Content{Type: &proto.Content_ToolCall{ToolCall: &proto.ToolCallContent{}}} + } + + t.Run("consecutive text chunks are concatenated", func(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + d := NewDisplay("test-id", &buf) + + d.Display(textContent("Hello ")) + d.Display(textContent("world")) + d.Display(textContent("!")) + + got := buf.String() + want := "Hello world!" + if got != want { + t.Errorf("got %q, want %q", got, want) + } + }) + + t.Run("tool call separates consecutive text blocks", func(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + d := NewDisplay("test-id", &buf) + + d.Display(textContent("...configured.")) + d.Display(toolCallContent()) + d.Display(textContent("I will list the contents.")) + + got := buf.String() + want := "...configured.\nI will list the contents." + if got != want { + t.Errorf("got %q, want %q", got, want) + } + }) + + t.Run("repeated tool calls do not add extra newlines", func(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + d := NewDisplay("test-id", &buf) + + d.Display(textContent("Done.")) + d.Display(toolCallContent()) + d.Display(toolCallContent()) + d.Display(textContent("Next.")) + + got := buf.String() + want := "Done.\nNext." + if got != want { + t.Errorf("got %q, want %q", got, want) + } + }) + + t.Run("consecutive thought chunks are concatenated with prefix", func(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + d := NewDisplay("test-id", &buf) + + d.Display(thoughtContent("thinking ")) + d.Display(thoughtContent("deeply")) + + got := buf.String() + want := "Thinking: thinking deeply" + if got != want { + t.Errorf("got %q, want %q", got, want) + } + }) + + t.Run("transition from thought to text adds newline", func(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + d := NewDisplay("test-id", &buf) + + d.Display(thoughtContent("thinking")) + d.Display(textContent("Hello")) + + got := buf.String() + want := "Thinking: thinking\nHello" + if got != want { + t.Errorf("got %q, want %q", got, want) + } + }) + + t.Run("transition from text to thought adds newline", func(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + d := NewDisplay("test-id", &buf) + + d.Display(textContent("Hello")) + d.Display(thoughtContent("thinking")) + + got := buf.String() + want := "Hello\nThinking: thinking" + if got != want { + t.Errorf("got %q, want %q", got, want) + } + }) + + t.Run("FinishOutput empty resets state and adds newlines", func(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + d := NewDisplay("test-id", &buf) + + d.Display(textContent("Hello")) + d.FinishOutput("") + + got := buf.String() + want := "Hello\n\n" + if got != want { + t.Errorf("got %q, want %q", got, want) + } + }) + + t.Run("FinishOutput with info prints info and resets state", func(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + d := NewDisplay("test-id", &buf) + + d.Display(textContent("Hello")) + d.FinishOutput("seq=1") + + got := buf.String() + if !bytes.HasPrefix([]byte(got), []byte("Hello\n")) { + t.Errorf("expected Hello to end with newline, got %q", got) + } + if !bytes.Contains([]byte(got), []byte("seq=1")) { + t.Errorf("expected output to contain seq=1, got %q", got) + } + }) + + t.Run("displaySystem resets state and prints newline", func(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + d := NewDisplay("test-id", &buf) + + d.Display(textContent("Hello")) + d.displaySystem("system message") + + got := buf.String() + want := "Hello\nsystem message\n" + if got != want { + t.Errorf("got %q, want %q", got, want) + } + }) + + t.Run("DisplayInput resets state and adds separation newlines", func(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + d := NewDisplay("test-id", &buf) + + d.Display(textContent("Hello")) + d.DisplayInput("prompt") + + got := buf.String() + if !bytes.HasPrefix([]byte(got), []byte("Hello\n")) { + t.Errorf("expected Hello to end with newline, got %q", got) + } + if !bytes.Contains([]byte(got), []byte("prompt")) { + t.Errorf("expected output to contain prompt, got %q", got) + } + }) +} diff --git a/cmd/ax/main.go b/cmd/ax/main.go new file mode 100644 index 00000000..ee65c4a2 --- /dev/null +++ b/cmd/ax/main.go @@ -0,0 +1,79 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// AX is a server for managing agent orchestrator tasks. +// It provides commands to execute tasks, resume from checkpoints, +// register agents, and run the controller server. +package main + +import ( + "errors" + "fmt" + "os" + + "github.com/google/ax/cmd/ax/internal/cliutil" + "github.com/spf13/cobra" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" +) + +func main() { + if err := rootCmd.Execute(); err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } +} + +var rootCmd = &cobra.Command{ + Use: "ax", + Short: "AX - Agent eXecutor", + Long: `ax provides a server and CLI tools for managing agent orchestrator tasks. +It provides commands to execute tasks, resume from checkpoints, +and run the controller server.`, +} + +func init() { + rootCmd.AddCommand(execCmd) + rootCmd.AddCommand(serveCmd) + + rootCmd.AddCommand(dashboardCmd) +} + +func connect(server string) (*grpc.ClientConn, error) { + conn, err := grpc.NewClient(server, + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + return nil, fmt.Errorf("failed to connect to server: %w", err) + } + return conn, nil +} + +const currentVersion = "v1alpha" + +func newConfig(cmd *cobra.Command, configFile string) (*cliutil.Config, error) { + cfg, err := cliutil.LoadFromFile(configFile) + if errors.Is(err, os.ErrNotExist) && !cmd.Flags().Changed("config") { + cfg := cliutil.DefaultConfig() + cfg.Version = currentVersion + return cfg, nil + } + if err != nil { + return nil, fmt.Errorf("error loading config file '%s': %w", configFile, err) + } + if cfg.Version != currentVersion { + return nil, fmt.Errorf("unsupported config version %q, must be %q", cfg.Version, currentVersion) + } + return cfg, nil +} diff --git a/cmd/ax/serve.go b/cmd/ax/serve.go new file mode 100644 index 00000000..d0b407e2 --- /dev/null +++ b/cmd/ax/serve.go @@ -0,0 +1,113 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "fmt" + "log/slog" + "os" + "os/signal" + "syscall" + + "github.com/google/ax/cmd/ax/internal/cliutil" + "github.com/google/ax/internal/server" + "github.com/google/ax/internal/telemetry" + "github.com/spf13/cobra" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc" +) + +var ( + serveConfigFile string +) + +var serveCmd = &cobra.Command{ + Use: "serve", + Short: "Run controller as a gRPC server", + Long: `Run the AX controller as a gRPC server. +Loads configuration from a YAML file (default: ax.yaml).`, + RunE: runServe, +} + +func init() { + serveCmd.Flags().StringVar(&serveConfigFile, "config", "ax.yaml", "Path to YAML configuration file") +} + +func runServe(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + + // Initialize structured logging (JSON to stdout) + slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, nil))) + + cfg, err := newConfig(cmd, serveConfigFile) + if err != nil { + return err + } + + // Validate configuration + if err := cfg.Validate(); err != nil { + return fmt.Errorf("invalid configuration: %w", err) + } + + // Initialize OpenTelemetry SDK + if cfg.Telemetry.OTLP.Enabled { + var opts []otlptracegrpc.Option + endpoint := cfg.Telemetry.OTLP.Endpoint + if endpoint == "" { + endpoint = defaultEndpoint + } + opts = append(opts, otlptracegrpc.WithEndpoint(endpoint)) + if gcpOpts, ok := gcpTelemetryOpts(endpoint); ok { + opts = append(opts, gcpOpts...) + } else { + opts = append(opts, otlptracegrpc.WithInsecure()) + } + + shutdown, err := telemetry.SetTraceProvider(ctx, "ax-server", opts...) + if err != nil { + return fmt.Errorf("failed to initialize telemetry: %w", err) + } + defer func() { + if err := shutdown(context.Background()); err != nil { + slog.Error("failed to shutdown telemetry", "error", err) + } + }() + } + + c, err := cliutil.NewControllerFromConfig(ctx, cfg) + if err != nil { + return fmt.Errorf("error creating controller: %w", err) + } + defer c.Close() + + // Create server + srv := server.New(c) + + // Setup signal handling for graceful shutdown + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + + go func() { + <-sigChan + slog.InfoContext(ctx, "Received interrupt, shutting down...") + srv.GracefulStop() + }() + slog.InfoContext(ctx, "Starting AX server", slog.String("address", cfg.Server.Address)) + if err := srv.Serve(cfg.Server.Address); err != nil { + return fmt.Errorf("error serving: %w", err) + } + + return nil +} diff --git a/cmd/ax/serve_gcp.go b/cmd/ax/serve_gcp.go new file mode 100644 index 00000000..63e559b3 --- /dev/null +++ b/cmd/ax/serve_gcp.go @@ -0,0 +1,32 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "github.com/google/ax/internal/telemetry" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc" +) + +const defaultEndpoint = "telemetry.googleapis.com" + +// gcpTelemetryOpts returns the OTLP options for Google Cloud Trace if the endpoint +// matches the Google Cloud Trace endpoint. It returns the options and a boolean +// indicating if the endpoint was matched. +func gcpTelemetryOpts(endpoint string) ([]otlptracegrpc.Option, bool) { + if endpoint == defaultEndpoint || endpoint == defaultEndpoint+":443" { + return []otlptracegrpc.Option{telemetry.WithGCPCredentials()}, true + } + return nil, false +} diff --git a/cmd/ax/web/index.html b/cmd/ax/web/index.html new file mode 100644 index 00000000..fc49046a --- /dev/null +++ b/cmd/ax/web/index.html @@ -0,0 +1,1242 @@ + + + + + + + +AX Dashboard + + + + + + + + +
+
+ +

AX Dashboard

+
+
+ + Connected +
+
+ +
+ +
+ +
+
+
+

Total

+
0
+
+
Σ
+
+
+
+

Active

+
0
+
+
⚡
+
+
+
+

Completed

+
0
+
+
✓
+
+
+
+

Failed

+
0
+
+
✗
+
+
+ + +
+
+ +
+ + +
+ + +
+ + + + + + + + + + + + + + +
Conversation IDAgentStatusLast SeqDurationActions
+ +
+
+ + +
+
+ +
+

Execution Trace

+

+
+
+ + +
+
Execution Timeline
+
+ +
+
+ + +
+ +
+
+
+ + + + diff --git a/cmd/gar/inspect.go b/cmd/gar/inspect.go deleted file mode 100644 index 8068f05f..00000000 --- a/cmd/gar/inspect.go +++ /dev/null @@ -1,74 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "context" - "fmt" - "time" - - "github.com/google/gar/proto" - "github.com/spf13/cobra" -) - -var ( - inspectSessionID string - inspectServerAddr string -) - -var inspectCmd = &cobra.Command{ - Use: "inspect", - Short: "Inspect a session", - Long: `Inspect a session to view its current state, step count, and other details.`, - RunE: runInspect, -} - -func init() { - inspectCmd.Flags().StringVar(&inspectSessionID, "session-id", "", "Session ID (required)") - inspectCmd.Flags().StringVar(&inspectServerAddr, "server", "localhost:8494", "gRPC controller server address (default: localhost:8494)") - inspectCmd.MarkFlagRequired("session-id") -} - -func runInspect(cmd *cobra.Command, args []string) error { - ctx := context.Background() - conn, err := connect(inspectServerAddr) - if err != nil { - return err - } - defer conn.Close() - - client := proto.NewGARServiceClient(conn) - - // Get session details - resp, err := client.GetSession(ctx, &proto.GetSessionRequest{ - SessionId: inspectSessionID, - }) - if err != nil { - return fmt.Errorf("error getting session: %w", err) - } - - session := resp.Session - - // Print session details - fmt.Println("\nSession Details:") - fmt.Printf(" ID: %s\n", inspectSessionID) - fmt.Printf(" State: %s\n", session.State) - fmt.Printf(" Created At: %s\n", session.CreatedAt.AsTime().Format(time.RFC3339)) - fmt.Printf(" Updated At: %s\n", session.UpdatedAt.AsTime().Format(time.RFC3339)) - fmt.Printf(" Checkpoints: %d\n", session.CheckpointCount) - fmt.Printf(" Active Agents: %v\n", session.ActiveAgents) - - return nil -} diff --git a/cmd/gar/main.go b/cmd/gar/main.go deleted file mode 100644 index 90ad3d51..00000000 --- a/cmd/gar/main.go +++ /dev/null @@ -1,58 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Gar is a CLI tool for managing agent orchestrator sessions. -// It provides commands to trigger sessions, resume from checkpoints, -// inspect session state, register agents, and run the controller server. -package main - -import ( - "fmt" - "os" - - "github.com/spf13/cobra" - "google.golang.org/grpc" - "google.golang.org/grpc/credentials/insecure" -) - -func main() { - if err := rootCmd.Execute(); err != nil { - fmt.Fprintf(os.Stderr, "Error: %v\n", err) - os.Exit(1) - } -} - -var rootCmd = &cobra.Command{ - Use: "gar", - Short: "GAR - Google Agent Runtime CLI", - Long: `Gar is a CLI tool for managing agent orchestrator sessions. -It provides commands to trigger sessions, resume from checkpoints, -inspect session state, register agents, and run the controller server.`, -} - -func init() { - // Add subcommands - rootCmd.AddCommand(triggerCmd) - rootCmd.AddCommand(inspectCmd) - rootCmd.AddCommand(registerCmd) - rootCmd.AddCommand(serveCmd) -} - -func connect(server string) (*grpc.ClientConn, error) { - conn, err := grpc.NewClient(server, grpc.WithTransportCredentials(insecure.NewCredentials())) - if err != nil { - return nil, fmt.Errorf("failed to connect to server: %w", err) - } - return conn, nil -} diff --git a/cmd/gar/register.go b/cmd/gar/register.go deleted file mode 100644 index 600c70f5..00000000 --- a/cmd/gar/register.go +++ /dev/null @@ -1,73 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "context" - "fmt" - - "github.com/google/gar/proto" - "github.com/spf13/cobra" -) - -var ( - registerAgentID string - registerAgentName string - registerAgentDesc string - registerAgentAddr string - registerServerAddr string -) - -var registerCmd = &cobra.Command{ - Use: "register", - Short: "Register a remote agent", - Long: `Register a remote agent with the controller so it can be used in sessions.`, - RunE: runRegister, -} - -func init() { - registerCmd.Flags().StringVar(®isterAgentID, "agent-id", "", "Agent ID (required)") - registerCmd.Flags().StringVar(®isterAgentName, "agent-name", "", "Agent name (required)") - registerCmd.Flags().StringVar(®isterAgentDesc, "agent-description", "", "Agent description (required)") - registerCmd.Flags().StringVar(®isterAgentAddr, "agent-addr", "", "Agent address (e.g., localhost:50051) (required)") - registerCmd.Flags().StringVar(®isterServerAddr, "server", "localhost:8494", "gRPC controller server address (default: localhost:8494)") - registerCmd.MarkFlagRequired("agent-id") - registerCmd.MarkFlagRequired("agent-name") - registerCmd.MarkFlagRequired("agent-description") - registerCmd.MarkFlagRequired("agent-addr") -} - -func runRegister(cmd *cobra.Command, args []string) error { - ctx := context.Background() - conn, err := connect(inspectServerAddr) - if err != nil { - return err - } - defer conn.Close() - - client := proto.NewGARServiceClient(conn) - - // Register remote agent - _, err = client.RegisterAgent(ctx, &proto.RegisterAgentRequest{ - AgentId: registerAgentID, - Name: registerAgentName, - Description: registerAgentDesc, - Address: registerAgentAddr, - }) - if err != nil { - return fmt.Errorf("error registering agent: %w", err) - } - return nil -} diff --git a/cmd/gar/serve.go b/cmd/gar/serve.go deleted file mode 100644 index 3eb82ec2..00000000 --- a/cmd/gar/serve.go +++ /dev/null @@ -1,121 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "context" - "fmt" - "os" - "os/signal" - "syscall" - - "github.com/google/gar/internal/config" - "github.com/google/gar/internal/controller" - "github.com/google/gar/internal/eventlog" - "github.com/google/gar/internal/server" - "github.com/spf13/cobra" -) - -var ( - serveConfigFile string -) - -var serveCmd = &cobra.Command{ - Use: "serve", - Short: "Run controller as a gRPC server", - Long: `Run the GAR controller as a gRPC server. -Loads configuration from a YAML file (default: gar.yaml).`, - RunE: runServe, -} - -func init() { - serveCmd.Flags().StringVar(&serveConfigFile, "config", "gar.yaml", "Path to YAML configuration file") -} - -func runServe(cmd *cobra.Command, args []string) error { - ctx := context.Background() - - // Load configuration from YAML file - cfg, err := config.LoadFromFile(serveConfigFile) - if err != nil { - return fmt.Errorf("error loading config file '%s': %w\nTip: Create a config file with 'gar serve --help' to see an example", serveConfigFile, err) - } - - // Validate configuration - if err := cfg.Validate(); err != nil { - return fmt.Errorf("invalid configuration: %w", err) - } - - fmt.Printf("Starting GAR server at %s...\n", cfg.Server.Address) - c, err := newControllerFromConfig(ctx, cfg) - if err != nil { - return fmt.Errorf("error creating controller: %w", err) - } - defer c.Close() - - // Create server - srv := server.New(c) - - // Setup signal handling for graceful shutdown - sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) - - go func() { - <-sigChan - fmt.Println("\nReceived interrupt, shutting down...") - os.Exit(0) - }() - - // Start serving - if err := srv.Serve(cfg.Server.Address); err != nil { - return fmt.Errorf("error serving: %w", err) - } - - return nil -} - -func newControllerFromConfig(ctx context.Context, cfg *config.Config) (*controller.Controller, error) { - // Create event log factory - eventLogFactory := func(sessionID string) (eventlog.EventLog, error) { - return eventlog.NewFileEventLog(eventlog.FileConfig{ - SessionID: sessionID, - Dir: cfg.EventLog.Dir, - }) - } - - // Build controller config - // The controller will create a default Gemini planner if PlanFunc is nil - // Gemini config can be customized via environment variables (GEMINI_API_KEY, GAR_GEMINI_MODEL) - controllerConfig := controller.Config{ - EventLogFactory: eventLogFactory, - MaxSteps: cfg.MaxSteps, - HealthCheckInterval: cfg.HealthCheckInterval, - } - - // Create controller - c, err := controller.New(ctx, controllerConfig) - if err != nil { - return nil, err - } - - // Register remote agents from config - for _, agentCfg := range cfg.RemoteAgents { - if err := c.Registry().RegisterRemote(agentCfg); err != nil { - return nil, fmt.Errorf("failed to register remote agent %s: %w", agentCfg.ID, err) - } - } - - return c, nil -} diff --git a/cmd/gar/trigger.go b/cmd/gar/trigger.go deleted file mode 100644 index 64244b91..00000000 --- a/cmd/gar/trigger.go +++ /dev/null @@ -1,119 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "context" - "fmt" - "io" - "os" - "os/signal" - "syscall" - - "github.com/google/gar/proto" - "github.com/google/uuid" - "github.com/spf13/cobra" -) - -var ( - triggerSessionID string - triggerInput string - triggerCheckpoint string - triggerServerAddr string -) - -var triggerCmd = &cobra.Command{ - Use: "trigger", - Short: "Trigger a new session or resume an existing one", - Long: `Trigger a new agentic session or resume an existing one. -If no session ID is provided, a new UUID will be generated. -Use --checkpoint to resume from a specific checkpoint.`, - RunE: runTrigger, -} - -func init() { - triggerCmd.Flags().StringVar(&triggerSessionID, "session-id", "", "Session ID (optional, generates UUID if not provided)") - triggerCmd.Flags().StringVar(&triggerInput, "input", "", "Input message to send (required)") - triggerCmd.Flags().StringVar(&triggerCheckpoint, "checkpoint", "", "Resume from specific checkpoint UUID (empty for latest)") - triggerCmd.Flags().StringVar(&triggerServerAddr, "server", "localhost:8494", "gRPC controller server address (default: localhost:8494)") - triggerCmd.MarkFlagRequired("input") -} - -// TODO(jbd): Add multimodal input flags, e.g. --input-image. - -func runTrigger(cmd *cobra.Command, args []string) error { - ctx := context.Background() - - // Generate UUID if no session ID provided - if triggerSessionID == "" { - triggerSessionID = uuid.New().String() - fmt.Printf("Generated session ID: %s\n", triggerSessionID) - } - - // Create input content - inputs := []*proto.Content{ - { - Role: "user", - Type: "text", - Mimetype: "text/plain", - Data: triggerInput, - }, - } - - // Setup signal handling for graceful shutdown - ctx, cancel := context.WithCancel(ctx) - defer cancel() - - sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) - - go func() { - <-sigChan - fmt.Println("\nReceived interrupt, shutting down...") - cancel() - }() - - conn, err := connect(inspectServerAddr) - if err != nil { - return err - } - defer conn.Close() - - client := proto.NewGARServiceClient(conn) - stream, err := client.TriggerSession(ctx, &proto.TriggerSessionRequest{ - SessionId: triggerSessionID, - Inputs: inputs, - CheckpointId: triggerCheckpoint, - }) - if err != nil { - return fmt.Errorf("error triggering session: %w", err) - } - - // Receive and print all responses - for { - resp, err := stream.Recv() - if err == io.EOF { - break - } - if err != nil { - return fmt.Errorf("error receiving response: %w", err) - } - - if resp.Output != nil { - fmt.Printf("[%s] %s\n", resp.State, resp.Output.Data) - } - } - return nil -} diff --git a/e2e b/e2e new file mode 100755 index 00000000..c1226113 Binary files /dev/null and b/e2e differ diff --git a/examples/local_agent/main.go b/examples/local_agent/main.go deleted file mode 100644 index b86a73e2..00000000 --- a/examples/local_agent/main.go +++ /dev/null @@ -1,133 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "context" - "flag" - "fmt" - "log" - "strings" - "time" - - "github.com/google/gar/agent" - "github.com/google/gar/internal/config" - "github.com/google/gar/internal/controller" - "github.com/google/gar/proto" - "github.com/google/uuid" -) - -var ( - sessionID string - input string -) - -func main() { - ctx := context.Background() - flag.StringVar(&input, "input", "Hello, uppercase this message!", "Input message to be processed by the agent") - flag.StringVar(&sessionID, "session_id", "", "Optional session ID") - flag.Parse() - - // Create a local echo agent - echoAgent, err := createEchoAgent() - if err != nil { - log.Fatalf("Error creating agent: %v\n", err) - } - - c, err := controller.New(ctx, controller.Config{ - MaxSteps: 10, - HealthCheckInterval: 30 * time.Second, - }) - if err != nil { - log.Fatalf("Error creating controller: %v\n", err) - } - defer c.Close() - - if err := c.Registry().RegisterLocal(config.LocalAgentConfig{ - ID: "local-echo-agent", - Name: "Echo Agent", - Description: "Simple echo agent that uppercases input", - Metadata: map[string]string{ - "version": "1.0", - }, - Agent: echoAgent, - }); err != nil { - log.Fatalf("Error registering agent: %v\n", err) - } - - inputs := []*proto.Content{ - { - Role: "user", - Type: "text", - Mimetype: "text/plain", - Data: input, - }, - } - - // Trigger a session. Alternatively, controller can be used - // with the server package to expose a gRPC server. - if sessionID == "" { - sessionID = uuid.New().String() - } - log.Printf("Session ID: %s\n", sessionID) - - handler := agent.OutputHandler(func(content *proto.Content) error { - fmt.Printf("Output received: %s\n", content.Data) - return nil - }) - if err := c.TriggerSession(ctx, sessionID, inputs, handler); err != nil { - log.Fatalf("Error triggering session: %v\n", err) - } -} - -// createEchoAgent creates a simple echo agent that repeats input with a prefix. -func createEchoAgent() (*agent.LocalAgent, error) { - processFunc := func(ctx context.Context, sessionID string, inputs []*proto.Content, handler agent.OutputHandler) error { - // Process each input and call handler with response - for _, content := range inputs { - // Echo the content back with a prefix - response := &proto.Content{ - Role: "assistant", - Type: content.Type, - Mimetype: content.Mimetype, - Data: strings.ToUpper(content.Data), - } - - // Call handler with the response - if err := handler(response); err != nil { - return err - } - - // Check for context cancellation - select { - case <-ctx.Done(): - return ctx.Err() - default: - } - } - - return nil - } - - healthCheckFunc := func(ctx context.Context) error { - // Always healthy - return nil - } - - return agent.NewLocalAgent(agent.LocalAgentConfig{ - ProcessFunc: processFunc, - HealthCheckFunc: healthCheckFunc, - }) -} diff --git a/examples/remote_agent/main.go b/examples/remote_agent/main.go deleted file mode 100644 index 87eb5c93..00000000 --- a/examples/remote_agent/main.go +++ /dev/null @@ -1,91 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "context" - "fmt" - "io" - "log" - "net" - "strings" - - "google.golang.org/grpc" - - "github.com/google/gar/proto" -) - -const port = ":50051" - -// server implements the AgentService gRPC server. -type server struct { - proto.UnimplementedAgentServiceServer -} - -// Process implements bidirectional streaming for content processing. -// This RPC is called by the gar controller when a session triggers this agent. -func (s *server) Process(stream proto.AgentService_ProcessServer) error { - for { - // Receive input content from gar controller - content, err := stream.Recv() - if err == io.EOF { - return nil - } - if err != nil { - return err - } - - // Process the content (simple uppercase echo) - response := &proto.Content{ - Role: "assistant", - Type: content.Type, - Mimetype: content.Mimetype, - Data: fmt.Sprintf("Remote Echo: %s", strings.ToUpper(content.Data)), - } - - // Send response back to gar controller - if err := stream.Send(response); err != nil { - return err - } - - log.Printf("Sent to controller: %s", response.Data) - } -} - -// HealthCheck checks if the agent is healthy. -func (s *server) HealthCheck(ctx context.Context, req *proto.HealthCheckRequest) (*proto.HealthCheckResponse, error) { - log.Println("Health check requested") - - return &proto.HealthCheckResponse{ - Healthy: true, - Message: "Agent is healthy", - }, nil -} - -func main() { - fmt.Printf("Listening on port: %s\n", port) - lis, err := net.Listen("tcp", port) - if err != nil { - log.Fatalf("Failed to listen: %v", err) - } - - grpcServer := grpc.NewServer() - proto.RegisterAgentServiceServer(grpcServer, &server{}) - - fmt.Println("\nAgent server is running...") - if err := grpcServer.Serve(lis); err != nil { - log.Fatalf("Failed to serve: %v", err) - } -} diff --git a/examples/skills/README.md b/examples/skills/README.md new file mode 100644 index 00000000..72769049 --- /dev/null +++ b/examples/skills/README.md @@ -0,0 +1,37 @@ +# Skills + +AX has built-in support for [Agent Skills](https://agentskills.io). + +## Configuration + +To enable skills with a custom directory, you can set the `skills_dir` in your `ax.yaml` file: + +```yaml +planner: + gemini: + model: "gemini-3.5-flash" + timeout: "60s" + skills_dir: "./examples/skills" +``` + +If you omit `skills_dir` in your `ax.yaml` file, AX will check the `SKILLS_DIR` environment variable first, and finally fall back to `~/.agents/skills`. + +## Example: Using the `emoji` Skill + +AX contains a sample `emoji` skill in `examples/skills/emoji`. If you run a task that matches the description of the skill, the planner will automatically activate it and execute its scripts! + +Run the following command in your terminal: + +```bash +ax exec --input "Give me an emoji for happy" +``` + +The planner will: +1. Discover the `emoji` skill. +2. Determine that your query matches its description. +3. Call `activate_skill` to load the instructions. +4. Execute the `emoji` skill via text generation or script execution based on its instructions. + +### Custom Skills + +Refer to the [Agent Skills](https://agentskills.io) documentation for more information on how to create custom skills. \ No newline at end of file diff --git a/examples/skills/emoji/SKILL.md b/examples/skills/emoji/SKILL.md new file mode 100644 index 00000000..4a502462 --- /dev/null +++ b/examples/skills/emoji/SKILL.md @@ -0,0 +1,9 @@ +--- +name: emoji +description: Returns a corresponding emoji for a given word (e.g. happy, magic etc). +--- + +# Instructions +When the user asks for an emoji for a specific word or concept, use this skill. +1. Call `run_skill_script` to execute `emoji.sh`, passing the word as the first argument. +2. Return the output emoji directly to the user. diff --git a/examples/skills/emoji/scripts/emoji.sh b/examples/skills/emoji/scripts/emoji.sh new file mode 100755 index 00000000..ff8f5cb7 --- /dev/null +++ b/examples/skills/emoji/scripts/emoji.sh @@ -0,0 +1,34 @@ +#!/bin/sh + +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Return an ASCII emoticon based on the input word +# Convert input to lowercase using tr for case-insensitive matching +word=$(echo "$1" | tr '[:upper:]' '[:lower:]') + +case "$word" in + "happy"|"smile"|"joy") echo ":-)" ;; + "sad"|"cry") echo ":-(" ;; + "shrug"|"dunno"|"whatever") echo "¯\_(ツ)_/¯" ;; + "flip"|"angry"|"rage") echo "(╯°□°）╯︵ ┻━┻" ;; + "table"|"putback") echo "┬─┬ノ( º _ ºノ)" ;; + "magic"|"sparkle") echo "(ﾉ◕ヮ◕)ﾉ*:･ﾟ✧" ;; + "sunglasses"|"cool") echo "(•_•) / ( •_•)>⌐■-■ / (⌐■_■)" ;; + "sundar"|"ceo") echo "👓(⌐■_■) 👍" ;; + "dance"|"party") echo "♪~ ᕕ(ᐛ)ᕗ" ;; + "wink") echo ";-)" ;; + "surprise"|"gasp") echo "(O_O)" ;; + *) echo "¯\_(ツ)_/¯ (unknown word)" ;; +esac diff --git a/examples/skills/lowercase/SKILL.md b/examples/skills/lowercase/SKILL.md new file mode 100644 index 00000000..47a2c77b --- /dev/null +++ b/examples/skills/lowercase/SKILL.md @@ -0,0 +1,9 @@ +--- +name: lowercase +description: Converts the provided input string to lowercase. +--- + +# Instructions +When the user asks to convert a text to lowercase, or if you need to lowercase a string for processing, use this skill. +1. Call `run_skill_script` to execute `lowercase.sh`, passing the text as the first argument. +2. Return the lowered text directly to the user. diff --git a/examples/skills/lowercase/scripts/lowercase.sh b/examples/skills/lowercase/scripts/lowercase.sh new file mode 100755 index 00000000..f28b33d8 --- /dev/null +++ b/examples/skills/lowercase/scripts/lowercase.sh @@ -0,0 +1,17 @@ +#!/bin/bash + +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Converts argument to lowercase +echo "$1" | tr '[:upper:]' '[:lower:]' diff --git a/gar.yaml b/gar.yaml deleted file mode 100644 index d26a05fa..00000000 --- a/gar.yaml +++ /dev/null @@ -1,67 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Server configuration -server: - # Address to listen on (host:port) - address: ":8494" - -# Event log configuration -eventlog: - # Directory where event logs are stored - dir: "eventlog" - -# Maximum number of steps per trigger -max_steps: 50 - -# Health check interval for agents (e.g., "30s", "1m") -health_check_interval: 30s - -# Gemini planner configuration (optional) -# The Gemini planner uses Google's Gemini models for intelligent agent selection. -# API key must be set via GEMINI_API_KEY environment variable (required). -# Model name can be set via GAR_GEMINI_MODEL environment variable (optional). -# -# gemini_planner: -# Set GEMINI_API_KEY for authentication. -# -# # Model name (e.g., "gemini-2.5-flash") -# model: "gemini-flash-latest" -# -# # Temperature for generation (0.0 to 1.0) -# temperature: 0.7 -# -# # Maximum output tokens -# max_tokens: 8192 -# -# # Request timeout (e.g., "30s", "1m") -# timeout: 60s -# -# # Number of recent messages to include in context -# context_window: 30 -# -# # Custom system prompt for the planner (optional) -# system_prompt: "" - -# Remote agents configuration (optional) -# List of remote agents to register on startup. -# Remote agents run as separate gRPC servers. -# -# remote_agents: -# - id: "python-agent" -# name: "Python Agent" -# description: "Python-based agent for text processing" -# address: "localhost:50051" -# metadata: -# version: "1.0" diff --git a/go.mod b/go.mod index bf426b27..73c022fa 100644 --- a/go.mod +++ b/go.mod @@ -1,34 +1,73 @@ -module github.com/google/gar +module github.com/google/ax -go 1.24.0 +go 1.26.3 require ( + charm.land/bubbles/v2 v2.0.0 + charm.land/bubbletea/v2 v2.0.2 + charm.land/huh/v2 v2.0.3 + charm.land/lipgloss/v2 v2.0.3 + cloud.google.com/go/compute/metadata v0.9.0 + github.com/agent-substrate/substrate v0.0.0-20260706222328-3cb7433bd8a8 github.com/google/uuid v1.6.0 + github.com/jackc/pgx/v5 v5.10.0 github.com/spf13/cobra v1.10.2 - google.golang.org/genai v1.43.0 - google.golang.org/grpc v1.78.0 - google.golang.org/protobuf v1.36.11 + go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0 + go.opentelemetry.io/otel v1.44.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.44.0 + go.opentelemetry.io/otel/sdk v1.44.0 + go.opentelemetry.io/otel/trace v1.44.0 + golang.org/x/oauth2 v0.36.0 + google.golang.org/grpc v1.81.1 + google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af gopkg.in/yaml.v3 v3.0.1 + modernc.org/sqlite v1.47.0 ) require ( - cloud.google.com/go v0.116.0 // indirect - cloud.google.com/go/auth v0.9.3 // indirect - cloud.google.com/go/compute/metadata v0.9.0 // indirect - github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect - github.com/google/go-cmp v0.7.0 // indirect - github.com/google/s2a-go v0.1.8 // indirect - github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect - github.com/gorilla/websocket v1.5.3 // indirect + github.com/atotto/clipboard v0.1.4 // indirect + github.com/catppuccin/go v0.3.0 // indirect + github.com/cenkalti/backoff/v5 v5.0.3 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/charmbracelet/colorprofile v0.4.3 // indirect + github.com/charmbracelet/ultraviolet v0.0.0-20260205113103-524a6607adb8 // indirect + github.com/charmbracelet/x/ansi v0.11.7 // indirect + github.com/charmbracelet/x/exp/ordered v0.1.0 // indirect + github.com/charmbracelet/x/exp/strings v0.0.0-20240722160745-212f7b056ed0 // indirect + github.com/charmbracelet/x/term v0.2.2 // indirect + github.com/charmbracelet/x/termios v0.1.1 // indirect + github.com/charmbracelet/x/windows v0.2.2 // indirect + github.com/clipperhouse/displaywidth v0.11.0 // indirect + github.com/clipperhouse/uax29/v2 v2.7.0 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/go-logr/logr v1.4.3 // indirect + github.com/go-logr/stdr v1.2.2 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect - github.com/kr/pretty v0.3.1 // indirect - github.com/rogpeppe/go-internal v1.14.1 // indirect + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect + github.com/lucasb-eyer/go-colorful v1.4.0 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-runewidth v0.0.23 // indirect + github.com/mitchellh/hashstructure/v2 v2.0.2 // indirect + github.com/muesli/cancelreader v0.2.2 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + github.com/rivo/uniseg v0.4.7 // indirect github.com/spf13/pflag v1.0.10 // indirect - go.opencensus.io v0.24.0 // indirect - golang.org/x/crypto v0.46.0 // indirect - golang.org/x/net v0.48.0 // indirect - golang.org/x/sys v0.40.0 // indirect - golang.org/x/text v0.32.0 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b // indirect - gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect + github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect + go.opentelemetry.io/auto/sdk v1.2.1 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.44.0 // indirect + go.opentelemetry.io/otel/metric v1.44.0 // indirect + go.opentelemetry.io/proto/otlp v1.10.0 // indirect + golang.org/x/net v0.55.0 // indirect + golang.org/x/sync v0.20.0 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/text v0.37.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect + modernc.org/libc v1.70.0 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect ) diff --git a/go.sum b/go.sum index 3f695972..a5bc044b 100644 --- a/go.sum +++ b/go.sum @@ -1,72 +1,122 @@ -cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= -cloud.google.com/go v0.116.0 h1:B3fRrSDkLRt5qSHWe40ERJvhvnQwdZiHu0bJOpldweE= -cloud.google.com/go v0.116.0/go.mod h1:cEPSRWPzZEswwdr9BxE6ChEn01dWlTaF05LiC2Xs70U= -cloud.google.com/go/auth v0.9.3 h1:VOEUIAADkkLtyfr3BLa3R8Ed/j6w1jTBmARx+wb5w5U= -cloud.google.com/go/auth v0.9.3/go.mod h1:7z6VY+7h3KUdRov5F1i8NDP5ZzWKYmEPO842BgCsmTk= +charm.land/bubbles/v2 v2.0.0 h1:tE3eK/pHjmtrDiRdoC9uGNLgpopOd8fjhEe31B/ai5s= +charm.land/bubbles/v2 v2.0.0/go.mod h1:rCHoleP2XhU8um45NTuOWBPNVHxnkXKTiZqcclL/qOI= +charm.land/bubbletea/v2 v2.0.2 h1:4CRtRnuZOdFDTWSff9r8QFt/9+z6Emubz3aDMnf/dx0= +charm.land/bubbletea/v2 v2.0.2/go.mod h1:3LRff2U4WIYXy7MTxfbAQ+AdfM3D8Xuvz2wbsOD9OHQ= +charm.land/huh/v2 v2.0.3 h1:2cJsMqEPwSywGHvdlKsJyQKPtSJLVnFKyFbsYZTlLkU= +charm.land/huh/v2 v2.0.3/go.mod h1:93eEveeeqn47MwiC3tf+2atZ2l7Is88rAtmZNZ8x9Wc= +charm.land/lipgloss/v2 v2.0.3 h1:yM2zJ4Cf5Y51b7RHIwioil4ApI/aypFXXVHSwlM6RzU= +charm.land/lipgloss/v2 v2.0.3/go.mod h1:7myLU9iG/3xluAWzpY/fSxYYHCgoKTie7laxk6ATwXA= cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs= cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10= -github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= -github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= -github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= -github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc= +github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ= +github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE= +github.com/agent-substrate/substrate v0.0.0-20260706222328-3cb7433bd8a8 h1:tUeBjLs9TJgIWfML2dlZ3UOFG3dacbwrIYkas2ctBgY= +github.com/agent-substrate/substrate v0.0.0-20260706222328-3cb7433bd8a8/go.mod h1:2cvSnnHPwZRAhkrC2LH1bQd1tQBfL3p+zU6pZiHBUkA= +github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= +github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= +github.com/aymanbagabas/go-udiff v0.4.1 h1:OEIrQ8maEeDBXQDoGCbbTTXYJMYRCRO1fnodZ12Gv5o= +github.com/aymanbagabas/go-udiff v0.4.1/go.mod h1:0L9PGwj20lrtmEMeyw4WKJ/TMyDtvAoK9bf2u/mNo3w= +github.com/catppuccin/go v0.3.0 h1:d+0/YicIq+hSTo5oPuRi5kOpqkVA5tAsU6dNhvRu+aY= +github.com/catppuccin/go v0.3.0/go.mod h1:8IHJuMGaUUjQM82qBrGNBv7LFq6JI3NnQCF6MOlZjpc= +github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM= +github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/charmbracelet/colorprofile v0.4.3 h1:QPa1IWkYI+AOB+fE+mg/5/4HRMZcaXex9t5KX76i20Q= +github.com/charmbracelet/colorprofile v0.4.3/go.mod h1:/zT4BhpD5aGFpqQQqw7a+VtHCzu+zrQtt1zhMt9mR4Q= +github.com/charmbracelet/ultraviolet v0.0.0-20260205113103-524a6607adb8 h1:eyFRbAmexyt43hVfeyBofiGSEmJ7krjLOYt/9CF5NKA= +github.com/charmbracelet/ultraviolet v0.0.0-20260205113103-524a6607adb8/go.mod h1:SQpCTRNBtzJkwku5ye4S3HEuthAlGy2n9VXZnWkEW98= +github.com/charmbracelet/x/ansi v0.11.7 h1:kzv1kJvjg2S3r9KHo8hDdHFQLEqn4RBCb39dAYC84jI= +github.com/charmbracelet/x/ansi v0.11.7/go.mod h1:9qGpnAVYz+8ACONkZBUWPtL7lulP9No6p1epAihUZwQ= +github.com/charmbracelet/x/conpty v0.1.1 h1:s1bUxjoi7EpqiXysVtC+a8RrvPPNcNvAjfi4jxsAuEs= +github.com/charmbracelet/x/conpty v0.1.1/go.mod h1:OmtR77VODEFbiTzGE9G1XiRJAga6011PIm4u5fTNZpk= +github.com/charmbracelet/x/errors v0.0.0-20240508181413-e8d8b6e2de86 h1:JSt3B+U9iqk37QUU2Rvb6DSBYRLtWqFqfxf8l5hOZUA= +github.com/charmbracelet/x/errors v0.0.0-20240508181413-e8d8b6e2de86/go.mod h1:2P0UgXMEa6TsToMSuFqKFQR+fZTO9CNGUNokkPatT/0= +github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f h1:pk6gmGpCE7F3FcjaOEKYriCvpmIN4+6OS/RD0vm4uIA= +github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f/go.mod h1:IfZAMTHB6XkZSeXUqriemErjAWCCzT0LwjKFYCZyw0I= +github.com/charmbracelet/x/exp/ordered v0.1.0 h1:55/qLwjIh0gL0Vni+QAWk7T/qRVP6sBf+2agPBgnOFE= +github.com/charmbracelet/x/exp/ordered v0.1.0/go.mod h1:5UHwmG+is5THxMyCJHNPCn2/ecI07aKNrW+LcResjJ8= +github.com/charmbracelet/x/exp/strings v0.0.0-20240722160745-212f7b056ed0 h1:qko3AQ4gK1MTS/de7F5hPGx6/k1u0w4TeYmBFwzYVP4= +github.com/charmbracelet/x/exp/strings v0.0.0-20240722160745-212f7b056ed0/go.mod h1:pBhA0ybfXv6hDjQUZ7hk1lVxBiUbupdw5R31yPUViVQ= +github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk= +github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI= +github.com/charmbracelet/x/termios v0.1.1 h1:o3Q2bT8eqzGnGPOYheoYS8eEleT5ZVNYNy8JawjaNZY= +github.com/charmbracelet/x/termios v0.1.1/go.mod h1:rB7fnv1TgOPOyyKRJ9o+AsTU/vK5WHJ2ivHeut/Pcwo= +github.com/charmbracelet/x/windows v0.2.2 h1:IofanmuvaxnKHuV04sC0eBy/smG6kIKrWG2/jYn2GuM= +github.com/charmbracelet/x/windows v0.2.2/go.mod h1:/8XtdKZzedat74NQFn0NGlGL4soHB0YQZrETF96h75k= +github.com/charmbracelet/x/xpty v0.1.3 h1:eGSitii4suhzrISYH50ZfufV3v085BXQwIytcOdFSsw= +github.com/charmbracelet/x/xpty v0.1.3/go.mod h1:poPYpWuLDBFCKmKLDnhBp51ATa0ooD8FhypRwEFtH3Y= +github.com/clipperhouse/displaywidth v0.11.0 h1:lBc6kY44VFw+TDx4I8opi/EtL9m20WSEFgwIwO+UVM8= +github.com/clipperhouse/displaywidth v0.11.0/go.mod h1:bkrFNkf81G8HyVqmKGxsPufD3JhNl3dSqnGhOoSD/o0= +github.com/clipperhouse/uax29/v2 v2.7.0 h1:+gs4oBZ2gPfVrKPthwbMzWZDaAFPGYK72F0NJv2v7Vk= +github.com/clipperhouse/uax29/v2 v2.7.0/go.mod h1:EFJ2TJMRUaplDxHKj1qAEhCtQPW2tJSwu5BF98AuoVM= +github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 h1:aBangftG7EVZoUb69Os8IaYg++6uMOdKK83QtkkvJik= +github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2/go.mod h1:qwXFYgsP6T7XnJtbKlf1HP8AjxZZyzxMmc+Lq5GjlU4= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= -github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s= +github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= -github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= -github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98= -github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/envoyproxy/go-control-plane v0.14.0 h1:hbG2kr4RuFj222B6+7T83thSPqLjwBIfQawTkC++2HA= +github.com/envoyproxy/go-control-plane/envoy v1.37.0 h1:u3riX6BoYRfF4Dr7dwSOroNfdSbEPe9Yyl09/B6wBrQ= +github.com/envoyproxy/go-control-plane/envoy v1.37.0/go.mod h1:DReE9MMrmecPy+YvQOAOHNYMALuowAnbjjEMkkWOi6A= +github.com/envoyproxy/protoc-gen-validate v1.3.3 h1:MVQghNeW+LZcmXe7SY1V36Z+WFMDjpqGAGacLe2T0ds= +github.com/envoyproxy/protoc-gen-validate v1.3.3/go.mod h1:TsndJ/ngyIdQRhMcVVGDDHINPLWB7C82oDArY51KfB0= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= -github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= -github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc= -github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE= -github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc= -github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A= -github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= -github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= -github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= -github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= -github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= -github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= -github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= -github.com/golang/protobuf v1.4.1/go.mod h1:U8fpvMrcmy5pZrNK1lt4xCsGvpyWQ/VVv6QDs8UjoX8= -github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= -github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= -github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= -github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= -github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= -github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= -github.com/google/go-cmp v0.5.3/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= -github.com/google/s2a-go v0.1.8 h1:zZDs9gcbt9ZPLV0ndSyQk6Kacx2g/X+SKYovpnz3SMM= -github.com/google/s2a-go v0.1.8/go.mod h1:6iNWHTpQ+nfNRN5E00MSdfDwVesa8hhS32PhPO8deJA= -github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/pprof v0.0.0-20250602020802-c6617b811d0e h1:FJta/0WsADCe1r9vQjdHbd3KuiLPu7Y9WlyLGwMUNyE= +github.com/google/pprof v0.0.0-20250602020802-c6617b811d0e/go.mod h1:5hDyRhoBCxViHszMt12TnOpEI4VVi+U8Gm9iphldiMA= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/googleapis/enterprise-certificate-proxy v0.3.4 h1:XYIDZApgAnrN1c855gTgghdIA6Stxb52D5RnLI1SLyw= -github.com/googleapis/enterprise-certificate-proxy v0.3.4/go.mod h1:YKe7cfqYXjKGpGvmSg28/fFvhNzinZQm8DGnaburhGA= -github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= -github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF28c5ZQfqCBQ5g2xfk= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= -github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= +github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= -github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= -github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= -github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= +github.com/lucasb-eyer/go-colorful v1.4.0 h1:UtrWVfLdarDgc44HcS7pYloGHJUjHV/4FwW4TvVgFr4= +github.com/lucasb-eyer/go-colorful v1.4.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-runewidth v0.0.23 h1:7ykA0T0jkPpzSvMS5i9uoNn2Xy3R383f9HDx3RybWcw= +github.com/mattn/go-runewidth v0.0.23/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= +github.com/mitchellh/hashstructure/v2 v2.0.2 h1:vGKWl0YJqUNxE8d+h8f6NJLcCJrgbhC4NcD46KavDd4= +github.com/mitchellh/hashstructure/v2 v2.0.2/go.mod h1:MG3aRVU/N29oo/V/IhBX8GR/zz4kQkprJgF2EVszyDE= +github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= +github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo= +github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= -github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= +github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= @@ -76,98 +126,93 @@ github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk= github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= -github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= -github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= -github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= -github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= -go.opencensus.io v0.24.0 h1:y73uSU6J157QMP2kn2r30vwW1A2W2WFwSCGnAVxeaD0= -go.opencensus.io v0.24.0/go.mod h1:vNK8G9p7aAivkbmorf4v+7Hgx+Zs0yY+0fOtgBfjQKo= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= -go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8= -go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM= -go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA= -go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI= -go.opentelemetry.io/otel/sdk v1.38.0 h1:l48sr5YbNf2hpCUj/FoGhW9yDkl+Ma+LrVl8qaM5b+E= -go.opentelemetry.io/otel/sdk v1.38.0/go.mod h1:ghmNdGlVemJI3+ZB5iDEuk4bWA3GkTpW+DOoZMYBVVg= -go.opentelemetry.io/otel/sdk/metric v1.38.0 h1:aSH66iL0aZqo//xXzQLYozmWrXxyFkBJ6qT5wthqPoM= -go.opentelemetry.io/otel/sdk/metric v1.38.0/go.mod h1:dg9PBnW9XdQ1Hd6ZnRz689CbtrUp0wMMs9iPcgT9EZA= -go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE= -go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0 h1:2yEATaop1/a1I4psnSLgWVPLWwCzkqWakgJy7xTDVy0= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0/go.mod h1:D7J12YRapIekYyPWgGPlA/23pRmpSEZC5xJC/TTLI9U= +go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= +go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.44.0 h1:4YsVu3B8+3qtWYYrsUYgn0OG78pN0rnNPRGX4SbokQI= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.44.0/go.mod h1:+wnlSn0mD1ADVMe3v9Z/WIaiz6q6gL2J/ejaAmdmv80= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.44.0 h1:qazEJlUOQzhCpzQpFETGby7EdqjI1wsd0W+6Gg1SCTU= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.44.0/go.mod h1:fOD2Yefuxixkx3ahVNf0O/PERb6r4OlbxfATVnYvzCo= +go.opentelemetry.io/otel/metric v1.44.0 h1:1w0gILTcHdr3YI+ixLyjemwrVnsMURbTZFrSYCdDdmc= +go.opentelemetry.io/otel/metric v1.44.0/go.mod h1:8O7hanEPBNgEMmybD3s2VBKcgWOCsA6tzHBPODAiquo= +go.opentelemetry.io/otel/sdk v1.44.0 h1:nHYwb9lK+fJPU/dnT6s7W7Z8itMWyqrnVfbheVYrZ58= +go.opentelemetry.io/otel/sdk v1.44.0/go.mod h1:Osuydd3Se74nqjAKxid74N5eC+jfEqfTegHRnq58oK0= +go.opentelemetry.io/otel/sdk/metric v1.44.0 h1:3LlKgI+VjbVsjNRFZJZAJ30WjXC5VkNRks6si09iEfI= +go.opentelemetry.io/otel/sdk/metric v1.44.0/go.mod h1:5B5pMARnXxKhltooO4xUuCBorl65a4EpnTalObqOigA= +go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/6gtIk= +go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE= +go.opentelemetry.io/proto/otlp v1.10.0 h1:IQRWgT5srOCYfiWnpqUYz9CVmbO8bFmKcwYxpuCSL2g= +go.opentelemetry.io/proto/otlp v1.10.0/go.mod h1:/CV4QoCR/S9yaPj8utp3lvQPoqMtxXdzn7ozvvozVqk= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= -golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= -golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= -golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU= -golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0= -golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= -golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= -golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= -golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= -golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= -golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= -golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= -golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= -golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= -golang.org/x/net v0.0.0-20201110031124-69a78807bb2b/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= -golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU= -golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY= -golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= -golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= -golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= -golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= -golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= -golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= -golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= -golang.org/x/sys v0.40.0 h1:DBZZqJ2Rkml6QMQsZywtnjnnGvHza6BTfYFWY9kjEWQ= -golang.org/x/sys v0.40.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= -golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= -golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= -golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU= -golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY= -golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= -golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= -golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY= -golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= -golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q= -golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk= -gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E= -google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= -google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= -google.golang.org/genai v1.43.0 h1:8vhqhzJNZu1U94e2m+KvDq/TUUjSmDrs1aKkvTa8SoM= -google.golang.org/genai v1.43.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk= -google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc= -google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98Agz4BDEuKkezgsaosCRResVns1a3J2ZsMNc= -google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo= -google.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b h1:Mv8VFug0MP9e5vUxfBcE3vUkV6CImK3cMNMIDFjmzxU= -google.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b/go.mod h1:j9x/tPzZkyxcgEFkiKEEGxfvyumM01BEtsW8xzOahRQ= -google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= -google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg= -google.golang.org/grpc v1.25.1/go.mod h1:c3i+UQWmh7LiEpx4sFZnkU36qjEYZ0imhYfXVyQciAY= -google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk= -google.golang.org/grpc v1.33.2/go.mod h1:JMHMWHQWaTccqQQlmk3MJZS+GWXOdAesneDmEnv2fbc= -google.golang.org/grpc v1.78.0 h1:K1XZG/yGDJnzMdd/uZHAkVqJE+xIDOcmdSFZkBUicNc= -google.golang.org/grpc v1.78.0/go.mod h1:I47qjTo4OKbMkjA/aOOwxDIiPSBofUtQUI5EfpWvW7U= -google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= -google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= -google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= -google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= -google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= -google.golang.org/protobuf v1.22.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= -google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= -google.golang.org/protobuf v1.23.1-0.20200526195155-81db48ad09cc/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= -google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlbajtzgsN7c= -google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= -google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY= +golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70= +golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM= +golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU= +golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= +golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= +golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= +golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= +golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= +golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= +golang.org/x/tools v0.44.0 h1:UP4ajHPIcuMjT1GqzDWRlalUEoY+uzoZKnhOjbIPD2c= +golang.org/x/tools v0.44.0/go.mod h1:KA0AfVErSdxRZIsOVipbv3rQhVXTnlU6UhKxHd1seDI= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa h1:Kjn0N0tCrDgiAFW+lGO4JZ3ck44CehvJQMAwj9QF0G8= +google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:q4lMZS6kskjT5HvCPrnnypcDPVJqT/f4nfxmkE7gryY= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa h1:mZHHdPZl0dbGHCflZgAq/Q468DWVFcU2whhB2KAo8fk= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.81.1 h1:VnnIIZ88UzOOKLukQi+ImGz8O1Wdp8nAGGnvOfEIWQQ= +google.golang.org/grpc v1.81.1/go.mod h1:xGH9GfzOyMTGIOXBJmXt+BX/V0kcdQbdcuwQ/zNw42I= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= -honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= -honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= +modernc.org/cc/v4 v4.27.1 h1:9W30zRlYrefrDV2JE2O8VDtJ1yPGownxciz5rrbQZis= +modernc.org/cc/v4 v4.27.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0= +modernc.org/ccgo/v4 v4.32.0 h1:hjG66bI/kqIPX1b2yT6fr/jt+QedtP2fqojG2VrFuVw= +modernc.org/ccgo/v4 v4.32.0/go.mod h1:6F08EBCx5uQc38kMGl+0Nm0oWczoo1c7cgpzEry7Uc0= +modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM= +modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo= +modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.70.0 h1:U58NawXqXbgpZ/dcdS9kMshu08aiA6b7gusEusqzNkw= +modernc.org/libc v1.70.0/go.mod h1:OVmxFGP1CI/Z4L3E0Q3Mf1PDE0BucwMkcXjjLntvHJo= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8= +modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.47.0 h1:R1XyaNpoW4Et9yly+I2EeX7pBza/w+pmYee/0HJDyKk= +modernc.org/sqlite v1.47.0/go.mod h1:hWjRO6Tj/5Ik8ieqxQybiEOUXy0NJFNp2tpvVpKlvig= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/hack/install-ax.sh b/hack/install-ax.sh new file mode 100755 index 00000000..78f2b630 --- /dev/null +++ b/hack/install-ax.sh @@ -0,0 +1,371 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e +set -u +set -o pipefail + +ROOT=$(git rev-parse --show-toplevel) +cd "${ROOT}" + +if [[ -n "${GOOGLE_CLOUD_PROJECT:-}" ]]; then + export AX_IMAGE_REPO="gcr.io/${GOOGLE_CLOUD_PROJECT}" + echo "Using AX_IMAGE_REPO: ${AX_IMAGE_REPO}" >&2 +fi + +export KO_DEFAULTPLATFORMS="${KO_DEFAULTPLATFORMS:-linux/amd64}" + +# ANSI color codes for prettier output +COLOR_CYAN='\033[1;36m' +COLOR_RESET='\033[0m' + +function log_step() { + local step_name="$1" + echo -e "${COLOR_CYAN}[step]: ${step_name}${COLOR_RESET}" +} + +# wait_with_spinner runs a blocking command while showing a simple spinner on an +# interactive terminal, then reports "done"/"failed" and returns the command's +# exit status. +wait_with_spinner() { + local msg="$1"; shift + if [[ ! -t 2 ]]; then + "$@" + return $? + fi + + local out; out="$(mktemp)" + "$@" >"${out}" 2>&1 & + local pid=$! frames='|/-\' i=0 + while kill -0 "${pid}" 2>/dev/null; do + i=$(( (i + 1) % ${#frames} )) + printf '\r%s %s' "${frames:${i}:1}" "${msg}" >&2 + sleep 0.1 + done + + local status=0 + wait "${pid}" || status=$? + if [[ "${status}" -eq 0 ]]; then + printf '\r\033[K%s... done\n' "${msg}" >&2 + else + printf '\r\033[K%s... failed\n' "${msg}" >&2 + cat "${out}" >&2 + fi + rm -f "${out}" + return "${status}" +} + +function usage() { + echo "Usage: $0 [options]" + echo "" + echo "Options:" + echo " --deploy-ax-server Build images and deploy AX server and components" + echo " --delete-ax-server Delete AX server and components, preserving the event-log database" + echo " --deploy-postgres With --deploy-ax-server: also deploy a bundled Postgres for testing (default: connect to an existing Postgres via AX_EVENTLOG_DSN)" + echo " -h, --help Show this help message" +} + +run_kubectl() { + kubectl \ + ${KUBECTL_CONTEXT:+--context=${KUBECTL_CONTEXT}} \ + "$@" +} + +# detect_container_engine selects the OCI build/push tool when CONTAINER_ENGINE +# is not set explicitly. It prefers a *working* docker (daemon reachable), then a +# working podman, so a docker CLI installed without a running daemon does not +# shadow a working podman. As a last resort it picks whichever CLI exists so the +# build step can surface an actionable daemon error. +detect_container_engine() { + if [[ -n "${CONTAINER_ENGINE:-}" ]]; then + return # Respect an explicit override; do not second-guess it. + fi + if docker info >/dev/null 2>&1; then + CONTAINER_ENGINE=docker + elif podman info >/dev/null 2>&1; then + CONTAINER_ENGINE=podman + elif command -v docker >/dev/null 2>&1; then + CONTAINER_ENGINE=docker + elif command -v podman >/dev/null 2>&1; then + CONTAINER_ENGINE=podman + else + CONTAINER_ENGINE=docker + fi +} + +# _build_and_push_ax builds one target of cmd/ax/Dockerfile into ${repo}:, +# pushes it, and echoes the digest-pinned reference on stdout. Args: +# $1 repo - fully-qualified image repo (e.g. gcr.io/proj/ax) +# $2 target - Dockerfile target stage (ax | antigravity) +# Requires AX_IMAGE_REPO and a container engine. +_build_and_push_ax() { + local repo="$1" target="$2" + if [[ -z "${AX_IMAGE_REPO:-}" ]]; then + echo "Error: AX_IMAGE_REPO environment variable must be set" >&2 + exit 1 + fi + detect_container_engine + if ! command -v "${CONTAINER_ENGINE}" >/dev/null 2>&1; then + echo "Error: container engine '${CONTAINER_ENGINE}' not found in PATH." >&2 + echo "Install it or set CONTAINER_ENGINE to an available builder." >&2 + exit 1 + fi + + local tag image digest + tag="$(git rev-parse --short HEAD)" + image="${repo}:${tag}" + + log_step "build ${target} -> ${image}" >&2 + "${CONTAINER_ENGINE}" build \ + --platform linux/amd64 \ + --target "${target}" \ + -f cmd/ax/Dockerfile \ + -t "${image}" \ + . 2>&1 \ + | awk '{ sub(/^\[[0-9]+\/[0-9]+\] /, ""); print; fflush() }' >&2 + + # Push the readable tag, then resolve the pushed manifest digest so the + # ActorTemplate can reference the image by digest (snapshot-safe). + if [[ "${CONTAINER_ENGINE}" == *podman* ]]; then + local digestfile + digestfile="$(mktemp)" + "${CONTAINER_ENGINE}" push --digestfile="${digestfile}" "${image}" >&2 + digest="$(cat "${digestfile}")" + rm -f "${digestfile}" + else + "${CONTAINER_ENGINE}" push "${image}" >&2 + local repo_digest + repo_digest="$("${CONTAINER_ENGINE}" image inspect --format '{{index .RepoDigests 0}}' "${image}")" + digest="${repo_digest##*@}" + fi + + if [[ "${digest}" != sha256:* ]]; then + echo "Error: could not resolve a sha256 digest for ${image} (got '${digest}')." >&2 + exit 1 + fi + + echo "${repo}@${digest}" +} + +# build_ax_image builds and pushes the Go-only ax image used by the ax-server and the +# Go interactions harness. +build_ax_image() { + _build_and_push_ax "${AX_IMAGE_REPO}/ax" ax +} + +# build_ax_antigravity_image builds and pushes the antigravity image used by the +# antigravity harness. +build_ax_antigravity_image() { + _build_and_push_ax "${AX_IMAGE_REPO}/ax-antigravity" antigravity +} + +build_ateom_image() { + if [[ -n "${ATEOM_IMAGE:-}" ]]; then + echo "${ATEOM_IMAGE}" + return + fi + if [[ -z "${AX_IMAGE_REPO:-}" ]]; then + echo "Error: AX_IMAGE_REPO environment variable must be set" >&2 + exit 1 + fi + + # Resolve the substrate source for the version AX is pinned to in go.mod. + go mod download github.com/agent-substrate/substrate + local sub_dir ateom_ref + sub_dir="$(go list -m -f '{{.Dir}}' github.com/agent-substrate/substrate)" + if [[ -z "${sub_dir}" ]]; then + echo "Error: could not locate the substrate module (go list -m)." >&2 + exit 1 + fi + + log_step "build_ateom_image (from ${sub_dir})" >&2 + ateom_ref="$(cd "${sub_dir}" && KO_DOCKER_REPO="${AX_IMAGE_REPO}" GOFLAGS="" ko build --platform=linux/amd64 -B ./cmd/ateom-gvisor)" + + if [[ "${ateom_ref}" != *@sha256:* ]]; then + echo "Error: ko did not return a digest-pinned ateom image (got '${ateom_ref}')." >&2 + exit 1 + fi + echo "${ateom_ref}" +} + +deploy_ax_server() { + log_step "deploy_ax_server" + + # Check dependencies + if [[ -z "${GEMINI_API_KEY:-}" ]]; then + echo "Error: GEMINI_API_KEY environment variable must be set" >&2 + exit 1 + fi + if [[ -z "${AX_SNAPSHOTS_BUCKET:-}" ]]; then + echo "Error: AX_SNAPSHOTS_BUCKET environment variable must be set" >&2 + exit 1 + fi + # The default (external Postgres) path needs a DSN; fail fast before building. + if [[ "${DEPLOY_POSTGRES}" != "true" && -z "${AX_EVENTLOG_DSN:-}" ]]; then + echo "Error: AX_EVENTLOG_DSN must be set to your Postgres DSN, or pass --deploy-postgres to deploy a bundled test Postgres." >&2 + exit 1 + fi + + echo "Using GCS Bucket: ${AX_SNAPSHOTS_BUCKET}" + + # Build and push the images, capturing their digest-pinned references. + local ax_image ax_antigravity_image ateom_image + ax_image=$(build_ax_image) + ax_antigravity_image=$(build_ax_antigravity_image) + ateom_image=$(build_ateom_image) + + # Resolve the event-log Postgres DSN. By default ax-server connects to an + # existing Postgres via AX_EVENTLOG_DSN; --deploy-postgres creates a bundled + # Postgres in-cluster (for testing) and derives the DSN from it. + local pg_dsn pg_password="" + if [[ "${DEPLOY_POSTGRES}" == "true" ]]; then + # Reuse the existing bundled-Postgres password if present, else POSTGRES_PASSWORD, + # else generate one. + local existing_pw + existing_pw="$(run_kubectl -n ax get secret ax-eventlog-postgres -o go-template='{{.data.password | base64decode}}' 2>/dev/null || true)" + if [[ -n "${existing_pw}" ]]; then + pg_password="${existing_pw}" + else + pg_password="${POSTGRES_PASSWORD:-$(openssl rand -hex 16)}" + fi + pg_dsn="postgres://axuser:${pg_password}@ax-eventlog-postgres.ax.svc:5432/axeventlog?sslmode=disable" + else + pg_dsn="${AX_EVENTLOG_DSN}" + echo "Using existing Postgres via AX_EVENTLOG_DSN." >&2 + fi + + # Base64 of the single-source config file (manifests/ax.yaml), injected into + # each ActorTemplate's AX_CONFIG_CONTENT env so substrate actors read their config + # without a mounted ConfigMap. + local ax_config_content + ax_config_content="$(base64 < manifests/ax.yaml | tr -d '\n')" + + # Common substitutions applied to every rendered manifest. + local render_sed=( + -e "s|\${GEMINI_API_KEY}|${GEMINI_API_KEY}|g" + -e "s|\${AX_SNAPSHOTS_BUCKET}|${AX_SNAPSHOTS_BUCKET}|g" + -e "s|\${AX_IMAGE}|${ax_image}|g" + -e "s|\${AX_ANTIGRAVITY_IMAGE}|${ax_antigravity_image}|g" + -e "s|\${ATEOM_IMAGE}|${ateom_image}|g" + -e "s|\${GOOGLE_CLOUD_PROJECT}|${GOOGLE_CLOUD_PROJECT:-}|g" + -e "s|\${AX_CONFIG_CONTENT}|${ax_config_content}|g" + ) + + # Render and apply the core manifest (namespace, harnesses, ax-server). + if ! sed "${render_sed[@]}" manifests/ax-deployment.yaml | run_kubectl apply -f -; then + echo >&2 + echo "Error: cluster rejected the manifest. An \"unknown field\" error usually means the" >&2 + echo "cluster's substrate is incompatible with AX's go.mod pin — see" >&2 + echo "manifests/README.md (\"Substrate compatibility\")." >&2 + exit 1 + fi + + # Create/update the controller ConfigMap from the single-source config file + # (manifests/ax.yaml) -- the same file that seeds the actors' AX_CONFIG_CONTENT. + # ax-server mounts it at /etc/ax/ax.yaml. + run_kubectl -n ax create configmap ax-server-config \ + --from-file=ax.yaml=manifests/ax.yaml \ + --dry-run=client -o yaml | run_kubectl apply -f - + + # Create/update the event-log Secret with the DSN (and, for the bundled Postgres, + # its password). ax-server reads AX_EVENTLOG_DSN from this Secret's dsn key. + local secret_args=(--from-literal=dsn="${pg_dsn}") + if [[ "${DEPLOY_POSTGRES}" == "true" ]]; then + secret_args+=(--from-literal=password="${pg_password}") + fi + run_kubectl -n ax create secret generic ax-eventlog-postgres "${secret_args[@]}" \ + --dry-run=client -o yaml | run_kubectl apply -f - + + # With --deploy-postgres, create the bundled Postgres and wait for it to be + # ready before ax-server relies on it. + if [[ "${DEPLOY_POSTGRES}" == "true" ]]; then + run_kubectl apply -f manifests/ax-postgres.yaml + log_step "wait for statefulset/ax-eventlog-postgres to be ready" + wait_with_spinner "waiting for postgres (timeout ${AX_WAIT_TIMEOUT:-5m})" \ + run_kubectl -n ax rollout status statefulset/ax-eventlog-postgres \ + --timeout="${AX_WAIT_TIMEOUT:-5m}" + fi + + # Wait for the antigravity ActorTemplate's golden snapshot to be ready. + log_step "wait for actortemplate/ax-harness-antigravity-template to be Ready" + wait_with_spinner "waiting for antigravity golden snapshot (timeout ${AX_WAIT_TIMEOUT:-5m})" \ + run_kubectl wait --for=condition=Ready actortemplate/ax-harness-antigravity-template \ + -n ax --timeout="${AX_WAIT_TIMEOUT:-5m}" + + # Wait for the interactions ActorTemplate's golden snapshot to be ready. + log_step "wait for actortemplate/ax-harness-interactions-template to be Ready" + wait_with_spinner "waiting for interactions golden snapshot (timeout ${AX_WAIT_TIMEOUT:-5m})" \ + run_kubectl wait --for=condition=Ready actortemplate/ax-harness-interactions-template \ + -n ax --timeout="${AX_WAIT_TIMEOUT:-5m}" + + echo "" + echo "Forward the AX server by running the following command (optional)" + echo "kubectl port-forward -n ax rs/ax-server 8494:8494" + echo "" + echo "Then, run: ax exec --server localhost:8494 [--harness antigravity|antigravity-interactions]" +} + +# delete_ax_server removes the AX server and harness resources but preserves the +# event-log database: it leaves the namespace and the Postgres subsystem +# (Service/Secret/StatefulSet and its PVC) intact so a later redeploy reuses the +# existing data. +delete_ax_server() { + log_step "delete_ax_server" + + run_kubectl -n ax delete --ignore-not-found \ + replicaset/ax-server \ + configmap/ax-server-config \ + actortemplate/ax-harness-antigravity-template \ + actortemplate/ax-harness-interactions-template \ + workerpool/ax-harness-workerpool +} + +if [ "$#" -eq 0 ]; then + usage + exit 1 +fi + +# Event-log Postgres: by default ax-server connects to an existing Postgres via +# the AX_EVENTLOG_DSN env var. --deploy-postgres additionally creates a bundled +# Postgres in-cluster (for testing on Substrate). +DEPLOY_POSTGRES=false + +# If -h or --help appears anywhere in the command line, print the usage and exit. +for arg in "$@"; do + case "$arg" in + -h|--help) + usage + exit 0 + ;; + --deploy-postgres) + DEPLOY_POSTGRES=true + ;; + esac +done + +while [[ "$#" -gt 0 ]]; do + case $1 in + --deploy-ax-server) deploy_ax_server ;; + --delete-ax-server) delete_ax_server ;; + --deploy-postgres) ;; # resolved in the pre-scan above + *) + echo "Error: unknown option: $1" >&2 + echo "" + usage + exit 1 + ;; + esac + shift +done diff --git a/internal/cmd/e2e/main.go b/internal/cmd/e2e/main.go new file mode 100644 index 00000000..0e48dc6e --- /dev/null +++ b/internal/cmd/e2e/main.go @@ -0,0 +1,133 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package main implements an end-to-end demonstration of the Antigravity harness +// integration with AX Controller V2. +// +// TO RUN THIS E2E DEMONSTRATION: +// +// Step 1: Start the Python gRPC Harness Server (in a separate terminal or background): +// +// PYTHONPATH=python:. /Users/anjalisridhar/.gemini/jetski/worktrees/harness-interface-3/implement-agy-sdk-streaming-20260528/.venv/bin/python python/antigravity/harness_server.py --port 50053 +// +// Step 2: Run this Go E2E client: +// +// go run internal/cmd/e2e/main.go +package main + +import ( + "context" + "fmt" + "log" + "net" + "os" + "time" + + "github.com/google/ax/internal/controller" + "github.com/google/ax/internal/controller/eventlog" + "github.com/google/ax/internal/controller/eventlog/eventlogtest" + "github.com/google/ax/internal/harness/antigravity" + "github.com/google/ax/proto" +) + +func main() { + ctx := context.Background() + fmt.Println("==================================================") + fmt.Println("AX Controller V2 - E2E Harness Demonstration") + fmt.Println("==================================================") + + // ------------------------------------------------------------------------- + // Demo 1: Unregistered Harness (Should fail) + // ------------------------------------------------------------------------- + fmt.Println("\n--- Demo 1: Unregistered Harness ---") + fmt.Println("Requesting 'unregistered-agent'. Exec should fail since no harness is registered.") + runDemo(ctx, "unregistered-agent", func(reg *controller.Registry) { + // Do not register any harness + }) + + // ------------------------------------------------------------------------- + // Demo 2: Antigravity Execution (Requires google-antigravity & GEMINI_API_KEY) + // ------------------------------------------------------------------------- + fmt.Println("\n--- Demo 2: Antigravity Execution ---") + fmt.Println("Registering 'antigravity' with real script. Attempting execution.") + if os.Getenv("GEMINI_API_KEY") == "" { + fmt.Println("WARNING: GEMINI_API_KEY is not set. Execution will likely fail if dependencies are missing, but we will try anyway.") + } + runDemo(ctx, "antigravity", func(reg *controller.Registry) { + // With the new stateful gRPC-based streaming harness, connectivity checks on the + // server address replace the build-time checks for local script file presence. + address := "localhost:50053" + conn, err := net.DialTimeout("tcp", address, 1*time.Second) + if err != nil { + log.Fatalf("Antigravity harness server not active at %s: %v", address, err) + } + 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) + reg.RegisterHarness("antigravity", h) + }) +} + +func runDemo(ctx context.Context, harnessID string, setupRegistry func(reg *controller.Registry)) { + reg := controller.NewRegistry() + setupRegistry(reg) + + log := &eventlogtest.MemoryEventLog{} + c, err := controller.New(ctx, controller.Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { + return log, nil + }, + }) + if err != nil { + fmt.Printf("Error creating controller: %v\n", err) + return + } + defer c.Close() + + handler := controller.ExecHandler(func(resp *proto.ExecResponse) error { + for _, out := range resp.Outputs { + if textContent := out.GetContent().GetText().GetText(); textContent != "" { + fmt.Printf("Agent Output: %s\n", textContent) + } else if toolCall := out.GetContent().GetToolCall(); toolCall != nil { + fmt.Printf("Agent Triggered Tool Call: %s (ID: %s)\n", toolCall.GetFunctionCall().Name, toolCall.Id) + } + } + return nil + }) + + inputs := []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{Text: "What is the weather in New York?"}, + }, + }, + }, + } + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: "e2e-conv", + Inputs: inputs, + HarnessId: harnessID, + }, handler) + + if err != nil { + fmt.Printf("Execution Failed (as expected if environment is not ready): %v\n", err) + } else { + fmt.Println("Execution Succeeded!") + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 737e43d0..58ec5008 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -12,26 +12,48 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Package config provides configuration structures for GAR server. +// Package config provides configuration for the controller server path. package config import ( "fmt" "os" - "time" + "path/filepath" - "github.com/google/gar/agent" + "github.com/google/ax/internal/harness" + "github.com/google/ax/internal/harness/substrate" "gopkg.in/yaml.v3" ) -// Config represents the main configuration for GAR server. +const ( + // The substrate namespace reserved for AX's built-in harnesses. + defaultNamespace = "ax" + // The port for harnesses running as substrate actors. Substrate's + // actor networking DNATs inbound workerPodIP:80 to the actor. + substrateDefaultPort = 80 + // Harness IDs reserved for AX's built-in harnesses. + AntigravityHarnessID = "antigravity" + AntigravityInteractionsHarnessID = "antigravity-interactions" + // AntigravityHarnessTemplate is the substrate ActorTemplate that runs the + // Antigravity harness. + AntigravityHarnessTemplate = "ax-harness-antigravity-template" + // AntigravityInteractionsTemplate is the substrate ActorTemplate that runs + // the Antigravity Interactions harness. + AntigravityInteractionsTemplate = "ax-harness-interactions-template" +) + +// Config represents the main configuration for the AX harness server. type Config struct { - Server ServerConfig `yaml:"server"` - EventLog EventLogConfig `yaml:"eventlog"` - MaxSteps int `yaml:"max_steps"` // Maximum steps per trigger - HealthCheckInterval time.Duration `yaml:"health_check_interval"` // Health check interval for agents - GeminiPlanner GeminiPlannerConfig `yaml:"gemini_planner,omitempty"` - RemoteAgents []RemoteAgentConfig `yaml:"remote_agents,omitempty"` // List of remote agents to register + Version string `yaml:"version"` + Server ServerConfig `yaml:"server"` + EventLog EventLogConfig `yaml:"eventlog"` + Harnesses HarnessesConfig `yaml:"harnesses,omitempty"` + // Skills sources skills from the Gemini Enterprise Skill Registry into on-disk folders + // before the harness starts. It is harness-agnostic: each actor runs exactly + // one harness, which consumes the materialized folder(s). Optional; disabled + // when no registry is enabled. + Skills SkillsConfig `yaml:"skills,omitempty"` + Telemetry TelemetryConfig `yaml:"telemetry,omitempty"` } // ServerConfig configures the gRPC server. @@ -39,44 +61,179 @@ type ServerConfig struct { Address string `yaml:"address"` // Server address to listen on (e.g., ":8494") } +// TelemetryConfig configures telemetry options. +type TelemetryConfig struct { + OTLP OTLPConfig `yaml:"otlp,omitempty"` +} + +// OTLPConfig configures the OTLP exporter. +type OTLPConfig struct { + Enabled bool `yaml:"enabled,omitempty"` + Endpoint string `yaml:"endpoint,omitempty"` // OTLP collector endpoint (e.g., "localhost:4317") +} + +// SQLiteConfig configures the SQLite event log file. +type SQLiteConfig struct { + Filename string `yaml:"filename"` // SQLite file for event log storage +} + +// PostgresConfig configures the Postgres event log. +type PostgresConfig struct { + DSN string `yaml:"dsn"` // Postgres connection DSN +} + // EventLogConfig configures the event log storage. type EventLogConfig struct { - Dir string `yaml:"dir"` // Directory for event log files + SQLiteConfig SQLiteConfig `yaml:"sqlite,omitempty"` + PostgresConfig PostgresConfig `yaml:"postgres,omitempty"` +} + +// HarnessesConfig groups harnesses to serve by type. There are two categories: +// - Built-in harnesses (e.g. Antigravity, AntigravityInteractions) whose +// implementation and container image are provided by AX. +// - Custom harnesses on substrate whose implementation and container image are +// provided by the user via their own ActorTemplate. +type HarnessesConfig struct { + Antigravity AntigravityHarnessConfig `yaml:"antigravity,omitempty"` + AntigravityInteractions AntigravityInteractionsHarnessConfig `yaml:"antigravity-interactions,omitempty"` + Substrate []SubstrateHarnessConfig `yaml:"substrate,omitempty"` +} + +// AntigravityHarnessConfig registers the built-in Antigravity harness. +type AntigravityHarnessConfig struct { + Default bool `yaml:"default,omitempty"` + Endpoint string `yaml:"endpoint,omitempty"` // HarnessService address +} + +// AntigravityInteractionsHarnessConfig registers the built-in Antigravity +// Interactions harness (over the Vertex GenAI Interactions API). +type AntigravityInteractionsHarnessConfig struct { + Default bool `yaml:"default,omitempty"` // Default harness or not + Agent string `yaml:"agent,omitempty"` // Interactions API agent (default: antigravityinteractions.DefaultAgent) + // SystemInstruction is a free-form system prompt sent on every turn. + SystemInstruction string `yaml:"system_instruction,omitempty"` } -// ControllerConfig configures the controller behavior. -type ControllerConfig struct { - MaxSteps int `yaml:"max_steps"` // Maximum steps per trigger - HealthCheckInterval time.Duration `yaml:"health_check_interval"` // Health check interval for agents +// SkillsConfig configures optional skill sources (top-level, harness-agnostic). +// Today the only source type is the Gemini Enterprise Skill Registry; it may source from more +// than one registry (e.g. a shared org-wide registry plus a team-specific one), +// each with its own project/location, selection, and target directory. +type SkillsConfig struct { + Registries []SkillsRegistryConfig `yaml:"registries,omitempty"` } -// GeminiPlannerConfig configures the Gemini-based planner. -// Note: API key is not configurable here for security reasons. -// Set GEMINI_API_KEY environment variable instead. -type GeminiPlannerConfig struct { - Model string `yaml:"model,omitempty"` // Model name - Temperature float32 `yaml:"temperature,omitempty"` - MaxTokens int32 `yaml:"max_tokens,omitempty"` - Timeout time.Duration `yaml:"timeout,omitempty"` - ContextWindow int `yaml:"context_window,omitempty"` - SystemPrompt string `yaml:"system_prompt,omitempty"` +// Validate checks the (top-level) skills config. +func (s SkillsConfig) Validate() error { + for i := range s.Registries { + if err := s.Registries[i].validate(i); err != nil { + return err + } + } + return nil +} + +// validate enforces that, when this registry source is enabled, exactly one +// selection mode is set (skills, query, or all) — a config-level "oneof". idx is +// the registry's index within the registries list, for error context. +func (rc SkillsRegistryConfig) validate(idx int) error { + if !rc.Enabled { + return nil + } + if rc.TargetDir == "" { + return fmt.Errorf("skills.registries[%d] requires target_dir (skills materialize to //)", idx) + } + modes := 0 + if len(rc.Skills) > 0 { + modes++ + } + if rc.Query != nil { + modes++ + } + if rc.All { + modes++ + } + switch { + case modes == 0: + return fmt.Errorf("skills.registries[%d] requires exactly one selection mode (set one of skills, query, or all)", idx) + case modes > 1: + return fmt.Errorf("skills.registries[%d] sets multiple selection modes; set exactly one of skills, query, or all", idx) + } + if rc.Query != nil && rc.Query.Text == "" { + return fmt.Errorf("skills.registries[%d].query requires a non-empty text", idx) + } + return nil } -// RemoteAgentConfig configures a remote agent to register on startup. -type RemoteAgentConfig struct { - ID string `yaml:"id"` // Unique agent identifier - Name string `yaml:"name"` // Human-readable name - Description string `yaml:"description"` // Description of agent capabilities - Address string `yaml:"address"` // gRPC address (e.g., "localhost:50051") - Metadata map[string]string `yaml:"metadata,omitempty"` // Optional metadata +// SkillsRegistryConfig sources agentskills.io skills from the Gemini Skill +// Registry. When Enabled, exactly one selection mode should be set (Skills, +// Query, or All); if none is set, all skills are materialized. +type SkillsRegistryConfig struct { + Enabled bool `yaml:"enabled,omitempty"` + // Project owns the skills (projects/{Project}/locations/{Location}/skills). + // Empty falls back to the GOOGLE_CLOUD_PROJECT environment variable. + Project string `yaml:"project,omitempty"` + // Location is the registry region, e.g. "us-central1". Empty falls back to + // GOOGLE_CLOUD_LOCATION, then a built-in default. + Location string `yaml:"location,omitempty"` + + // --- selection (choose one) --- + + // Skills is an explicit allowlist of skills, each optionally pinned to a + // revision. Takes precedence over Query and All. + Skills []SkillRefConfig `yaml:"skills,omitempty"` + // Query is a semantic search selection; its top matches are materialized. + // TopK lives inside it because it only has meaning for a query. + Query *SkillsQueryConfig `yaml:"query,omitempty"` + // All materializes every skill in the project/location (used when neither + // Skills nor Query is set; can also be set explicitly). + All bool `yaml:"all,omitempty"` + + // TargetDir is the base directory skills are materialized into: each skill + // is written to //. Required when Enabled. + TargetDir string `yaml:"target_dir,omitempty"` } -type LocalAgentConfig struct { - ID string // Unique agent identifier - Name string // Human-readable name - Description string // Description of agent capabilities - Metadata map[string]string // Optional metadata - Agent agent.Agent // In-process agent instance (not in YAML) +// SkillsQueryConfig selects skills by semantic search. +type SkillsQueryConfig struct { + // Text is the semantic search string (required for query selection). + Text string `yaml:"text"` + // TopK bounds the number of matches (<=0 uses the server default). + TopK int `yaml:"top_k,omitempty"` +} + +// SkillRefConfig identifies a skill to materialize, optionally pinned. +type SkillRefConfig struct { + ID string `yaml:"id"` + Revision string `yaml:"revision,omitempty"` +} + +// SubstrateHarnessConfig registers a custom harness deployed on substrate +// from a user-provided container image. +type SubstrateHarnessConfig struct { + ID string `yaml:"id"` // Unique harness identifier + Namespace string `yaml:"namespace"` // ActorTemplate namespace (user-owned, not "ax") + Template string `yaml:"template"` // ActorTemplate name + Port int `yaml:"port,omitempty"` // HarnessService port + Default bool `yaml:"default,omitempty"` // Default harness or not +} + +// NewHarness builds the custom harness. Custom harnesses always run as substrate +// actors from the user's own ActorTemplate. +func (c SubstrateHarnessConfig) NewHarness(endpoint string) (harness.Harness, error) { + port := c.Port + if port == 0 { + port = substrateDefaultPort + } + return newSubstrateHarness(c.ID, endpoint, c.Namespace, c.Template, port) +} + +// newSubstrateHarness brings up a harness that is deployed as a substrate actor. +func newSubstrateHarness(harnessID, endpoint, namespace, template string, port int) (harness.Harness, error) { + sh, err := substrate.New(harnessID, endpoint, namespace, template, port) + if err != nil { + return nil, err + } + return sh, nil } // LoadFromFile loads configuration from a YAML file. @@ -85,36 +242,35 @@ func LoadFromFile(path string) (*Config, error) { if err != nil { return nil, fmt.Errorf("failed to read config file: %w", err) } + return LoadFromBytes(data) +} +// LoadFromBytes parses configuration from YAML bytes and applies defaults. +func LoadFromBytes(data []byte) (*Config, error) { var cfg Config if err := yaml.Unmarshal(data, &cfg); err != nil { - return nil, fmt.Errorf("failed to parse config file: %w", err) + return nil, fmt.Errorf("failed to parse config: %w", err) } - // Set defaults cfg.setDefaults() return &cfg, nil } +// DefaultConfig returns a configuration with default values set. +func DefaultConfig() *Config { + var cfg Config + cfg.setDefaults() + return &cfg +} + // setDefaults sets default values for optional fields. func (c *Config) setDefaults() { - // Server defaults if c.Server.Address == "" { c.Server.Address = ":8494" } - - // EventLog defaults - if c.EventLog.Dir == "" { - c.EventLog.Dir = "eventlog" - } - - // Controller defaults - if c.MaxSteps == 0 { - c.MaxSteps = 100 - } - if c.HealthCheckInterval == 0 { - c.HealthCheckInterval = 30 * time.Second + if c.EventLog.SQLiteConfig.Filename == "" { + c.EventLog.SQLiteConfig.Filename = "eventlog/log.sqlite" } } @@ -123,14 +279,54 @@ func (c *Config) Validate() error { if c.Server.Address == "" { return fmt.Errorf("server.address is required") } - if c.EventLog.Dir == "" { - return fmt.Errorf("eventlog.dir is required") + if c.EventLog.PostgresConfig.DSN == "" && c.EventLog.SQLiteConfig.Filename == "" { + return fmt.Errorf("eventlog requires either postgres.dsn or sqlite.filename") + } + + var defaultCount int + if c.Harnesses.Antigravity.Default { + defaultCount++ + } + if c.Harnesses.AntigravityInteractions.Default { + defaultCount++ } - if c.MaxSteps <= 0 { - return fmt.Errorf("max_steps must be positive") + + for _, sc := range c.Harnesses.Substrate { + if sc.ID == "" { + return fmt.Errorf("substrate harness id is required") + } + if sc.ID == AntigravityHarnessID || sc.ID == AntigravityInteractionsHarnessID { + return fmt.Errorf("substrate harness id %q is reserved for built-in harnesses", sc.ID) + } + if sc.Namespace == "" { + return fmt.Errorf("substrate harness %q: namespace is required", sc.ID) + } + if sc.Namespace == defaultNamespace { + return fmt.Errorf("substrate harness %q: namespace %q is reserved for built-in harnesses", sc.ID, defaultNamespace) + } + if sc.Template == "" { + return fmt.Errorf("substrate harness %q: template is required", sc.ID) + } + if sc.Default { + defaultCount++ + } } - if c.HealthCheckInterval <= 0 { - return fmt.Errorf("health_check_interval must be positive") + + if defaultCount > 1 { + return fmt.Errorf("multiple harnesses marked as default") } + + if err := c.Skills.Validate(); err != nil { + return err + } + return nil } + +func AXAssetsDir() (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("resolving home directory: %w", err) + } + return filepath.Join(home, ".ax"), nil +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go new file mode 100644 index 00000000..baebb037 --- /dev/null +++ b/internal/config/config_test.go @@ -0,0 +1,347 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package config + +import ( + "strings" + "testing" + + "gopkg.in/yaml.v3" +) + +func TestSubstrateNewHarness(t *testing.T) { + h, err := SubstrateHarnessConfig{ID: "c", Namespace: "team-ns", Template: "custom-template"}.NewHarness("api.ate-system.svc:443") + if err != nil { + t.Fatalf("NewHarness: %v", err) + } + if h == nil { + t.Fatal("expected non-nil harness") + } +} + +// validConfig returns a config that passes Validate, that tests can mutate. +func validConfig() *Config { + c := DefaultConfig() + c.Harnesses = HarnessesConfig{ + Antigravity: AntigravityHarnessConfig{Default: true}, + Substrate: []SubstrateHarnessConfig{ + {ID: "custom", Namespace: "team-ns", Template: "custom-template"}, + }, + } + return c +} + +func TestValidate_ValidConfig(t *testing.T) { + if err := validConfig().Validate(); err != nil { + t.Fatalf("Validate() = %v, want nil", err) + } +} + +func TestValidate_CustomIDRequired(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].ID = "" + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "substrate harness id") { + t.Fatalf("Validate() = %v, want substrate id error", err) + } +} + +func TestValidate_CustomIDReserved(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].ID = "antigravity" + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "reserved") { + t.Fatalf("Validate() = %v, want reserved id error", err) + } +} + +func TestValidate_CustomNamespaceRequired(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].Namespace = "" + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "namespace is required") { + t.Fatalf("Validate() = %v, want namespace-required error", err) + } +} + +func TestValidate_CustomNamespaceReserved(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].Namespace = defaultNamespace + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "reserved") { + t.Fatalf("Validate() = %v, want reserved-namespace error", err) + } +} + +func TestValidate_CustomTemplateRequired(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].Template = "" + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "template is required") { + t.Fatalf("Validate() = %v, want template-required error", err) + } +} + +func TestValidate_MultipleDefaults(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].Default = true + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "multiple harnesses marked as default") { + t.Fatalf("Validate() = %v, want multiple defaults error", err) + } +} + +func TestValidate_InteractionsIDReserved(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].ID = "antigravity-interactions" + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "reserved") { + t.Fatalf("Validate() = %v, want reserved id error", err) + } +} + +func TestValidate_InteractionsValid(t *testing.T) { + c := validConfig() + c.Harnesses.AntigravityInteractions = AntigravityInteractionsHarnessConfig{ + Agent: "projects/p/locations/global/agents/a", + } + if err := c.Validate(); err != nil { + t.Fatalf("Validate() = %v, want nil", err) + } +} + +func TestLoadFromFile_Version(t *testing.T) { + data := ` +version: "1.2.3" +server: + address: ":8080" +eventlog: + sqlite: + filename: "test.db" +` + var cfg Config + if err := yaml.Unmarshal([]byte(data), &cfg); err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + if cfg.Version != "1.2.3" { + t.Errorf("cfg.Version = %q, want %q", cfg.Version, "1.2.3") + } +} + +func TestLoadFromBytes(t *testing.T) { + cfg, err := LoadFromBytes([]byte(` +version: v1alpha +harnesses: + antigravity: + default: true +`)) + if err != nil { + t.Fatalf("LoadFromBytes: %v", err) + } + if cfg.Version != "v1alpha" { + t.Errorf("Version = %q, want %q", cfg.Version, "v1alpha") + } + if !cfg.Harnesses.Antigravity.Default { + t.Error("Harnesses.Antigravity.Default = false, want true") + } + // setDefaults must run (same as LoadFromFile). + if got, want := cfg.Server.Address, ":8494"; got != want { + t.Errorf("Server.Address = %q, want default %q", got, want) + } +} + +// TestLoadFromBytes_Invalid returns an error on malformed YAML. +func TestLoadFromBytes_Invalid(t *testing.T) { + if _, err := LoadFromBytes([]byte("harnesses: [unterminated")); err == nil { + t.Fatal("LoadFromBytes(invalid): got nil error, want error") + } +} + +func TestParse_SkillsByID(t *testing.T) { + data := ` +skills: + registries: + - enabled: true + project: my-proj + location: us-central1 + target_dir: /tmp/ax-skills + skills: + - id: emoji + - id: lowercase + revision: rev-3 +` + var cfg Config + if err := yaml.Unmarshal([]byte(data), &cfg); err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + if len(cfg.Skills.Registries) != 1 { + t.Fatalf("registries = %d, want 1", len(cfg.Skills.Registries)) + } + reg := cfg.Skills.Registries[0] + if !reg.Enabled || reg.Project != "my-proj" || reg.Location != "us-central1" || reg.TargetDir != "/tmp/ax-skills" { + t.Errorf("registry = %+v, want enabled my-proj/us-central1 /tmp/ax-skills", reg) + } + if len(reg.Skills) != 2 { + t.Fatalf("skills = %d, want 2", len(reg.Skills)) + } + if reg.Skills[0].ID != "emoji" || reg.Skills[0].Revision != "" { + t.Errorf("skills[0] = %+v, want {emoji }", reg.Skills[0]) + } + if reg.Skills[1].ID != "lowercase" || reg.Skills[1].Revision != "rev-3" { + t.Errorf("skills[1] = %+v, want {lowercase rev-3}", reg.Skills[1]) + } + if reg.Query != nil { + t.Errorf("Query = %+v, want nil in by-id mode", reg.Query) + } +} + +func TestParse_SkillsByQuery(t *testing.T) { + data := ` +skills: + registries: + - enabled: true + project: my-proj + target_dir: /tmp/ax-skills + query: + text: "find gcp skills" + top_k: 5 +` + var cfg Config + if err := yaml.Unmarshal([]byte(data), &cfg); err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + regs := cfg.Skills.Registries + if len(regs) != 1 || regs[0].Query == nil { + t.Fatalf("registries = %+v, want one with a query block", regs) + } + if regs[0].Query.Text != "find gcp skills" || regs[0].Query.TopK != 5 { + t.Errorf("Query = %+v, want {find gcp skills 5}", *regs[0].Query) + } +} + +func TestParse_MultipleRegistries(t *testing.T) { + data := ` +skills: + registries: + - enabled: true + project: org-proj + target_dir: /tmp/org + all: true + - enabled: true + project: team-proj + target_dir: /tmp/team + skills: + - id: teamskill +` + var cfg Config + if err := yaml.Unmarshal([]byte(data), &cfg); err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + regs := cfg.Skills.Registries + if len(regs) != 2 { + t.Fatalf("registries = %d, want 2", len(regs)) + } + if regs[0].Project != "org-proj" || regs[0].TargetDir != "/tmp/org" || !regs[0].All { + t.Errorf("registries[0] = %+v, want org-proj /tmp/org all", regs[0]) + } + if regs[1].Project != "team-proj" || regs[1].TargetDir != "/tmp/team" || len(regs[1].Skills) != 1 { + t.Errorf("registries[1] = %+v, want team-proj /tmp/team [teamskill]", regs[1]) + } + // Each registry has exactly one selection mode + target_dir, so validation passes. + if err := cfg.Skills.Validate(); err != nil { + t.Errorf("Skills.Validate() = %v, want nil", err) + } +} + +func TestValidate_SkillsSelectionOneof(t *testing.T) { + // withRegistry returns a valid config carrying a single top-level registry. + // It fills TargetDir (a required field) unless the caller already set one, so + // selection-mode assertions aren't masked by the target_dir check. + withRegistry := func(rc SkillsRegistryConfig) *Config { + if rc.Enabled && rc.TargetDir == "" { + rc.TargetDir = "/tmp/skills" + } + c := validConfig() + c.Skills.Registries = []SkillsRegistryConfig{rc} + return c + } + + t.Run("disabled skips validation", func(t *testing.T) { + if err := withRegistry(SkillsRegistryConfig{Enabled: false}).Validate(); err != nil { + t.Fatalf("Validate() = %v, want nil", err) + } + }) + + t.Run("enabled without target_dir is an error", func(t *testing.T) { + c := validConfig() + c.Skills.Registries = []SkillsRegistryConfig{ + {Enabled: true, Project: "p", All: true}, // valid mode, but no target_dir + } + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "target_dir") { + t.Fatalf("Validate() = %v, want target_dir error", err) + } + }) + + t.Run("zero selection modes is an error", func(t *testing.T) { + err := withRegistry(SkillsRegistryConfig{Enabled: true, Project: "p"}).Validate() + if err == nil || !strings.Contains(err.Error(), "exactly one selection mode") { + t.Fatalf("Validate() = %v, want exactly-one error", err) + } + }) + + t.Run("multiple selection modes is an error", func(t *testing.T) { + err := withRegistry(SkillsRegistryConfig{ + Enabled: true, Project: "p", + Skills: []SkillRefConfig{{ID: "emoji"}}, + All: true, + }).Validate() + if err == nil || !strings.Contains(err.Error(), "multiple selection modes") { + t.Fatalf("Validate() = %v, want multiple-modes error", err) + } + }) + + t.Run("exactly one is valid", func(t *testing.T) { + err := withRegistry(SkillsRegistryConfig{ + Enabled: true, Project: "p", + Skills: []SkillRefConfig{{ID: "emoji"}}, + }).Validate() + if err != nil { + t.Fatalf("Validate() = %v, want nil", err) + } + }) + + t.Run("query with empty text is an error", func(t *testing.T) { + err := withRegistry(SkillsRegistryConfig{ + Enabled: true, Project: "p", + Query: &SkillsQueryConfig{Text: ""}, + }).Validate() + if err == nil || !strings.Contains(err.Error(), "non-empty text") { + t.Fatalf("Validate() = %v, want empty-text error", err) + } + }) + + t.Run("second registry invalid is caught", func(t *testing.T) { + c := validConfig() + c.Skills.Registries = []SkillsRegistryConfig{ + {Enabled: true, Project: "p", TargetDir: "/tmp/a", All: true}, + {Enabled: true, Project: "p", TargetDir: "/tmp/b"}, // no selection mode + } + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "registries[1]") { + t.Fatalf("Validate() = %v, want error citing registries[1]", err) + } + }) +} diff --git a/internal/controller/controller.go b/internal/controller/controller.go index 2034ac4c..a3775284 100644 --- a/internal/controller/controller.go +++ b/internal/controller/controller.go @@ -13,152 +13,180 @@ // limitations under the License. // Package controller implements the single-writer orchestrator that coordinates -// agentic loops, manages sessions, and communicates with local and remote agents. +// agentic loops, manages executions, and communicates with local and remote agents. package controller import ( "context" "fmt" - "sync" - "time" + "log/slog" - "github.com/google/gar/agent" - "github.com/google/gar/internal/eventlog" - "github.com/google/gar/proto" + "github.com/google/ax/internal/controller/eventlog" + "github.com/google/ax/proto" + "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/types/known/structpb" ) +type ExecHandler func(resp *proto.ExecResponse) error + // Controller is the main controller that coordinates all components. // It acts as a single-writer system for managing agentic loops. type Controller struct { - inFlightSessionsMu sync.Mutex - inFlightSessions map[string]struct{} - - sessionManager *SessionManager - registry *Registry - loopExecutor *LoopExecutor + registry *Registry + eventLog eventlog.EventLog } // Config configures the controller. type Config struct { - EventLogFactory eventlog.EventLogFactory - PlanFunc PlanFunc + Registry *Registry + EventLogBuilder eventlog.EventLogBuilder +} - // TODO(jbd): Add CompactionFunc. +// New creates a new controller instance. +func New(ctx context.Context, cfg Config) (*Controller, error) { + if cfg.Registry == nil { + return nil, fmt.Errorf("registry is required") + } + if cfg.EventLogBuilder == nil { + return nil, fmt.Errorf("event log builder is required") + } + eventLog, err := cfg.EventLogBuilder() + if err != nil { + return nil, fmt.Errorf("failed to create event log: %w", err) + } - HealthCheckInterval time.Duration - MaxSteps int + return &Controller{ + registry: cfg.Registry, + eventLog: eventLog, + }, nil } -// New creates a new controller instance. -func New(ctx context.Context, config Config) (*Controller, error) { - if config.MaxSteps == 0 { - config.MaxSteps = 100 - } - if config.HealthCheckInterval == 0 { - config.HealthCheckInterval = 30 * time.Second - } - if config.EventLogFactory == nil { - config.EventLogFactory = func(sessionID string) (eventlog.EventLog, error) { - return eventlog.NewFileEventLog(eventlog.FileConfig{ - SessionID: sessionID, - Dir: "eventlog", - }) - } +// Exec executes a new agentic loop execution or resumes an existing one. +// If id is empty, a UUID will be generated. +// If the execution already exists, it will be resumed with optional new inputs. +func (d *Controller) Exec(ctx context.Context, req *proto.ExecRequest, handler ExecHandler) error { + if req.ConversationId == "" { + return fmt.Errorf("conversation_id is required") } - // Initialize session manager with file-based event logs - sessionManager := NewSessionManager(config.EventLogFactory) + // TODO(jbd): Resume an incomplete execution if there exists one. + // TODO(jbd): Enable bringing a remote harness that implements HarnessService. + // TODO(anj): We need to consolidate agents and harness registration. + // Adding harness registration support temporarily. + l := newLogger(d.eventLog, req.ConversationId, req.HarnessId) + state, storedHarnessID, err := l.ResumptionState(ctx) + if err != nil { + return fmt.Errorf("failed to check resumption state: %w", err) + } - // Initialize agent registry - registry := NewRegistry(config.HealthCheckInterval) + // On resume, use the conversation's recorded harness. Using a different harness + // for the same conversation is not allowed. + if req.HarnessId != "" && storedHarnessID != "" && req.HarnessId != storedHarnessID { + return fmt.Errorf("resumption not allowed: harness ID changed from %s to %s", storedHarnessID, req.HarnessId) + } + harnessID := req.HarnessId + // Use the conversations's stored harness if no harness is specified. + if harnessID == "" { + harnessID = storedHarnessID + } + // For new conversations, use the default harness if no harness is specified. + if harnessID == "" { + harnessID = d.registry.defaultHarness + } + l.harnessID = harnessID - // Initialize loop executor - loopExecutor, err := NewLoopExecutor(ctx, LoopConfig{ - Registry: registry, - SessionManager: sessionManager, - MaxSteps: config.MaxSteps, - PlanFunc: config.PlanFunc, - }) + h, err := d.registry.Harness(harnessID) if err != nil { - return nil, fmt.Errorf("failed to create loop executor: %w", err) + return fmt.Errorf("failed to get harness %q: %w", harnessID, err) } - return &Controller{ - inFlightSessions: make(map[string]struct{}), - sessionManager: sessionManager, - registry: registry, - loopExecutor: loopExecutor, - }, nil -} + hhandler := &harnessHandler{ + logger: l, + execHandler: handler, + } -// TriggerSession triggers a new agentic loop session or resumes an existing one. -// If sessionID is empty, a UUID will be generated. -// If the session already exists, it will be resumed with optional new inputs. -// If checkpointID is provided, resumes from that specific checkpoint instead of the latest. -func (d *Controller) TriggerSession(ctx context.Context, sessionID string, inputs []*proto.Content, handler agent.OutputHandler) error { - if sessionID == "" { - return fmt.Errorf("session_id is required") - } - - d.inFlightSessionsMu.Lock() - _, ok := d.inFlightSessions[sessionID] - d.inFlightSessionsMu.Unlock() - - if ok { - return fmt.Errorf("session is already in flight") - } - defer func() { - d.inFlightSessionsMu.Lock() - delete(d.inFlightSessions, sessionID) - d.inFlightSessionsMu.Unlock() - }() - - // Check if session already exists - sess, err := d.sessionManager.LoadSession(ctx, sessionID) - if err == nil && sess == nil { - // Session doesn't exist - create new session - // Checkpoint ID is ignored for new sessions - sess, err = d.sessionManager.NewSession(sessionID) + if state == proto.State_STATE_PENDING { + // If the state is pending, first try to resume the + // pending execution. If the state is COMPLETED or FAILED, start + // a new execution. + exec, err := h.Start(ctx, req.ConversationId, req.HarnessConfig) if err != nil { - return fmt.Errorf("failed to create session: %w", err) + return fmt.Errorf("failed to start harness session: %w", err) + } + defer exec.Close(ctx) + + if err := exec.Run(ctx, hhandler); err != nil { + return fmt.Errorf("harness execution failed: %w", err) } } - if sess.State() == proto.State_STATE_FAILED { - return fmt.Errorf("session has failed and cannot continue") + if len(req.Inputs) == 0 { + // No more inputs, just return. + return nil } - // Write input content to session - for _, content := range inputs { - if _, err := sess.WriteContentIn(ctx, content); err != nil { - return fmt.Errorf("failed to write input content: %w", err) - } + exec, err := h.Start(ctx, req.ConversationId, req.HarnessConfig) + if err != nil { + return fmt.Errorf("failed to start harness session: %w", err) } + defer exec.Close(ctx) - if err := d.loopExecutor.Execute(ctx, sess, handler); err != nil { - if err := sess.SetState(ctx, proto.State_STATE_FAILED); err != nil { - return fmt.Errorf("failed to set session state: %w", err) - } - return fmt.Errorf("loop execution failed: %w", err) + if err := exec.Queue(ctx, req.Inputs...); err != nil { + return fmt.Errorf("failed to queue inputs: %w", err) + } + // Log inputs before running harness + if _, err := l.LogInputs(ctx, req.Inputs, req.HarnessConfig); err != nil { + return fmt.Errorf("failed to log inputs: %w", err) + } + if err := exec.Run(ctx, hhandler); err != nil { + return fmt.Errorf("harness execution failed: %w", err) } return nil } -func (d *Controller) TriggerForkedSession(ctx context.Context, sessionID string, checkpointID string, inputs []*proto.Content, handler agent.OutputHandler) error { - if sessionID == "" { - return fmt.Errorf("session_id is required") +type harnessHandler struct { + logger *logger + execHandler ExecHandler +} + +func (a *harnessHandler) OnMessage(ctx context.Context, execID string, msg *proto.Message) error { + // Log every response received from the harness + // TODO(anj): The harness should send the full input sent to get this particular response. + seq, err := a.logger.LogOutputs(ctx, []*proto.Message{msg}, proto.State_STATE_PENDING) + if err != nil { + slog.WarnContext(ctx, "Failed to log streamed message to event log", + slog.String("conversation_id", a.logger.conversationID), + slog.Any("error", err), + ) + } + + if a.execHandler == nil { + return nil } - if checkpointID == "" { - return fmt.Errorf("checkpoint_id is required") + return a.execHandler(&proto.ExecResponse{ + Outputs: []*proto.Message{msg}, + Seq: seq, + }) +} + +func (a *harnessHandler) OnComplete(ctx context.Context, execID string) error { + // Mark the execution turn as completed in the conversation log + if _, err := a.logger.LogOutputs(ctx, nil, proto.State_STATE_COMPLETED); err != nil { + slog.WarnContext(ctx, "Failed to log completion event to event log", + slog.String("conversation_id", a.logger.conversationID), + slog.Any("error", err), + ) } - // TODO(jbd): Fork a new session by copying all content to - // the provided checkpoint ID, and execute the loop executor. - panic("not yet implemented") + return nil } -// LoadSession loads a session from event log. -func (d *Controller) LoadSession(ctx context.Context, sessionID string) (*Session, error) { - return d.sessionManager.LoadSession(ctx, sessionID) +// Delete deletes all events for a specific conversation ID. +func (d *Controller) Delete(ctx context.Context, conversationID string) error { + if conversationID == "" { + return fmt.Errorf("conversation_id is required") + } + + return d.eventLog.DeleteAll(ctx, conversationID) } // Registry returns the agent registry. @@ -168,9 +196,85 @@ func (d *Controller) Registry() *Registry { // Close gracefully shuts down the controller. func (d *Controller) Close() error { + if err := d.eventLog.Close(); err != nil { + return fmt.Errorf("failed to close event log: %w", err) + } if err := d.registry.Close(); err != nil { return fmt.Errorf("failed to close registry: %w", err) } - d.sessionManager.CloseAll() return nil } + +func newLogger( + el eventlog.EventLog, + conversationID string, + harnessID string) *logger { + return &logger{ + el: el, + conversationID: conversationID, + harnessID: harnessID, + } +} + +type logger struct { + conversationID string + execID string + el eventlog.EventLog + harnessID string +} + +// ResumptionState returns the conversation's current state and the harness it used. +func (l *logger) ResumptionState(ctx context.Context) (proto.State, string, error) { + events, err := l.el.Events(ctx, l.conversationID) + if err != nil { + return proto.State_STATE_UNSPECIFIED, "", err + } + + var state proto.State + var harnessID string + for _, ev := range events { + if harnessID == "" && ev.HarnessId != "" { + harnessID = ev.HarnessId + } + if l.execID == "" || ev.ExecId == l.execID { + if ev.State != proto.State_STATE_UNSPECIFIED { + state = ev.State + } + } + } + return state, harnessID, nil +} + +func (l *logger) LogInputs(ctx context.Context, inputs []*proto.Message, harnessConfig []byte) (int32, error) { + // Parse the harness config into a human-readable struct for logging. + var cfg *structpb.Struct + if len(harnessConfig) > 0 { + cfg = &structpb.Struct{} + if err := protojson.Unmarshal(harnessConfig, cfg); err != nil { + slog.WarnContext(ctx, "Failed to parse harness config for logging", + slog.String("conversation_id", l.conversationID), + slog.Any("error", err), + ) + cfg = nil + } + } + ev := &proto.ConversationEvent{ + ConversationId: l.conversationID, + ExecId: l.execID, + HarnessId: l.harnessID, + HarnessConfig: cfg, + Messages: inputs, + State: proto.State_STATE_PENDING, + } + return l.el.Append(ctx, ev) +} + +func (l *logger) LogOutputs(ctx context.Context, outputs []*proto.Message, state proto.State) (int32, error) { + ev := &proto.ConversationEvent{ + ConversationId: l.conversationID, + ExecId: l.execID, + Messages: outputs, + State: state, + } + return l.el.Append(ctx, ev) +} diff --git a/internal/controller/controller_test.go b/internal/controller/controller_test.go new file mode 100644 index 00000000..24ab7559 --- /dev/null +++ b/internal/controller/controller_test.go @@ -0,0 +1,665 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controller + +import ( + "context" + "fmt" + "strings" + "testing" + + "github.com/google/ax/internal/controller/eventlog" + "github.com/google/ax/internal/controller/eventlog/eventlogtest" + "github.com/google/ax/internal/harness" + "github.com/google/ax/proto" +) + +type fakeHarness struct{} + +func (f *fakeHarness) Start(ctx context.Context, conversationID string, harnessConfig []byte) (harness.Execution, error) { + return &fakeExecution{id: "fake-exec-id"}, nil +} + +type fakeExecution struct { + id string + queued []*proto.Message +} + +func (f *fakeExecution) ID() string { + return f.id +} + +func (f *fakeExecution) Queue(ctx context.Context, msg ...*proto.Message) error { + f.queued = append(f.queued, msg...) + return nil +} + +func (f *fakeExecution) Run(ctx context.Context, handler harness.Handler) error { + msg := &proto.Message{ + Role: "assistant", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{ + Text: "Hello world", + }, + }, + }, + } + if err := handler.OnMessage(ctx, f.id, msg); err != nil { + return err + } + return handler.OnComplete(ctx, f.id) +} + +func (f *fakeExecution) Close(ctx context.Context) error { + return nil +} + +func TestController2_ExecHelloWorld(t *testing.T) { + ctx := context.Background() + cid := "test-conversation-id" + + log := &eventlogtest.MemoryEventLog{} + reg := NewRegistry() + if err := reg.RegisterHarness("default-harness", &fakeHarness{}); err != nil { + t.Fatal(err) + } + if err := reg.SetDefaultHarness("default-harness"); err != nil { + t.Fatal(err) + } + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { + return log, nil + }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + var outputs []*proto.Message + handler := ExecHandler(func(resp *proto.ExecResponse) error { + outputs = append(outputs, resp.Outputs...) + return nil + }) + + inputs := []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{Text: "Trigger prompt"}, + }, + }, + }, + } + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + Inputs: inputs, + }, handler) + if err != nil { + t.Fatalf("Controller2.Exec failed: %v", err) + } + + if len(outputs) != 1 { + t.Fatalf("expected exactly 1 output message, got %d", len(outputs)) + } + + gotText := outputs[0].GetContent().GetText().GetText() + if gotText != "Hello world" { + t.Errorf("expected 'Hello world' output text response, got %q", gotText) + } + + // Verify that events were logged correctly in Conversation Log + events, err := log.Events(ctx, cid) + if err != nil { + t.Fatalf("failed to retrieve logged events: %v", err) + } + + if len(events) != 3 { + t.Fatalf("expected 3 logged events, got %d", len(events)) + } + + // 1. First event should be inputs + if len(events[0].Messages) != 1 { + t.Errorf("expected 1 message in first event, got %d", len(events[0].Messages)) + } else { + gotInputText := events[0].Messages[0].GetContent().GetText().GetText() + if gotInputText != "Trigger prompt" { + t.Errorf("expected 'Trigger prompt' in logged input, got %q", gotInputText) + } + } + if events[0].State != proto.State_STATE_PENDING { + t.Errorf("expected first event state to be PENDING, got %v", events[0].State) + } + + // 2. Second event should be output + if len(events[1].Messages) != 1 { + t.Errorf("expected 1 message in second event, got %d", len(events[1].Messages)) + } else { + gotOutputText := events[1].Messages[0].GetContent().GetText().GetText() + if gotOutputText != "Hello world" { + t.Errorf("expected 'Hello world' in logged output, got %q", gotOutputText) + } + } + if events[1].State != proto.State_STATE_PENDING { + t.Errorf("expected second event state to be PENDING, got %v", events[1].State) + } + + // 3. Third event should be completion + if len(events[2].Messages) != 0 { + t.Errorf("expected 0 messages in third event, got %d", len(events[2].Messages)) + } + if events[2].State != proto.State_STATE_COMPLETED { + t.Errorf("expected third event state to be COMPLETED, got %v", events[2].State) + } + +} + +func TestController2_ExecWithAgentID(t *testing.T) { + ctx := context.Background() + cid := "test-conversation-id" + + log := &eventlogtest.MemoryEventLog{} + reg := NewRegistry() + if err := reg.RegisterHarness("my-agent", &fakeHarness{}); err != nil { + t.Fatal(err) + } + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { + return log, nil + }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + var outputs []*proto.Message + handler := ExecHandler(func(resp *proto.ExecResponse) error { + outputs = append(outputs, resp.Outputs...) + return nil + }) + + inputs := []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{Text: "Trigger prompt"}, + }, + }, + }, + } + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + HarnessId: "my-agent", + Inputs: inputs, + }, handler) + if err != nil { + t.Fatalf("Controller2.Exec failed: %v", err) + } + + if len(outputs) != 1 { + t.Fatalf("expected exactly 1 output message, got %d", len(outputs)) + } + + gotText := outputs[0].GetContent().GetText().GetText() + if gotText != "Hello world" { + t.Errorf("expected 'Hello world' output text response, got %q", gotText) + } +} + +func TestController2_ExecHarnessNotFound(t *testing.T) { + ctx := context.Background() + cid := "test-conversation-id" + + log := &eventlogtest.MemoryEventLog{} + reg := NewRegistry() // Empty registry, will force error for any requested agent + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { + return log, nil + }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + handler := ExecHandler(func(resp *proto.ExecResponse) error { + return nil + }) + + inputs := []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{Text: "Trigger prompt"}, + }, + }, + }, + } + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + Inputs: inputs, + HarnessId: "antigravity", + }, handler) + if err == nil { + t.Fatal("expected error requesting unregistered agent, got nil") + } +} + +type testHarness struct { + startCalls int + startFunc func(ctx context.Context, conversationID string) (harness.Execution, error) +} + +func (c *testHarness) Start(ctx context.Context, conversationID string, harnessConfig []byte) (harness.Execution, error) { + c.startCalls++ + return c.startFunc(ctx, conversationID) +} + +type testExecution struct { + id string + queueCalls int + runCalls int + closeCalls int + queued []*proto.Message + runFunc func(ctx context.Context, execID string, handler harness.Handler) error +} + +func (c *testExecution) ID() string { + return c.id +} + +func (c *testExecution) Queue(ctx context.Context, msg ...*proto.Message) error { + c.queueCalls++ + c.queued = append(c.queued, msg...) + return nil +} + +func (c *testExecution) Run(ctx context.Context, handler harness.Handler) error { + c.runCalls++ + if c.runFunc != nil { + return c.runFunc(ctx, c.id, handler) + } + return nil +} + +func (c *testExecution) Close(ctx context.Context) error { + c.closeCalls++ + return nil +} + +func TestController2_ExecResumptionFlow(t *testing.T) { + // Subtest 1: New Execution with Inputs + t.Run("NewExecutionWithInputs", func(t *testing.T) { + ctx := context.Background() + cid := "new-conv" + + log := &eventlogtest.MemoryEventLog{} + reg := NewRegistry() + + var exec *testExecution + h := &testHarness{ + startFunc: func(ctx context.Context, conversationID string) (harness.Execution, error) { + exec = &testExecution{ + id: "exec-new", + runFunc: func(ctx context.Context, execID string, handler harness.Handler) error { + return handler.OnComplete(ctx, execID) + }, + } + return exec, nil + }, + } + if err := reg.RegisterHarness("test-agent", h); err != nil { + t.Fatal(err) + } + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { return log, nil }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + HarnessId: "test-agent", + Inputs: []*proto.Message{ + {Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "Hello"}}}}, + }, + }, func(resp *proto.ExecResponse) error { return nil }) + if err != nil { + t.Fatal(err) + } + + if h.startCalls != 1 { + t.Errorf("expected 1 Start call, got %d", h.startCalls) + } + if exec.queueCalls != 1 { + t.Errorf("expected 1 Queue call, got %d", exec.queueCalls) + } + if exec.runCalls != 1 { + t.Errorf("expected 1 Run call, got %d", exec.runCalls) + } + }) + + // Subtest 2: Pending Execution with NO New Inputs + t.Run("PendingExecutionWithoutNewInputs", func(t *testing.T) { + ctx := context.Background() + cid := "pending-no-inputs" + + log := &eventlogtest.MemoryEventLog{} + // Seed the event log with a pending event + _, err := log.Append(ctx, &proto.ConversationEvent{ + ConversationId: cid, + HarnessId: "test-agent", + State: proto.State_STATE_PENDING, + Messages: []*proto.Message{ + {Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "Initial"}}}}, + }, + }) + if err != nil { + t.Fatal(err) + } + + reg := NewRegistry() + + var exec *testExecution + h := &testHarness{ + startFunc: func(ctx context.Context, conversationID string) (harness.Execution, error) { + exec = &testExecution{ + id: "exec-pending", + runFunc: func(ctx context.Context, execID string, handler harness.Handler) error { + return handler.OnComplete(ctx, execID) + }, + } + return exec, nil + }, + } + if err := reg.RegisterHarness("test-agent", h); err != nil { + t.Fatal(err) + } + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { return log, nil }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + HarnessId: "test-agent", + Inputs: nil, // NO new inputs + }, func(resp *proto.ExecResponse) error { return nil }) + if err != nil { + t.Fatal(err) + } + + if h.startCalls != 1 { + t.Errorf("expected 1 Start call, got %d", h.startCalls) + } + if exec.queueCalls != 0 { + t.Errorf("expected 0 Queue calls, got %d", exec.queueCalls) + } + if exec.runCalls != 1 { + t.Errorf("expected 1 Run call, got %d", exec.runCalls) + } + }) + + // Subtest 3: Pending Execution WITH New Inputs + t.Run("PendingExecutionWithNewInputs", func(t *testing.T) { + ctx := context.Background() + cid := "pending-with-inputs" + + log := &eventlogtest.MemoryEventLog{} + // Seed the event log with a pending event + _, err := log.Append(ctx, &proto.ConversationEvent{ + ConversationId: cid, + HarnessId: "test-agent", + State: proto.State_STATE_PENDING, + Messages: []*proto.Message{ + {Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "Initial"}}}}, + }, + }) + if err != nil { + t.Fatal(err) + } + + reg := NewRegistry() + + var execs []*testExecution + h := &testHarness{ + startFunc: func(ctx context.Context, conversationID string) (harness.Execution, error) { + exec := &testExecution{ + id: fmt.Sprintf("exec-%d", len(execs)+1), + runFunc: func(ctx context.Context, execID string, handler harness.Handler) error { + return handler.OnComplete(ctx, execID) + }, + } + execs = append(execs, exec) + return exec, nil + }, + } + if err := reg.RegisterHarness("test-agent", h); err != nil { + t.Fatal(err) + } + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { return log, nil }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + HarnessId: "test-agent", + Inputs: []*proto.Message{ + {Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "New input"}}}}, + }, + }, func(resp *proto.ExecResponse) error { return nil }) + if err != nil { + t.Fatal(err) + } + + // It should start the harness twice: once for resumption, once for new inputs. + if h.startCalls != 2 { + t.Errorf("expected 2 Start calls, got %d", h.startCalls) + } + if len(execs) != 2 { + t.Fatalf("expected 2 execution sessions, got %d", len(execs)) + } + + // The first session (resumption) should NOT have queued inputs and should run. + if execs[0].queueCalls != 0 { + t.Errorf("expected first session to have 0 Queue calls, got %d", execs[0].queueCalls) + } + if execs[0].runCalls != 1 { + t.Errorf("expected first session to have 1 Run call, got %d", execs[0].runCalls) + } + + // The second session (new inputs) should have queued inputs and run. + if execs[1].queueCalls != 1 { + t.Errorf("expected second session to have 1 Queue call, got %d", execs[1].queueCalls) + } + if execs[1].runCalls != 1 { + t.Errorf("expected second session to have 1 Run call, got %d", execs[1].runCalls) + } + }) +} + +// A conversation started on one harness must resume on that harness when the +// harness id is empty, not fall back to the registry default. +func TestExec_ResumeEmptyHarnessUsesStored(t *testing.T) { + ctx := context.Background() + cid := "resume-empty" + log := &eventlogtest.MemoryEventLog{} + reg := NewRegistry() + + def := &testHarness{startFunc: func(context.Context, string) (harness.Execution, error) { + return &testExecution{}, nil + }} + stored := &testHarness{startFunc: func(context.Context, string) (harness.Execution, error) { + return &testExecution{}, nil + }} + if err := reg.RegisterHarness("harness-a", def); err != nil { + t.Fatal(err) + } + if err := reg.RegisterHarness("harness-b", stored); err != nil { + t.Fatal(err) + } + if err := reg.SetDefaultHarness("harness-a"); err != nil { + t.Fatal(err) + } + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { return log, nil }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + noop := ExecHandler(func(*proto.ExecResponse) error { return nil }) + + // Turn 1: explicitly run the NON-default harness. + if err := c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + HarnessId: "harness-b", + Inputs: []*proto.Message{{Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "hi"}}}}}, + }, noop); err != nil { + t.Fatalf("turn 1: %v", err) + } + + // Turn 2: resume WITHOUT a harness id. Must reuse harness-b, not the default. + before := stored.startCalls + if err := c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + Inputs: []*proto.Message{{Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "more"}}}}}, + }, noop); err != nil { + t.Fatalf("turn 2 (resume, empty harness): %v", err) + } + if stored.startCalls <= before { + t.Errorf("stored harness not used on resume (startCalls stayed %d)", stored.startCalls) + } + if def.startCalls != 0 { + t.Errorf("default harness used on resume (startCalls=%d), want the stored harness", def.startCalls) + } +} + +// Verifies that an explicit harness change during resume is rejected. +func TestExec_ResumeExplicitDifferentHarnessRejected(t *testing.T) { + ctx := context.Background() + cid := "resume-mismatch" + log := &eventlogtest.MemoryEventLog{} + reg := NewRegistry() + if err := reg.RegisterHarness("harness-a", &testHarness{startFunc: func(context.Context, string) (harness.Execution, error) { + return &testExecution{}, nil + }}); err != nil { + t.Fatal(err) + } + if err := reg.RegisterHarness("harness-b", &testHarness{startFunc: func(context.Context, string) (harness.Execution, error) { + return &testExecution{}, nil + }}); err != nil { + t.Fatal(err) + } + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { return log, nil }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + noop := ExecHandler(func(*proto.ExecResponse) error { return nil }) + + if err := c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + HarnessId: "harness-a", + Inputs: []*proto.Message{{Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "hi"}}}}}, + }, noop); err != nil { + t.Fatalf("turn 1: %v", err) + } + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + HarnessId: "harness-b", + Inputs: []*proto.Message{{Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "more"}}}}}, + }, noop) + if err == nil || !strings.Contains(err.Error(), "harness ID changed from harness-a to harness-b") { + t.Fatalf("resume with a different harness: got %v, want error 'harness ID changed from harness-a to harness-b'", err) + } +} + +// Verifies that a new conversation started without a harness id records the +// default harness's canonical id, not "". +func TestExec_NewConversationLogsCanonicalDefault(t *testing.T) { + ctx := context.Background() + cid := "canonical-log" + log := &eventlogtest.MemoryEventLog{} + reg := NewRegistry() + if err := reg.RegisterHarness("harness-a", &testHarness{startFunc: func(context.Context, string) (harness.Execution, error) { + return &testExecution{}, nil + }}); err != nil { + t.Fatal(err) + } + if err := reg.SetDefaultHarness("harness-a"); err != nil { + t.Fatal(err) + } + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { return log, nil }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + if err := c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + Inputs: []*proto.Message{{Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "hi"}}}}}, + }, ExecHandler(func(*proto.ExecResponse) error { return nil })); err != nil { + t.Fatalf("exec: %v", err) + } + _, stored, err := newLogger(log, cid, "").ResumptionState(ctx) + if err != nil { + t.Fatalf("ResumptionState: %v", err) + } + if stored != "harness-a" { + t.Errorf("logged harness id = %q, want canonical %q (not empty)", stored, "harness-a") + } +} diff --git a/internal/controller/eventlog/eventlog.go b/internal/controller/eventlog/eventlog.go new file mode 100644 index 00000000..4e5e1883 --- /dev/null +++ b/internal/controller/eventlog/eventlog.go @@ -0,0 +1,46 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package eventlog contains an AX event log. +package eventlog + +import ( + "context" + + "github.com/google/ax/proto" + "google.golang.org/protobuf/encoding/protojson" +) + +// EventLogBuilder builds new event logs. +type EventLogBuilder func() (EventLog, error) + +// EventLog is the persistent, append-only record of all actions taken in an +// exec. Every entry is an atomic step: replaying the log in order brings +// the executor back to a consistent state from which execution can resume. +type EventLog interface { + // Append adds a conversation event to the end of the log. + Append(ctx context.Context, event *proto.ConversationEvent) (int32, error) + + // Events returns all events for the conversation. + Events(ctx context.Context, conversationID string) ([]*proto.ConversationEvent, error) + + // DeleteAll deletes all events for a specific conversation ID. + DeleteAll(ctx context.Context, conversationID string) error + + // Close releases the underlying resources and closes the log. + Close() error +} + +var marshalOpts = protojson.MarshalOptions{UseProtoNames: true} +var unmarshalOpts = protojson.UnmarshalOptions{DiscardUnknown: true} diff --git a/internal/controller/eventlog/eventlogtest/eventlog.go b/internal/controller/eventlog/eventlogtest/eventlog.go new file mode 100644 index 00000000..f24abe5b --- /dev/null +++ b/internal/controller/eventlog/eventlogtest/eventlog.go @@ -0,0 +1,95 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package eventlogtest + +import ( + "context" + "sync" + + "github.com/google/ax/proto" +) + +// MemoryEventLog is an in-memory EventLog useful for testing and short-lived +// executions. It does not survive process restarts. +type MemoryEventLog struct { + mu sync.Mutex + AllEvents []*proto.ConversationEvent +} + +func (m *MemoryEventLog) Append(_ context.Context, event *proto.ConversationEvent) (int32, error) { + m.mu.Lock() + defer m.mu.Unlock() + + seq := event.Seq + if seq == 0 { + maxSeq := int32(0) + for _, ev := range m.AllEvents { + if ev.ConversationId == event.ConversationId && ev.Seq > maxSeq { + maxSeq = ev.Seq + } + } + seq = maxSeq + 1 + event.Seq = seq + } + m.AllEvents = append(m.AllEvents, event) + return seq, nil +} + +func (m *MemoryEventLog) Events(_ context.Context, conversationID string) ([]*proto.ConversationEvent, error) { + m.mu.Lock() + defer m.mu.Unlock() + + out := make([]*proto.ConversationEvent, 0) + for _, ev := range m.AllEvents { + if ev.ConversationId == conversationID { + out = append(out, ev) + } + } + return out, nil +} + +// Drop removes every event for which drop returns true. +// It is provided for testing and crash-simulation purposes. +func (m *MemoryEventLog) Drop(drop func(*proto.ConversationEvent) bool) { + m.mu.Lock() + defer m.mu.Unlock() + + kept := m.AllEvents[:0] + for _, ev := range m.AllEvents { + if !drop(ev) { + kept = append(kept, ev) + } + } + m.AllEvents = kept +} + +func (m *MemoryEventLog) DeleteAll(_ context.Context, conversationID string) error { + m.mu.Lock() + defer m.mu.Unlock() + + var keptEvents []*proto.ConversationEvent + for _, ev := range m.AllEvents { + if ev.ConversationId != conversationID { + keptEvents = append(keptEvents, ev) + } + } + m.AllEvents = keptEvents + + return nil +} + +func (m *MemoryEventLog) Close() error { + return nil +} diff --git a/internal/controller/eventlog/postgres.go b/internal/controller/eventlog/postgres.go new file mode 100644 index 00000000..c99c103d --- /dev/null +++ b/internal/controller/eventlog/postgres.go @@ -0,0 +1,57 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package eventlog + +import ( + "context" + "database/sql" + "fmt" + "time" + + _ "github.com/jackc/pgx/v5/stdlib" +) + +// OpenPostgresEventLog connects to the PostgreSQL database described by dsn and +// initializes the event log schema. Caller is responsible to ensure it is safe +// for concurrent use. +func OpenPostgresEventLog(dsn string) (EventLog, error) { + db, err := sql.Open("pgx", dsn) + if err != nil { + return nil, fmt.Errorf("postgres_eventlog: open: %w", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + if err := db.PingContext(ctx); err != nil { + db.Close() + return nil, fmt.Errorf("postgres_eventlog: ping: %w", err) + } + + // Create tables if they don't exist. + if _, err := db.ExecContext(ctx, ` + CREATE TABLE IF NOT EXISTS conversation_log ( + conversation_id TEXT NOT NULL, + seq INTEGER NOT NULL, + payload TEXT NOT NULL, + PRIMARY KEY (conversation_id, seq) + )`); err != nil { + db.Close() + return nil, fmt.Errorf("postgres_eventlog: create conversation_log table: %w", err) + } + + + return &sqlEventLog{db: db}, nil +} diff --git a/internal/controller/eventlog/sql.go b/internal/controller/eventlog/sql.go new file mode 100644 index 00000000..522b6717 --- /dev/null +++ b/internal/controller/eventlog/sql.go @@ -0,0 +1,133 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package eventlog + +import ( + "context" + "database/sql" + "fmt" + + "github.com/google/ax/proto" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" +) + +// sqlEventLog is a database backed EventLog shared by the SQLite and +// PostgreSQL implementations. +type sqlEventLog struct { + db *sql.DB +} + +// Append serializes the event to JSON and inserts it into the database. +func (l *sqlEventLog) Append(ctx context.Context, event *proto.ConversationEvent) (seq int32, err error) { + ctx, endSpan := l.startSpan(ctx, "Append", event.ConversationId) + defer func() { endSpan(err) }() + + tx, err := l.db.BeginTx(ctx, nil) + if err != nil { + return 0, fmt.Errorf("eventlog: begin tx: %w", err) + } + defer tx.Rollback() + + seq = event.Seq + if seq == 0 { + if err := tx.QueryRowContext(ctx, "SELECT COALESCE(MAX(seq), 0) + 1 FROM conversation_log WHERE conversation_id = $1", event.ConversationId).Scan(&seq); err != nil { + return 0, fmt.Errorf("eventlog: compute seq: %w", err) + } + event.Seq = seq + } + + payload, err := marshalOpts.Marshal(event) + if err != nil { + return 0, fmt.Errorf("eventlog: marshal event: %w", err) + } + + if _, err := tx.ExecContext(ctx, + "INSERT INTO conversation_log (conversation_id, seq, payload) VALUES ($1, $2, $3)", + event.ConversationId, event.Seq, string(payload)); err != nil { + return 0, fmt.Errorf("eventlog: insert conversation: %w", err) + } + + if err := tx.Commit(); err != nil { + return 0, fmt.Errorf("eventlog: commit: %w", err) + } + + return seq, nil +} + +// Events retrieves all events from the database for a conversation, ordered by seq. +func (l *sqlEventLog) Events(ctx context.Context, conversationID string) (events []*proto.ConversationEvent, err error) { + ctx, endSpan := l.startSpan(ctx, "Events", conversationID) + defer func() { endSpan(err) }() + + rows, err := l.db.QueryContext(ctx, "SELECT payload FROM conversation_log WHERE conversation_id = $1 ORDER BY seq", conversationID) + if err != nil { + return nil, fmt.Errorf("eventlog: query conversation: %w", err) + } + defer rows.Close() + + for rows.Next() { + var payload string + if err := rows.Scan(&payload); err != nil { + return nil, fmt.Errorf("eventlog: scan conversation: %w", err) + } + + ev := &proto.ConversationEvent{} + if err := unmarshalOpts.Unmarshal([]byte(payload), ev); err != nil { + return nil, fmt.Errorf("eventlog: unmarshal event: %w", err) + } + events = append(events, ev) + } + + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("eventlog: iterate conversation: %w", err) + } + + return events, nil +} + +// DeleteAll deletes all events for a specific conversation ID. +func (l *sqlEventLog) DeleteAll(ctx context.Context, conversationID string) (err error) { + ctx, endSpan := l.startSpan(ctx, "DeleteAll", conversationID) + defer func() { endSpan(err) }() + + if _, err := l.db.ExecContext(ctx, "DELETE FROM conversation_log WHERE conversation_id = $1", conversationID); err != nil { + return fmt.Errorf("eventlog: delete conversation: %w", err) + } + return nil +} + +// Close releases the database connection. +func (l *sqlEventLog) Close() error { + return l.db.Close() +} + +func (l *sqlEventLog) startSpan(ctx context.Context, name string, conversationID string) (context.Context, func(err error)) { + const tracerName = "eventlog.sql" + tracer := otel.Tracer(tracerName) + + ctx, span := tracer.Start(ctx, tracerName+"/"+name, trace.WithAttributes( + attribute.String("conversation_id", conversationID), + )) + return ctx, func(err error) { + if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + } + span.End() + } +} diff --git a/internal/controller/eventlog/sql_test.go b/internal/controller/eventlog/sql_test.go new file mode 100644 index 00000000..86eb4cf3 --- /dev/null +++ b/internal/controller/eventlog/sql_test.go @@ -0,0 +1,205 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package eventlog + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/google/ax/proto" +) + +// testEventLog runs the EventLog contract against a backend. newLog returns a +// fresh, empty event log for each subtest, using the subtest's *testing.T for +// temp dirs and cleanup. +// +// Subtests derive their IDs from t.Name() (which includes the backend and +// subtest name) and clear them with DeleteAll before and after running, so they +// are safe against the shared Postgres database and harmless against SQLite's +// per-test file. +func testEventLog(t *testing.T, newLog func(t *testing.T) EventLog) { + t.Run("AppendAndEvents", func(t *testing.T) { + ctx := context.Background() + log := newLog(t) + + conv := t.Name() + "-conv-1" + task1 := t.Name() + "-task-1" + task2 := t.Name() + "-task-2" + _ = log.DeleteAll(ctx, conv) + t.Cleanup(func() { _ = log.DeleteAll(ctx, conv) }) + + // 1. Conversation log. + cev1 := &proto.ConversationEvent{ConversationId: conv, Seq: 1, ExecId: task1} + cev2 := &proto.ConversationEvent{ConversationId: conv, Seq: 2, ExecId: task2} + if _, err := log.Append(ctx, cev1); err != nil { + t.Fatalf("failed to append cev1: %v", err) + } + if _, err := log.Append(ctx, cev2); err != nil { + t.Fatalf("failed to append cev2: %v", err) + } + + cEvents, err := log.Events(ctx, conv) + if err != nil { + t.Fatalf("failed to read conversation events: %v", err) + } + if len(cEvents) != 2 { + t.Fatalf("expected 2 conversation events, got %d", len(cEvents)) + } + if cEvents[0].Seq != 1 || cEvents[1].Seq != 2 { + t.Errorf("conversation events out of order: %d, %d", cEvents[0].Seq, cEvents[1].Seq) + } + if cEvents[0].ExecId != task1 || cEvents[1].ExecId != task2 { + t.Errorf("conversation events mismatch: %q, %q", cEvents[0].ExecId, cEvents[1].ExecId) + } + + + }) + + t.Run("Empty", func(t *testing.T) { + ctx := context.Background() + log := newLog(t) + + events, err := log.Events(ctx, t.Name()+"-none") + if err != nil { + t.Fatalf("failed to read events: %v", err) + } + if len(events) != 0 { + t.Fatalf("expected 0 events, got %d", len(events)) + } + }) + + t.Run("DeleteAll", func(t *testing.T) { + ctx := context.Background() + log := newLog(t) + + conv1 := t.Name() + "-conv-1" + conv2 := t.Name() + "-conv-2" + task1 := t.Name() + "-task-1" + task3 := t.Name() + "-task-3" + _ = log.DeleteAll(ctx, conv1) + _ = log.DeleteAll(ctx, conv2) + t.Cleanup(func() { + _ = log.DeleteAll(ctx, conv1) + _ = log.DeleteAll(ctx, conv2) + }) + + if _, err := log.Append(ctx, &proto.ConversationEvent{ConversationId: conv1, Seq: 1, ExecId: task1}); err != nil { + t.Fatalf("append: %v", err) + } + if _, err := log.Append(ctx, &proto.ConversationEvent{ConversationId: conv2, Seq: 1, ExecId: task3}); err != nil { + t.Fatalf("append: %v", err) + } + + + if err := log.DeleteAll(ctx, conv1); err != nil { + t.Fatalf("failed to delete events: %v", err) + } + + if ev, _ := log.Events(ctx, conv1); len(ev) != 0 { + t.Errorf("expected 0 events for conv1, got %d", len(ev)) + } + if ev, _ := log.Events(ctx, conv2); len(ev) != 1 { + t.Errorf("expected 1 event for conv2, got %d", len(ev)) + } + + }) + + // AutoSeq exercises the seq==0 auto-assignment path: appends with Seq unset + // receive sequential numbers starting at 1. + t.Run("AutoSeq", func(t *testing.T) { + ctx := context.Background() + log := newLog(t) + + conv := t.Name() + "-conv" + _ = log.DeleteAll(ctx, conv) + t.Cleanup(func() { _ = log.DeleteAll(ctx, conv) }) + + const n = 3 + for i := int32(1); i <= n; i++ { + seq, err := log.Append(ctx, &proto.ConversationEvent{ConversationId: conv, ExecId: "t"}) + if err != nil { + t.Fatalf("auto-seq append failed: %v", err) + } + if seq != i { + t.Errorf("expected seq %d, got %d", i, seq) + } + } + + events, err := log.Events(ctx, conv) + if err != nil { + t.Fatalf("failed to read events: %v", err) + } + if len(events) != n { + t.Fatalf("expected %d events, got %d", n, len(events)) + } + for i, e := range events { + if e.Seq != int32(i+1) { + t.Errorf("event %d: expected seq %d, got %d", i, i+1, e.Seq) + } + } + }) +} + +// TestSQLiteEventLog runs the EventLog contract against the SQLite backend. +func TestSQLiteEventLog(t *testing.T) { + testEventLog(t, func(t *testing.T) EventLog { + t.Helper() + log, err := OpenSQLiteEventLog(filepath.Join(t.TempDir(), "test.db")) + if err != nil { + t.Fatalf("failed to open sqlite event log: %v", err) + } + t.Cleanup(func() { log.Close() }) + return log + }) +} + +// TestPostgresEventLog runs the EventLog contract against the Postgres backend +// described by PG_TEST_DSN, skipping when that variable is not set. +func TestPostgresEventLog(t *testing.T) { + dsn := os.Getenv("PG_TEST_DSN") + if dsn == "" { + t.Skip("PG_TEST_DSN not set; skipping Postgres event log tests") + } + testEventLog(t, func(t *testing.T) EventLog { + t.Helper() + log, err := OpenPostgresEventLog(dsn) + if err != nil { + t.Fatalf("failed to open postgres event log: %v", err) + } + t.Cleanup(func() { log.Close() }) + return log + }) +} + +// TestSQLiteEventLog_CreatesParentDirectory is SQLite-specific: OpenSQLiteEventLog +// must create the database file's parent directory if it does not exist. +func TestSQLiteEventLog_CreatesParentDirectory(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "newdir", "test.db") + + log, err := OpenSQLiteEventLog(dbPath) + if err != nil { + t.Fatalf("failed to open sqlite event log and create directory: %v", err) + } + defer log.Close() + + if _, err := os.Stat(filepath.Dir(dbPath)); os.IsNotExist(err) { + t.Fatalf("expected parent directory to be created, but it does not exist") + } + if _, err := os.Stat(dbPath); os.IsNotExist(err) { + t.Fatalf("expected database file to be created, but it does not exist") + } +} diff --git a/internal/controller/eventlog/sqlite.go b/internal/controller/eventlog/sqlite.go new file mode 100644 index 00000000..b29ea5a6 --- /dev/null +++ b/internal/controller/eventlog/sqlite.go @@ -0,0 +1,56 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package eventlog + +import ( + "database/sql" + "fmt" + "os" + "path/filepath" + "time" + + _ "modernc.org/sqlite" +) + +const sqliteBusyTimeout = 10 * time.Second + +// OpenSQLiteEventLog opens (or creates) a SQLite database at path and initializes the event log schema. +func OpenSQLiteEventLog(path string) (EventLog, error) { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0755); err != nil { + return nil, fmt.Errorf("sqlite_eventlog: mkdir %s: %w", dir, err) + } + + dsn := fmt.Sprintf("%s?_pragma=busy_timeout(%d)&_txlock=immediate", path, sqliteBusyTimeout.Milliseconds()) + db, err := sql.Open("sqlite", dsn) + if err != nil { + return nil, fmt.Errorf("sqlite_eventlog: open %s: %w", dsn, err) + } + + // Create tables if they don't exist. + if _, err := db.Exec(` + CREATE TABLE IF NOT EXISTS conversation_log ( + conversation_id TEXT NOT NULL, + seq INTEGER NOT NULL, + payload TEXT NOT NULL, + PRIMARY KEY (conversation_id, seq) + )`); err != nil { + db.Close() + return nil, fmt.Errorf("sqlite_eventlog: create conversation_log table: %w", err) + } + + + return &sqlEventLog{db: db}, nil +} diff --git a/internal/controller/gemini.go b/internal/controller/gemini.go deleted file mode 100644 index 8de3d82d..00000000 --- a/internal/controller/gemini.go +++ /dev/null @@ -1,184 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package controller - -import ( - "context" - "fmt" - "os" - "time" - - "github.com/google/gar/proto" - "google.golang.org/genai" -) - -// GeminiPlannerConfig configures the Gemini-based planner. -type GeminiPlannerConfig struct { - APIKey string // Google AI API key (for programmatic use only; if empty, uses GEMINI_API_KEY env var - recommended) - Model string // Model name (default: gemini-2.5-flash) - MaxTokens int32 // Max output tokens (default: 8192) - Timeout time.Duration // Request timeout (default: 60s) - SystemPrompt string // Custom system prompt (optional) -} - -// NewGeminiPlanFunc creates a planning function that uses Gemini for intelligent agent selection. -// It converts available agents into function declarations and lets Gemini decide which agent -// to invoke based on the session context and agent capabilities. -func NewGeminiPlanFunc(ctx context.Context, registry *Registry, config GeminiPlannerConfig) (PlanFunc, error) { - // Set defaults - if config.Timeout == 0 { - config.Timeout = 60 * time.Second - } - if config.Model == "" { - config.Model = os.Getenv("GAR_GEMINI_MODEL") - if config.Model == "" { - config.Model = "gemini-2.5-flash" - } - } - if config.APIKey == "" { - config.APIKey = os.Getenv("GEMINI_API_KEY") - if config.APIKey == "" { - return nil, fmt.Errorf("GEMINI_API_KEY not set and no API key provided in config") - } - } - - // Default system prompt - if config.SystemPrompt == "" { - config.SystemPrompt = `You are an intelligent agent orchestrator. Your role is to analyze the conversation history and user requests, then select the most appropriate agent to handle the task. - -Available agents have been provided to you as function tools. Each agent has: -- A unique ID -- A description of its capabilities - -Your job is to: -1. Analyze the current conversation context and understand what needs to be done -2. Select the best agent for the task by calling the appropriate function -3. If enough work is done, stop to indicate completion - -Guidelines: -- Choose agents based on their capabilities and the user's needs -- If no suitable agent exists, stop. -- Keep the conversation context in mind when selecting agents` - } - - client, err := genai.NewClient(ctx, &genai.ClientConfig{ - APIKey: config.APIKey, - }) - if err != nil { - return nil, fmt.Errorf("failed to create Gemini client: %w", err) - } - - // Return the plan function - return func(ctx context.Context, inputs []*proto.Content) (*Task, error) { - // Create a context with timeout - ctx, cancel := context.WithTimeout(ctx, config.Timeout) - defer cancel() - - // Convert agents to Gemini function declarations - tools, err := agentsToTools(registry) - if err != nil { - return nil, fmt.Errorf("failed to convert agents to tools: %w", err) - } - - // Convert session to conversation history - contents := protoToContents(inputs) - resp, err := client.Models.GenerateContent(ctx, config.Model, contents, &genai.GenerateContentConfig{ - Tools: tools, - SystemInstruction: genai.Text(config.SystemPrompt)[0], - MaxOutputTokens: config.MaxTokens, - CandidateCount: 1, - }) - - if err != nil { - return nil, fmt.Errorf("failed to generate in planner: %w", err) - } - if len(resp.Candidates) == 0 { - return nil, fmt.Errorf("no candidates from Gemini in planner") - } - candidate := resp.Candidates[0] - if candidate.Content == nil || candidate.Content.Parts == nil { - if candidate.FinishReason == genai.FinishReasonStop { - return nil, nil // No more tasks - } - return nil, fmt.Errorf("no content in candidates from Gemini") - } - - // Look for function calls in the response - for _, part := range candidate.Content.Parts { - if part == nil { - continue - } - - if fc := part.FunctionCall; fc != nil { - return &Task{ - AgentID: fc.Name, - Inputs: inputs, - }, nil - } - } - return nil, nil - }, nil -} - -// agentsToTools converts registry agents to Gemini function declarations. -func agentsToTools(registry *Registry) ([]*genai.Tool, error) { - healthyAgents := registry.ListHealthy() - if len(healthyAgents) == 0 { - return nil, fmt.Errorf("no healthy agents available") - } - - var tools []*genai.Tool - for _, id := range healthyAgents { - info, err := registry.GetInfo(id) - if err != nil { - continue // Skip agents we can't get info for - } - - // Create a function declaration for this agent - funcDecl := &genai.FunctionDeclaration{ - Name: id, // Use agent ID as function name - Description: fmt.Sprintf("%s, %s", info.Name, info.Description), - } - - tools = append(tools, &genai.Tool{ - FunctionDeclarations: []*genai.FunctionDeclaration{funcDecl}, - }) - } - return tools, nil -} - -// protoToContents converts session message history to Gemini conversation format. -func protoToContents(inputs []*proto.Content) []*genai.Content { - var contents []*genai.Content - - // Convert each message to Gemini format - for _, msg := range inputs { - role := msg.Role - if role != "user" { - role = "model" - } - contents = append(contents, &genai.Content{ - Role: role, - Parts: []*genai.Part{ - { - Text: msg.Data, - // TODO(jbd): Handle other content types (e.g., images, files) - }, - }, - }) - } - - return contents -} diff --git a/internal/controller/loop.go b/internal/controller/loop.go deleted file mode 100644 index ce1c4ecd..00000000 --- a/internal/controller/loop.go +++ /dev/null @@ -1,143 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package controller - -import ( - "context" - "fmt" - - "github.com/google/gar/agent" - "github.com/google/gar/proto" -) - -// Task represents a task to be executed by an agent. -type Task struct { - AgentID string - Inputs []*proto.Content -} - -// LoopExecutor orchestrates the agentic loop workflow. -// It implements the plan-execute-evaluate-checkpoint cycle. -type LoopExecutor struct { - registry *Registry - sessionManager *SessionManager - maxSteps int - planFunc PlanFunc -} - -// PlanFunc determines the next agent task to execute. -// It receives the current session state and returns the next task. -type PlanFunc func(ctx context.Context, inputs []*proto.Content) (*Task, error) - -// LoopConfig configures the loop executor. -type LoopConfig struct { - Registry *Registry - SessionManager *SessionManager - MaxSteps int - PlanFunc PlanFunc -} - -// NewLoopExecutor creates a new loop executor. -func NewLoopExecutor(ctx context.Context, config LoopConfig) (*LoopExecutor, error) { - if config.Registry == nil { - return nil, fmt.Errorf("registry cannot be nil") - } - if config.SessionManager == nil { - return nil, fmt.Errorf("session manager cannot be nil") - } - if config.MaxSteps == 0 { - config.MaxSteps = 10 // Default max steps per trigger - } - - // Provide default plan function if not specified - if config.PlanFunc == nil { - // Use Gemini planner by default - geminiPlanFunc, err := NewGeminiPlanFunc(ctx, config.Registry, GeminiPlannerConfig{}) - if err != nil { - return nil, fmt.Errorf("failed to initialize default Gemini planner: %w (set GEMINI_API_KEY env var or provide custom PlanFunc)", err) - } - config.PlanFunc = geminiPlanFunc - } - - return &LoopExecutor{ - registry: config.Registry, - sessionManager: config.SessionManager, - maxSteps: config.MaxSteps, - planFunc: config.PlanFunc, - }, nil -} - -// Execute starts a new agentic loop execution for the given session. -func (e *LoopExecutor) Execute(ctx context.Context, session *Session, handler agent.OutputHandler) error { - return e.runLoop(ctx, session, handler) -} - -// runLoop executes the main agentic loop. -// It runs up to maxSteps iterations per trigger/resume invocation. -func (e *LoopExecutor) runLoop(ctx context.Context, session *Session, handler agent.OutputHandler) error { - steps := 0 - - for steps < e.maxSteps { - // Check context cancellation - select { - case <-ctx.Done(): - return ctx.Err() - default: - } - - task, err := e.planFunc(ctx, session.History()) - if err != nil { - return fmt.Errorf("planning failed: %w", err) - } - - if task == nil { - // No more tasks to execute; loop is complete. - return nil - } - - taskOutputHandler := func(content *proto.Content) error { - if _, err := session.WriteContentOut(ctx, content); err != nil { - return fmt.Errorf("failed to write output content: %w", err) - } - if err := handler(content); err != nil { - return fmt.Errorf("output handler error: %w", err) - } - return nil - } - if err := e.executeTask(ctx, session.ID(), task, taskOutputHandler); err != nil { - return err - } - - steps++ - } - - // Can be resumed later with another trigger - return fmt.Errorf("max steps per trigger (%d) reached", e.maxSteps) -} - -// executeTask sends input to an agent and collects output. -func (e *LoopExecutor) executeTask(ctx context.Context, sessionID string, task *Task, handler agent.OutputHandler) error { - // Get the agent from registry - ag, err := e.registry.Get(task.AgentID) - if err != nil { - return fmt.Errorf("failed to get agent: %w", err) - } - - // Process inputs with the agent - if err := ag.Process(ctx, sessionID, task.Inputs, handler); err != nil { - return fmt.Errorf("agent process failed: %w", err) - } - return nil -} diff --git a/internal/controller/registry.go b/internal/controller/registry.go index 8122e145..7c1ab57b 100644 --- a/internal/controller/registry.go +++ b/internal/controller/registry.go @@ -15,259 +15,65 @@ package controller import ( - "context" "fmt" "sync" - "time" - "github.com/google/gar/agent" - "github.com/google/gar/internal/config" + "github.com/google/ax/internal/harness" ) -// AgentType represents the type of agent (local or remote). -type AgentType string - -const ( - AgentTypeLocal AgentType = "local" - AgentTypeRemote AgentType = "remote" -) - -// AgentInfo contains metadata about a registered agent. -type AgentInfo struct { - ID string - Name string - Description string - Type AgentType - Healthy bool - LastHealthCheck time.Time - Metadata map[string]string -} - -// Registry manages a collection of local and remote agents. -// It provides agent discovery, health monitoring, and load balancing. +// Registry manages a collection of harnesses. type Registry struct { - mu sync.RWMutex - agents map[string]agent.Agent - agentInfo map[string]*AgentInfo - healthCheckTicker *time.Ticker - stopHealthCheck chan struct{} + mu sync.RWMutex + harnesses map[string]harness.Harness + defaultHarness string } -// NewRegistry creates a new agent registry. -func NewRegistry(healthCheckInterval time.Duration) *Registry { - if healthCheckInterval == 0 { - healthCheckInterval = 30 * time.Second +// NewRegistry creates a new harness registry. +func NewRegistry() *Registry { + return &Registry{ + harnesses: make(map[string]harness.Harness), } - - r := &Registry{ - agents: make(map[string]agent.Agent), - agentInfo: make(map[string]*AgentInfo), - stopHealthCheck: make(chan struct{}), - } - - // Start background health check if interval is positive - if healthCheckInterval > 0 { - r.healthCheckTicker = time.NewTicker(healthCheckInterval) - go r.runHealthChecks() - } - - return r } -// RegisterLocal registers a local (in-process) agent. -func (r *Registry) RegisterLocal(cfg config.LocalAgentConfig) error { - r.mu.Lock() - defer r.mu.Unlock() - - if _, ok := r.agents[cfg.ID]; ok { - return fmt.Errorf("agent %s already registered", cfg.ID) - } - - r.agents[cfg.ID] = cfg.Agent - r.agentInfo[cfg.ID] = &AgentInfo{ - ID: cfg.ID, - Name: cfg.Name, - Description: cfg.Description, - Type: AgentTypeLocal, - Healthy: true, - LastHealthCheck: time.Now(), - Metadata: cfg.Metadata, - } - - return nil -} - -// RegisterRemote registers a remote agent by creating a remote agent client. -func (r *Registry) RegisterRemote(cfg config.RemoteAgentConfig) error { - r.mu.Lock() - defer r.mu.Unlock() - - if _, ok := r.agents[cfg.ID]; ok { - return fmt.Errorf("agent %s already registered", cfg.ID) - } - - // Create remote agent client - remoteAgent, err := agent.NewRemoteAgent(agent.RemoteAgentConfig{ - Address: cfg.Address, - Reconnect: true, - MaxRetries: 3, - }) - if err != nil { - return fmt.Errorf("failed to create remote agent: %w", err) - } - - r.agents[cfg.ID] = remoteAgent - r.agentInfo[cfg.ID] = &AgentInfo{ - ID: cfg.ID, - Name: cfg.Name, - Description: cfg.Description, - Type: AgentTypeRemote, - Healthy: false, // Will be checked by health monitor - LastHealthCheck: time.Time{}, - Metadata: cfg.Metadata, +// RegisterHarness registers a harness under the given id. +func (r *Registry) RegisterHarness(id string, h harness.Harness) error { + if err := validateID(id); err != nil { + return err } - return nil -} - -// Unregister removes an agent from the registry. -func (r *Registry) Unregister(id string) error { r.mu.Lock() defer r.mu.Unlock() - agent, ok := r.agents[id] - if !ok { - return fmt.Errorf("agent %s not found", id) - } - - // Close the agent - if err := agent.Close(); err != nil { - return fmt.Errorf("failed to close agent: %w", err) + if _, ok := r.harnesses[id]; ok { + return fmt.Errorf("harness %q already registered", id) } - - delete(r.agents, id) - delete(r.agentInfo, id) - + r.harnesses[id] = h return nil } -// Get retrieves an agent by ID. -func (r *Registry) Get(id string) (agent.Agent, error) { - r.mu.RLock() - defer r.mu.RUnlock() - - a, exists := r.agents[id] - if !exists { - return nil, fmt.Errorf("agent %s not found", id) - } - - return a, nil -} - -// GetInfo retrieves agent metadata by ID. -func (r *Registry) GetInfo(id string) (*AgentInfo, error) { +// Harness retrieves a harness by id. +func (r *Registry) Harness(id string) (harness.Harness, error) { r.mu.RLock() defer r.mu.RUnlock() - - info, exists := r.agentInfo[id] - if !exists { - return nil, fmt.Errorf("agent %s not found", id) - } - - return info, nil -} - -// List returns all registered agent IDs. -func (r *Registry) List() []string { - r.mu.RLock() - defer r.mu.RUnlock() - - ids := make([]string, 0, len(r.agents)) - for id := range r.agents { - ids = append(ids, id) - } - - return ids -} - -// ListHealthy returns all healthy agent IDs. -func (r *Registry) ListHealthy() []string { - r.mu.RLock() - defer r.mu.RUnlock() - - ids := make([]string, 0) - for id, info := range r.agentInfo { - if info.Healthy { - ids = append(ids, id) - } + h, ok := r.harnesses[id] + if !ok { + return nil, fmt.Errorf("harness %s not found", id) } - - return ids + return h, nil } -// healthCheck performs a health check on a specific agent. -func (r *Registry) healthCheck(id string) error { - a, err := r.Get(id) - if err != nil { - return err - } - - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - - err = a.HealthCheck(ctx) - +// SetDefaultHarness marks a registered harness id as the default. +func (r *Registry) SetDefaultHarness(id string) error { r.mu.Lock() defer r.mu.Unlock() - - info, exists := r.agentInfo[id] - if exists { - info.Healthy = (err == nil) - info.LastHealthCheck = time.Now() - } - - return err -} - -// runHealthChecks runs periodic health checks on all agents. -func (r *Registry) runHealthChecks() { - for { - select { - case <-r.healthCheckTicker.C: - r.performHealthChecks() - case <-r.stopHealthCheck: - return - } - } -} - -// performHealthChecks checks the health of all registered agents. -func (r *Registry) performHealthChecks() { - ids := r.List() - for _, id := range ids { - // Run health check (ignore errors, status is updated in HealthCheck method) - _ = r.healthCheck(id) + if _, ok := r.harnesses[id]; !ok { + return fmt.Errorf("harness %q not found", id) } + r.defaultHarness = id + return nil } -// Close stops the registry and closes all agents. +// Close releases resources held by the registry. func (r *Registry) Close() error { - r.mu.Lock() - defer r.mu.Unlock() - - // Stop health checks - if r.healthCheckTicker != nil { - r.healthCheckTicker.Stop() - close(r.stopHealthCheck) - } - - // Close all agents - var firstErr error - for id, a := range r.agents { - if err := a.Close(); err != nil && firstErr == nil { - firstErr = fmt.Errorf("failed to close agent %s: %w", id, err) - } - } - - return firstErr + return nil } diff --git a/internal/controller/registry_test.go b/internal/controller/registry_test.go new file mode 100644 index 00000000..c3eb7516 --- /dev/null +++ b/internal/controller/registry_test.go @@ -0,0 +1,87 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controller + +import ( + "context" + "testing" + + "github.com/google/ax/internal/harness" +) + +type dummyHarness struct{} + +func (d *dummyHarness) Start(ctx context.Context, conversationID string, harnessConfig []byte) (harness.Execution, error) { + return nil, nil +} + +func TestRegistry_RegisterHarness(t *testing.T) { + r := NewRegistry() + h := &dummyHarness{} + + if err := r.RegisterHarness("antigravity", h); err != nil { + t.Fatalf("RegisterHarness(valid id): %v", err) + } + + // Duplicate id is rejected. + if err := r.RegisterHarness("antigravity", h); err == nil { + t.Error("expected error registering duplicate id, got nil") + } + + // Invalid id is rejected. + if err := r.RegisterHarness("bad id", h); err == nil { + t.Error("expected error registering invalid id, got nil") + } + + // Empty id is reserved for the default harness. + if err := r.RegisterHarness("", h); err == nil { + t.Error("expected error registering empty id, got nil") + } +} + +func TestRegistry_FindHarness(t *testing.T) { + r := NewRegistry() + h := &dummyHarness{} + if err := r.RegisterHarness("antigravity", h); err != nil { + t.Fatalf("RegisterHarness: %v", err) + } + + if _, err := r.Harness("antigravity"); err != nil { + t.Errorf("Harness(antigravity): %v", err) + } + if _, err := r.Harness("missing"); err == nil { + t.Error("expected error looking up missing harness, got nil") + } +} + +func TestRegistry_SetDefaultHarness(t *testing.T) { + r := NewRegistry() + if err := r.RegisterHarness("antigravity", &dummyHarness{}); err != nil { + t.Fatalf("RegisterHarness: %v", err) + } + + // An unregistered id is rejected. + if err := r.SetDefaultHarness("missing"); err == nil { + t.Error("SetDefaultHarness(missing): expected error, got nil") + } + + // A registered id becomes the default. + if err := r.SetDefaultHarness("antigravity"); err != nil { + t.Fatalf("SetDefaultHarness(antigravity): %v", err) + } + if r.defaultHarness != "antigravity" { + t.Errorf("defaultHarness = %q, want %q", r.defaultHarness, "antigravity") + } +} diff --git a/internal/controller/session.go b/internal/controller/session.go deleted file mode 100644 index cb2fa285..00000000 --- a/internal/controller/session.go +++ /dev/null @@ -1,344 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package controller - -import ( - "context" - "fmt" - "sync" - "time" - - "github.com/google/gar/internal/eventlog" - "github.com/google/gar/proto" - "github.com/google/uuid" -) - -// Session represents an agentic loop execution session. -// It maintains in-memory state and uses event log for durability. -type Session struct { - id string - - mu sync.RWMutex - eventLog eventlog.EventLog - currentStep int - state proto.State - activeAgents []string - messageHistory []*proto.Content - checkpointIDs []string // Ordered list of checkpoint UUIDs - - createdAt time.Time - updatedAt time.Time -} - -// SessionManager manages multiple sessions. -type SessionManager struct { - mu sync.RWMutex - sessions map[string]*Session - eventLogFactory eventlog.EventLogFactory -} - -// NewSessionManager creates a new session manager with a custom EventLog factory. -func NewSessionManager(factory eventlog.EventLogFactory) *SessionManager { - return &SessionManager{ - sessions: make(map[string]*Session), - eventLogFactory: factory, - } -} - -// NewSession creates a new session with the given ID. -func (sm *SessionManager) NewSession(sessionID string) (*Session, error) { - sm.mu.Lock() - defer sm.mu.Unlock() - - // Check if session already ok - if _, ok := sm.sessions[sessionID]; ok { - return nil, fmt.Errorf("session %s already exists", sessionID) - } - - // Create event log for this session using the factory - el, err := sm.eventLogFactory(sessionID) - if err != nil { - return nil, fmt.Errorf("failed to create event log: %w", err) - } - - now := time.Now() - session := &Session{ - id: sessionID, - state: proto.State_STATE_UNSPECIFIED, - currentStep: 0, - activeAgents: []string{}, - messageHistory: []*proto.Content{}, - checkpointIDs: []string{}, - createdAt: now, - updatedAt: now, - eventLog: el, - } - - sm.sessions[sessionID] = session - return session, nil -} - -// LoadSession loads an existing session from event log. -func (sm *SessionManager) LoadSession(ctx context.Context, sessionID string) (*Session, error) { - return sm.LoadSessionFromCheckpoint(ctx, sessionID, "") -} - -// LoadSessionFromCheckpoint loads an existing session from event log up to a specific checkpoint. -// If checkpointID is empty, loads to the latest state. -// If checkpointID is provided, loads up to and including that checkpoint UUID. -func (sm *SessionManager) LoadSessionFromCheckpoint(ctx context.Context, sessionID string, checkpointID string) (*Session, error) { - sm.mu.Lock() - defer sm.mu.Unlock() - - // Check if already loaded - remove it to reload fresh from checkpoint - delete(sm.sessions, sessionID) - - // Open event log for replay using the factory - replayLog, err := sm.eventLogFactory(sessionID) - if err != nil { - return nil, fmt.Errorf("failed to open event log for replay: %w", err) - } - defer replayLog.Close() - - entries, state, err := replayLog.RetrieveEntries(ctx) - if err != nil { - return nil, fmt.Errorf("failed to get event log entries: %w", err) - } - - // Reconstruct session state from event log - session := &Session{ - id: sessionID, - state: state, - currentStep: 0, - activeAgents: []string{}, - messageHistory: []*proto.Content{}, - checkpointIDs: []string{}, - createdAt: time.Now(), - updatedAt: time.Now(), - } - - targetReached := false - checkpointFound := false - - // Replay entries to rebuild state - for _, entry := range entries { - // If we've reached the target checkpoint, stop processing - if targetReached { - break - } - - switch entry.Type { - case eventlog.EventTypeContentIn, eventlog.EventTypeContentOut: - content := &proto.Content{ - Role: getStringFromData(entry.Data, "role"), - Type: getStringFromData(entry.Data, "type"), - Mimetype: getStringFromData(entry.Data, "mimetype"), - Data: getStringFromData(entry.Data, "data"), - } - session.messageHistory = append(session.messageHistory, content) - - // Track checkpoint ID if present - if entry.CheckpointID != "" { - session.checkpointIDs = append(session.checkpointIDs, entry.CheckpointID) - - // Check if this is the target checkpoint - if checkpointID != "" && entry.CheckpointID == checkpointID { - targetReached = true - checkpointFound = true - } - } - } - - session.updatedAt = entry.Timestamp - } - - // Validate checkpoint ID if provided - if checkpointID != "" && !checkpointFound { - return nil, fmt.Errorf("checkpoint ID %s not found in session", checkpointID) - } - - // Reopen event log for appending using the factory - el, err := sm.eventLogFactory(sessionID) - if err != nil { - return nil, fmt.Errorf("failed to reopen event log: %w", err) - } - session.eventLog = el - - sm.sessions[sessionID] = session - return session, nil -} - -// GetSession retrieves a session by ID. -func (sm *SessionManager) GetSession(sessionID string) (*Session, error) { - sm.mu.RLock() - defer sm.mu.RUnlock() - - session, ok := sm.sessions[sessionID] - if !ok { - return nil, fmt.Errorf("session %s not found", sessionID) - } - - return session, nil -} - -// CloseSession closes a session and its event log. -func (sm *SessionManager) CloseSession(sessionID string) error { - sm.mu.Lock() - defer sm.mu.Unlock() - - session, ok := sm.sessions[sessionID] - if !ok { - return fmt.Errorf("session %s not found", sessionID) - } - - if err := session.eventLog.Close(); err != nil { - return fmt.Errorf("failed to close event log: %w", err) - } - - delete(sm.sessions, sessionID) - return nil -} - -// CloseAll closes all active sessions. -func (sm *SessionManager) CloseAll() { - sm.mu.Lock() - defer sm.mu.Unlock() - - for sessionID, session := range sm.sessions { - if err := session.eventLog.Close(); err != nil { - // Log error but continue closing other sessions - _ = err - } - delete(sm.sessions, sessionID) - } -} - -// WriteContentIn appends an incoming content message to the session. -// Creates a checkpoint only if checkpoint_id is provided in the content. -func (s *Session) WriteContentIn(ctx context.Context, content *proto.Content) (string, error) { - s.mu.Lock() - defer s.mu.Unlock() - - // Use checkpoint_id from content if provided - checkpointID := content.CheckpointId - - if checkpointID != "" { - // TODO(jbd): Optimize the lookup. - for _, existingID := range s.checkpointIDs { - if existingID == checkpointID { - return "", fmt.Errorf("checkpoint %s already exists", checkpointID) - } - } - } - - if err := s.eventLog.AppendContent(ctx, eventlog.EventTypeContentIn, checkpointID, content); err != nil { - return "", err - } - - s.messageHistory = append(s.messageHistory, content) - if checkpointID != "" { - s.checkpointIDs = append(s.checkpointIDs, checkpointID) - } - s.updatedAt = time.Now() - return checkpointID, nil -} - -// WriteContentOut appends an outgoing content message to the session with a new checkpoint. -func (s *Session) WriteContentOut(ctx context.Context, content *proto.Content) (string, error) { - s.mu.Lock() - defer s.mu.Unlock() - - // Generate a new checkpoint UUID - checkpointID := uuid.New().String() - - if err := s.eventLog.AppendContent(ctx, eventlog.EventTypeContentOut, checkpointID, content); err != nil { - return "", err - } - - s.messageHistory = append(s.messageHistory, content) - s.checkpointIDs = append(s.checkpointIDs, checkpointID) - s.updatedAt = time.Now() - return checkpointID, nil -} - -// SetState updates the session state. -func (s *Session) SetState(ctx context.Context, state proto.State) error { - s.mu.Lock() - defer s.mu.Unlock() - - if err := s.eventLog.AppendState(ctx, state); err != nil { - return fmt.Errorf("failed to append state: %w", err) - } - - s.state = state - s.updatedAt = time.Now() - return nil -} - -func (s *Session) ID() string { - return s.id -} - -func (s *Session) State() proto.State { - s.mu.RLock() - defer s.mu.RUnlock() - - return s.state -} - -func (s *Session) ActiveAgents() []string { - s.mu.RLock() - defer s.mu.RUnlock() - - return s.activeAgents -} - -func (s *Session) History() []*proto.Content { - s.mu.RLock() - defer s.mu.RUnlock() - - return s.messageHistory -} - -func (s *Session) CheckpointIDs() []string { - s.mu.RLock() - defer s.mu.RUnlock() - - return s.checkpointIDs -} - -func (s *Session) CreatedAt() time.Time { - s.mu.RLock() - defer s.mu.RUnlock() - - return s.createdAt -} - -func (s *Session) UpdatedAt() time.Time { - s.mu.RLock() - defer s.mu.RUnlock() - - return s.updatedAt -} - -// Helper function to extract string from map[string]any -func getStringFromData(data map[string]any, key string) string { - if val, ok := data[key]; ok { - if str, ok := val.(string); ok { - return str - } - } - return "" -} diff --git a/internal/controller/validation.go b/internal/controller/validation.go new file mode 100644 index 00000000..729b2435 --- /dev/null +++ b/internal/controller/validation.go @@ -0,0 +1,36 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controller + +import ( + "errors" + "fmt" + "regexp" +) + +var validIDRegex = regexp.MustCompile(`^[A-Za-z0-9\-_]+$`) + +// validateID checks if a harness ID contains allowed characters. +func validateID(id string) error { + if id == "" { + return errors.New("empty ID") + } + + if !validIDRegex.MatchString(id) { + return fmt.Errorf("invalid harness ID %q: must only contain A-Z, a-z, 0-9, -, and _", id) + } + + return nil +} diff --git a/internal/controller/validation_test.go b/internal/controller/validation_test.go new file mode 100644 index 00000000..4b1dde13 --- /dev/null +++ b/internal/controller/validation_test.go @@ -0,0 +1,72 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controller + +import ( + "testing" +) + +func TestValidateID(t *testing.T) { + tests := []struct { + name string + id string + wantErr bool + }{ + { + name: "valid lowercase", + id: "task123", + wantErr: false, + }, + { + name: "valid mixed", + id: "Task-ID_123", + wantErr: false, + }, + { + name: "valid simple", + id: "Task-ID", + wantErr: false, + }, + { + name: "valid underscore", + id: "task_id", + wantErr: false, + }, + { + name: "invalid space", + id: "task id", + wantErr: true, + }, + { + name: "invalid char", + id: "task!", + wantErr: true, + }, + { + name: "empty", + id: "", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateID(tt.id) + if (err != nil) != tt.wantErr { + t.Errorf("validateID(%q) error = %v, wantErr %v", tt.id, err, tt.wantErr) + } + }) + } +} diff --git a/internal/eventlog/eventlog.go b/internal/eventlog/eventlog.go deleted file mode 100644 index 495d0f43..00000000 --- a/internal/eventlog/eventlog.go +++ /dev/null @@ -1,66 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Package eventlog implements event logging for session durability. -// Event log entries use JSON Lines format with one log file per session. -package eventlog - -import ( - "context" - "time" - - "github.com/google/gar/proto" -) - -// EventType represents the type of event stored in the event log. -type EventType string - -const ( - EventTypeState EventType = "STATE" - EventTypeContentIn EventType = "CONTENT_IN" - EventTypeContentOut EventType = "CONTENT_OUT" - // TODO(jbd): Add EventTypeCompaction. -) - -// Entry represents a single entry in the event log. -type Entry struct { - SessionID string `json:"session_id"` - CheckpointID string `json:"checkpoint_id,omitempty"` // UUID for checkpoint tracking - Sequence int64 `json:"seq"` // Monotonic sequence number - Timestamp time.Time `json:"timestamp"` - Type EventType `json:"type"` - Data map[string]any `json:"data"` -} - -// EventLog is the interface that all event log implementations must satisfy. -// It provides methods for appending events, reading entries, and managing the log lifecycle. -type EventLog interface { - // AppendContent appends a content message to the event log with a checkpoint UUID. - AppendContent(ctx context.Context, t EventType, checkpointID string, content *proto.Content) error - - // AppendState appends a state to the event log. - AppendState(ctx context.Context, s proto.State) error - - // RetrieveEntries returns all entries from the event log in order. - RetrieveEntries(ctx context.Context) ([]Entry, proto.State, error) - - // Close closes the event log and releases any resources. - Close() error - - // SessionID returns the session ID for this event log. - SessionID() string -} - -// EventLogFactory is a function that creates EventLog instances for sessions. -type EventLogFactory func(sessionID string) (EventLog, error) diff --git a/internal/eventlog/file.go b/internal/eventlog/file.go deleted file mode 100644 index 4c2c8584..00000000 --- a/internal/eventlog/file.go +++ /dev/null @@ -1,246 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package eventlog - -import ( - "bufio" - "context" - "encoding/json" - "fmt" - "os" - "path/filepath" - "sync" - "time" - - "github.com/google/gar/proto" -) - -// FileEventLog represents a file-based event log for a single session. -// It provides append-only durability with JSON Lines format. -type FileEventLog struct { - mu sync.Mutex - sessionID string - filePath string - file *os.File - writer *bufio.Writer - sequence int64 // Monotonic sequence counter -} - -// FileConfig configures a FileEventLog instance. -type FileConfig struct { - SessionID string - Dir string // Directory where event log files are stored -} - -// NewFileEventLog creates a new file-based EventLog instance for the given session. -// It creates the event log directory and file if they don't exist. -func NewFileEventLog(config FileConfig) (*FileEventLog, error) { - if config.SessionID == "" { - return nil, fmt.Errorf("session ID cannot be empty") - } - if config.Dir == "" { - config.Dir = "eventlog" - } - - // Create event log directory if it doesn't exist - if err := os.MkdirAll(config.Dir, 0755); err != nil { - return nil, fmt.Errorf("failed to create event log directory: %w", err) - } - - filePath := filepath.Join(config.Dir, fmt.Sprintf("%s.log", config.SessionID)) - - // Open file in append mode, create if doesn't exist - file, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) - if err != nil { - return nil, fmt.Errorf("failed to open event log file: %w", err) - } - - // Read existing entries to determine the current sequence number - var sequence int64 - if fileInfo, err := os.Stat(filePath); err == nil && fileInfo.Size() > 0 { - readFile, err := os.Open(filePath) - if err != nil { - file.Close() - return nil, fmt.Errorf("failed to read existing log for sequence: %w", err) - } - scanner := bufio.NewScanner(readFile) - for scanner.Scan() { - var entry Entry - if err := json.Unmarshal(scanner.Bytes(), &entry); err == nil { - if entry.Sequence > sequence { - sequence = entry.Sequence - } - } - } - readFile.Close() - } - - return &FileEventLog{ - sessionID: config.SessionID, - filePath: filePath, - file: file, - writer: bufio.NewWriter(file), - sequence: sequence, - }, nil -} - -// NewFileConfig creates a new file-based EventLog instance for the given session. -// This is a convenience function that wraps NewFileEventLog. -// Deprecated: Use NewFileEventLog for clarity, or implement your own EventLog. -func NewFileConfig(config FileConfig) (EventLog, error) { - return NewFileEventLog(config) -} - -// Append writes an entry to the event log. -// Entries are written in JSON Lines format with atomic appends. -func (e *FileEventLog) Append(entryType EventType, checkpointID string, data map[string]any) error { - e.mu.Lock() - defer e.mu.Unlock() - - entry := Entry{ - SessionID: e.sessionID, - Timestamp: time.Now(), - Sequence: e.sequence, - Type: entryType, - CheckpointID: checkpointID, - Data: data, - } - - // Increment sequence for next entry - e.sequence++ - - jsonData, err := json.Marshal(entry) - if err != nil { - return fmt.Errorf("failed to marshal entry: %w", err) - } - if _, err := e.writer.Write(jsonData); err != nil { - return fmt.Errorf("failed to write entry: %w", err) - } - if _, err := e.writer.WriteString("\n"); err != nil { - return fmt.Errorf("failed to write newline: %w", err) - } - - // Flush to ensure durability - if err := e.writer.Flush(); err != nil { - return fmt.Errorf("failed to flush: %w", err) - } - if err := e.file.Sync(); err != nil { - return fmt.Errorf("failed to sync: %w", err) - } - return nil -} - -// AppendContent writes a content message to the event log with a checkpoint UUID. -func (e *FileEventLog) AppendContent(ctx context.Context, direction EventType, checkpointID string, content *proto.Content) error { - data := map[string]any{ - "role": content.Role, - "type": content.Type, - "mimetype": content.Mimetype, - "data": content.Data, - } - return e.Append(direction, checkpointID, data) -} - -func (e *FileEventLog) AppendState(ctx context.Context, state proto.State) error { - data := map[string]any{ - "state": state, - } - return e.Append(EventTypeState, "", data) -} - -// Close closes the event log file. -func (e *FileEventLog) Close() error { - e.mu.Lock() - defer e.mu.Unlock() - - if err := e.writer.Flush(); err != nil { - return fmt.Errorf("failed to flush on close: %w", err) - } - if err := e.file.Close(); err != nil { - return fmt.Errorf("failed to close file: %w", err) - } - return nil -} - -// FilePath returns the path to the event log file. -// This method is specific to FileEventLog and not part of the EventLog interface. -func (e *FileEventLog) FilePath() string { - return e.filePath -} - -// SessionID returns the session ID for this event log. -func (e *FileEventLog) SessionID() string { - return e.sessionID -} - -// RetrieveEntries reads and returns all entries from the event log file. -// Returns entries in order. -func (e *FileEventLog) RetrieveEntries(ctx context.Context) ([]Entry, proto.State, error) { - e.mu.Lock() - defer e.mu.Unlock() - - var state proto.State - - // Flush any pending writes first - if err := e.writer.Flush(); err != nil { - return nil, 0, fmt.Errorf("failed to flush before reading: %w", err) - } - - // Check if file exists - if _, err := os.Stat(e.filePath); os.IsNotExist(err) { - return nil, 0, fmt.Errorf("event log file does not exist") - } - - // Open file for reading - readFile, err := os.Open(e.filePath) - if err != nil { - return nil, 0, fmt.Errorf("failed to open event log file for reading: %w", err) - } - defer readFile.Close() - - var entries []Entry - scanner := bufio.NewScanner(readFile) - - lineNum := 0 - var expectedSequence int64 = 0 - for scanner.Scan() { - lineNum++ - line := scanner.Bytes() - - var entry Entry - if err := json.Unmarshal(line, &entry); err != nil { - return nil, 0, fmt.Errorf("failed to unmarshal entry at line %d: %w", lineNum, err) - } - - // Validate sequence number - if entry.Sequence != expectedSequence { - return nil, 0, fmt.Errorf("sequence number mismatch at line %d: expected %d, got %d", lineNum, expectedSequence, entry.Sequence) - } - - expectedSequence++ - if entry.Type == EventTypeState { - if s, ok := entry.Data["state"].(proto.State); ok { - state = s - } - } - entries = append(entries, entry) - } - - if err := scanner.Err(); err != nil { - return nil, 0, fmt.Errorf("error reading event log file: %w", err) - } - - return entries, state, nil -} diff --git a/internal/harness/antigravity/antigravity.go b/internal/harness/antigravity/antigravity.go new file mode 100644 index 00000000..1185dad3 --- /dev/null +++ b/internal/harness/antigravity/antigravity.go @@ -0,0 +1,216 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravity + +import ( + "context" + "fmt" + "io" + "net" + "os" + "os/signal" + "path/filepath" + "sync" + "syscall" + + "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" + "go.opentelemetry.io/otel" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + + "github.com/google/ax/internal/config" + "github.com/google/ax/internal/harness" + "github.com/google/ax/internal/pythonsidecar" + "github.com/google/ax/proto" + "github.com/google/ax/python" + "github.com/google/uuid" +) + +// Compile-time interface assertions. +var _ harness.Harness = (*AntigravityHarness)(nil) +var _ harness.Execution = (*antigravityExecution)(nil) + +// DefaultStateDir returns the default trajectory storage directory, +// ~/.ax/antigravity/conversations. AX owns this path so it stays a harness +// implementation detail rather than a user-facing ax.yaml knob. It lives +// outside the agent's working directory on purpose: the working directory is +// the agent's operating surface (it reads and edits files there), so AX's +// internal state is kept separate to avoid the agent seeing or clobbering it. +func DefaultStateDir() (string, error) { + axDir, err := config.AXAssetsDir() + if err != nil { + return "", err + } + return filepath.Join(axDir, "antigravity", "conversations"), nil +} + +// AntigravityHarness implements the Harness interface by connecting to the +// Antigravity Python agent server over gRPC. +type AntigravityHarness struct { + address string + sidecar *pythonsidecar.Sidecar // non-nil when this harness forked the process +} + +// New creates a new AntigravityHarness. When autoStart is true, New forks the +// 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" + } + h := &AntigravityHarness{address: address} + if !autoStart { + return h, nil + } + host, port, err := net.SplitHostPort(address) + 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: args, + ReadyFunc: pythonsidecar.TCPReady(address), + } + sidecar := pythonsidecar.New(cfg) + path, err := pythonsidecar.Setup(ctx, pythonsidecar.SetupOptions{ + FS: python.FS, + }) + if err != nil { + return nil, fmt.Errorf("failed to setup antigravity harness server assets: %w", err) + } + if err := sidecar.Start(ctx, path); err != nil { + return nil, fmt.Errorf("failed to start antigravity harness server: %w", err) + } + h.sidecar = sidecar + // Own the sidecar lifecycle here (no registry->controller->antigravity + // teardown chain): on ctrl-C / SIGTERM, stop the Python process + // directly. Mirrors cmd/ax/harness.go. + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + go func() { + <-sigChan + _ = sidecar.Stop() + }() + return h, nil +} + +// Start implements Harness.Start. +func (h *AntigravityHarness) Start(ctx context.Context, conversationID string, harnessConfig []byte) (harness.Execution, error) { + return &antigravityExecution{ + harness: h, + conversationID: conversationID, + id: uuid.NewString(), + harnessConfig: harnessConfig, + }, nil +} + +// antigravityExecution implements the Execution interface. +type antigravityExecution struct { + harness *AntigravityHarness + conversationID string + id string + harnessConfig []byte + + mu sync.Mutex + queued []*proto.Message + closed bool +} + +// ID implements Execution.ID. +func (e *antigravityExecution) ID() string { + return e.id +} + +// Queue implements Execution.Queue. +func (e *antigravityExecution) Queue(ctx context.Context, msg ...*proto.Message) error { + e.mu.Lock() + defer e.mu.Unlock() + if e.closed { + return fmt.Errorf("execution session already closed") + } + e.queued = append(e.queued, msg...) + return nil +} + +// Run executes the turn over gRPC bidirectional streaming and forwards events to the handler. +func (e *antigravityExecution) Run(ctx context.Context, handler harness.Handler) error { + ctx, span := otel.Tracer("antigravity-harness").Start(ctx, "Run") + defer span.End() + + e.mu.Lock() + if e.closed { + e.mu.Unlock() + return fmt.Errorf("execution session already closed") + } + // Retrieve queued inputs. + inputs := e.queued + e.queued = nil + e.mu.Unlock() + + // 1. Connect to the gRPC server + dialOpts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())} + dialOpts = append(dialOpts, grpc.WithStatsHandler(otelgrpc.NewClientHandler())) + conn, err := grpc.DialContext(ctx, e.harness.address, dialOpts...) + if err != nil { + return fmt.Errorf("failed to connect to gRPC harness server at %s: %w", e.harness.address, err) + } + defer conn.Close() + + // 2. Create HarnessService client. + client := proto.NewHarnessServiceClient(conn) + + // 3. Build standard HarnessRequest. + start := &proto.HarnessRequest{ + ConversationId: e.conversationID, + HarnessId: "antigravity", + Type: &proto.HarnessRequest_Start{ + Start: &proto.HarnessStart{ + HarnessConfig: e.harnessConfig, + Messages: inputs, + }, + }, + } + + // 4. Call Connect to start bidirectional streaming + stream, err := client.Connect(ctx) + if err != nil { + return fmt.Errorf("failed to call gRPC HarnessService.Connect: %w", err) + } + // A server that fails before reading the start frame makes Send/CloseSend + // report io.EOF; the real status is surfaced by DrainStream's Recv below, so + // only treat non-EOF errors as send failures. + if err := stream.Send(start); err != nil && err != io.EOF { + return fmt.Errorf("failed to send harness start: %w", err) + } + if err := stream.CloseSend(); err != nil && err != io.EOF { + return fmt.Errorf("failed to close stream send direction: %w", err) + } + + // 5. Stream responses and trigger callbacks + return harness.DrainStream(ctx, stream, e.id, handler) +} + +// Close implements Execution.Close. +func (e *antigravityExecution) Close(ctx context.Context) error { + e.mu.Lock() + defer e.mu.Unlock() + e.closed = true + return nil +} diff --git a/internal/harness/antigravity/antigravity_test.go b/internal/harness/antigravity/antigravity_test.go new file mode 100644 index 00000000..c01ade27 --- /dev/null +++ b/internal/harness/antigravity/antigravity_test.go @@ -0,0 +1,124 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravity + +import ( + "bytes" + "context" + "path/filepath" + "strings" + "testing" + + "github.com/google/ax/internal/harness/harnesstest" + "github.com/google/ax/proto" +) + +var antigravityHarnessConfig = []byte(`{"system_instructions":"be terse"}`) + +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) + if err != nil { + t.Fatalf("New: %v", err) + } + + exec, err := harnessClient.Start(context.Background(), "conv-test", antigravityHarnessConfig) + if err != nil { + t.Fatalf("failed to start execution: %v", err) + } + defer exec.Close(context.Background()) + + if err := exec.Queue(context.Background(), harnesstest.UserText("Hi")); err != nil { + t.Fatalf("failed to queue message: %v", err) + } + + handler := &harnesstest.MockHandler{} + if err := exec.Run(context.Background(), handler); err != nil { + t.Fatalf("Run failed: %v", err) + } + + if !handler.IsDone() { + t.Error("expected OnComplete to be called") + } + msgs := handler.Collected() + if len(msgs) != 2 { + t.Fatalf("expected 2 messages, got %d", len(msgs)) + } + if got := msgs[0].GetContent().GetThought().GetSummary()[0].GetText().GetText(); got != "Analyzing" { + t.Errorf("expected 'Analyzing', got %q", got) + } + if got := msgs[1].GetContent().GetText().GetText(); got != "Hello world" { + t.Errorf("expected 'Hello world', got %q", got) + } + // The harness propagated the conversation id and config to the server. + convID, _, harnessConfig, _ := srv.Received() + if convID != "conv-test" { + t.Errorf("server got convID=%q, want conv-test", convID) + } + if !bytes.Equal(harnessConfig, antigravityHarnessConfig) { + t.Errorf("server got harnessConfig=%q, want %q", harnessConfig, antigravityHarnessConfig) + } +} + +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) + if err != nil { + t.Fatalf("New: %v", err) + } + + exec, _ := harnessClient.Start(context.Background(), "conv-test", antigravityHarnessConfig) + defer exec.Close(context.Background()) + + if err := exec.Queue(context.Background(), harnesstest.UserText("Hi")); err != nil { + t.Fatalf("failed to queue message: %v", err) + } + + err = exec.Run(context.Background(), &harnesstest.MockHandler{}) + if err == nil { + t.Fatal("expected error from Run(), got nil") + } + if !strings.Contains(err.Error(), "internal mock server crash") { + t.Errorf("unexpected error message: %v", err) + } +} + +// 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) + if err != nil { + t.Fatalf("New: %v", err) + } + if h.sidecar != nil { + t.Errorf("expected sidecar to be nil, got %v", h.sidecar) + } +} + +// TestDefaultStateDir returns ~/.ax/antigravity/conversations under the user's +// home directory. +func TestDefaultStateDir(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + + got, err := DefaultStateDir() + if err != nil { + t.Fatalf("DefaultStateDir: %v", err) + } + if want := filepath.Join(home, ".ax", "antigravity", "conversations"); got != want { + t.Errorf("DefaultStateDir() = %q, want %q", got, want) + } +} diff --git a/internal/harness/antigravityinteractions/antigravityinteractions.go b/internal/harness/antigravityinteractions/antigravityinteractions.go new file mode 100644 index 00000000..cb84faf8 --- /dev/null +++ b/internal/harness/antigravityinteractions/antigravityinteractions.go @@ -0,0 +1,975 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +// AntigravityInteractionsHarness drives an Antigravity agent through the Vertex +// GenAI Interactions API over HTTPS + Server-Sent Events, using the steps-based +// ("step_list") request format. It implements the Harness interface for the +// client-side ("local") environment: the agent runs server-side as the brain +// while the harness, on the client, drives the turn loop and executes every tool +// the agent yields. +// +// Tool execution (all internal to the harness): +// +// - Built-in environment tools (view_file, run_command, list_dir, ...) are +// executed against the local filesystem/shell. +// - Third-party / MCP tools are executed via a ThirdPartyExecutor, the seam +// through which the controller injects the caller's tool implementations. If +// no executor is configured, no third-party tools are advertised. +// +// Neither kind of tool call is surfaced to the caller: Run drives the whole +// interaction loop to completion (initial turn -> continuation turn -> ... -> +// final answer) and only forwards the agent's text output via Handler.OnMessage. +// Each turn after the first is chained to the previous one via the Interactions +// API's previous_interaction_id; this is an implementation detail of the loop, +// not an AX resume. +// +// Queue carries human input only -- the initial prompt and "steering" messages +// injected mid-run. It never carries tool results (the harness produces those +// itself). Queued input is drained at each interaction gap, which is the only +// place the harness can influence an otherwise atomic interaction. + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "math/rand" + "net/http" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + + "github.com/google/ax/internal/config" + "github.com/google/ax/internal/harness" + "github.com/google/ax/proto" + "github.com/google/uuid" + "golang.org/x/oauth2" + "golang.org/x/oauth2/google" +) + +const ( + // interactionsEndpoint is the public Vertex GenAI dataplane endpoint. + interactionsEndpoint = "https://aiplatform.googleapis.com" + // interactionsAPIVersion is the Interactions API version this harness targets. + interactionsAPIVersion = "v1beta1" + // cloudPlatformScope is the OAuth2 scope required to call Vertex AI. + cloudPlatformScope = "https://www.googleapis.com/auth/cloud-platform" + + // Cloud project and location are read from these standard environment + // variables (see https://github.com/google/ax#authentication). + envCloudProject = "GOOGLE_CLOUD_PROJECT" + envCloudLocation = "GOOGLE_CLOUD_LOCATION" + // defaultLocation is used when GOOGLE_CLOUD_LOCATION is unset. + defaultLocation = "global" + + // DefaultAgent is the Interactions API agent used when the harness is + // registered without an explicit agent override. + DefaultAgent = "antigravity-preview-05-2026" +) + +// Compile-time interface assertions. +var _ harness.Harness = (*AntigravityInteractionsHarness)(nil) +var _ harness.Execution = (*antigravityInteractionsExecution)(nil) + +// AntigravityInteractionsConfig configures an AntigravityInteractionsHarness. +// Use New, which fills sensible defaults. +// +// Cloud project and location come from the standard GOOGLE_CLOUD_PROJECT and +// GOOGLE_CLOUD_LOCATION environment variables. +type AntigravityInteractionsConfig struct { + // --- Required --- + + // Agent is the Interactions API agent name to run. Required: the API rejects a + // request with no agent. + Agent string + + // StateDir is the directory where each conversation's resume cursor is + // persisted, so a conversation can resume after a restart. Required: New + // returns an error if it is empty. + // + // Correctness relies on a single writer per conversation: writes are + // last-write-wins with no compare-and-swap. This is an expectation the caller + // must satisfy (e.g. by routing a conversation id to one worker, or resuming + // a worker from that conversation's snapshot) -- it is not currently enforced + // by the controller. + StateDir string + + // --- Optional --- + + // SystemInstruction, if set, is sent as the interaction's system_instruction + // (a free-form system prompt prepended to the agent's own instructions). It + // is sent on every turn of the interaction loop so it persists across them. + SystemInstruction string + // MaxTurns caps the number of interaction turns the harness will drive within + // a single Run before giving up. Defaults to 100. + MaxTurns int + // Debug, if true, logs concise per-conversation tool activity to stderr: a + // line for each function call (FC) the agent yields and each function result + // (FR) the harness produces. Useful for observing the FC/FR exchange that is + // otherwise internal to the harness. + Debug bool + // ThirdPartyExecutor executes third-party (non-built-in) function tool calls + // and declares them to the agent. It is the seam for the controller to inject + // the caller's tool implementations. If nil, the harness advertises no + // third-party tools and any non-built-in call the agent attempts yields an + // error result. + ThirdPartyExecutor ThirdPartyExecutor + // TokenSource overrides how the bearer token is obtained. If nil, the harness + // builds an auto-refreshing source from Application Default Credentials. + TokenSource oauth2.TokenSource +} + +func (c *AntigravityInteractionsConfig) withDefaults() { + if c.MaxTurns == 0 { + c.MaxTurns = 100 + } +} + +// cloudProject returns the Cloud project id from GOOGLE_CLOUD_PROJECT. +func cloudProject() string { + return os.Getenv(envCloudProject) +} + +// cloudLocation returns the Cloud location from GOOGLE_CLOUD_LOCATION, falling +// back to the default ("global"). +func cloudLocation() string { + if loc := os.Getenv(envCloudLocation); loc != "" { + return loc + } + return defaultLocation +} + +// DefaultStateDir returns the default resume-cursor directory, ~/.ax/antigravityinteractions/cursors, +// used when a caller does not set StateDir explicitly. It lives outside the +// agent's working directory on purpose: the working directory is the agent's +// operating surface (it reads and edits files there), so AX's internal state is +// kept separate to avoid the agent seeing or clobbering it. New still requires a +// non-empty StateDir; callers apply this default. +func DefaultStateDir() (string, error) { + axDir, err := config.AXAssetsDir() + if err != nil { + return "", err + } + return filepath.Join(axDir, "antigravityinteractions", "cursors"), nil +} + +// AntigravityInteractionsHarness implements Harness by talking to the public +// Vertex GenAI Interactions API. +type AntigravityInteractionsHarness struct { + cfg AntigravityInteractionsConfig + httpClient *http.Client + + // cursors persists each conversation's resume cursor to disk so a conversation + // can resume after a restart. It is always non-nil (the constructor requires a + // usable state directory). + cursors *cursorStore + + // tsOnce guards lazy initialization of ts, the resolved OAuth2 token source. + // It is resolved on first use (rather than in the constructor) so credential + // errors surface to the caller of Run instead of at construction time. + tsOnce sync.Once + ts oauth2.TokenSource + tsErr error +} + +// New creates a harness from the given config, filling in defaults for unset +// fields. It returns an error if cfg.StateDir is empty or the cursor store cannot +// be created: resume-cursor persistence is required, so a usable state directory +// must be provided. +func New(cfg AntigravityInteractionsConfig) (*AntigravityInteractionsHarness, error) { + return newWithHTTPClient(cfg, nil) +} + +// newWithHTTPClient is New with an injectable HTTP client. If hc is nil it builds +// the default client. It exists so tests can inject a fake transport; the public +// config intentionally does not expose an HTTP client override. +func newWithHTTPClient(cfg AntigravityInteractionsConfig, hc *http.Client) (*AntigravityInteractionsHarness, error) { + cfg.withDefaults() + if hc == nil { + // Disable HTTP keep-alives so no TCP/TLS connection is reused across + // requests. On substrate the actor is suspended after a turn and resumed + // for the next (with a new routable IP), which leaves any pooled keep-alive + // connection stale; reusing it makes the next turn's request fail. Opening a + // fresh connection per request avoids that at the cost of an extra handshake. + // + // Clone DefaultTransport (rather than a bare &http.Transport{}) so we inherit + // its production defaults -- proxy from the environment, dial/TLS/handshake + // timeouts, and HTTP/2 -- and only override keep-alives. + tr := http.DefaultTransport.(*http.Transport).Clone() + tr.DisableKeepAlives = true + hc = &http.Client{Timeout: 10 * time.Minute, Transport: tr} + } + if cfg.StateDir == "" { + return nil, errors.New("AntigravityInteractionsConfig.StateDir must be set") + } + cursors, err := newCursorStore(cfg.StateDir) + if err != nil { + return nil, fmt.Errorf("creating cursor store: %w", err) + } + return &AntigravityInteractionsHarness{cfg: cfg, httpClient: hc, cursors: cursors}, nil +} + +// Start implements Harness.Start. It loads any previously persisted resume +// cursor for conversationID so the returned Execution resumes the existing +// interaction chain instead of starting a new one. +func (h *AntigravityInteractionsHarness) Start(ctx context.Context, conversationID string, harnessConfig []byte) (harness.Execution, error) { + e := &antigravityInteractionsExecution{ + harness: h, + conversationID: conversationID, + id: uuid.NewString(), + harnessConfig: harnessConfig, + } + + cur, found, err := h.cursors.load(conversationID) + if err != nil { + // A real load failure must not be silently treated as "new", which would + // lose an existing conversation's history. + return nil, fmt.Errorf("loading resume cursor for %q: %w", conversationID, err) + } + if found { + // A persisted cursor is only written after a successful turn, so a + // non-empty interaction id means the conversation has already started; + // the first-turn check in Run derives that from prevInteractionID. + e.prevInteractionID = cur.PrevInteractionID + } + return e, nil +} + +// antigravityInteractionsExecution implements Execution. It is long-lived +// across Run calls and owns the interaction chain (prevInteractionID). Queue +// may be called concurrently with Run to inject steering input, which is +// drained at each interaction gap. +type antigravityInteractionsExecution struct { + harness *AntigravityInteractionsHarness + conversationID string + id string + harnessConfig []byte + + mu sync.Mutex + queued []*proto.Message + closed bool + + // prevInteractionID chains the turns of the interaction loop (the interaction + // chain this Execution owns), and is the value persisted for cross-Execution + // resume. It is empty until the first turn completes successfully; an empty + // value therefore means "no turn has succeeded yet" (the first turn). + prevInteractionID string +} + +// ID implements Execution.ID. +func (e *antigravityInteractionsExecution) ID() string { return e.id } + +// Queue implements Execution.Queue. It carries human input only: the initial +// prompt, or steering messages injected mid-run. Tool results are NOT queued by +// the caller -- the harness executes all tools itself. Queued messages are +// drained at the next interaction gap within Run. +func (e *antigravityInteractionsExecution) Queue(ctx context.Context, msg ...*proto.Message) error { + e.mu.Lock() + defer e.mu.Unlock() + if e.closed { + return fmt.Errorf("execution session already closed") + } + e.queued = append(e.queued, msg...) + return nil +} + +// Close implements Execution.Close. +func (e *antigravityInteractionsExecution) Close(ctx context.Context) error { + e.mu.Lock() + defer e.mu.Unlock() + e.closed = true + return nil +} + +// drainQueue removes and returns any queued human input, converting it to +// user_input steps. Called at each interaction gap so newly-queued steering +// messages are folded into the next interaction. +func (e *antigravityInteractionsExecution) drainQueue() []any { + e.mu.Lock() + msgs := e.queued + e.queued = nil + e.mu.Unlock() + return messagesToInputSteps(msgs) +} + +// setPrevID records the latest interaction id (in memory) and durably persists +// the resume cursor so the conversation can resume after a restart. A persistence +// failure is returned to the caller so the run can decide how to handle it rather +// than silently losing durability. +func (e *antigravityInteractionsExecution) setPrevID(ctx context.Context, id string) error { + if id == "" { + return nil + } + e.mu.Lock() + e.prevInteractionID = id + e.mu.Unlock() + + if err := e.harness.cursors.save(e.conversationID, resumeCursor{PrevInteractionID: id}); err != nil { + return fmt.Errorf("persisting resume cursor: %w", err) + } + return nil +} + +// Run implements Execution.Run. It drives the interaction to completion, +// executing every tool the agent yields internally (built-in env tools locally, +// third-party tools via the configured executor) and resuming until the agent +// produces a final answer (a turn with no tool calls) or the turn cap is hit. +// +// The agent's text output is forwarded via handler.OnMessage; handler.OnComplete +// is called once when the conversation finishes. At each interaction gap, any +// human input queued via Queue (steering) is folded into the next turn. +func (e *antigravityInteractionsExecution) Run(ctx context.Context, handler harness.Handler) error { + e.mu.Lock() + if e.closed { + e.mu.Unlock() + return fmt.Errorf("execution session already closed") + } + prevID := e.prevInteractionID + e.mu.Unlock() + + token, err := e.harness.token(ctx) + if err != nil { + return fmt.Errorf("obtaining access token: %w", err) + } + + // Initial input for this Run: drain whatever is queued (the prompt on the + // first Run, or steering input on later Runs). An empty prevID means no turn + // has completed yet, so this is the conversation's first turn. (A first turn + // that failed leaves prevID empty, so a retried Run is correctly treated as + // the first turn again.) + input := e.drainQueue() + if len(input) == 0 { + if prevID == "" { + return fmt.Errorf("no input messages queued for the initial turn") + } + return fmt.Errorf("Run called with no queued input and no work pending") + } + + // Set the working directory for this execution. Substrate resets the + // process cwd to the OCI spec's root directory, "/" after gVisor + // restore. We need to explicitly set it for each execution. + if workDir := os.Getenv("AX_HARNESS_WORKDIR"); workDir != "" { + if err := os.Chdir(workDir); err != nil { + return fmt.Errorf("set working directory %q: %w", workDir, err) + } + } + + res, err := e.harness.postTurn(ctx, token, e.harness.newRequest(input, prevID)) + if err != nil { + return fmt.Errorf("interaction turn failed: %w", err) + } + prevID = res.interactionID + if err := e.setPrevID(ctx, prevID); err != nil { + return err + } + + for turn := 0; turn < e.harness.cfg.MaxTurns; turn++ { + e.harness.debugTurn(e.conversationID, turn+1, len(res.toolCalls)) + if err := emitText(ctx, handler, e.id, res.modelText); err != nil { + return err + } + + // Execute every tool the agent yielded (built-in env tools and + // third-party tools alike are executed internally) and collect the + // results to send back. + var next []any + for i, call := range res.toolCalls { + e.harness.debugCall(e.conversationID, i+1, len(res.toolCalls), call) + out := e.harness.executeTool(ctx, call) + e.harness.debugResult(e.conversationID, call.name, out) + next = append(next, toolResultStep{ + Type: "function_result", + CallID: call.callID, + Name: call.name, + Result: out, + }) + } + + // Fold in any human input queued (steering) since the last gap. + next = append(next, e.drainQueue()...) + + // If the agent yielded no tool calls and there is no queued input to send, + // the conversation is complete. + if len(next) == 0 { + return handler.OnComplete(ctx, e.id) + } + + res, err = e.harness.postTurn(ctx, token, e.harness.newRequest(next, prevID)) + if err != nil { + return fmt.Errorf("continuation turn failed: %w", err) + } + prevID = res.interactionID + if err := e.setPrevID(ctx, prevID); err != nil { + return err + } + } + + // Hit the turn cap while still driving tools. + return handler.OnComplete(ctx, e.id) +} + +// emitText forwards non-empty model text to the handler as a Message. +func emitText(ctx context.Context, handler harness.Handler, execID, text string) error { + if strings.TrimSpace(text) == "" { + return nil + } + return handler.OnMessage(ctx, execID, &proto.Message{ + Role: "assistant", + Content: &proto.Content{ + Type: &proto.Content_Text{Text: &proto.TextContent{Text: text}}, + }, + }) +} + +// messagesToInputSteps converts queued ax Messages (human input) into user_input +// steps. Only text content is supported as input today. +func messagesToInputSteps(msgs []*proto.Message) []any { + var steps []any + for _, m := range msgs { + if t, ok := m.GetContent().GetType().(*proto.Content_Text); ok && t.Text.GetText() != "" { + steps = append(steps, userInputStep{ + Type: "user_input", + Content: []textPart{{Type: "text", Text: t.Text.GetText()}}, + }) + } + } + return steps +} + +// --------------------------------------------------------------------------- +// Wire types and helpers (steps-based Interactions API). +// --------------------------------------------------------------------------- + +// textPart is a single Content part within a user_input step. +type textPart struct { + Type string `json:"type"` + Text string `json:"text"` +} + +// userInputStep is the steps-based representation of a user prompt. +type userInputStep struct { + Type string `json:"type"` // "user_input" + Content []textPart `json:"content"` +} + +// toolResultStep is one tool result returned on a continuation turn: a flat Step with +// "type":"function_result", the matching call_id, the tool name, and the result +// object under "result". +type toolResultStep struct { + Type string `json:"type"` // "function_result" + CallID string `json:"call_id"` + Name string `json:"name"` + Result any `json:"result"` +} + +// FunctionTool is a client-declared function tool declaration. It is the element +// type a ThirdPartyExecutor returns from Declarations. +type FunctionTool struct { + Type string `json:"type"` // "function" + Name string `json:"name"` + Desc string `json:"description,omitempty"` + Params any `json:"parameters,omitempty"` +} + +// environmentConfig is the JSON shape of the Interaction.environment oneof. This +// harness always uses {"type":"local"} (the client-side environment). +type environmentConfig struct { + Type string `json:"type"` +} + +type interactionRequest struct { + Stream bool `json:"stream"` + Background bool `json:"background"` + Store bool `json:"store"` + Agent string `json:"agent"` + SystemInstruction string `json:"system_instruction,omitempty"` + Environment *environmentConfig `json:"environment,omitempty"` + PreviousInteractionID string `json:"previous_interaction_id,omitempty"` + Input []any `json:"input,omitempty"` + Tools []FunctionTool `json:"tools,omitempty"` +} + +// capturedToolCall is a tool call the agent yielded during a turn. +type capturedToolCall struct { + callID string + name string + arguments map[string]any +} + +// turnResult is what we extracted from streaming one turn. +type turnResult struct { + interactionID string + toolCalls []capturedToolCall + modelText string +} + +// Debug output layout. When Debug is enabled, each Run turn is logged as a +// hierarchical block so the function-call / function-result (FC/FR) exchange is +// easy to follow: +// +// [harness:e2e-conv] ==== turn 1 (2 function calls) ============ +// [harness:e2e-conv] FC 1/2 list_dir +// [harness:e2e-conv] DirectoryPath: . +// [harness:e2e-conv] explanation: Listing the current directory. +// [harness:e2e-conv] FR list_dir +// [harness:e2e-conv] results: [...] +// [harness:e2e-conv] FC 2/2 run_command +// [harness:e2e-conv] CommandLine: pwd +// [harness:e2e-conv] FR run_command +// [harness:e2e-conv] ExitCode: 0 +// [harness:e2e-conv] Output: /home/user/project +const ( + debugIndentCall = " " // FC / FR lines, under the turn header + debugIndentDetail = " " // params and result values, under FC / FR +) + +// debugln writes a single debug line (prefixed with the conversation id) to +// stderr when Debug is enabled. +func (h *AntigravityInteractionsHarness) debugln(conversationID, line string) { + if !h.cfg.Debug { + return + } + fmt.Fprintf(os.Stderr, "[harness:%s] %s\n", conversationID, line) +} + +// debugTurn logs the separator header that begins a turn. +func (h *AntigravityInteractionsHarness) debugTurn(conversationID string, turn, numCalls int) { + if !h.cfg.Debug { + return + } + plural := "s" + if numCalls == 1 { + plural = "" + } + header := fmt.Sprintf("==== turn %d (%d function call%s) ", turn, numCalls, plural) + // Pad the header with '=' to a fixed width for an easy visual break. + if pad := 56 - len(header); pad > 0 { + header += strings.Repeat("=", pad) + } + h.debugln(conversationID, header) +} + +// debugCall logs a function call (FC): a header line plus one indented line per +// argument (keys sorted for stable output). +func (h *AntigravityInteractionsHarness) debugCall(conversationID string, idx, total int, call capturedToolCall) { + if !h.cfg.Debug { + return + } + h.debugln(conversationID, fmt.Sprintf("%sFC %d/%d %s", debugIndentCall, idx, total, call.name)) + h.debugKeyValues(conversationID, call.arguments) +} + +// debugResult logs a function result (FR): a header line plus the indented +// result value(s). +func (h *AntigravityInteractionsHarness) debugResult(conversationID, name string, result any) { + if !h.cfg.Debug { + return + } + h.debugln(conversationID, fmt.Sprintf("%sFR %s", debugIndentCall, name)) + if m, ok := result.(map[string]any); ok { + h.debugKeyValues(conversationID, m) + return + } + h.debugln(conversationID, debugIndentDetail+debugTruncate(fmt.Sprintf("%v", result))) +} + +// debugKeyValues logs a map as one "key: value" line per entry, indented under +// an FC/FR header, with keys sorted for stable output. +func (h *AntigravityInteractionsHarness) debugKeyValues(conversationID string, m map[string]any) { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + for _, k := range keys { + h.debugln(conversationID, fmt.Sprintf("%s%s: %s", debugIndentDetail, k, debugTruncate(fmt.Sprintf("%v", m[k])))) + } +} + +// debugTruncate collapses newlines and caps very long values so a single debug +// line stays readable. +func debugTruncate(s string) string { + s = strings.ReplaceAll(s, "\n", "\\n") + const max = 200 + if len(s) > max { + return s[:max] + "…" + } + return s +} + +func (h *AntigravityInteractionsHarness) interactionsURL() string { + return fmt.Sprintf("%s/%s/projects/%s/locations/%s/interactions", + interactionsEndpoint, interactionsAPIVersion, cloudProject(), cloudLocation()) +} + +// Endpoint returns the Interactions API base endpoint this harness targets +// (e.g. the prod or autopush host). Exposed so callers/tools can log which +// backend is in use, since the endpoint is a compile-time constant. +func Endpoint() string { return interactionsEndpoint } + +// APIVersion returns the Interactions API version this harness targets. +func APIVersion() string { return interactionsAPIVersion } + +// token returns a bearer access token from the harness's OAuth2 token source. +// The source is resolved once (lazily) and auto-refreshes thereafter. +func (h *AntigravityInteractionsHarness) token(ctx context.Context) (string, error) { + h.tsOnce.Do(func() { + if h.cfg.TokenSource != nil { + h.ts = h.cfg.TokenSource + return + } + h.ts, h.tsErr = newTokenSource(ctx) + }) + if h.tsErr != nil { + return "", h.tsErr + } + tok, err := h.ts.Token() + if err != nil { + return "", fmt.Errorf("obtaining access token: %w", err) + } + return tok.AccessToken, nil +} + +// newTokenSource builds an auto-refreshing OAuth2 token source for Vertex AI +// from Application Default Credentials. +func newTokenSource(ctx context.Context) (oauth2.TokenSource, error) { + creds, err := google.FindDefaultCredentials(ctx, cloudPlatformScope) + if err != nil { + return nil, fmt.Errorf("finding application default credentials: %w", err) + } + return creds.TokenSource, nil +} + +// newRequest builds an interactionRequest common to every turn. The environment +// is always the client-side ("local") environment -- this harness exists to +// execute the agent's built-in env tools locally. Tools are re-declared on every +// turn so they stay known to the agent across the turns of the interaction loop. +func (h *AntigravityInteractionsHarness) newRequest(input []any, previousID string) interactionRequest { + var tools []FunctionTool + if h.cfg.ThirdPartyExecutor != nil { + tools = h.cfg.ThirdPartyExecutor.Declarations() + } + return interactionRequest{ + Stream: true, + Background: true, + Store: true, + Agent: h.cfg.Agent, + SystemInstruction: h.cfg.SystemInstruction, + Environment: &environmentConfig{Type: "local"}, + PreviousInteractionID: previousID, + Input: input, + Tools: tools, + } +} + +// retryableHTTPError marks a transient HTTP response that should be retried. +// retryAfter is the server-suggested delay (parsed from the Retry-After +// header), or 0 if none was provided. +type retryableHTTPError struct { + status int + retryAfter time.Duration + body string +} + +func (e *retryableHTTPError) Error() string { + return fmt.Sprintf("HTTP %d: %s", e.status, e.body) +} + +// Retry tuning for transient rate-limit (429) responses. The stateful +// interaction quota is per-minute, so the backoff must be able to span ~a +// minute; with these values the cumulative wait reaches ~1-2 minutes before +// giving up. +const ( + rateLimitMaxRetries = 6 + rateLimitBaseDelay = 2 * time.Second + rateLimitMaxDelay = 32 * time.Second +) + +// postTurn POSTs one turn, retrying on HTTP 429 (rate limit / quota) with +// exponential backoff and jitter, honoring a Retry-After header when present. +// Creating an interaction is not idempotent, so only 429 -- which is rejected +// before any interaction is created -- is retried here. +func (h *AntigravityInteractionsHarness) postTurn(ctx context.Context, token string, reqBody interactionRequest) (*turnResult, error) { + delay := rateLimitBaseDelay + for attempt := 0; ; attempt++ { + res, err := h.postTurnOnce(ctx, token, reqBody) + if err == nil { + return res, nil + } + var re *retryableHTTPError + if !errors.As(err, &re) || attempt >= rateLimitMaxRetries { + return nil, err + } + + wait := re.retryAfter + if wait <= 0 { + wait = backoffWithJitter(delay) + } + fmt.Fprintf(os.Stderr, "[harness] rate limited (HTTP %d); retrying in %s (attempt %d/%d)\n", + re.status, wait.Round(time.Millisecond), attempt+1, rateLimitMaxRetries) + + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(wait): + } + if delay *= 2; delay > rateLimitMaxDelay { + delay = rateLimitMaxDelay + } + } +} + +// backoffWithJitter returns a randomized delay in [delay/2, delay] to avoid +// synchronized retries. +func backoffWithJitter(delay time.Duration) time.Duration { + half := delay / 2 + return half + time.Duration(rand.Int63n(int64(half)+1)) +} + +// parseRetryAfter parses a Retry-After header value, which may be either an +// integer number of seconds or an HTTP date. Returns 0 if absent or unparseable. +func parseRetryAfter(v string) time.Duration { + v = strings.TrimSpace(v) + if v == "" { + return 0 + } + if secs, err := strconv.Atoi(v); err == nil { + if secs <= 0 { + return 0 + } + return time.Duration(secs) * time.Second + } + if t, err := http.ParseTime(v); err == nil { + if d := time.Until(t); d > 0 { + return d + } + } + return 0 +} + +// postTurnOnce POSTs the request and streams the SSE response for one turn. +func (h *AntigravityInteractionsHarness) postTurnOnce(ctx context.Context, token string, reqBody interactionRequest) (*turnResult, error) { + body, err := json.Marshal(reqBody) + if err != nil { + return nil, err + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, h.interactionsURL(), bytes.NewReader(body)) + if err != nil { + return nil, err + } + httpReq.Header.Set("Authorization", "Bearer "+token) + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("X-Goog-User-Project", cloudProject()) + + resp, err := h.httpClient.Do(httpReq) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + b := new(bytes.Buffer) + if _, err := b.ReadFrom(resp.Body); err != nil { + return nil, fmt.Errorf("HTTP %d (failed to read error body: %v)", resp.StatusCode, err) + } + msg := b.String() + // 429 means the request was rejected by quota before any interaction was + // created, so it is safe to retry (unlike a partial 5xx). Mark it. + if resp.StatusCode == http.StatusTooManyRequests { + return nil, &retryableHTTPError{ + status: resp.StatusCode, + retryAfter: parseRetryAfter(resp.Header.Get("Retry-After")), + body: msg, + } + } + return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, msg) + } + return h.parseStreamedTurn(resp.Body) +} + +// serverErrorMessage extracts a human-readable message from an SSE "error" +// event. The error payload is an object with a "message" field (and often a +// "status"/"code"); e.g. {"error":{"message":"...","status":"INVALID_ARGUMENT"}}. +// Falls back to a generic string when the shape is unexpected so the caller +// always gets a non-empty error. +func serverErrorMessage(event map[string]any) string { + errObj, ok := event["error"].(map[string]any) + if !ok { + return "unknown server error" + } + msg, _ := errObj["message"].(string) + status, _ := errObj["status"].(string) + switch { + case msg != "" && status != "": + return fmt.Sprintf("%s (%s)", msg, status) + case msg != "": + return msg + case status != "": + return status + default: + return "unknown server error" + } +} + +// parseStreamedTurn parses the event:/data: Server-Sent Events stream for one +// interaction turn and extracts the tool calls the agent yielded and the model's +// text output. +// +// Function-call streaming works per step index: a step.start announces the call +// (id, name, sometimes inline arguments), then step.delta "arguments_delta" +// events carry the arguments as a JSON STRING. The server re-emits the SAME call +// (same id) at several consecutive indices while it streams, and each emission +// carries the COMPLETE arguments JSON (a snapshot, not an incremental fragment). +// So we track state per step index, always REPLACING the latest arguments +// snapshot, then dedupe by id at the end -- concatenating snapshots would +// corrupt the JSON. +func (h *AntigravityInteractionsHarness) parseStreamedTurn(body io.Reader) (*turnResult, error) { + turn := &turnResult{} + + // pendingCall accumulates one function call as it streams in. + type pendingCall struct { + id string + name string + argsJSON string // latest complete arguments JSON snapshot + args map[string]any // inline arguments, if provided directly + } + callsByStepIndex := map[int]*pendingCall{} + var stepIndexOrder []int + + scanner := bufio.NewScanner(body) + scanner.Buffer(make([]byte, 64*1024), 64*1024*1024) + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "data:") { + continue + } + eventJSON := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if eventJSON == "[DONE]" { + break + } + var event map[string]any + if err := json.Unmarshal([]byte(eventJSON), &event); err != nil { + continue + } + switch event["event_type"] { + case "error": + // The server reports a turn-level failure (e.g. INVALID_ARGUMENT for a + // malformed client tool result) as an SSE "error" event. Surface it as + // a real error so the caller aborts the turn instead of treating the + // stream as a completed-but-empty turn. Silently dropping it made an + // empty directory listing look like a blank "no response, seq=0" turn + // AND caused the failing turn's interaction id to be persisted as a + // poisoned resume cursor. + return nil, fmt.Errorf( + "antigravity interactions: server error event: %s", + serverErrorMessage(event), + ) + + case "interaction.created", "interaction.completed", "interaction.status_update": + if interaction, ok := event["interaction"].(map[string]any); ok { + if id, ok := interaction["id"].(string); ok && id != "" { + turn.interactionID = id + } + } + if id, ok := event["interaction_id"].(string); ok && id != "" { + turn.interactionID = id + } + + case "step.start": + // encoding/json decodes JSON numbers as float64; convert the step + // index to int at this boundary so the maps are keyed by a plain int. + stepIndexF, _ := event["index"].(float64) + stepIndex := int(stepIndexF) + if step, ok := event["step"].(map[string]any); ok && step["type"] == "function_call" { + call := callsByStepIndex[stepIndex] + if call == nil { + call = &pendingCall{} + callsByStepIndex[stepIndex] = call + stepIndexOrder = append(stepIndexOrder, stepIndex) + } + if id, ok := step["id"].(string); ok && id != "" { + call.id = id + } + if name, ok := step["name"].(string); ok && name != "" { + call.name = name + } + if inlineArgs, ok := step["arguments"].(map[string]any); ok && len(inlineArgs) > 0 { + call.args = inlineArgs + } + } + case "step.delta": + stepIndexF, _ := event["index"].(float64) + stepIndex := int(stepIndexF) + if delta, ok := event["delta"].(map[string]any); ok { + switch delta["type"] { + case "arguments_delta": + if call := callsByStepIndex[stepIndex]; call != nil { + if argsSnapshot, ok := delta["arguments"].(string); ok && argsSnapshot != "" { + call.argsJSON = argsSnapshot // snapshot: replace, not append + } + } + case "text": + if textChunk, ok := delta["text"].(string); ok { + // Text deltas are mostly incremental, but the server + // periodically re-sends a full cumulative restatement of the + // text so far (e.g. when finalizing a turn before a tool + // call). Detect that case and replace rather than append, so + // the restated text is not duplicated. + if strings.HasPrefix(textChunk, turn.modelText) && turn.modelText != "" { + turn.modelText = textChunk + } else { + turn.modelText += textChunk + } + } + } + } + } + } + if err := scanner.Err(); err != nil { + return nil, err + } + + // Collapse the per-step-index pending calls into unique tool calls, in + // first-seen order, deduped by id. A call missing a name or id is an + // incomplete streaming artifact and is skipped. + seenCallIDs := map[string]bool{} + for _, stepIndex := range stepIndexOrder { + call := callsByStepIndex[stepIndex] + if call.name == "" || call.id == "" || seenCallIDs[call.id] { + continue + } + seenCallIDs[call.id] = true + if call.args == nil && call.argsJSON != "" { + var parsedArgs map[string]any + if err := json.Unmarshal([]byte(call.argsJSON), &parsedArgs); err == nil { + call.args = parsedArgs + } + } + if call.args == nil { + call.args = map[string]any{} + } + turn.toolCalls = append(turn.toolCalls, capturedToolCall{ + callID: call.id, name: call.name, arguments: call.args, + }) + } + return turn, nil +} diff --git a/internal/harness/antigravityinteractions/antigravityinteractions_test.go b/internal/harness/antigravityinteractions/antigravityinteractions_test.go new file mode 100644 index 00000000..fe4e17b7 --- /dev/null +++ b/internal/harness/antigravityinteractions/antigravityinteractions_test.go @@ -0,0 +1,302 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/google/ax/internal/harness/harnesstest" + "golang.org/x/oauth2" +) + +// fakeInteractions is a fake Interactions API: an http.RoundTripper that records +// the decoded request body of every POST and replies with a canned SSE stream. +// It lets the harness's real Start/Run/cursorStore code run end to end while the +// network is faked, so we can assert exactly what previous_interaction_id the +// harness sends on each turn. +type fakeInteractions struct { + mu sync.Mutex + // requests holds the decoded body of each interaction request, in order. + requests []interactionRequest + // interactionIDs are returned (in order) as the completed interaction id for + // successive turns; the i-th request gets interactionIDs[i]. + interactionIDs []string +} + +func (f *fakeInteractions) RoundTrip(req *http.Request) (*http.Response, error) { + body, err := io.ReadAll(req.Body) + if err != nil { + return nil, err + } + var decoded interactionRequest + if err := json.Unmarshal(body, &decoded); err != nil { + return nil, fmt.Errorf("fake server: decoding request: %w", err) + } + + f.mu.Lock() + idx := len(f.requests) + f.requests = append(f.requests, decoded) + id := fmt.Sprintf("INT-%d", idx+1) + if idx < len(f.interactionIDs) { + id = f.interactionIDs[idx] + } + f.mu.Unlock() + + // A minimal completed turn: no tool calls, so Run finishes in one turn. + sse := "" + + "event: interaction.created\n" + + `data: {"interaction":{"id":"` + id + `","status":"in_progress"},"event_type":"interaction.created"}` + "\n\n" + + "event: interaction.completed\n" + + `data: {"interaction":{"id":"` + id + `","status":"completed"},"event_type":"interaction.completed"}` + "\n\n" + + "event: done\n" + + "data: [DONE]\n\n" + + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(sse)), + Request: req, + }, nil +} + +func (f *fakeInteractions) recorded() []interactionRequest { + f.mu.Lock() + defer f.mu.Unlock() + return append([]interactionRequest(nil), f.requests...) +} + +// newTestHarness builds a harness wired to the fake server, a static token (no +// ADC), and the given state dir. It also sets the project env so the request URL +// and X-Goog-User-Project header are well-formed. +func newTestHarness(t *testing.T, fake *fakeInteractions, stateDir string) *AntigravityInteractionsHarness { + t.Helper() + t.Setenv(envCloudProject, "test-project") + h, err := newWithHTTPClient(AntigravityInteractionsConfig{ + Agent: "test-agent", + StateDir: stateDir, + TokenSource: oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "fake-token"}), + }, &http.Client{Transport: fake}) + if err != nil { + t.Fatalf("newWithHTTPClient: %v", err) + } + return h +} + +// runOneTurn starts an Execution for conversationID, queues a single user +// message, and runs it to completion. +func runOneTurn(t *testing.T, h *AntigravityInteractionsHarness, conversationID, prompt string) { + t.Helper() + ctx := context.Background() + exec, err := h.Start(ctx, conversationID, nil) + if err != nil { + t.Fatalf("Start(%q): %v", conversationID, err) + } + if err := exec.Queue(ctx, harnesstest.UserText(prompt)); err != nil { + t.Fatalf("Queue: %v", err) + } + if err := exec.Run(ctx, &harnesstest.MockHandler{}); err != nil { + t.Fatalf("Run: %v", err) + } + if err := exec.Close(ctx); err != nil { + t.Fatalf("Close: %v", err) + } +} + +// TestResumeAcrossRestart is the core CUJ: a first Execution starts a new +// interaction chain (empty previous_interaction_id) and persists the returned +// interaction id; then a brand-new harness over the SAME state dir (a simulated +// restart / snapshot restore) loads that cursor and sends it as +// previous_interaction_id on the next request. +func TestResumeAcrossRestart(t *testing.T) { + fake := &fakeInteractions{interactionIDs: []string{"INT-1", "INT-2"}} + stateDir := t.TempDir() + + // First Execution: starts the chain. + h1 := newTestHarness(t, fake, stateDir) + runOneTurn(t, h1, "conv-1", "hello") + + // Simulated restart: a brand-new harness over the same state dir, so any + // resumed state must come from disk, not h1's memory. + h2 := newTestHarness(t, fake, stateDir) + runOneTurn(t, h2, "conv-1", "again") + + reqs := fake.recorded() + if len(reqs) != 2 { + t.Fatalf("expected 2 requests, got %d", len(reqs)) + } + if reqs[0].PreviousInteractionID != "" { + t.Errorf("turn 1: previous_interaction_id = %q, want empty (new chain)", reqs[0].PreviousInteractionID) + } + if got, want := reqs[1].PreviousInteractionID, "INT-1"; got != want { + t.Errorf("turn 2 (after restart): previous_interaction_id = %q, want %q (resumed from persisted cursor)", got, want) + } +} + +// TestNewRequiresStateDir verifies that the constructor rejects an empty +// StateDir: resume-cursor persistence is required. +func TestNewRequiresStateDir(t *testing.T) { + t.Setenv(envCloudProject, "test-project") + _, err := New(AntigravityInteractionsConfig{ + Agent: "test-agent", + StateDir: "", // missing + TokenSource: oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "fake-token"}), + }) + if err == nil { + t.Fatal("New with empty StateDir: got nil error, want error") + } +} + +// TestSameHarnessSecondTurnResumes checks that even without a "restart", a second +// Execution on the same harness/conversation continues the chain via the cursor. +func TestSameHarnessSecondTurnResumes(t *testing.T) { + fake := &fakeInteractions{interactionIDs: []string{"INT-1", "INT-2"}} + h := newTestHarness(t, fake, t.TempDir()) + + runOneTurn(t, h, "conv-1", "hello") + runOneTurn(t, h, "conv-1", "again") + + reqs := fake.recorded() + if len(reqs) != 2 { + t.Fatalf("expected 2 requests, got %d", len(reqs)) + } + if got, want := reqs[1].PreviousInteractionID, "INT-1"; got != want { + t.Errorf("turn 2: previous_interaction_id = %q, want %q", got, want) + } +} + +// TestCursorStoreLoadSave is a focused unit test of the harness-local cursor +// store round-trip. +func TestCursorStoreLoadSave(t *testing.T) { + cs, err := newCursorStore(t.TempDir()) + if err != nil { + t.Fatalf("newCursorStore: %v", err) + } + + // Missing key: found is false, no error. + if _, found, err := cs.load("missing"); err != nil || found { + t.Fatalf("load(missing) = found=%v err=%v, want found=false err=nil", found, err) + } + + // Round-trip. + if err := cs.save("conv-1", resumeCursor{PrevInteractionID: "INT-7"}); err != nil { + t.Fatalf("save: %v", err) + } + cur, found, err := cs.load("conv-1") + if err != nil || !found { + t.Fatalf("load(conv-1) = found=%v err=%v, want found=true err=nil", found, err) + } + if cur.PrevInteractionID != "INT-7" { + t.Errorf("loaded PrevInteractionID = %q, want %q", cur.PrevInteractionID, "INT-7") + } + + // Last-write-wins overwrite. + if err := cs.save("conv-1", resumeCursor{PrevInteractionID: "INT-8"}); err != nil { + t.Fatalf("save (overwrite): %v", err) + } + cur, _, err = cs.load("conv-1") + if err != nil { + t.Fatalf("load after overwrite: %v", err) + } + if cur.PrevInteractionID != "INT-8" { + t.Errorf("after overwrite PrevInteractionID = %q, want %q", cur.PrevInteractionID, "INT-8") + } +} + +// TestDefaultStateDir returns ~/.ax/antigravityinteractions/cursors under the +// user's home directory. +func TestDefaultStateDir(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + + got, err := DefaultStateDir() + if err != nil { + t.Fatalf("DefaultStateDir: %v", err) + } + if want := filepath.Join(home, ".ax", "antigravityinteractions", "cursors"); got != want { + t.Errorf("DefaultStateDir() = %q, want %q", got, want) + } +} + +// TestParseStreamedTurnSurfacesErrorEvent verifies that a server-emitted SSE +// "error" event (e.g. INVALID_ARGUMENT for a malformed client tool result) is +// surfaced as a real error, instead of being silently dropped and returned as a +// completed-but-empty turn. Dropping it made an empty directory listing look +// like a blank "no response" turn and poisoned the resume cursor with the +// failing interaction's id. +func TestParseStreamedTurnSurfacesErrorEvent(t *testing.T) { + sse := "" + + "event: interaction.created\n" + + `data: {"interaction":{"id":"INT-1","status":"in_progress"},"event_type":"interaction.created"}` + "\n\n" + + "event: error\n" + + `data: {"event_type":"error","error":{"message":"field 'results' must be a list_value, got unset","status":"INVALID_ARGUMENT"}}` + "\n\n" + + "event: done\n" + + "data: [DONE]\n\n" + + h := &AntigravityInteractionsHarness{} + turn, err := h.parseStreamedTurn(strings.NewReader(sse)) + if err == nil { + t.Fatalf("parseStreamedTurn() = %+v, nil error; want an error for the SSE error event", turn) + } + if !strings.Contains(err.Error(), "INVALID_ARGUMENT") { + t.Errorf("error = %q, want it to include the server status INVALID_ARGUMENT", err) + } + if !strings.Contains(err.Error(), "must be a list_value") { + t.Errorf("error = %q, want it to include the server message", err) + } +} + +func TestServerErrorMessage(t *testing.T) { + tests := []struct { + name string + event map[string]any + want string + }{ + { + name: "message and status", + event: map[string]any{"error": map[string]any{"message": "boom", "status": "INVALID_ARGUMENT"}}, + want: "boom (INVALID_ARGUMENT)", + }, + { + name: "message only", + event: map[string]any{"error": map[string]any{"message": "boom"}}, + want: "boom", + }, + { + name: "status only", + event: map[string]any{"error": map[string]any{"status": "UNAVAILABLE"}}, + want: "UNAVAILABLE", + }, + { + name: "no error object", + event: map[string]any{"event_type": "error"}, + want: "unknown server error", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := serverErrorMessage(tt.event); got != tt.want { + t.Errorf("serverErrorMessage() = %q, want %q", got, tt.want) + } + }) + } +} diff --git a/internal/harness/antigravityinteractions/antigravityinteractions_tools.go b/internal/harness/antigravityinteractions/antigravityinteractions_tools.go new file mode 100644 index 00000000..ad4b979f --- /dev/null +++ b/internal/harness/antigravityinteractions/antigravityinteractions_tools.go @@ -0,0 +1,521 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +import ( + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "sort" + "strconv" + "strings" + "time" + "unicode/utf8" +) + +// executeTool runs a tool call the agent yielded and returns the result value to +// send back as a function_result step. The built-in environment tools enumerated +// in the switch below (file reads, command execution, and the file-mutation +// family) are executed internally against the local filesystem/shell; every +// other tool is dispatched to the configured ThirdPartyExecutor, or, if none is +// configured, returned as an error result. All execution is internal to the +// harness -- no tool call is surfaced to the caller. +// +// Argument names match the agent's tool schema (PascalCase). Mutation tools +// (move, delete_dir, create_file, edit_file, multi_edit_file, delete_file) +// require no success payload -- on success they return an empty result; on +// failure they return {"error": }, which marks the step as failed. +func (h *AntigravityInteractionsHarness) executeTool(ctx context.Context, call capturedToolCall) any { + switch call.name { + case "view_file": + return execViewFile(call) + case "run_command": + return execRunCommand(ctx, call) + case "list_dir", "list_directory": + return execListDir(call) + case "move": + return execMove(call) + case "delete_dir": + return execDeleteDir(call) + case "create_file": + return execCreateFile(call) + case "edit_file": + return execEditFile(call) + case "multi_edit_file": + return execMultiEditFile(call) + case "delete_file": + return execDeleteFile(call) + default: + if h.cfg.ThirdPartyExecutor == nil { + return map[string]any{"error": fmt.Sprintf("no executor configured for third-party tool %q", call.name)} + } + return h.cfg.ThirdPartyExecutor.Execute(ctx, call.name, call.arguments) + } +} + +// --------------------------------------------------------------------------- +// Built-in environment tools: real local implementations. +// +// These run against the actual local filesystem/shell because they ARE the +// client-side environment. Argument names follow the agent's tool schema +// (view_file -> AbsolutePath, run_command -> CommandLine, list_dir -> +// DirectoryPath), and result field names follow what the server maps back into +// the step's own output (content, Output/ExitCode, results). +// --------------------------------------------------------------------------- + +// view_file result caps. Mirrors the Antigravity view_file contract: +// StartLine/EndLine are 1-indexed inclusive, at most viewFileMaxLines lines are +// returned per view, and content is byte-capped with a ContentOffset +// continuation so a large file can't blob a multi-hundred-KB tool result that +// stalls the turn / blows past API context limits. +const ( + viewFileMaxLines = 2000 + viewFileMaxBytes = 256 * 1024 // 256 KiB +) + +func execViewFile(call capturedToolCall) any { + path := stringArg(call.arguments, "AbsolutePath") + if path == "" { + return map[string]any{"error": "view_file: missing required argument 'AbsolutePath'"} + } + data, err := os.ReadFile(path) + if err != nil { + return map[string]any{"error": fmt.Sprintf("view_file: %v", err)} + } + + // Resolve the 1-indexed inclusive line window using slice notation (see + // resolveLineWindow), then apply the byte cap, honoring ContentOffset as the + // read position within the windowed content. The caps bound the result size so + // a large file can't blob. + start, startSet := intArgOK(call.arguments, "StartLine") + end, endSet := intArgOK(call.arguments, "EndLine") + offset := intArg(call.arguments, "ContentOffset") + + whole := string(data) + windowed, lo, hi, totalLines := resolveLineWindow(whole, start, startSet, end, endSet) + content := applyByteWindow(windowed, offset) + + // Return the metadata the server needs to distinguish a complete read from a + // paginated/byte-truncated one: + // - start_line/end_line: 0-indexed inclusive served line range (the result + // is 0-indexed; the tool-call StartLine/EndLine args are 1-indexed). + // - content_offset: byte offset within the line-range content this slice + // starts at (for pagination continuation). + // - line_range_bytes: total bytes of the selected line range *before* byte + // truncation; the server compares content_offset+len(content) against + // this to detect byte truncation. + // - num_lines/num_bytes: whole-file totals. + result := map[string]any{ + "content": content, + "content_offset": offset, + "line_range_bytes": len(windowed), + "num_lines": totalLines, + "num_bytes": len(whole), + } + // Only report a concrete served range when there was content to serve; + // otherwise leave the (0,0) default. + if lo > 0 && hi >= lo { + result["start_line"] = lo - 1 // 1-indexed -> 0-indexed + result["end_line"] = hi - 1 + } + return result +} + +// resolveLineWindow returns the requested line window of content, following the +// Antigravity view_file slice notation over 1-indexed inclusive [start, end]: +// +// - neither set: the first viewFileMaxLines lines (or the whole file if smaller) +// - start only: viewFileMaxLines lines starting at start (forward window) +// - end only: viewFileMaxLines lines ending at end (backward window) +// - both set: lines [start, end], capped to viewFileMaxLines from start +// +// It returns the joined window text, the served 1-indexed inclusive [lo, hi] +// range (0, 0 when nothing is served), and the whole-file total line count. +func resolveLineWindow(content string, start int, startSet bool, end int, endSet bool) (windowed string, lo, hi, total int) { + lines := strings.Split(content, "\n") + // strings.Split on a trailing newline yields a final empty element; drop it so + // line counts match the file's logical lines. + if n := len(lines); n > 0 && lines[n-1] == "" { + lines = lines[:n-1] + } + total = len(lines) + if total == 0 { + return "", 0, 0, 0 + } + + // Compute a 1-indexed inclusive [lo, hi] window per slice notation. + switch { + case !startSet && !endSet: + lo, hi = 1, viewFileMaxLines + case startSet && !endSet: + lo = start + hi = start + viewFileMaxLines - 1 + case !startSet && endSet: + hi = end + lo = end - viewFileMaxLines + 1 + default: // both set + lo, hi = start, end + if hi-lo+1 > viewFileMaxLines { + hi = lo + viewFileMaxLines - 1 + } + } + + // Clamp to the file bounds. + if lo < 1 { + lo = 1 + } + if hi > total { + hi = total + } + if lo > total || hi < lo { + // Window is entirely past EOF (or inverted after clamping): nothing to show. + return "", 0, 0, total + } + + return strings.Join(lines[lo-1:hi], "\n"), lo, hi, total +} + +// applyByteWindow returns the slice of content starting at offset (the agent's +// ContentOffset), capped to viewFileMaxBytes so a large window can't blob. When +// it truncates, the cut is backed off to the last complete UTF-8 rune so the +// returned string is always valid UTF-8 (never a split multi-byte character). +// +// Callers detect truncation and the resume point from the returned slice: the +// window was truncated when offset+len(out) < len(content), and the agent +// resumes by re-reading with ContentOffset == offset+len(out). +func applyByteWindow(content string, offset int) string { + if offset < 0 { + offset = 0 + } + if offset > len(content) { + offset = len(content) + } + remaining := content[offset:] + if len(remaining) <= viewFileMaxBytes { + return remaining + } + + // Cap at viewFileMaxBytes, then back off to a valid UTF-8 boundary so we + // don't split a multi-byte rune across the cut. + cut := viewFileMaxBytes + for cut > 0 && !utf8.RuneStart(remaining[cut]) { + cut-- + } + // remaining[cut] is now the start of a rune (or cut == 0). Guard the + // pathological case of a single rune larger than the cap: emit at least that + // rune so we always make forward progress. + if cut == 0 { + _, size := utf8.DecodeRuneInString(remaining) + cut = size + } + return remaining[:cut] +} + +// runCommandTimeout bounds how long a single run_command may take, so a runaway +// command (e.g. `find /`, or `ping` without a count) cannot wedge the harness. +const runCommandTimeout = 60 * time.Second + +func execRunCommand(ctx context.Context, call capturedToolCall) any { + cmdLine := stringArg(call.arguments, "CommandLine") + if cmdLine == "" { + return map[string]any{"error": "run_command: missing required argument 'CommandLine'"} + } + + // Bound the command's runtime so it cannot hang the harness. + runCtx, cancel := context.WithTimeout(ctx, runCommandTimeout) + defer cancel() + + cmd := exec.CommandContext(runCtx, "/bin/sh", "-c", cmdLine) + if cwd := stringArg(call.arguments, "Cwd"); cwd != "" { + cmd.Dir = cwd + } + + out, err := cmd.CombinedOutput() + + // Timed out: report a clear, non-zero result rather than blocking. + if runCtx.Err() == context.DeadlineExceeded { + return map[string]any{ + "Output": fmt.Sprintf("%scommand timed out after %s", out, runCommandTimeout), + "ExitCode": 124, // conventional timeout exit code + } + } + + exitCode := 0 + if err != nil { + if ee, ok := err.(*exec.ExitError); ok { + exitCode = ee.ExitCode() + } else { + // Failed to start: surface as a non-zero exit + output. + return map[string]any{"Output": fmt.Sprintf("%s%v", out, err), "ExitCode": 1} + } + } + return map[string]any{"Output": string(out), "ExitCode": exitCode} +} + +func execListDir(call capturedToolCall) any { + dir := stringArg(call.arguments, "DirectoryPath") + if dir == "" { + return map[string]any{"error": "list_dir: missing required argument 'DirectoryPath'"} + } + entries, err := os.ReadDir(dir) + if err != nil { + return map[string]any{"error": fmt.Sprintf("list_dir: %v", err)} + } + results := make([]map[string]any, 0, len(entries)) + for _, e := range entries { + var sizeBytes int64 + if info, err := e.Info(); err == nil { + sizeBytes = info.Size() + } + results = append(results, map[string]any{ + "name": e.Name(), + "is_dir": e.IsDir(), + "size_bytes": sizeBytes, + }) + } + sort.Slice(results, func(i, j int) bool { + return results[i]["name"].(string) < results[j]["name"].(string) + }) + return map[string]any{"results": results} +} + +// --------------------------------------------------------------------------- +// Built-in file mutation tools. +// +// These mutate the real local filesystem. By contract, a successful mutation +// needs no result payload (the server renders success from the call's input +// args); a failure is reported as {"error": }, which marks the step +// ERROR. Argument names follow the agent's tool schema (PascalCase). +// --------------------------------------------------------------------------- + +// mutationOK is the empty success result for a mutation tool. +func mutationOK() any { return map[string]any{} } + +func mutationErr(tool string, err error) any { + return map[string]any{"error": fmt.Sprintf("%s: %v", tool, err)} +} + +// execMove implements the "move" tool: SourcePath -> DestinationPath. +func execMove(call capturedToolCall) any { + src := stringArg(call.arguments, "SourcePath") + dst := stringArg(call.arguments, "DestinationPath") + if src == "" || dst == "" { + return mutationErr("move", fmt.Errorf("missing required argument 'SourcePath' and/or 'DestinationPath'")) + } + if err := os.Rename(src, dst); err != nil { + return mutationErr("move", err) + } + return mutationOK() +} + +// execDeleteDir implements the "delete_dir" tool: removes DirectoryPath. Force +// allows recursive removal; otherwise only an empty directory is removed. +func execDeleteDir(call capturedToolCall) any { + dir := stringArg(call.arguments, "DirectoryPath") + if dir == "" { + return mutationErr("delete_dir", fmt.Errorf("missing required argument 'DirectoryPath'")) + } + var err error + if boolArg(call.arguments, "Force") { + err = os.RemoveAll(dir) + } else { + err = os.Remove(dir) // fails if not empty / not a dir + } + if err != nil { + return mutationErr("delete_dir", err) + } + return mutationOK() +} + +// execCreateFile implements the "create_file" tool: writes Content to TargetFile. +// Overwrite must be true to replace an existing file. +func execCreateFile(call capturedToolCall) any { + path := stringArg(call.arguments, "TargetFile") + if path == "" { + return mutationErr("create_file", fmt.Errorf("missing required argument 'TargetFile'")) + } + if !boolArg(call.arguments, "Overwrite") { + if _, err := os.Stat(path); err == nil { + return mutationErr("create_file", fmt.Errorf("file %q already exists and Overwrite is false", path)) + } + } + if dir := filepath.Dir(path); dir != "" { + if err := os.MkdirAll(dir, 0o755); err != nil { + return mutationErr("create_file", err) + } + } + if err := os.WriteFile(path, []byte(stringArg(call.arguments, "Content")), 0o644); err != nil { + return mutationErr("create_file", err) + } + return mutationOK() +} + +// execEditFile implements the "edit_file" tool: replaces the first occurrence of +// TargetContent with ReplacementContent in TargetFile. +func execEditFile(call capturedToolCall) any { + path := stringArg(call.arguments, "TargetFile") + target := stringArg(call.arguments, "TargetContent") + repl := stringArg(call.arguments, "ReplacementContent") + if path == "" || target == "" { + return mutationErr("edit_file", fmt.Errorf("missing required argument 'TargetFile' and/or 'TargetContent'")) + } + return applyReplacements(path, "edit_file", []replacement{{target: target, repl: repl}}) +} + +// execMultiEditFile implements the "multi_edit_file" tool: applies a list of +// search/replace ReplacementChunks to TargetFile, in order. +func execMultiEditFile(call capturedToolCall) any { + path := stringArg(call.arguments, "TargetFile") + if path == "" { + return mutationErr("multi_edit_file", fmt.Errorf("missing required argument 'TargetFile'")) + } + chunks, ok := call.arguments["ReplacementChunks"].([]any) + if !ok { + return mutationErr("multi_edit_file", fmt.Errorf("missing or invalid 'ReplacementChunks' array")) + } + reps := make([]replacement, 0, len(chunks)) + for _, c := range chunks { + m, ok := c.(map[string]any) + if !ok { + return mutationErr("multi_edit_file", fmt.Errorf("each ReplacementChunk must be an object")) + } + reps = append(reps, replacement{ + target: stringArg(m, "TargetContent"), + repl: stringArg(m, "ReplacementContent"), + }) + } + return applyReplacements(path, "multi_edit_file", reps) +} + +// execDeleteFile implements the "delete_file" tool: removes TargetFile. +func execDeleteFile(call capturedToolCall) any { + path := stringArg(call.arguments, "TargetFile") + if path == "" { + return mutationErr("delete_file", fmt.Errorf("missing required argument 'TargetFile'")) + } + if err := os.Remove(path); err != nil { + return mutationErr("delete_file", err) + } + return mutationOK() +} + +// replacement is a single search/replace edit. +type replacement struct { + target string + repl string +} + +// applyReplacements reads path, applies each search/replace (first occurrence, +// in order), and writes the result back. Each target must be found. +func applyReplacements(path, tool string, reps []replacement) any { + data, err := os.ReadFile(path) + if err != nil { + return mutationErr(tool, err) + } + content := string(data) + for i, r := range reps { + if r.target == "" { + return mutationErr(tool, fmt.Errorf("chunk %d has empty TargetContent", i)) + } + idx := strings.Index(content, r.target) + if idx < 0 { + return mutationErr(tool, fmt.Errorf("chunk %d: TargetContent not found in %q", i, path)) + } + content = content[:idx] + r.repl + content[idx+len(r.target):] + } + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + return mutationErr(tool, err) + } + return mutationOK() +} + +// --------------------------------------------------------------------------- +// Third-party / MCP function tools. +// +// Third-party tools are executed INTERNALLY by the harness, via a +// ThirdPartyExecutor. The executor both declares the tools (advertised to the +// agent in the request's "tools" field) and executes calls to them. +// +// The executor is the seam for the controller to inject the caller's own tool +// implementations (and their declarations). If no executor is configured, the +// harness advertises no third-party tools and reports an error result for any +// non-built-in call the agent attempts. +// --------------------------------------------------------------------------- + +// ThirdPartyExecutor declares and executes third-party (non-built-in) function +// tools. The harness owns when it is called; an implementation just needs to +// describe its tools and produce a result for a given call. +type ThirdPartyExecutor interface { + // Declarations returns the tool declarations advertised to the agent. + Declarations() []FunctionTool + // Execute runs the named tool with the given arguments and returns the result + // value to send back to the agent (wrapped into the function_result step). + Execute(ctx context.Context, name string, args map[string]any) any +} + +// --------------------------------------------------------------------------- +// Argument helpers. +// --------------------------------------------------------------------------- + +func stringArg(args map[string]any, name string) string { + if args == nil { + return "" + } + if v, ok := args[name].(string); ok { + return v + } + return "" +} + +func boolArg(args map[string]any, name string) bool { + if args == nil { + return false + } + if v, ok := args[name].(bool); ok { + return v + } + return false +} + +// intArg reads an integer argument. JSON numbers decode to float64, but tolerate +// int and numeric strings too. Returns 0 when absent or unparseable. +func intArg(args map[string]any, name string) int { + n, _ := intArgOK(args, name) + return n +} + +// intArgOK is like intArg but also reports whether the argument was present and +// parseable. Callers use the ok flag to distinguish "unset" from an explicit 0 +// (needed for view_file slice notation, where start-only vs end-only differ). +func intArgOK(args map[string]any, name string) (int, bool) { + if args == nil { + return 0, false + } + switch v := args[name].(type) { + case float64: + return int(v), true + case int: + return v, true + case int64: + return int(v), true + case string: + if n, err := strconv.Atoi(strings.TrimSpace(v)); err == nil { + return n, true + } + } + return 0, false +} diff --git a/internal/harness/antigravityinteractions/antigravityinteractions_tools_test.go b/internal/harness/antigravityinteractions/antigravityinteractions_tools_test.go new file mode 100644 index 00000000..e38f3a36 --- /dev/null +++ b/internal/harness/antigravityinteractions/antigravityinteractions_tools_test.go @@ -0,0 +1,343 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +import ( + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "unicode/utf8" +) + +// numberedLines returns "1\n2\n...\nN\n". +func numberedLines(n int) string { + var b strings.Builder + for i := 1; i <= n; i++ { + fmt.Fprintf(&b, "%d\n", i) + } + return b.String() +} + +func TestResolveLineWindow(t *testing.T) { + content := "l1\nl2\nl3\nl4\nl5\n" + tests := []struct { + name string + start int + startSet bool + end int + endSet bool + want string + }{ + {"neither set (small file)", 0, false, 0, false, "l1\nl2\nl3\nl4\nl5"}, + {"both set middle", 2, true, 4, true, "l2\nl3\nl4"}, + {"both set whole", 1, true, 5, true, "l1\nl2\nl3\nl4\nl5"}, + {"start only", 3, true, 0, false, "l3\nl4\nl5"}, + {"end only", 0, false, 3, true, "l1\nl2\nl3"}, + {"end past EOF clamps", 3, true, 999, true, "l3\nl4\nl5"}, + {"start past EOF empty", 99, true, 0, false, ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, _, _, _ := resolveLineWindow(content, tt.start, tt.startSet, tt.end, tt.endSet) + if got != tt.want { + t.Errorf("content = %q, want %q", got, tt.want) + } + }) + } +} + +func TestResolveLineWindow_StartOnlyForwardWindow(t *testing.T) { + // start only => viewFileMaxLines lines starting at start (forward window). + content := numberedLines(viewFileMaxLines + 500) + got, _, _, _ := resolveLineWindow(content, 10, true, 0, false) + lines := strings.Split(got, "\n") + if len(lines) != viewFileMaxLines { + t.Fatalf("got %d lines, want forward window of %d", len(lines), viewFileMaxLines) + } + if lines[0] != "10" { + t.Errorf("first line = %q, want %q (window starts at start)", lines[0], "10") + } +} + +func TestResolveLineWindow_EndOnlyBackwardWindow(t *testing.T) { + // end only => viewFileMaxLines lines ending at end (backward window). + total := viewFileMaxLines + 500 + content := numberedLines(total) + end := total // last line + got, _, _, _ := resolveLineWindow(content, 0, false, end, true) + lines := strings.Split(got, "\n") + if len(lines) != viewFileMaxLines { + t.Fatalf("got %d lines, want backward window of %d", len(lines), viewFileMaxLines) + } + if lines[len(lines)-1] != fmt.Sprintf("%d", end) { + t.Errorf("last line = %q, want %q (window ends at end)", lines[len(lines)-1], fmt.Sprintf("%d", end)) + } +} + +func TestResolveLineWindow_NeitherSetCapsToMaxLines(t *testing.T) { + content := numberedLines(viewFileMaxLines + 100) + got, _, _, _ := resolveLineWindow(content, 0, false, 0, false) + if n := len(strings.Split(got, "\n")); n != viewFileMaxLines { + t.Errorf("got %d lines, want cap %d", n, viewFileMaxLines) + } +} + +func TestApplyByteWindow(t *testing.T) { + // truncated reports whether applyByteWindow(content, offset)==out capped the + // window, i.e. there is still content past the returned slice. Callers derive + // this the same way from the wire fields (offset+len(content) < + // line_range_bytes). + truncated := func(content string, offset, gotLen int) bool { + return offset+gotLen < len(content) + } + + t.Run("under cap: not truncated", func(t *testing.T) { + got := applyByteWindow("hello", 0) + if got != "hello" || truncated("hello", 0, len(got)) { + t.Errorf("got %q (truncated=%v), want hello (false)", got, truncated("hello", 0, len(got))) + } + }) + t.Run("over cap truncates to viewFileMaxBytes + resumes at next", func(t *testing.T) { + big := strings.Repeat("x", viewFileMaxBytes+100) + got := applyByteWindow(big, 0) + if len(got) != viewFileMaxBytes || !truncated(big, 0, len(got)) { + t.Errorf("got len %d (truncated=%v), want %d (true)", len(got), truncated(big, 0, len(got)), viewFileMaxBytes) + } + // Resume offset is offset+len(got); the remainder is the trailing 100 bytes. + if next := len(got); next != viewFileMaxBytes { + t.Errorf("resume offset = %d, want %d", next, viewFileMaxBytes) + } + }) + t.Run("resume via offset returns remainder", func(t *testing.T) { + big := strings.Repeat("x", viewFileMaxBytes+100) + got := applyByteWindow(big, viewFileMaxBytes) + if len(got) != 100 || truncated(big, viewFileMaxBytes, len(got)) { + t.Errorf("got len %d (truncated=%v), want 100 (false)", len(got), truncated(big, viewFileMaxBytes, len(got))) + } + }) + t.Run("offset past end is empty", func(t *testing.T) { + got := applyByteWindow("abc", 999) + if got != "" { + t.Errorf("got %q, want empty", got) + } + }) + t.Run("caps at valid UTF-8 boundary (no split rune)", func(t *testing.T) { + // Fill just under the cap with ASCII, then place a 3-byte rune ("€" is + // U+20AC, 3 bytes) straddling the cap so a naive byte cut would split it. + const r = "€" // 3 bytes: E2 82 AC + pad := viewFileMaxBytes - 1 // cap lands 1 byte into the rune + content := strings.Repeat("a", pad) + r + "tail" + got := applyByteWindow(content, 0) + if !truncated(content, 0, len(got)) { + t.Fatal("expected truncation") + } + if !utf8.ValidString(got) { + t.Errorf("returned content is not valid UTF-8 (split rune): last bytes %x", got[len(got)-3:]) + } + // The rune must not be partially included: length backs off to pad. + if len(got) != pad { + t.Errorf("got len %d, want %d (backed off before the straddling rune)", len(got), pad) + } + // Resuming from offset+len(got) yields the rune intact at the front. + rest := applyByteWindow(content, len(got)) + if !strings.HasPrefix(rest, r) { + t.Errorf("resume did not start at the rune boundary: %.6q", rest) + } + }) +} + +func TestExecViewFile_HonorsRange(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "f.txt") + if err := os.WriteFile(path, []byte("a\nb\nc\nd\ne\n"), 0o644); err != nil { + t.Fatal(err) + } + res := execViewFile(capturedToolCall{arguments: map[string]any{ + "AbsolutePath": path, + "StartLine": float64(2), // JSON numbers arrive as float64 + "EndLine": float64(3), + }}) + m := res.(map[string]any) + if m["content"] != "b\nc" { + t.Errorf("content = %q, want %q", m["content"], "b\nc") + } +} + +func TestExecViewFile_LargeFileIsCapped(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "big.csv") + var b strings.Builder + for i := 0; i < 5000; i++ { + fmt.Fprintf(&b, "row-%d-with-some-padding\n", i) + } + if err := os.WriteFile(path, []byte(b.String()), 0o644); err != nil { + t.Fatal(err) + } + // No range requested (the regression: previously returned the whole file). + res := execViewFile(capturedToolCall{arguments: map[string]any{"AbsolutePath": path}}) + content := res.(map[string]any)["content"].(string) + if len(content) > viewFileMaxBytes { + t.Errorf("content = %d bytes, exceeds cap %d", len(content), viewFileMaxBytes) + } + if lines := len(strings.Split(content, "\n")); lines > viewFileMaxLines { + t.Errorf("content = %d lines, exceeds cap %d", lines, viewFileMaxLines) + } +} + +func TestExecViewFile_PaginationMetadataAndResume(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "wide.txt") + totalBytes := viewFileMaxBytes + 50 + // One long single line exceeding the byte cap. + if err := os.WriteFile(path, []byte(strings.Repeat("z", totalBytes)), 0o644); err != nil { + t.Fatal(err) + } + // First read: capped at viewFileMaxBytes; reports pagination metadata so the + // converter can detect byte truncation (content_offset+len(content) < + // line_range_bytes). + first := execViewFile(capturedToolCall{arguments: map[string]any{"AbsolutePath": path}}).(map[string]any) + if got := first["content"].(string); len(got) != viewFileMaxBytes { + t.Fatalf("first read = %d bytes, want cap %d", len(got), viewFileMaxBytes) + } + if first["content_offset"] != 0 { + t.Errorf("content_offset = %v, want 0", first["content_offset"]) + } + if first["line_range_bytes"] != totalBytes { + t.Errorf("line_range_bytes = %v, want %d (pre-truncation total)", first["line_range_bytes"], totalBytes) + } + // Byte-truncation detectable: offset + len(content) < line_range_bytes. + if got := first["content_offset"].(int) + len(first["content"].(string)); got >= first["line_range_bytes"].(int) { + t.Errorf("offset+len = %d, want < line_range_bytes %d (should look truncated)", got, first["line_range_bytes"]) + } + + // Resume from where the first slice ended: returns the remaining 50 bytes. + resumeOffset := len(first["content"].(string)) // 0 + viewFileMaxBytes + second := execViewFile(capturedToolCall{arguments: map[string]any{ + "AbsolutePath": path, + "ContentOffset": float64(resumeOffset), + }}).(map[string]any) + if got := second["content"].(string); len(got) != 50 { + t.Errorf("resume read = %d bytes, want 50", len(got)) + } + // Now complete: offset + len == line_range_bytes. + if got := second["content_offset"].(int) + len(second["content"].(string)); got != second["line_range_bytes"].(int) { + t.Errorf("resume offset+len = %d, want == line_range_bytes %d (complete)", got, second["line_range_bytes"]) + } +} + +func TestExecViewFile_MultiByteRuneStraddlesByteCap(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "wide_utf8.txt") + // One long single line: ASCII padding so the byte cap lands 1 byte into a + // 3-byte rune ("€" is U+20AC, E2 82 AC), followed by more content so the read + // is byte-truncated and must continue. + const r = "€" + pad := viewFileMaxBytes - 1 + content := strings.Repeat("a", pad) + r + strings.Repeat("b", 40) + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatal(err) + } + + // First page: capped, backed off before the straddling rune -> valid UTF-8. + first := execViewFile(capturedToolCall{arguments: map[string]any{"AbsolutePath": path}}).(map[string]any) + firstContent := first["content"].(string) + if !utf8.ValidString(firstContent) { + t.Errorf("first page is not valid UTF-8 (split rune): last bytes %x", firstContent[len(firstContent)-3:]) + } + if len(firstContent) != pad { + t.Errorf("first page = %d bytes, want %d (backed off before the straddling rune)", len(firstContent), pad) + } + // Byte-truncation detectable: content_offset+len(content) < line_range_bytes. + if got := first["content_offset"].(int) + len(firstContent); got >= first["line_range_bytes"].(int) { + t.Errorf("offset+len = %d, want < line_range_bytes %d (should look truncated)", got, first["line_range_bytes"]) + } + + // Second page: resume from where the first ended. The rune must appear intact + // at the front and the page must be valid UTF-8. + resumeOffset := first["content_offset"].(int) + len(firstContent) + second := execViewFile(capturedToolCall{arguments: map[string]any{ + "AbsolutePath": path, + "ContentOffset": float64(resumeOffset), + }}).(map[string]any) + secondContent := second["content"].(string) + if !utf8.ValidString(secondContent) { + t.Errorf("resume page is not valid UTF-8: first bytes %x", secondContent[:3]) + } + if !strings.HasPrefix(secondContent, r) { + t.Errorf("resume page did not start at the rune boundary: %.6q", secondContent) + } + // Now complete: offset + len == line_range_bytes. + if got := second["content_offset"].(int) + len(secondContent); got != second["line_range_bytes"].(int) { + t.Errorf("resume offset+len = %d, want == line_range_bytes %d (complete)", got, second["line_range_bytes"]) + } +} + +func TestExecViewFile_CompleteReadMetadata(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "small.txt") + if err := os.WriteFile(path, []byte("a\nb\nc\n"), 0o644); err != nil { + t.Fatal(err) + } + res := execViewFile(capturedToolCall{arguments: map[string]any{"AbsolutePath": path}}).(map[string]any) + // 3 lines, served whole: start_line=0, end_line=2 (0-indexed inclusive). + if res["start_line"] != 0 || res["end_line"] != 2 { + t.Errorf("served range = (%v,%v), want (0,2)", res["start_line"], res["end_line"]) + } + if res["num_lines"] != 3 { + t.Errorf("num_lines = %v, want 3", res["num_lines"]) + } + // Complete (not byte-truncated): content_offset+len(content) == line_range_bytes. + if got := res["content_offset"].(int) + len(res["content"].(string)); got != res["line_range_bytes"].(int) { + t.Errorf("offset+len = %d, want == line_range_bytes %d (complete read)", got, res["line_range_bytes"]) + } +} + +func TestExecViewFile_MissingPath(t *testing.T) { + res := execViewFile(capturedToolCall{arguments: map[string]any{}}) + m := res.(map[string]any) + if _, ok := m["error"]; !ok { + t.Errorf("expected error for missing AbsolutePath, got %+v", m) + } +} + +func TestIntArgOK(t *testing.T) { + args := map[string]any{ + "f": float64(7), + "i": 9, + "s": "12", + "z": float64(0), + "bad": "nope", + } + cases := []struct { + name string + wantN int + wantOK bool + }{ + {"f", 7, true}, + {"i", 9, true}, + {"s", 12, true}, + {"z", 0, true}, // explicit 0 is present + {"bad", 0, false}, + {"absent", 0, false}, + } + for _, c := range cases { + n, ok := intArgOK(args, c.name) + if n != c.wantN || ok != c.wantOK { + t.Errorf("intArgOK(%q) = (%d,%v), want (%d,%v)", c.name, n, ok, c.wantN, c.wantOK) + } + } +} diff --git a/internal/harness/antigravityinteractions/cursorstore.go b/internal/harness/antigravityinteractions/cursorstore.go new file mode 100644 index 00000000..9ebfc99a --- /dev/null +++ b/internal/harness/antigravityinteractions/cursorstore.go @@ -0,0 +1,119 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" +) + +// resumeCursor is the small per-conversation state persisted so a conversation +// can resume across restarts/replicas. It records the tail of the server-side +// interaction chain (PrevInteractionID). +// +// It is a struct (rather than a bare string) so it can grow to hold richer +// resume state later, e.g. partial function-call results for mid-tool-loop +// recovery, without changing the on-disk shape's identity. +type resumeCursor struct { + PrevInteractionID string `json:"prev_interaction_id"` +} + +// cursorStore is a minimal filesystem-backed store for resume cursors, local to +// the Antigravity Interactions harness. Each conversation maps to one file whose +// contents are the JSON-encoded cursor. +// +// It assumes a single writer per conversation (the controller guarantees at most +// one Execution per conversation), so writes are last-write-wins with no +// compare-and-swap. Writes are atomic (temp file + rename) so a reader never +// observes a torn value. +type cursorStore struct { + dir string +} + +// newCursorStore creates a cursorStore rooted at dir, creating it if needed. +func newCursorStore(dir string) (*cursorStore, error) { + if dir == "" { + return nil, errors.New("cursorStore dir must not be empty") + } + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, fmt.Errorf("creating cursor dir %q: %w", dir, err) + } + return &cursorStore{dir: dir}, nil +} + +// path maps a conversation id to its cursor file. The id is hashed so arbitrary +// id strings map to a safe, fixed-length filename. +func (s *cursorStore) path(conversationID string) string { + sum := sha256.Sum256([]byte(conversationID)) + return filepath.Join(s.dir, hex.EncodeToString(sum[:])+".json") +} + +// load returns the stored cursor for conversationID. found is false if no cursor +// has been persisted yet; a non-nil error means the lookup itself failed (which +// callers must not treat as "no cursor"). +func (s *cursorStore) load(conversationID string) (cur resumeCursor, found bool, err error) { + blob, err := os.ReadFile(s.path(conversationID)) + if errors.Is(err, os.ErrNotExist) { + return resumeCursor{}, false, nil + } + if err != nil { + return resumeCursor{}, false, fmt.Errorf("reading cursor: %w", err) + } + if err := json.Unmarshal(blob, &cur); err != nil { + return resumeCursor{}, false, fmt.Errorf("decoding cursor: %w", err) + } + return cur, true, nil +} + +// save durably writes the cursor for conversationID (last-write-wins). +func (s *cursorStore) save(conversationID string, cur resumeCursor) error { + blob, err := json.Marshal(cur) + if err != nil { + return fmt.Errorf("encoding cursor: %w", err) + } + return s.atomicWrite(s.path(conversationID), blob) +} + +// atomicWrite writes value to a temp file, fsyncs it, and renames it into place +// so a reader never observes a partial write. +func (s *cursorStore) atomicWrite(path string, value []byte) error { + tmp, err := os.CreateTemp(s.dir, ".tmp-*") + if err != nil { + return fmt.Errorf("creating temp file: %w", err) + } + tmpName := tmp.Name() + defer os.Remove(tmpName) // no-op if the rename succeeded + + if _, err := tmp.Write(value); err != nil { + tmp.Close() + return fmt.Errorf("writing temp file: %w", err) + } + if err := tmp.Sync(); err != nil { + tmp.Close() + return fmt.Errorf("syncing temp file: %w", err) + } + if err := tmp.Close(); err != nil { + return fmt.Errorf("closing temp file: %w", err) + } + if err := os.Rename(tmpName, path); err != nil { + return fmt.Errorf("renaming temp file into place: %w", err) + } + return nil +} diff --git a/internal/harness/antigravityinteractions/server.go b/internal/harness/antigravityinteractions/server.go new file mode 100644 index 00000000..7d77a438 --- /dev/null +++ b/internal/harness/antigravityinteractions/server.go @@ -0,0 +1,225 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +import ( + "context" + "errors" + "fmt" + "log" + "net" + "net/http" + "os" + "os/signal" + "strconv" + "sync" + "syscall" + + "github.com/google/ax/internal/harness" + "github.com/google/ax/proto" + "google.golang.org/grpc" + "google.golang.org/grpc/health" + healthpb "google.golang.org/grpc/health/grpc_health_v1" +) + +// server adapts the Antigravity Interactions harness (a Go harness.Harness) onto +// the gRPC HarnessService that the AX controller connects to. One server serves +// one actor's lifetime; substrate creates the actor per request. +type server struct { + proto.UnimplementedHarnessServiceServer + h harness.Harness +} + +// Serve builds the Antigravity Interactions harness from cfg and serves the +// HarnessService (plus gRPC health) on host:port, and an HTTP /readyz endpoint on +// readyzPort that reflects this server's own serving state. It blocks until ctx +// is cancelled or a termination signal (SIGINT/SIGTERM) is received, then shuts +// down gracefully. +func Serve(ctx context.Context, cfg AntigravityInteractionsConfig, host string, port, readyzPort int) error { + h, err := New(cfg) + if err != nil { + return fmt.Errorf("creating antigravity interactions harness: %w", err) + } + + lis, err := net.Listen("tcp", net.JoinHostPort(host, strconv.Itoa(port))) + if err != nil { + return fmt.Errorf("listen on %s:%d: %w", host, port, err) + } + + grpcServer := grpc.NewServer() + proto.RegisterHarnessServiceServer(grpcServer, &server{h: h}) + + // gRPC health: report SERVING for the default service so the controller's + // health checks pass. + hs := health.NewServer() + hs.SetServingStatus("", healthpb.HealthCheckResponse_SERVING) + healthpb.RegisterHealthServer(grpcServer, hs) + + // Go-based readiness: /readyz is OK once this server is up and serving. + readyzSrv := serveReadyz(host, readyzPort) + + // Shut down on ctx cancellation or termination signals. + ctx, stop := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM) + defer stop() + go func() { + <-ctx.Done() + log.Printf("antigravity interactions harness: shutting down") + if readyzSrv != nil { + _ = readyzSrv.Close() + } + grpcServer.GracefulStop() + }() + + log.Printf("antigravity interactions harness serving on %s:%d (readyz :%d)", host, port, readyzPort) + if err := grpcServer.Serve(lis); err != nil { + return fmt.Errorf("harness gRPC server: %w", err) + } + return nil +} + +// serveReadyz starts an HTTP server exposing /readyz that returns 200 as long as +// this process's gRPC server is running. Readiness is intrinsic to this server. +func serveReadyz(host string, readyzPort int) *http.Server { + mux := http.NewServeMux() + mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok\n")) + }) + srv := &http.Server{ + Addr: net.JoinHostPort(host, strconv.Itoa(readyzPort)), + Handler: mux, + } + go func() { + if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { + log.Printf("readyz server on %s exited: %v", srv.Addr, err) + } + }() + return srv +} + +// Connect drives one harness execution. Per the HarnessService contract, the +// client sends HarnessRequest{start} (and may send one HarnessRequest{cancel} +// mid-stream); the server streams zero or more HarnessResponse{outputs} frames +// terminated by exactly one HarnessResponse{end}. +func (s *server) Connect(stream proto.HarnessService_ConnectServer) error { + ctx, cancel := context.WithCancel(stream.Context()) + defer cancel() + + // The first frame must be a start. + req, err := stream.Recv() + if err != nil { + return err + } + start := req.GetStart() + if start == nil { + return fmt.Errorf("first HarnessRequest must be a start frame") + } + convID := req.GetConversationId() + + exec, err := s.h.Start(ctx, convID, start.GetHarnessConfig()) + if err != nil { + return sendEnd(stream, convID, proto.State_STATE_FAILED, err) + } + defer func() { _ = exec.Close(context.WithoutCancel(ctx)) }() + + if len(start.GetMessages()) > 0 { + if err := exec.Queue(ctx, start.GetMessages()...); err != nil { + return sendEnd(stream, convID, proto.State_STATE_FAILED, err) + } + } + + // Watch for a mid-stream cancel frame; it cancels the run's context. + // + // Note: mid-run steering (injecting extra human input while a turn is running) + // is intentionally NOT supported here. The execution can accept steering via + // Queue, but the HarnessRequest oneof only defines {start, cancel} -- there is + // no wire frame to carry additional input after start -- so a client cannot + // deliver steering over Connect. Multi-turn conversation is instead expressed + // as separate Connect executions that resume via the persisted cursor. Any + // non-cancel frame received here is therefore ignored. + go func() { + for { + r, err := stream.Recv() + if err != nil { + return // stream closed or errored; run ends by other means + } + if r.GetCancel() != nil { + cancel() + return + } + } + }() + + // Run the turn, forwarding each agent message as an outputs frame. + handler := &streamHandler{stream: stream, convID: convID} + if err := exec.Run(ctx, handler); err != nil { + if ctx.Err() != nil { + return sendEnd(stream, convID, proto.State_STATE_CANCELED, ctx.Err()) + } + return sendEnd(stream, convID, proto.State_STATE_FAILED, err) + } + if serr := handler.err(); serr != nil { + return serr // failed while sending an outputs frame; stream is broken + } + return sendEnd(stream, convID, proto.State_STATE_COMPLETED, nil) +} + +// streamHandler forwards each message the harness produces during a turn as a +// HarnessResponse{outputs} frame. +type streamHandler struct { + stream proto.HarnessService_ConnectServer + convID string + + mu sync.Mutex + sendErr error +} + +var _ harness.Handler = (*streamHandler)(nil) + +func (h *streamHandler) OnMessage(_ context.Context, _ string, msg *proto.Message) error { + err := h.stream.Send(&proto.HarnessResponse{ + ConversationId: h.convID, + Type: &proto.HarnessResponse_Outputs{ + Outputs: &proto.HarnessOutputs{Messages: []*proto.Message{msg}}, + }, + }) + if err != nil { + h.mu.Lock() + h.sendErr = err + h.mu.Unlock() + } + return err +} + +func (h *streamHandler) OnComplete(_ context.Context, _ string) error { return nil } + +func (h *streamHandler) err() error { + h.mu.Lock() + defer h.mu.Unlock() + return h.sendErr +} + +// sendEnd sends the single terminal HarnessResponse{end} frame. If err is +// non-nil, it is attached as the terminal Error. +func sendEnd(stream proto.HarnessService_ConnectServer, convID string, state proto.State, err error) error { + end := &proto.HarnessEnd{State: state} + if err != nil { + end.Error = &proto.Error{Description: err.Error()} + } + return stream.Send(&proto.HarnessResponse{ + ConversationId: convID, + Type: &proto.HarnessResponse_End{End: end}, + }) +} diff --git a/internal/harness/antigravityinteractions/server_test.go b/internal/harness/antigravityinteractions/server_test.go new file mode 100644 index 00000000..e0e454b2 --- /dev/null +++ b/internal/harness/antigravityinteractions/server_test.go @@ -0,0 +1,142 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +import ( + "context" + "errors" + "io" + "net" + "testing" + + "github.com/google/ax/internal/harness/harnesstest" + "github.com/google/ax/proto" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/test/bufconn" +) + +// startTestServer serves the HarnessService (backed by h) over an in-memory +// bufconn and returns a connected HarnessService client. +func startTestServer(t *testing.T, h *AntigravityInteractionsHarness) proto.HarnessServiceClient { + t.Helper() + lis := bufconn.Listen(1024 * 1024) + s := grpc.NewServer() + proto.RegisterHarnessServiceServer(s, &server{h: h}) + go func() { _ = s.Serve(lis) }() + t.Cleanup(s.Stop) + + conn, err := grpc.NewClient( + "passthrough:///bufnet", + grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) { + return lis.DialContext(ctx) + }), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + t.Fatalf("dial bufconn: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + return proto.NewHarnessServiceClient(conn) +} + +// TestConnect_StartToEnd verifies the Connect wire contract: a start frame drives +// one harness turn, and the stream terminates with exactly one HarnessEnd +// (COMPLETED) against the fake Interactions API. +func TestConnect_StartToEnd(t *testing.T) { + fake := &fakeInteractions{interactionIDs: []string{"INT-1"}} + h := newTestHarness(t, fake, t.TempDir()) + client := startTestServer(t, h) + + stream, err := client.Connect(context.Background()) + if err != nil { + t.Fatalf("Connect: %v", err) + } + if err := stream.Send(&proto.HarnessRequest{ + ConversationId: "conv-1", + HarnessId: "antigravity-interactions", + Type: &proto.HarnessRequest_Start{ + Start: &proto.HarnessStart{Messages: []*proto.Message{harnesstest.UserText("hello")}}, + }, + }); err != nil { + t.Fatalf("send start: %v", err) + } + _ = stream.CloseSend() + + var gotEnd *proto.HarnessEnd + for { + resp, err := stream.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("recv: %v", err) + } + if end := resp.GetEnd(); end != nil { + if gotEnd != nil { + t.Fatalf("received more than one HarnessEnd") + } + gotEnd = end + } + } + + if gotEnd == nil { + t.Fatal("stream ended without a HarnessEnd frame") + } + if gotEnd.GetState() != proto.State_STATE_COMPLETED { + t.Errorf("terminal state = %v, want COMPLETED (error: %v)", gotEnd.GetState(), gotEnd.GetError()) + } + + // The fake recorded the request: first turn has no previous_interaction_id. + reqs := fake.recorded() + if len(reqs) != 1 { + t.Fatalf("fake received %d requests, want 1", len(reqs)) + } + if reqs[0].PreviousInteractionID != "" { + t.Errorf("previous_interaction_id = %q, want empty", reqs[0].PreviousInteractionID) + } +} + +// TestConnect_FirstFrameMustBeStart verifies that a non-start first frame is +// rejected. +func TestConnect_FirstFrameMustBeStart(t *testing.T) { + fake := &fakeInteractions{} + h := newTestHarness(t, fake, t.TempDir()) + client := startTestServer(t, h) + + stream, err := client.Connect(context.Background()) + if err != nil { + t.Fatalf("Connect: %v", err) + } + if err := stream.Send(&proto.HarnessRequest{ + ConversationId: "conv-1", + Type: &proto.HarnessRequest_Cancel{Cancel: &proto.HarnessCancel{}}, + }); err != nil { + t.Fatalf("send cancel: %v", err) + } + _ = stream.CloseSend() + + // Draining the stream should surface an error (Connect returns it). + for { + _, err := stream.Recv() + if err == nil { + continue + } + if errors.Is(err, io.EOF) { + t.Fatal("expected an error for a non-start first frame, got clean EOF") + } + break // got the expected non-EOF error + } +} diff --git a/internal/harness/antigravityinteractions/skills.go b/internal/harness/antigravityinteractions/skills.go new file mode 100644 index 00000000..5e0d57f8 --- /dev/null +++ b/internal/harness/antigravityinteractions/skills.go @@ -0,0 +1,61 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +import ( + "fmt" + "strings" + + "github.com/google/ax/internal/skills/geminienterprise" +) + +// SkillsSystemInstruction builds a system-instruction pointer telling the agent +// where its materialized skills live and lists them. +// +// This is discovery logic specific to the Antigravity Interactions harness: it +// has no SKILLS_DIR concept, so it must be told where to find skills via its +// system instruction (its built-in file tools then read that directory). +// Harnesses that auto-discover a skills directory (e.g. the Antigravity SDK +// harness via SKILLS_DIR) do not use this. +func SkillsSystemInstruction(res geminienterprise.Result) string { + if res.Empty() { + return "" + } + var b strings.Builder + for _, w := range res.Written { + if b.Len() > 0 { + b.WriteString("\n") + } + fmt.Fprintf(&b, "Agent skills are available under %s. Available skills:", w.Dir) + for _, s := range w.Skills { + fmt.Fprintf(&b, " %s", s.SkillID) + } + b.WriteString(". Read a skill's SKILL.md before using it.") + } + return b.String() +} + +// JoinSystemInstruction combines a user-configured system instruction with an +// optional skills pointer (from SkillsSystemInstruction), dropping empties. +func JoinSystemInstruction(base, pointer string) string { + parts := make([]string, 0, 2) + if strings.TrimSpace(base) != "" { + parts = append(parts, base) + } + if strings.TrimSpace(pointer) != "" { + parts = append(parts, pointer) + } + return strings.Join(parts, "\n\n") +} diff --git a/internal/harness/antigravityinteractions/skills_test.go b/internal/harness/antigravityinteractions/skills_test.go new file mode 100644 index 00000000..151e993b --- /dev/null +++ b/internal/harness/antigravityinteractions/skills_test.go @@ -0,0 +1,68 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package antigravityinteractions + +import ( + "strings" + "testing" + + "github.com/google/ax/internal/skills/geminienterprise" +) + +func TestSkillsSystemInstruction(t *testing.T) { + t.Run("empty result yields empty pointer", func(t *testing.T) { + if got := SkillsSystemInstruction(geminienterprise.Result{}); got != "" { + t.Errorf("empty result pointer = %q, want empty", got) + } + }) + + t.Run("mentions dir and skill ids", func(t *testing.T) { + res := geminienterprise.Result{Written: []geminienterprise.Written{{ + Dir: "/workspace", + Skills: []geminienterprise.MaterializedSkill{{SkillID: "emoji"}, {SkillID: "lowercase"}}, + }}} + got := SkillsSystemInstruction(res) + if !strings.Contains(got, "/workspace") || !strings.Contains(got, "emoji") || !strings.Contains(got, "lowercase") { + t.Errorf("pointer = %q, want it to mention dir and skill ids", got) + } + }) + + t.Run("multiple registries produce multiple lines", func(t *testing.T) { + res := geminienterprise.Result{Written: []geminienterprise.Written{ + {Dir: "/a", Skills: []geminienterprise.MaterializedSkill{{SkillID: "s1"}}}, + {Dir: "/b", Skills: []geminienterprise.MaterializedSkill{{SkillID: "s2"}}}, + }} + got := SkillsSystemInstruction(res) + if !strings.Contains(got, "/a") || !strings.Contains(got, "/b") || !strings.Contains(got, "\n") { + t.Errorf("pointer = %q, want both dirs on separate lines", got) + } + }) +} + +func TestJoinSystemInstruction(t *testing.T) { + cases := []struct { + base, ptr, want string + }{ + {"", "", ""}, + {"base only", "", "base only"}, + {"", "ptr only", "ptr only"}, + {"base", "ptr", "base\n\nptr"}, + } + for _, c := range cases { + if got := JoinSystemInstruction(c.base, c.ptr); got != c.want { + t.Errorf("join(%q,%q) = %q, want %q", c.base, c.ptr, got, c.want) + } + } +} diff --git a/internal/harness/harness.go b/internal/harness/harness.go new file mode 100644 index 00000000..716f0165 --- /dev/null +++ b/internal/harness/harness.go @@ -0,0 +1,63 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package harness defines a common execution boundary and lifecycle hooks +// for interacting with agents, planners, or external gRPC services. +package harness + +import ( + "context" + + "github.com/google/ax/proto" +) + +// Handler defines the streaming event hook callbacks for an execution turn. +type Handler interface { + // OnMessage is invoked when the agent generates output content during its turn. + OnMessage(ctx context.Context, execID string, msg *proto.Message) error + + // OnComplete is invoked when the agent finishes its current execution turn. + OnComplete(ctx context.Context, execID string) error +} + +// Harness represents a service capable of starting execution sessions. +// +// Single-writer expectation: the controller must ensure that at most one +// Execution exists per conversation id at a time. Harness implementations rely +// on this invariant -- for example, a harness that durably persists +// per-conversation state may use a last-write-wins store without +// compare-and-swap, which is correct only because there is a single writer per +// conversation. +type Harness interface { + // Start initializes a new Execution session for a conversation. harnessConfig + // carries optional per-request harness configuration; it is opaque to the + // controller and interpreted by the harness implementation. + Start(ctx context.Context, conversationID string, harnessConfig []byte) (Execution, error) +} + +// Execution represents an active interactive session with an agent or planner. +type Execution interface { + // Run executes the session and streams events to the provided Handler. + // It blocks until the current turn completes or fails. + Run(ctx context.Context, handler Handler) error + + // Queue enqueues new input messages to be processed in the next turn. + Queue(ctx context.Context, msg ...*proto.Message) error + + // ID returns the unique execution session ID. + ID() string + + // Close cleanly releases all resources associated with the execution session. + Close(ctx context.Context) error +} diff --git a/internal/harness/harnesstest/harnesstest.go b/internal/harness/harnesstest/harnesstest.go new file mode 100644 index 00000000..6328b063 --- /dev/null +++ b/internal/harness/harnesstest/harnesstest.go @@ -0,0 +1,294 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package harnesstest + +// Shared in-process mocks for the harness tests: a mock Substrate Control server +// (the substrate control plane), a mock HarnessService server (the harness +// inside an actor), a recording Handler, and message builders. + +import ( + "context" + "net" + "sync" + "testing" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/google/ax/internal/harness" + "github.com/google/ax/proto" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/health" + "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" +) + +// mockControlServer is an in-process ateapipb.ControlServer that records the +// actor lifecycle calls SubstrateHarness makes and lets tests steer the +// CreateActor/ResumeActor responses. Only the three RPCs SubstrateHarness uses +// are implemented; the rest come from the embedded Unimplemented server. +type MockControlServer struct { + ateapipb.UnimplementedControlServer + + mu sync.Mutex + createCalls []string + resumeCalls []string + suspendCalls []string + + CreateErr error // returned from CreateActor when non-nil + ResumeIP string // AteomPodIp returned from ResumeActor + ResumeNilActor bool // when true, ResumeActor returns a nil Actor +} + +func (f *MockControlServer) CreateAtespace(_ context.Context, req *ateapipb.CreateAtespaceRequest) (*ateapipb.CreateAtespaceResponse, error) { + return &ateapipb.CreateAtespaceResponse{Atespace: &ateapipb.Atespace{Name: req.GetName()}}, nil +} + +func (f *MockControlServer) CreateActor(_ context.Context, req *ateapipb.CreateActorRequest) (*ateapipb.CreateActorResponse, error) { + f.mu.Lock() + f.createCalls = append(f.createCalls, req.GetActorRef().GetName()) + f.mu.Unlock() + if f.CreateErr != nil { + return nil, f.CreateErr + } + return &ateapipb.CreateActorResponse{Actor: &ateapipb.Actor{ActorId: req.GetActorRef().GetName()}}, nil +} + +func (f *MockControlServer) ResumeActor(_ context.Context, req *ateapipb.ResumeActorRequest) (*ateapipb.ResumeActorResponse, error) { + f.mu.Lock() + f.resumeCalls = append(f.resumeCalls, req.GetActorRef().GetName()) + f.mu.Unlock() + if f.ResumeNilActor { + return &ateapipb.ResumeActorResponse{}, nil + } + return &ateapipb.ResumeActorResponse{Actor: &ateapipb.Actor{ActorId: req.GetActorRef().GetName(), AteomPodIp: f.ResumeIP}}, nil +} + +func (f *MockControlServer) SuspendActor(_ context.Context, req *ateapipb.SuspendActorRequest) (*ateapipb.SuspendActorResponse, error) { + f.mu.Lock() + f.suspendCalls = append(f.suspendCalls, req.GetActorRef().GetName()) + f.mu.Unlock() + return &ateapipb.SuspendActorResponse{}, nil +} + +// Calls returns copies of the recorded call lists. +func (f *MockControlServer) Calls() (create, resume, suspend []string) { + f.mu.Lock() + defer f.mu.Unlock() + return append([]string(nil), f.createCalls...), + append([]string(nil), f.resumeCalls...), + append([]string(nil), f.suspendCalls...) +} + +// mockHarnessServer is an in-process proto.HarnessServiceServer standing in for +// the harness running inside an actor (substrate) or a local subprocess +// (antigravity). It records the start frame and emits its configured outputs +// followed by a terminal HarnessEnd. +type MockHarnessServer struct { + proto.UnimplementedHarnessServiceServer + + // Outputs are the messages emitted (in a single Outputs frame) before the + // terminal HarnessEnd. When nil, each input is echoed as "ack: ". + Outputs []*proto.Message + // FailConnect makes Connect return an RPC error before any frame. + FailConnect bool + // FailFrame makes Connect terminate the turn with HarnessEnd{STATE_FAILED}. + FailFrame bool + // ErrCode is the error code used by FailFrame. + ErrCode int32 + // ErrMessage is the error text used by FailConnect/FailFrame. + ErrMessage string + + mu sync.Mutex + gotConvID string + gotHarnessID string + gotHarnessConfig []byte + gotInputs []string +} + +func (s *MockHarnessServer) Connect(stream proto.HarnessService_ConnectServer) error { + if s.FailConnect { + return status.Error(codes.Internal, s.ErrMessage) + } + + req, err := stream.Recv() + if err != nil { + return err + } + + var inputs []string + for _, m := range req.GetStart().GetMessages() { + if text := m.GetContent().GetText().GetText(); text != "" { + inputs = append(inputs, text) + } + } + s.mu.Lock() + s.gotConvID = req.GetConversationId() + s.gotHarnessID = req.GetHarnessId() + s.gotHarnessConfig = req.GetStart().GetHarnessConfig() + s.gotInputs = inputs + s.mu.Unlock() + + convID := req.GetConversationId() + if s.FailFrame { + return stream.Send(&proto.HarnessResponse{ + ConversationId: convID, + Type: &proto.HarnessResponse_End{ + End: &proto.HarnessEnd{ + State: proto.State_STATE_FAILED, + Error: &proto.Error{ + Code: s.ErrCode, + Description: s.ErrMessage, + }, + }, + }, + }) + } + + msgs := s.Outputs + if msgs == nil { + for _, in := range inputs { + msgs = append(msgs, AssistantText("ack: "+in)) + } + } + if len(msgs) > 0 { + if err := stream.Send(&proto.HarnessResponse{ + ConversationId: convID, + Type: &proto.HarnessResponse_Outputs{ + Outputs: &proto.HarnessOutputs{Messages: msgs}, + }, + }); err != nil { + return err + } + } + return stream.Send(&proto.HarnessResponse{ + ConversationId: convID, + Type: &proto.HarnessResponse_End{End: &proto.HarnessEnd{State: proto.State_STATE_COMPLETED}}, + }) +} + +// Received returns a copy of the start frame the server received. +func (s *MockHarnessServer) Received() (convID, harnessID string, harnessConfig []byte, inputs []string) { + s.mu.Lock() + defer s.mu.Unlock() + return s.gotConvID, s.gotHarnessID, append([]byte(nil), s.gotHarnessConfig...), append([]string(nil), s.gotInputs...) +} + +// mockHandler records the messages and completion streamed during a turn. +type MockHandler struct { + mu sync.Mutex + messages []*proto.Message + complete bool +} + +var _ harness.Handler = (*MockHandler)(nil) + +func (h *MockHandler) OnMessage(_ context.Context, _ string, msg *proto.Message) error { + h.mu.Lock() + defer h.mu.Unlock() + h.messages = append(h.messages, msg) + return nil +} + +func (h *MockHandler) OnComplete(_ context.Context, _ string) error { + h.mu.Lock() + defer h.mu.Unlock() + h.complete = true + return nil +} + +func (h *MockHandler) IsDone() bool { + h.mu.Lock() + defer h.mu.Unlock() + return h.complete +} + +// Collected returns a copy of the messages received via OnMessage. +func (h *MockHandler) Collected() []*proto.Message { + h.mu.Lock() + defer h.mu.Unlock() + return append([]*proto.Message(nil), h.messages...) +} + +// Texts returns the text content of each received message, in order. +func (h *MockHandler) Texts() []string { + h.mu.Lock() + defer h.mu.Unlock() + var out []string + for _, m := range h.messages { + out = append(out, m.GetContent().GetText().GetText()) + } + return out +} + +func AssistantText(text string) *proto.Message { + return &proto.Message{ + Role: "assistant", + Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: text}}}, + } +} + +func UserText(text string) *proto.Message { + return &proto.Message{ + Role: "user", + Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: text}}}, + } +} + +func ThoughtText(summary string) *proto.Message { + return &proto.Message{ + Role: "model", + Content: &proto.Content{ + Type: &proto.Content_Thought{ + Thought: &proto.ThoughtContent{ + Summary: []*proto.ThoughtSummaryContent{ + {Type: &proto.ThoughtSummaryContent_Text{Text: &proto.TextContent{Text: summary}}}, + }, + }, + }, + }, + } +} + +// StartHarnessServer starts a HarnessService + health server (status SERVING) +// on a random local port and returns its address. +func StartHarnessServer(t *testing.T, srv *MockHarnessServer) string { + t.Helper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + s := grpc.NewServer() + proto.RegisterHarnessServiceServer(s, srv) + hs := health.NewServer() + hs.SetServingStatus("", grpc_health_v1.HealthCheckResponse_SERVING) + grpc_health_v1.RegisterHealthServer(s, hs) + go func() { _ = s.Serve(lis) }() + t.Cleanup(s.Stop) + return lis.Addr().String() +} + +// StartControlServer starts a mock Substrate Control server on a random local port. +func StartControlServer(t *testing.T, srv *MockControlServer) string { + t.Helper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + s := grpc.NewServer() + ateapipb.RegisterControlServer(s, srv) + go func() { _ = s.Serve(lis) }() + t.Cleanup(s.Stop) + return lis.Addr().String() +} diff --git a/internal/harness/stream.go b/internal/harness/stream.go new file mode 100644 index 00000000..7551da85 --- /dev/null +++ b/internal/harness/stream.go @@ -0,0 +1,68 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package harness + +import ( + "context" + "fmt" + "io" + + "github.com/google/ax/proto" +) + +// DrainStream reads from the harness gRPC stream until io.EOF, dispatching messages +// to the handler, and returns the final execution status. +func DrainStream(ctx context.Context, stream proto.HarnessService_ConnectClient, execID string, handler Handler) error { + var endState proto.State + var endErr error + hasEnd := false + + for { + resp, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + return fmt.Errorf("gRPC harness streaming failure: %w", err) + } + + switch payload := resp.Type.(type) { + case *proto.HarnessResponse_Outputs: + for _, outMsg := range payload.Outputs.Messages { + if err := handler.OnMessage(ctx, execID, outMsg); err != nil { + return fmt.Errorf("failed to dispatch streamed output: %w", err) + } + } + case *proto.HarnessResponse_End: + hasEnd = true + endState = payload.End.GetState() + if endState == proto.State_STATE_FAILED { + if errDetail := payload.End.GetError(); errDetail != nil { + endErr = fmt.Errorf("harness failed: [%d] %s", errDetail.GetCode(), errDetail.GetDescription()) + } else { + endErr = fmt.Errorf("harness failed with no error details") + } + } + } + } + + if !hasEnd { + return fmt.Errorf("harness stream ended without HarnessEnd frame") + } + if endState == proto.State_STATE_FAILED { + return endErr + } + return handler.OnComplete(ctx, execID) +} diff --git a/internal/harness/substrate/substrate.go b/internal/harness/substrate/substrate.go new file mode 100644 index 00000000..b1e84030 --- /dev/null +++ b/internal/harness/substrate/substrate.go @@ -0,0 +1,256 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package substrate + +import ( + "context" + "crypto/tls" + "errors" + "fmt" + "io" + "log/slog" + "sync" + "time" + + "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" + "go.opentelemetry.io/otel" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" + + "github.com/google/ax/internal/harness" + "github.com/google/ax/internal/k8s/ate" + "github.com/google/ax/proto" + "github.com/google/uuid" +) + +// Compile-time interface assertions. +var _ harness.Harness = (*SubstrateHarness)(nil) +var _ harness.Execution = (*substrateExecution)(nil) + +// healthCheckTimeout defines the maximum time Start waits for a freshly +// created/resumed actor's harness to become reachable and ready. +const healthCheckTimeout = 60 * time.Second + +// SubstrateHarness manages execution in a SubstrATE sandboxed actor over gRPC HarnessService. +type SubstrateHarness struct { + harnessID string + ateClient *ate.Client + port int + dialOpts []grpc.DialOption +} + +// New creates a new SubstrateHarness. +func New(harnessID string, endpoint string, namespace string, template string, port int, opts ...grpc.DialOption) (*SubstrateHarness, error) { + if port == 0 { + port = 50053 // Default HarnessService port + } + if namespace == "" { + namespace = "ax" + } + if template == "" { + template = "ax-harness-antigravity-template" + } + controlCreds := grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{InsecureSkipVerify: true})) + client, err := ate.NewClient(namespace, template, endpoint, controlCreds) + if err != nil { + return nil, fmt.Errorf("failed to create ATE client: %w", err) + } + if len(opts) == 0 { + opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials())) + } + opts = append(opts, grpc.WithStatsHandler(otelgrpc.NewClientHandler())) + return &SubstrateHarness{ + harnessID: harnessID, + ateClient: client, + port: port, + dialOpts: opts, + }, nil +} + +// Start implements Harness interface. It creates/resumes the target actor. +func (h *SubstrateHarness) Start(ctx context.Context, conversationID string, harnessConfig []byte) (harness.Execution, error) { + if conversationID == "" { + return nil, errors.New("SubstrateHarness needs valid conversationID") + } + + // CreateActor is idempotent here: on follow-up turns the actor was created + // (and suspended) on a previous turn, so AlreadyExists is expected and fine. + if _, err := h.ateClient.CreateActor(ctx, conversationID); err != nil && status.Code(err) != codes.AlreadyExists { + return nil, fmt.Errorf("failed to create substrate actor %s: %w", conversationID, err) + } + + // Resume the actor so it is scheduled onto a worker and gets a routable IP. + resumeResp, err := h.ateClient.ResumeActor(ctx, conversationID) + if err != nil { + return nil, fmt.Errorf("failed to resume substrate actor %s: %w", conversationID, err) + } + actor := resumeResp.Actor + if actor == nil { + return nil, fmt.Errorf("received nil actor in response for %s", conversationID) + } + if actor.AteomPodIp == "" { + return nil, fmt.Errorf("actor %s has no active worker IP address", conversationID) + } + + // Establish connection to the actor's worker IP + workerAddr := fmt.Sprintf("%s:%d", actor.AteomPodIp, h.port) + conn, err := grpc.NewClient(workerAddr, h.dialOpts...) + if err != nil { + return nil, fmt.Errorf("failed to dial remote harness service at %s: %w", workerAddr, err) + } + + // Wait for the harness to be reachable and ready before handing back the + // execution. + if err := waitForHealthy(ctx, conn, healthCheckTimeout); err != nil { + conn.Close() + return nil, fmt.Errorf("harness for %s not ready at %s: %w", conversationID, workerAddr, err) + } + + return &substrateExecution{ + harness: h, + conversationID: conversationID, + execID: uuid.NewString(), + conn: conn, + client: proto.NewHarnessServiceClient(conn), + harnessConfig: harnessConfig, + }, nil +} + +// waitForHealthy blocks until the harness behind conn reports SERVING via the +// standard gRPC health protocol until timeout. A harness that is reachable +// but does not implement the health service (Unimplemented) is treated as +// ready; connection failures (Unavailable) and NOT_SERVING are retried. +func waitForHealthy(ctx context.Context, conn *grpc.ClientConn, timeout time.Duration) error { + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + client := grpc_health_v1.NewHealthClient(conn) + const maxBackoff = 2 * time.Second + backoff := 100 * time.Millisecond + for { + resp, err := client.Check(ctx, &grpc_health_v1.HealthCheckRequest{Service: ""}) + if err == nil && resp.GetStatus() == grpc_health_v1.HealthCheckResponse_SERVING { + return nil + } + if status.Code(err) == codes.Unimplemented { + // Reachable but no health service: the port is up, proceed. + return nil + } + + select { + case <-ctx.Done(): + if err != nil { + return fmt.Errorf("harness not healthy within %s: %w", timeout, err) + } + return fmt.Errorf("harness not healthy within %s (last status: %s)", timeout, resp.GetStatus()) + case <-time.After(backoff): + } + backoff *= 2 + if backoff > maxBackoff { + backoff = maxBackoff + } + } +} + +type substrateExecution struct { + harness *SubstrateHarness + conversationID string + execID string + conn *grpc.ClientConn + client proto.HarnessServiceClient + harnessConfig []byte + + mu sync.Mutex + pending []*proto.Message +} + +func (e *substrateExecution) ID() string { + return e.execID +} + +func (e *substrateExecution) Queue(ctx context.Context, msg ...*proto.Message) error { + e.mu.Lock() + defer e.mu.Unlock() + e.pending = append(e.pending, msg...) + return nil +} + +func (e *substrateExecution) Run(ctx context.Context, handler harness.Handler) error { + ctx, span := otel.Tracer("substrate-harness").Start(ctx, "Run") + defer span.End() + + e.mu.Lock() + inputs := e.pending + e.pending = nil + e.mu.Unlock() + + stream, err := e.client.Connect(ctx) + if err != nil { + return fmt.Errorf("failed to open harness service stream: %w", err) + } + + // Send a HarnessRequest to initiate the turn. + start := &proto.HarnessRequest{ + ConversationId: e.conversationID, + HarnessId: e.harness.harnessID, + Type: &proto.HarnessRequest_Start{ + Start: &proto.HarnessStart{ + HarnessConfig: e.harnessConfig, + Messages: inputs, + }, + }, + } + // A server that fails before reading the start frame makes Send/CloseSend + // report io.EOF; the real status is surfaced by DrainStream's Recv below, so + // only treat non-EOF errors as send failures. + if err := stream.Send(start); err != nil && err != io.EOF { + return fmt.Errorf("failed to send harness start: %w", err) + } + + // Close send direction to trigger server processing. + if err := stream.CloseSend(); err != nil && err != io.EOF { + return fmt.Errorf("failed to close stream send direction: %w", err) + } + + // Drain HarnessResponse frames until the terminal HarnessEnd. + return harness.DrainStream(ctx, stream, e.execID, handler) +} + +func (e *substrateExecution) Close(ctx context.Context) error { + // Close connection + if e.conn != nil { + e.conn.Close() + } + + // Suspend actor to return resource to standard standby pool + slog.InfoContext(ctx, "Suspending SubstrATE actor", + slog.String("conversation_id", e.conversationID), + slog.String("exec_id", e.execID), + ) + suspendCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if _, err := e.harness.ateClient.SuspendActor(suspendCtx, e.conversationID); err != nil { + slog.ErrorContext(ctx, "Failed to suspend SubstrATE actor", + slog.String("conversation_id", e.conversationID), + slog.Any("error", err), + ) + } + + return nil +} diff --git a/internal/harness/substrate/substrate_test.go b/internal/harness/substrate/substrate_test.go new file mode 100644 index 00000000..111e0cc3 --- /dev/null +++ b/internal/harness/substrate/substrate_test.go @@ -0,0 +1,290 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package substrate + +import ( + "bytes" + "context" + "net" + "slices" + "strconv" + "strings" + "testing" + "time" + + "github.com/google/ax/internal/harness/harnesstest" + "github.com/google/ax/internal/k8s/ate" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/health" + "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" +) + +var substrateHarnessConfig = []byte(`{"model":"gemini-2.5-pro"}`) + +// startHealthTestServer starts a gRPC server on a random local port. If hs is +// non-nil the standard health service is registered. Returns the listen address. +func startHealthTestServer(t *testing.T, hs *health.Server) string { + t.Helper() + lis, err := net.Listen("tcp", "localhost:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + s := grpc.NewServer() + if hs != nil { + grpc_health_v1.RegisterHealthServer(s, hs) + } + go func() { + _ = s.Serve(lis) + }() + t.Cleanup(s.Stop) + return lis.Addr().String() +} + +func dialTestConn(t *testing.T, addr string) *grpc.ClientConn { + t.Helper() + conn, err := grpc.NewClient(addr, grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("failed to dial %s: %v", addr, err) + } + t.Cleanup(func() { conn.Close() }) + return conn +} + +func TestWaitForHealthy_Serving(t *testing.T) { + hs := health.NewServer() + hs.SetServingStatus("", grpc_health_v1.HealthCheckResponse_SERVING) + conn := dialTestConn(t, startHealthTestServer(t, hs)) + + if err := waitForHealthy(context.Background(), conn, 5*time.Second); err != nil { + t.Fatalf("expected healthy, got %v", err) + } +} + +func TestWaitForHealthy_UnimplementedProceeds(t *testing.T) { + // Server is up but does not register the health service. + conn := dialTestConn(t, startHealthTestServer(t, nil)) + + if err := waitForHealthy(context.Background(), conn, 5*time.Second); err != nil { + t.Fatalf("expected to proceed when health is unimplemented, got %v", err) + } +} + +func TestWaitForHealthy_TimesOut(t *testing.T) { + hs := health.NewServer() + hs.SetServingStatus("", grpc_health_v1.HealthCheckResponse_NOT_SERVING) + conn := dialTestConn(t, startHealthTestServer(t, hs)) + + if err := waitForHealthy(context.Background(), conn, 500*time.Millisecond); err == nil { + t.Fatal("expected timeout error while NOT_SERVING, got nil") + } +} + +func TestWaitForHealthy_StatusChange(t *testing.T) { + hs := health.NewServer() + hs.SetServingStatus("", grpc_health_v1.HealthCheckResponse_NOT_SERVING) + conn := dialTestConn(t, startHealthTestServer(t, hs)) + + go func() { + time.Sleep(150 * time.Millisecond) + hs.SetServingStatus("", grpc_health_v1.HealthCheckResponse_SERVING) + }() + + if err := waitForHealthy(context.Background(), conn, 5*time.Second); err != nil { + t.Fatalf("expected healthy after status flip, got %v", err) + } +} + +func TestWaitForHealthy_ServerDown(t *testing.T) { + // Reserve a port then release it so nothing is listening there. + lis, err := net.Listen("tcp", "localhost:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + addr := lis.Addr().String() + lis.Close() + conn := dialTestConn(t, addr) + + if err := waitForHealthy(context.Background(), conn, 500*time.Millisecond); err == nil { + t.Fatal("expected timeout error when server is down, got nil") + } +} + +// newTestSubstrateHarness builds a SubstrateHarness wired to the mock control +// server and the mock harness server. It constructs the struct directly (rather +// than via NewSubstrateHarness) so the control client can use insecure +// credentials instead of the TLS that NewSubstrateHarness hard-codes. +func newTestSubstrateHarness(t *testing.T, ctrlAddr, harnessAddr string) *SubstrateHarness { + t.Helper() + _, portStr, err := net.SplitHostPort(harnessAddr) + if err != nil { + t.Fatalf("bad harness addr %q: %v", harnessAddr, err) + } + port, err := strconv.Atoi(portStr) + if err != nil { + t.Fatalf("bad harness port %q: %v", portStr, err) + } + client, err := ate.NewClient("ax", "antigravity-template", ctrlAddr, + grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("failed to create ate client: %v", err) + } + return &SubstrateHarness{ + harnessID: "antigravity", + ateClient: client, + port: port, + dialOpts: []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}, + } +} + +// Test full SubstrateHarness Start -> Run -> Close flow against the shared +// in-process mocks (see mocks_test.go). +// They lock in the wiring that a substrate bump or an ax-side change could silently +// break: create/resume idempotency, worker-IP extraction, the health gate, the +// Connect streaming protocol, and suspend-on-close. +func TestSubstrateHarness_EndToEnd(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"} + srv := &harnesstest.MockHarnessServer{} + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv)) + + ctx := context.Background() + exec, err := h.Start(ctx, "conv-1", substrateHarnessConfig) + if err != nil { + t.Fatalf("Start: %v", err) + } + if err := exec.Queue(ctx, harnesstest.UserText("hi")); err != nil { + t.Fatalf("Queue: %v", err) + } + handler := &harnesstest.MockHandler{} + if err := exec.Run(ctx, handler); err != nil { + t.Fatalf("Run: %v", err) + } + + // The harness server received the start frame with the right identifiers. + convID, harnessID, harnessConfig, inputs := srv.Received() + if convID != "conv-1" || harnessID != "antigravity" { + t.Errorf("server got convID=%q harnessID=%q, want conv-1/antigravity", convID, harnessID) + } + if !bytes.Equal(harnessConfig, substrateHarnessConfig) { + t.Errorf("server got harnessConfig=%q, want %q", harnessConfig, substrateHarnessConfig) + } + if !slices.Equal(inputs, []string{"hi"}) { + t.Errorf("server got inputs=%v, want [hi]", inputs) + } + + // The handler streamed the output and completed. + if !handler.IsDone() { + t.Error("handler did not complete") + } + if got := handler.Texts(); !slices.Equal(got, []string{"ack: hi"}) { + t.Errorf("handler messages=%v, want [ack: hi]", got) + } + + // CreateActor then ResumeActor ran for the conversation; no suspend yet. + create, resume, suspend := ctrl.Calls() + want := []string{"conv-1"} + if !slices.Equal(create, want) || !slices.Equal(resume, want) { + t.Errorf("create=%v resume=%v, want %v each", create, resume, want) + } + if len(suspend) != 0 { + t.Errorf("suspend called before Close: %v", suspend) + } + + // Close suspends the actor. + if err := exec.Close(ctx); err != nil { + t.Fatalf("Close: %v", err) + } + if _, _, suspend = ctrl.Calls(); !slices.Equal(suspend, want) { + t.Errorf("suspend=%v, want %v", suspend, want) + } +} + +func TestSubstrateHarness_CreateAlreadyExistsTolerated(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ + ResumeIP: "127.0.0.1", + CreateErr: status.Error(codes.AlreadyExists, "exists"), + } + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, &harnesstest.MockHarnessServer{})) + + ctx := context.Background() + exec, err := h.Start(ctx, "conv-1", substrateHarnessConfig) + if err != nil { + t.Fatalf("Start should tolerate AlreadyExists: %v", err) + } + t.Cleanup(func() { _ = exec.Close(ctx) }) + + if err := exec.Queue(ctx, harnesstest.UserText("hi")); err != nil { + t.Fatalf("Queue: %v", err) + } + handler := &harnesstest.MockHandler{} + if err := exec.Run(ctx, handler); err != nil { + t.Fatalf("Run: %v", err) + } + if !handler.IsDone() { + t.Error("handler did not complete") + } + if _, resume, _ := ctrl.Calls(); !slices.Equal(resume, []string{"conv-1"}) { + t.Errorf("resume=%v, want [conv-1]", resume) + } +} + +func TestSubstrateHarness_ResumeNoWorkerIP(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: ""} // empty AteomPodIp + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, &harnesstest.MockHarnessServer{})) + + _, err := h.Start(context.Background(), "conv-1", substrateHarnessConfig) + if err == nil { + t.Fatal("expected error for empty worker IP, got nil") + } + if !strings.Contains(err.Error(), "no active worker IP") { + t.Errorf("error = %v, want it to mention 'no active worker IP'", err) + } +} + +func TestSubstrateHarness_ResumeNilActor(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeNilActor: true} + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, &harnesstest.MockHarnessServer{})) + + _, err := h.Start(context.Background(), "conv-1", substrateHarnessConfig) + if err == nil { + t.Fatal("expected error for nil actor, got nil") + } + if !strings.Contains(err.Error(), "nil actor") { + t.Errorf("error = %v, want it to mention 'nil actor'", err) + } +} + +func TestSubstrateHarness_HarnessFailedFrame(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"} + srv := &harnesstest.MockHarnessServer{FailFrame: true, ErrCode: 13, ErrMessage: "boom"} + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv)) + + ctx := context.Background() + exec, err := h.Start(ctx, "conv-1", substrateHarnessConfig) + if err != nil { + t.Fatalf("Start: %v", err) + } + t.Cleanup(func() { _ = exec.Close(ctx) }) + if err := exec.Queue(ctx, harnesstest.UserText("hi")); err != nil { + t.Fatalf("Queue: %v", err) + } + if err := exec.Run(ctx, &harnesstest.MockHandler{}); err == nil { + t.Fatal("expected error from failed harness frame, got nil") + } else if !strings.Contains(err.Error(), "harness failed") { + t.Errorf("error = %v, want it to mention 'harness failed'", err) + } +} diff --git a/internal/historyutil/confirmations.go b/internal/historyutil/confirmations.go new file mode 100644 index 00000000..65730006 --- /dev/null +++ b/internal/historyutil/confirmations.go @@ -0,0 +1,66 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package historyutil + +import "github.com/google/ax/proto" + +// WaitsForConfirmation returns true if the last message in the history +// is a confirmation question waiting for user input. +func WaitsForConfirmation(history []*proto.Message) *proto.Message { + if len(history) == 0 { + return nil + } + last := history[len(history)-1] + if last.GetContent().GetConfirmation() != nil && last.GetContent().GetConfirmation().Question != "" { + return last + } + return nil +} + +// HasConfirmationAnswer returns true if the last message in the history +// is a confirmation question waiting for user input. +func HasConfirmationAnswer(history []*proto.Message) (approved bool, conf *proto.ConfirmationContent) { + if len(history) == 0 { + return false, nil + } + last := history[len(history)-1] + if last.GetContent().GetConfirmation() == nil { + return false, nil + } + if last.GetContent().GetConfirmation().GetApproval() != nil { + conf = last.GetContent().GetConfirmation() + approved = true + } + if last.GetContent().GetConfirmation().GetDecline() != nil { + conf = last.GetContent().GetConfirmation() + approved = false + } + return approved, conf +} + +func WaitsForUser(history []*proto.Message) *proto.Message { + if len(history) == 0 { + return nil + } + if msg := WaitsForConfirmation(history); msg != nil { + return msg + } + + last := history[len(history)-1] + if last.GetContent().GetConfirmation() != nil && last.GetContent().GetConfirmation().GetDecline() != nil { + return last + } + return nil +} diff --git a/internal/k8s/ate/client.go b/internal/k8s/ate/client.go new file mode 100644 index 00000000..81a5c061 --- /dev/null +++ b/internal/k8s/ate/client.go @@ -0,0 +1,109 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package ate provides a client for the Agent Substrate Control API. +package ate + +import ( + "context" + "fmt" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/status" +) + +type Client struct { + namespace string + template string + conn *grpc.ClientConn +} + +// NewClient creates a new actor client. +func NewClient(ns, template, target string, opts ...grpc.DialOption) (*Client, error) { + if ns == "" { + return nil, fmt.Errorf("namespace cannot be empty") + } + if template == "" { + return nil, fmt.Errorf("template cannot be empty") + } + if target == "" { + target = "api.ate-system.svc:443" + } + if len(opts) == 0 { + opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials())) + } + conn, err := grpc.NewClient(target, opts...) + if err != nil { + return nil, fmt.Errorf("error when creating Control client: %w", err) + } + return &Client{ + namespace: ns, + template: template, + conn: conn, + }, nil +} + +// CreateActor creates a new actor. +func (c *Client) CreateActor(ctx context.Context, id string) (*ateapipb.CreateActorResponse, error) { + client := ateapipb.NewControlClient(c.conn) + // TODO(wjjclaud): Configure atespace in manifests instead of reusing the namespace. + if _, err := client.CreateAtespace(ctx, &ateapipb.CreateAtespaceRequest{Name: c.namespace}); err != nil && status.Code(err) != codes.AlreadyExists { + return nil, fmt.Errorf("error when calling Control.CreateAtespace: %w", err) + } + resp, err := client.CreateActor(ctx, &ateapipb.CreateActorRequest{ + ActorRef: &ateapipb.ActorRef{Atespace: c.namespace, Name: id}, + ActorTemplateNamespace: c.namespace, + ActorTemplateName: c.template, + }) + if err != nil { + return nil, fmt.Errorf("error when calling Control.CreateActor: %w", err) + } + return resp, nil +} + +// ResumeActor resumes the actor, scheduling it onto a worker. The returned +// actor carries the worker IP once it is running. +func (c *Client) ResumeActor(ctx context.Context, id string) (*ateapipb.ResumeActorResponse, error) { + client := ateapipb.NewControlClient(c.conn) + resp, err := client.ResumeActor(ctx, &ateapipb.ResumeActorRequest{ + ActorRef: &ateapipb.ActorRef{Atespace: c.namespace, Name: id}, + }) + if err != nil { + return nil, fmt.Errorf("error when calling Control.ResumeActor: %w", err) + } + return resp, nil +} + +// SuspendActor suspends the actor. +func (c *Client) SuspendActor(ctx context.Context, id string) (*ateapipb.SuspendActorResponse, error) { + client := ateapipb.NewControlClient(c.conn) + resp, err := client.SuspendActor(ctx, &ateapipb.SuspendActorRequest{ + ActorRef: &ateapipb.ActorRef{Atespace: c.namespace, Name: id}, + }) + if err != nil { + return nil, fmt.Errorf("error when calling Control.SuspendActor: %w", err) + } + return resp, nil +} + +// Close closes the gRPC connection. +func (c *Client) Close() error { + if c.conn != nil { + return c.conn.Close() + } + return nil +} diff --git a/internal/pythonsidecar/setup.go b/internal/pythonsidecar/setup.go new file mode 100644 index 00000000..344189dd --- /dev/null +++ b/internal/pythonsidecar/setup.go @@ -0,0 +1,143 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pythonsidecar + +import ( + "context" + "fmt" + "io" + "io/fs" + "os" + "os/exec" + "path/filepath" + + "github.com/google/ax/internal/config" +) + +// SetupOptions configures asset extraction and environment setup for a Python sidecar. +type SetupOptions struct { + // FS is the embedded filesystem containing Python assets (e.g., python.FS). (Required) + FS fs.FS +} + +// Setup extracts the embedded filesystem assets to TargetDir. +// It returns TargetDir, which can be set as PythonPath in Config. +func Setup(ctx context.Context, opts SetupOptions) (string, error) { + if opts.FS == nil { + return "", fmt.Errorf("SetupOptions.FS cannot be nil") + } + axDir, err := config.AXAssetsDir() + if err != nil { + return "", err + } + extractDir := filepath.Join(axDir, "python") + reqPath := filepath.Join(extractDir, "antigravity", "requirements.txt") + + pythonPath := axDir + string(os.PathListSeparator) + extractDir + updated, err := extractFS(ctx, opts.FS, extractDir) + if err != nil { + return "", fmt.Errorf("failed to extract embedded assets: %w", err) + } + + pkgPath, err := install(ctx, reqPath, updated) + if err != nil { + return "", err + } + return pythonPath + string(os.PathListSeparator) + pkgPath, nil +} + +func extractFS(ctx context.Context, filesystem fs.FS, destDir string) (bool, error) { + var updated bool + err := fs.WalkDir(filesystem, ".", func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if ctx.Err() != nil { + return ctx.Err() + } + + destPath := filepath.Join(destDir, filepath.FromSlash(path)) + if d.IsDir() { + return os.MkdirAll(destPath, 0755) + } + + info, err := d.Info() + if err != nil { + return err + } + + stat, err := os.Stat(destPath) + // If the file already exists on disk with the exact same size, + // skip copying to preserve its timestamp and avoid false re-installations. + if err == nil { + if stat.Size() == info.Size() { + return nil + } + } + + if err := os.MkdirAll(filepath.Dir(destPath), 0755); err != nil { + return fmt.Errorf("creating directory for %s: %w", destPath, err) + } + + // There is at least, one file updated. + updated = true + src, err := filesystem.Open(path) + if err != nil { + return fmt.Errorf("opening embedded file %s: %w", path, err) + } + defer src.Close() + + mode := os.FileMode(0644) + if info.Mode().Perm() != 0 { + mode = info.Mode().Perm() | 0200 + } + _ = os.Chmod(destPath, 0644) + dst, err := os.OpenFile(destPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode) + if err != nil { + return fmt.Errorf("creating destination file %s: %w", destPath, err) + } + + if _, err := io.Copy(dst, src); err != nil { + _ = dst.Close() + return fmt.Errorf("writing file %s: %w", destPath, err) + } + if err := dst.Close(); err != nil { + return fmt.Errorf("closing file %s: %w", destPath, err) + } + return nil + }) + return updated, err +} + +func install(ctx context.Context, reqPath string, install bool) (string, error) { + pkgDir := filepath.Join(filepath.Dir(reqPath), "site-packages") + if !install { + // Make sure that pip install ran once and we have + // the dependencies installed under site-packages. + if _, err := os.Stat(pkgDir); err == nil { + return pkgDir, nil + } + } + fmt.Println("Setting up Antigravity SDK, this may take a while...") + + // Some corp machines are overriding the indexes, ensure that + // simple index is included where Antigravity SDK is distributed from. + cmd := exec.CommandContext(ctx, "python3", "-m", "pip", "install", "--extra-index-url", "https://pypi.org/simple", "--target", pkgDir, "-r", reqPath) + out, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("pip install failed for %s: %w\nOutput:\n%s", reqPath, err, string(out)) + } + return pkgDir, nil +} diff --git a/internal/pythonsidecar/setup_test.go b/internal/pythonsidecar/setup_test.go new file mode 100644 index 00000000..cfcf3da3 --- /dev/null +++ b/internal/pythonsidecar/setup_test.go @@ -0,0 +1,112 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pythonsidecar_test + +import ( + "bytes" + "context" + "io/fs" + "os" + "path/filepath" + "strings" + "testing" + "testing/fstest" + "time" + + "github.com/google/ax/internal/pythonsidecar" + "github.com/google/ax/python" +) + +func TestSetup_EmbeddedFS(t *testing.T) { + if _, err := fs.Stat(python.FS, "antigravity/__pycache__"); err == nil { + t.Errorf("expected antigravity/__pycache__ to be ignored when embedding, but it was found") + } + + testFS := fstest.MapFS{ + "antigravity/harness_server.py": &fstest.MapFile{Data: []byte("print('hello')")}, + "antigravity/requirements.txt": &fstest.MapFile{Data: []byte("# empty requirements for test\n")}, + "proto/ax_pb2.py": &fstest.MapFile{Data: []byte("print('proto')")}, + } + + targetDir := filepath.Join(t.TempDir(), "home") + axDir := filepath.Join(targetDir, ".ax") + t.Setenv("HOME", targetDir) + + opts := pythonsidecar.SetupOptions{ + FS: testFS, + } + gotDir, err := pythonsidecar.Setup(context.Background(), opts) + if err != nil { + t.Fatalf("Setup() failed: %v", err) + } + if !strings.Contains(gotDir, "python") { + t.Errorf("expected gotDir to contain extracted python directory in PYTHONPATH, got %q", gotDir) + } + + // Verify files were extracted under TargetDir/python + harnessPath := filepath.Join(axDir, "python", "antigravity", "harness_server.py") + if _, err := os.Stat(harnessPath); err != nil { + t.Errorf("expected file %s to exist, got stat error: %v", harnessPath, err) + } + protoPath := filepath.Join(axDir, "python", "proto", "ax_pb2.py") + if _, err := os.Stat(protoPath); err != nil { + t.Errorf("expected file %s to exist, got stat error: %v", protoPath, err) + } + + // Verify that subsequent Setup calls when TargetDir exists succeed without re-extracting + if _, err := pythonsidecar.Setup(context.Background(), opts); err != nil { + t.Fatalf("subsequent Setup() failed when TargetDir exists: %v", err) + } +} + +func TestSidecar_PythonPath(t *testing.T) { + tmpDir := t.TempDir() + customPath := filepath.Join(tmpDir, "custom_python_path") + if err := os.MkdirAll(customPath, 0755); err != nil { + t.Fatalf("failed to create custom path: %v", err) + } + + modulePath := filepath.Join(customPath, "path_check.py") + moduleContent := ` +import sys, os +print("SYSPATH:" + str(sys.path)) +sys.exit(0) +` + if err := os.WriteFile(modulePath, []byte(moduleContent), 0644); err != nil { + t.Fatalf("failed to write path_check module: %v", err) + } + + var stdout bytes.Buffer + cfg := pythonsidecar.Config{ + Module: "path_check", + Stdout: &stdout, + } + + s := pythonsidecar.New(cfg) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if err := s.Start(ctx, customPath); err != nil { + t.Fatalf("Start() failed: %v", err) + } + if err := s.Wait(); err != nil { + t.Fatalf("Wait() failed: %v", err) + } + + out := stdout.String() + if !strings.Contains(out, customPath) { + t.Errorf("expected sys.path to contain customPath, got output:\n%s", out) + } +} diff --git a/internal/pythonsidecar/sidecar.go b/internal/pythonsidecar/sidecar.go new file mode 100644 index 00000000..d02799ba --- /dev/null +++ b/internal/pythonsidecar/sidecar.go @@ -0,0 +1,269 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package pythonsidecar provides a mechanism to manage the lifecycle +// of a Python process as a sidecar component in a Go application. +package pythonsidecar + +import ( + "context" + "fmt" + "io" + "net" + "os" + "os/exec" + "strings" + "sync" + "syscall" + "time" +) + +// Config holds the configuration parameters for the sidecar lifecycle. +type Config struct { + // Module is the Python module name to run using the "-m" flag. (Required) + Module string + // Args contains any additional arguments to pass to the module. (Optional) + Args []string + // Stdout redirects the sidecar's standard output. (Optional) + Stdout io.Writer + // Stderr redirects the sidecar's standard error. (Optional) + Stderr io.Writer + // ReadyFunc is an optional function to check if the server is ready to accept requests. + // When provided, Start will poll ReadyFunc until it returns nil or the context expires. (Optional) + ReadyFunc func(ctx context.Context) error +} + +// TODO: Use /var/ax_agy_harness_service for communication instead of TCP. + +// Sidecar manages the lifecycle of the underlying Python process. +type Sidecar struct { + cfg Config + + mu sync.Mutex + cmd *exec.Cmd + running bool + stopping bool + exitErr error + doneChan chan struct{} +} + +// New creates a new Sidecar instance using the provided configuration struct. +func New(cfg Config) *Sidecar { + return &Sidecar{ + cfg: cfg, + } +} + +// Start launches the Python process and monitors its lifecycle in the background. +// If ReadyFunc is configured, Start blocks until the server is ready or the context expires. +// If the process fails to start or become ready, an error is returned immediately. +func (s *Sidecar) Start(ctx context.Context, pythonPath string) error { + s.mu.Lock() + if s.running { + s.mu.Unlock() + return fmt.Errorf("sidecar is already running") + } + + if s.cfg.Module == "" { + s.mu.Unlock() + return fmt.Errorf("Module cannot be empty") + } + + // Prepare arguments: python -u -m module [args...] + // -u forces unbuffered stdout/stderr so logs stream to Go instantly + fullArgs := append([]string{"-u", "-m", s.cfg.Module}, s.cfg.Args...) + + cmd := exec.CommandContext(ctx, "python3", fullArgs...) + if pythonPath != "" { + env := append([]string(nil), os.Environ()...) + var found bool + for i, kv := range env { + if strings.HasPrefix(kv, "PYTHONPATH=") { + existing := strings.TrimPrefix(kv, "PYTHONPATH=") + if existing != "" { + env[i] = "PYTHONPATH=" + pythonPath + string(os.PathListSeparator) + existing + } else { + env[i] = "PYTHONPATH=" + pythonPath + } + found = true + break + } + } + if !found { + env = append(env, "PYTHONPATH="+pythonPath) + } + cmd.Env = env + } + + if s.cfg.Stdout != nil { + cmd.Stdout = s.cfg.Stdout + } + if s.cfg.Stderr != nil { + cmd.Stderr = s.cfg.Stderr + } + + if err := cmd.Start(); err != nil { + s.mu.Unlock() + return fmt.Errorf("failed to start python process: %w", err) + } + + s.cmd = cmd + s.running = true + s.stopping = false + s.exitErr = nil + s.doneChan = make(chan struct{}) + s.mu.Unlock() + + // Monitor lifecycle asynchronously + go s.monitor() + + // If a readiness probe is configured, wait for the server to become ready + if s.cfg.ReadyFunc != nil { + if err := s.WaitUntilReady(ctx); err != nil { + _ = s.Stop() + return fmt.Errorf("server failed to become ready: %w", err) + } + } + + return nil +} + +// monitor waits for the process to exit and records its exit status. +func (s *Sidecar) monitor() { + err := s.cmd.Wait() + s.mu.Lock() + s.running = false + if err != nil && !s.stopping { + s.exitErr = fmt.Errorf("python process exited with error: %w", err) + } else { + s.exitErr = nil + } + s.mu.Unlock() + close(s.doneChan) +} + +// IsRunning returns true if the Python process is currently active. +func (s *Sidecar) IsRunning() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.running +} + +// Pid returns the process ID of the running sidecar, or 0 if not running. +func (s *Sidecar) Pid() int { + s.mu.Lock() + defer s.mu.Unlock() + if !s.running || s.cmd == nil || s.cmd.Process == nil { + return 0 + } + return s.cmd.Process.Pid +} + +// WaitUntilReady blocks until ReadyFunc returns nil, the context is canceled, or the process exits prematurely. +// If no ReadyFunc is configured in Config, this method returns nil immediately. +func (s *Sidecar) WaitUntilReady(ctx context.Context) error { + if s.cfg.ReadyFunc == nil { + return nil + } + + ticker := time.NewTicker(50 * time.Millisecond) + defer ticker.Stop() + + for { + // 1. Check if the process exited prematurely + s.mu.Lock() + running := s.running + exitErr := s.exitErr + s.mu.Unlock() + if !running { + if exitErr != nil { + return fmt.Errorf("process exited before becoming ready: %w", exitErr) + } + return fmt.Errorf("process exited unexpectedly before becoming ready") + } + + // 2. Try the readiness check + if err := s.cfg.ReadyFunc(ctx); err == nil { + return nil + } + + // 3. Wait for the next ticker or context cancellation + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + continue + } + } +} + +// Wait blocks until the sidecar exits or crashes, returning the exit error if any. +func (s *Sidecar) Wait() error { + s.mu.Lock() + done := s.doneChan + s.mu.Unlock() + if done != nil { + <-done + } + s.mu.Lock() + defer s.mu.Unlock() + return s.exitErr +} + +// Stop gracefully terminates the Python process using SIGTERM, falling back to SIGKILL if necessary. +func (s *Sidecar) Stop() error { + s.mu.Lock() + if !s.running || s.cmd == nil || s.cmd.Process == nil { + s.mu.Unlock() + return nil + } + s.stopping = true + done := s.doneChan + s.mu.Unlock() + + // 1. Send graceful SIGTERM + if err := s.cmd.Process.Signal(syscall.SIGTERM); err != nil { + _ = s.cmd.Process.Kill() + if done != nil { + <-done + } + return nil + } + + // 2. Give it a small window to exit gracefully before killing it + select { + case <-done: + return nil + case <-time.After(3 * time.Second): + // Fallback to force kill + _ = s.cmd.Process.Kill() + if done != nil { + <-done + } + return nil + } +} + +// TCPReady returns a ReadyFunc that attempts to establish a TCP connection to addr (e.g., "127.0.0.1:50053"). +func TCPReady(addr string) func(ctx context.Context) error { + return func(ctx context.Context) error { + var d net.Dialer + conn, err := d.DialContext(ctx, "tcp", addr) + if err != nil { + return err + } + _ = conn.Close() + return nil + } +} diff --git a/internal/pythonsidecar/sidecar_test.go b/internal/pythonsidecar/sidecar_test.go new file mode 100644 index 00000000..3052f414 --- /dev/null +++ b/internal/pythonsidecar/sidecar_test.go @@ -0,0 +1,158 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pythonsidecar_test + +import ( + "bytes" + "context" + "fmt" + "net" + "os" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "github.com/google/ax/internal/pythonsidecar" +) + +func getFreePort(t *testing.T) int { + t.Helper() + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to get free port: %v", err) + } + port := l.Addr().(*net.TCPAddr).Port + _ = l.Close() + return port +} + +func TestSidecar_ConfigValidation(t *testing.T) { + ctx := context.Background() + + t.Run("empty Module", func(t *testing.T) { + s := pythonsidecar.New(pythonsidecar.Config{}) + err := s.Start(ctx, "") + if err == nil || !strings.Contains(err.Error(), "Module cannot be empty") { + t.Fatalf("expected error about empty Module, got %v", err) + } + }) +} + +func TestSidecar_ModuleExecution(t *testing.T) { + tmpDir := t.TempDir() + modulePath := filepath.Join(tmpDir, "test_module.py") + moduleContent := ` +import sys +print("hello stdout") +print("hello stderr", file=sys.stderr) +sys.exit(0) +` + if err := os.WriteFile(modulePath, []byte(moduleContent), 0644); err != nil { + t.Fatalf("failed to write module: %v", err) + } + + var stdout, stderr bytes.Buffer + cfg := pythonsidecar.Config{ + Module: "test_module", + Stdout: &stdout, + Stderr: &stderr, + } + + s := pythonsidecar.New(cfg) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if err := s.Start(ctx, tmpDir); err != nil { + t.Fatalf("Start() failed: %v", err) + } + + if err := s.Wait(); err != nil { + t.Fatalf("Wait() failed: %v", err) + } + + if !strings.Contains(stdout.String(), "hello stdout") { + t.Errorf("stdout expected 'hello stdout', got %q", stdout.String()) + } + if !strings.Contains(stderr.String(), "hello stderr") { + t.Errorf("stderr expected 'hello stderr', got %q", stderr.String()) + } + if s.IsRunning() { + t.Errorf("expected IsRunning() to be false after exit") + } +} + +func TestSidecar_ModuleServerWithTCPReady(t *testing.T) { + port := getFreePort(t) + addr := fmt.Sprintf("127.0.0.1:%d", port) + + cfg := pythonsidecar.Config{ + Module: "http.server", + Args: []string{strconv.Itoa(port), "--bind", "127.0.0.1"}, + ReadyFunc: pythonsidecar.TCPReady(addr), + } + + s := pythonsidecar.New(cfg) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + if err := s.Start(ctx, ""); err != nil { + t.Fatalf("Start() with TCPReady failed: %v", err) + } + + if !s.IsRunning() { + t.Fatalf("expected server sidecar to be running") + } + if s.Pid() <= 0 { + t.Fatalf("expected valid PID, got %d", s.Pid()) + } + + if err := s.Stop(); err != nil { + t.Fatalf("Stop() failed: %v", err) + } +} + +func TestSidecar_ReadinessFailureOnPrematureExit(t *testing.T) { + tmpDir := t.TempDir() + modulePath := filepath.Join(tmpDir, "crash.py") + if err := os.WriteFile(modulePath, []byte("import sys; sys.exit(1)\n"), 0644); err != nil { + t.Fatalf("failed to write script: %v", err) + } + + t.Setenv("PYTHONPATH", tmpDir) + port := getFreePort(t) + addr := fmt.Sprintf("127.0.0.1:%d", port) + + cfg := pythonsidecar.Config{ + Module: "crash", + ReadyFunc: pythonsidecar.TCPReady(addr), + } + + s := pythonsidecar.New(cfg) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + err := s.Start(ctx, "") + if err == nil { + t.Fatalf("expected Start() to fail when process exits prematurely") + } + if !strings.Contains(err.Error(), "exited before becoming ready") { + t.Fatalf("expected 'exited before becoming ready' in error, got: %v", err) + } + if s.IsRunning() { + t.Fatalf("expected IsRunning() to be false") + } +} diff --git a/internal/server/interceptors.go b/internal/server/interceptors.go new file mode 100644 index 00000000..45d7732e --- /dev/null +++ b/internal/server/interceptors.go @@ -0,0 +1,79 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "context" + "log/slog" + "time" + + "google.golang.org/grpc" +) + +type conversationer interface { + GetConversationId() string +} + +// LoggingInterceptor logs unary RPC details, including latency, errors, and conversation ID if present. +func LoggingInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + start := time.Now() + logger := slog.With(slog.String("method", info.FullMethod)) + + // Extract conversation_id if the request supports it + if conv, ok := req.(conversationer); ok { + if id := conv.GetConversationId(); id != "" { + logger = logger.With(slog.String("conversation_id", id)) + } + } + + logger.InfoContext(ctx, "Handling unary request") + resp, err := handler(ctx, req) + + duration := time.Since(start) + if err != nil { + logger.ErrorContext(ctx, "Request failed", + slog.Duration("duration", duration), + slog.Any("error", err), + ) + } else { + logger.InfoContext(ctx, "Request completed", + slog.Duration("duration", duration), + ) + } + return resp, err +} + +// StreamLoggingInterceptor logs stream RPC connection details and latency. +func StreamLoggingInterceptor(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + start := time.Now() + ctx := ss.Context() + logger := slog.With(slog.String("method", info.FullMethod)) + + logger.InfoContext(ctx, "Handling stream request") + err := handler(srv, ss) + + duration := time.Since(start) + if err != nil { + logger.ErrorContext(ctx, "Stream failed", + slog.Duration("duration", duration), + slog.Any("error", err), + ) + } else { + logger.InfoContext(ctx, "Stream completed", + slog.Duration("duration", duration), + ) + } + return err +} diff --git a/internal/server/interceptors_test.go b/internal/server/interceptors_test.go new file mode 100644 index 00000000..0d02e8ee --- /dev/null +++ b/internal/server/interceptors_test.go @@ -0,0 +1,264 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "bytes" + "context" + "encoding/json" + "io" + "log/slog" + "net" + "strings" + "testing" + + "github.com/google/ax/proto" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" +) + +// Mock servers for testing interceptors +type mockConversationServer struct { + proto.UnimplementedConversationServiceServer +} + +func (m *mockConversationServer) DeleteConversation(ctx context.Context, req *proto.DeleteConversationRequest) (*proto.DeleteConversationResponse, error) { + if req.ConversationId == "fail" { + return nil, status.Error(codes.InvalidArgument, "mock error") + } + return &proto.DeleteConversationResponse{}, nil +} + +type mockExecutionServer struct { + proto.UnimplementedExecutionServiceServer +} + +func (m *mockExecutionServer) Exec(req *proto.ExecRequest, stream proto.ExecutionService_ExecServer) error { + if req.ConversationId == "fail" { + return status.Error(codes.InvalidArgument, "mock error") + } + return stream.Send(&proto.ExecResponse{}) +} + +const bufSize = 1024 * 1024 + +func setupTestServer(t *testing.T) (*grpc.ClientConn, func()) { + lis := bufconn.Listen(bufSize) + s := grpc.NewServer( + grpc.ChainUnaryInterceptor(LoggingInterceptor), + grpc.ChainStreamInterceptor(StreamLoggingInterceptor), + ) + + proto.RegisterConversationServiceServer(s, &mockConversationServer{}) + proto.RegisterExecutionServiceServer(s, &mockExecutionServer{}) + + go func() { + if err := s.Serve(lis); err != nil && err != grpc.ErrServerStopped { + t.Errorf("Server exited with error: %v", err) + } + }() + + conn, err := grpc.NewClient("passthrough:///bufnet", + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { + return lis.Dial() + }), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + t.Fatalf("Failed to dial bufnet: %v", err) + } + + cleanup := func() { + conn.Close() + s.GracefulStop() + lis.Close() + } + + return conn, cleanup +} + +type logEntry struct { + Time string `json:"time"` + Level string `json:"level"` + Msg string `json:"msg"` + Method string `json:"method"` + ConversationID string `json:"conversation_id,omitempty"` + Duration int64 `json:"duration,omitempty"` // in nanoseconds + Error string `json:"error,omitempty"` +} + +func TestLoggingInterceptors(t *testing.T) { + conn, cleanup := setupTestServer(t) + defer cleanup() + + // Capture slog output to a buffer + var logBuf bytes.Buffer + testLogger := slog.New(slog.NewJSONHandler(&logBuf, nil)) + oldLogger := slog.Default() + slog.SetDefault(testLogger) + defer slog.SetDefault(oldLogger) + + convClient := proto.NewConversationServiceClient(conn) + executionClient := proto.NewExecutionServiceClient(conn) + + t.Run("Unary Success", func(t *testing.T) { + logBuf.Reset() + ctx := context.Background() + _, err := convClient.DeleteConversation(ctx, &proto.DeleteConversationRequest{ConversationId: "conv-123"}) + if err != nil { + t.Fatalf("DeleteConversation failed: %v", err) + } + + entries := parseLogs(t, &logBuf) + if len(entries) != 2 { + t.Fatalf("Expected 2 log entries, got %d", len(entries)) + } + + // Verify Start Log + if entries[0].Msg != "Handling unary request" { + t.Errorf("Expected start log msg 'Handling unary request', got %q", entries[0].Msg) + } + if entries[0].Method != "/ax.ConversationService/DeleteConversation" { + t.Errorf("Expected method '/ax.ConversationService/DeleteConversation', got %q", entries[0].Method) + } + if entries[0].ConversationID != "conv-123" { + t.Errorf("Expected conversation_id 'conv-123', got %q", entries[0].ConversationID) + } + + // Verify End Log + if entries[1].Msg != "Request completed" { + t.Errorf("Expected end log msg 'Request completed', got %q", entries[1].Msg) + } + if entries[1].Level != "INFO" { + t.Errorf("Expected Level INFO, got %s", entries[1].Level) + } + }) + + t.Run("Unary Failure", func(t *testing.T) { + logBuf.Reset() + ctx := context.Background() + _, err := convClient.DeleteConversation(ctx, &proto.DeleteConversationRequest{ConversationId: "fail"}) + if err == nil { + t.Fatal("Expected DeleteConversation to fail") + } + + entries := parseLogs(t, &logBuf) + if len(entries) != 2 { + t.Fatalf("Expected 2 log entries, got %d", len(entries)) + } + + // Verify End Log + if entries[1].Msg != "Request failed" { + t.Errorf("Expected end log msg 'Request failed', got %q", entries[1].Msg) + } + if entries[1].Level != "ERROR" { + t.Errorf("Expected Level ERROR, got %s", entries[1].Level) + } + if !strings.Contains(entries[1].Error, "mock error") { + t.Errorf("Expected error details to contain 'mock error', got %q", entries[1].Error) + } + }) + + t.Run("Stream Success", func(t *testing.T) { + logBuf.Reset() + ctx := context.Background() + stream, err := executionClient.Exec(ctx, &proto.ExecRequest{ConversationId: "conv-456"}) + if err != nil { + t.Fatalf("Exec stream init failed: %v", err) + } + + // Consume stream + for { + _, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + t.Fatalf("Error receiving from stream: %v", err) + } + } + + entries := parseLogs(t, &logBuf) + if len(entries) != 2 { + t.Fatalf("Expected 2 log entries, got %d", len(entries)) + } + + // Verify Start Log + if entries[0].Msg != "Handling stream request" { + t.Errorf("Expected start log msg 'Handling stream request', got %q", entries[0].Msg) + } + if entries[0].Method != "/ax.ExecutionService/Exec" { + t.Errorf("Expected method '/ax.ExecutionService/Exec', got %q", entries[0].Method) + } + + // Verify End Log + if entries[1].Msg != "Stream completed" { + t.Errorf("Expected end log msg 'Stream completed', got %q", entries[1].Msg) + } + if entries[1].Level != "INFO" { + t.Errorf("Expected Level INFO, got %s", entries[1].Level) + } + }) + + t.Run("Stream Failure", func(t *testing.T) { + logBuf.Reset() + ctx := context.Background() + stream, err := executionClient.Exec(ctx, &proto.ExecRequest{ConversationId: "fail"}) + if err != nil { + t.Fatalf("Exec stream init failed: %v", err) + } + + // Consume stream (should fail) + _, err = stream.Recv() + if err == nil || err == io.EOF { + t.Fatal("Expected stream read to fail") + } + + entries := parseLogs(t, &logBuf) + if len(entries) != 2 { + t.Fatalf("Expected 2 log entries, got %d", len(entries)) + } + + // Verify End Log + if entries[1].Msg != "Stream failed" { + t.Errorf("Expected end log msg 'Stream failed', got %q", entries[1].Msg) + } + if entries[1].Level != "ERROR" { + t.Errorf("Expected Level ERROR, got %s", entries[1].Level) + } + if !strings.Contains(entries[1].Error, "mock error") { + t.Errorf("Expected error details to contain 'mock error', got %q", entries[1].Error) + } + }) +} + +func parseLogs(t *testing.T, buf *bytes.Buffer) []logEntry { + t.Helper() + var entries []logEntry + decoder := json.NewDecoder(buf) + for { + var entry logEntry + if err := decoder.Decode(&entry); err == io.EOF { + break + } else if err != nil { + t.Fatalf("Failed to decode log JSON: %v. Raw buffer: %s", err, buf.String()) + } + entries = append(entries, entry) + } + return entries +} diff --git a/internal/server/server.go b/internal/server/server.go index f9d070b2..267460ff 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -12,160 +12,139 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Package server implements the gRPC server for GARService, -// exposing session management and agent registration APIs. +// Package server implements the gRPC server for AXService, +// exposing execution management and agent registration APIs. + package server import ( "context" "fmt" + "log/slog" "net" + "sync" + "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" "google.golang.org/grpc" - "google.golang.org/protobuf/types/known/timestamppb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/health" + "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" - "github.com/google/gar/agent" - "github.com/google/gar/internal/config" - "github.com/google/gar/internal/controller" - "github.com/google/gar/proto" + "github.com/google/ax/internal/controller" + "github.com/google/ax/proto" ) -// Server implements the GARService gRPC service. +// Server implements the AXService gRPC service. type Server struct { - proto.UnimplementedGARServiceServer + proto.UnimplementedExecutionServiceServer + proto.UnimplementedConversationServiceServer controller *controller.Controller + grpcServer *grpc.Server + inFlight map[string]struct{} + inFlightMu sync.Mutex } // New creates a new controller server. func New(c *controller.Controller) *Server { return &Server{ controller: c, + inFlight: make(map[string]struct{}), } } -// TriggerSession triggers a new agentic loop session with streaming responses. -func (s *Server) TriggerSession(req *proto.TriggerSessionRequest, stream grpc.ServerStreamingServer[proto.TriggerSessionResponse]) error { - sessionID := req.SessionId - inputs := req.Inputs - checkpointID := req.CheckpointId - - if sessionID == "" { - return fmt.Errorf("session_id is required") - } +// Exec executes a new agentic task with streaming responses. +func (s *Server) Exec(req *proto.ExecRequest, stream grpc.ServerStreamingServer[proto.ExecResponse]) error { + ctx := stream.Context() + slog.InfoContext(ctx, "Executing request", + slog.String("request", req.String()), + ) - // TODO(jbd): This state update should be sent directly from the controller. - if err := stream.Send(&proto.TriggerSessionResponse{ - State: proto.State_STATE_RUNNING, - SessionId: sessionID, - }); err != nil { - return err + inFlight, cleanup := s.markInFlight(req.ConversationId) + if inFlight { + return status.Errorf(codes.FailedPrecondition, "conversation %q is already in flight", req.ConversationId) } + defer cleanup() - // Create output handler to stream outputs back to client - outputHandler := agent.OutputHandler(func(content *proto.Content) error { - return stream.Send(&proto.TriggerSessionResponse{ - SessionId: sessionID, - State: proto.State_STATE_RUNNING, - Output: content, - }) + outputHandler := controller.ExecHandler(func(resp *proto.ExecResponse) error { + return stream.Send(resp) }) - - if checkpointID == "" { - return s.controller.TriggerSession(stream.Context(), sessionID, inputs, outputHandler) - } - return s.controller.TriggerForkedSession(stream.Context(), sessionID, checkpointID, inputs, outputHandler) + return s.controller.Exec(ctx, req, outputHandler) } -// GetSession retrieves session details. -func (s *Server) GetSession(ctx context.Context, req *proto.GetSessionRequest) (*proto.GetSessionResponse, error) { - if req.SessionId == "" { - return nil, fmt.Errorf("session_id is required") - } - - // Load session if not already loaded - session, err := s.controller.LoadSession(ctx, req.SessionId) - if err != nil { - return nil, fmt.Errorf("error loading session: %w", err) - } - return &proto.GetSessionResponse{ - Session: &proto.SessionInfo{ - State: session.State(), - ActiveAgents: session.ActiveAgents(), - CreatedAt: timestamppb.New(session.CreatedAt()), - UpdatedAt: timestamppb.New(session.UpdatedAt()), - CheckpointCount: int32(len(session.CheckpointIDs())), - }, - }, nil -} +func (s *Server) DeleteConversation(ctx context.Context, req *proto.DeleteConversationRequest) (*proto.DeleteConversationResponse, error) { + slog.InfoContext(ctx, "Deleting conversation...", + slog.String("conversation_id", req.ConversationId)) -// RegisterAgent registers a new remote agent with the controller. -func (s *Server) RegisterAgent(ctx context.Context, req *proto.RegisterAgentRequest) (*proto.RegisterAgentResponse, error) { - if req.AgentId == "" { - return nil, fmt.Errorf("agent_id is required") + if req.ConversationId == "" { + return nil, status.Errorf(codes.InvalidArgument, "conversation_id is required") } - - if req.Address == "" { - return nil, fmt.Errorf("address is required for remote agents") + inFlight, cleanup := s.markInFlight(req.ConversationId) + if inFlight { + return nil, status.Errorf(codes.FailedPrecondition, "conversation %q is already in flight", req.ConversationId) } + defer cleanup() - registry := s.controller.Registry() - - // All registered agents are remote - err := registry.RegisterRemote(config.RemoteAgentConfig{ - ID: req.AgentId, - Name: req.Name, - Description: req.Description, - Address: req.Address, - Metadata: req.Metadata, - }) - if err != nil { - return nil, fmt.Errorf("failed to register agent: %w", err) + if err := s.controller.Delete(ctx, req.ConversationId); err != nil { + return nil, status.Errorf(codes.Internal, "failed to delete conversation: %v", err) } - - return &proto.RegisterAgentResponse{}, nil + return &proto.DeleteConversationResponse{}, nil } -// UnregisterAgent removes a remote agent from the controller. -// Local agents cannot be unregistered via this API. -func (s *Server) UnregisterAgent(ctx context.Context, req *proto.UnregisterAgentRequest) (*proto.UnregisterAgentResponse, error) { - if req.AgentId == "" { - return nil, fmt.Errorf("agent_id is required") +// Serve starts the gRPC server on the specified address. +func (s *Server) Serve(address string, opts ...grpc.ServerOption) error { + lis, err := net.Listen("tcp", address) + if err != nil { + return fmt.Errorf("failed to listen: %w", err) } - registry := s.controller.Registry() + opts = append(opts, + grpc.ChainUnaryInterceptor(LoggingInterceptor), + grpc.ChainStreamInterceptor(StreamLoggingInterceptor), + grpc.StatsHandler(otelgrpc.NewServerHandler()), + ) - // Check if the agent is local - info, err := registry.GetInfo(req.AgentId) - if err != nil { - return nil, fmt.Errorf("agent not found: %w", err) - } + s.grpcServer = grpc.NewServer(opts...) + proto.RegisterExecutionServiceServer(s.grpcServer, s) + proto.RegisterConversationServiceServer(s.grpcServer, s) - if info.Type == controller.AgentTypeLocal { - return nil, fmt.Errorf("cannot unregister local agents via API") - } + // Register standard gRPC Health Check server. + hs := health.NewServer() + hs.SetServingStatus("AX", grpc_health_v1.HealthCheckResponse_SERVING) + grpc_health_v1.RegisterHealthServer(s.grpcServer, hs) - if err := registry.Unregister(req.AgentId); err != nil { - return nil, fmt.Errorf("failed to unregister agent: %w", err) + if err := s.grpcServer.Serve(lis); err != nil { + return fmt.Errorf("failed to serve: %w", err) } - - return &proto.UnregisterAgentResponse{}, nil + return nil } -// Serve starts the gRPC server on the specified address. -func (s *Server) Serve(address string, opts ...grpc.ServerOption) error { - lis, err := net.Listen("tcp", address) - if err != nil { - return fmt.Errorf("failed to listen: %w", err) +// GracefulStop stops the gRPC server gracefully. +func (s *Server) GracefulStop() { + slog.Info("Stopping server gracefully...") + if s.controller != nil { + s.controller.Close() + } + if s.grpcServer != nil { + s.grpcServer.GracefulStop() } +} - grpcServer := grpc.NewServer(opts...) - proto.RegisterGARServiceServer(grpcServer, s) +func (s *Server) markInFlight(id string) (exists bool, cleanup func()) { + s.inFlightMu.Lock() + defer s.inFlightMu.Unlock() - if err := grpcServer.Serve(lis); err != nil { - return fmt.Errorf("failed to serve: %w", err) + _, ok := s.inFlight[id] + if ok { + return true, func() {} } + s.inFlight[id] = struct{}{} - return nil + return false, func() { + s.inFlightMu.Lock() + delete(s.inFlight, id) + s.inFlightMu.Unlock() + } } diff --git a/internal/skills/geminienterprise/client.go b/internal/skills/geminienterprise/client.go new file mode 100644 index 00000000..93e4f726 --- /dev/null +++ b/internal/skills/geminienterprise/client.go @@ -0,0 +1,400 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// This file is the internal Gemini Enterprise Skill Registry transport client: it talks to the +// Vertex AI v1beta1 REST API (ListSkills / GetSkill / GetSkillRevision / +// skills:retrieve), fetches skill payloads, and safe-unzips them to disk. The +// public, config-driven orchestration lives in materialize.go. + +package geminienterprise + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "time" + + "golang.org/x/oauth2" + "golang.org/x/oauth2/google" +) + +// cloudPlatformScope is the OAuth2 scope required to call Vertex AI (read path +// needs roles/aiplatform.viewer). +const cloudPlatformScope = "https://www.googleapis.com/auth/cloud-platform" + +// apiVersion is the Vertex AI API version the Skill Registry is exposed under. +const apiVersion = "v1beta1" + +// defaultHTTPTimeout bounds a single registry HTTP call. +const defaultHTTPTimeout = 60 * time.Second + +// client fetches a selected set of skills from the Gemini Enterprise Skill Registry and +// writes them into a target directory as agentskills.io skill folders. +type client struct { + baseURL string // .../v1beta1/projects/{project}/locations/{location}/skills + http *http.Client + ts oauth2.TokenSource + caps unzipCaps +} + +// clientOptions configures a client. newClient fills defaults for optional fields. +type clientOptions struct { + // --- Required --- + + // Project is the Google Cloud project that owns the skills. + Project string + // Location is the registry region, e.g. "us-central1". Selects the Vertex + // host (https://{Location}-aiplatform.googleapis.com) unless Endpoint is set. + Location string + + // --- Optional (newClient fills defaults) --- + + // Endpoint overrides the API host (scheme+host, no trailing slash), e.g. a + // sandbox host or an httptest.Server URL in unit tests. + Endpoint string + // HTTPClient is the client used for all registry calls. Defaults to + // http.Client{Timeout: 60s}. + HTTPClient *http.Client + // TokenSource provides the OAuth2 bearer token. If nil, ADC is used + // (roles/aiplatform.viewer). Set to a static source in tests. + TokenSource oauth2.TokenSource + // Caps bounds a defensive local unzip of untrusted payloads. + Caps unzipCaps +} + +// newClient constructs a client from options. It validates required fields and +// fills defaults; it performs no network I/O. +func newClient(opts clientOptions) (*client, error) { + if opts.Project == "" { + return nil, errors.New("registry: Project is required") + } + if opts.Location == "" { + return nil, errors.New("registry: Location is required") + } + + host := opts.Endpoint + if host == "" { + host = fmt.Sprintf("https://%s-aiplatform.googleapis.com", opts.Location) + } + host = strings.TrimSuffix(host, "/") + baseURL := fmt.Sprintf("%s/%s/projects/%s/locations/%s/skills", + host, apiVersion, opts.Project, opts.Location) + + httpClient := opts.HTTPClient + if httpClient == nil { + httpClient = &http.Client{Timeout: defaultHTTPTimeout} + } + + ts := opts.TokenSource + if ts == nil { + creds, err := google.FindDefaultCredentials(context.Background(), cloudPlatformScope) + if err != nil { + return nil, fmt.Errorf("registry: finding application default credentials: %w", err) + } + ts = creds.TokenSource + } + + return &client{ + baseURL: baseURL, + http: httpClient, + ts: ts, + caps: opts.Caps.withDefaults(), + }, nil +} + +// selection describes WHICH skills to fetch. Exactly one of the three modes +// should be set (all / by-id / by-query). +type selection struct { + // All fetches every skill returned by ListSkills for the project/location. + All bool + // SkillRefs is an explicit allowlist, each optionally pinned to a revision. + SkillRefs []skillRef + // Query is a semantic search string; the top matches are fetched. TopK bounds + // the result count (<=0 means the server default). + Query string + TopK int +} + +// skillRef identifies a single skill to fetch, optionally pinned. +type skillRef struct { + SkillID string + Revision string // empty => latest default revision +} + +// fetchResult reports the outcome of a fetch call. fetch is fail-safe: a skill +// that cannot be fetched or unzipped is recorded in skipped. +type fetchResult struct { + materialized []MaterializedSkill + skipped []skippedSkill +} + +// skippedSkill records a skill that was intentionally not written, and why. +type skippedSkill struct { + skillID string + reason error +} + +// fetch resolves the selection, fetches each skill's payload, and safe-unzips it +// into targetDir as //... targetDir is created if missing. +// +// claimed enforces first-wins across the whole operation: a skill id whose +// (targetDir, id) was already materialized (by this registry's own selection or +// an earlier registry sharing the dir) is skipped with a warning rather than +// overwriting the winner. regIdx is the registry's index, for log context. +// +// It is fail-safe: individual skill failures are collected in fetchResult.skipped; +// fetch returns a non-nil error only for whole-operation failures. +func (c *client) fetch(ctx context.Context, sel selection, targetDir string, claimed *claimSet, regIdx int) (*fetchResult, error) { + refs, err := c.resolveSelection(ctx, sel) + if err != nil { + return nil, fmt.Errorf("registry: resolving selection: %w", err) + } + if err := os.MkdirAll(targetDir, 0o755); err != nil { + return nil, fmt.Errorf("registry: creating target dir %q: %w", targetDir, err) + } + + res := &fetchResult{} + for _, ref := range refs { + // First-wins: skip (with a warning) any id already materialized into this + // dir, rather than overwriting it. + if won, byRegistry := claimed.claim(targetDir, ref.SkillID, regIdx); !won { + log.Printf("skills: registries[%d] skill %q already materialized into %s by registries[%d]; keeping the first, skipping this one", + regIdx, ref.SkillID, targetDir, byRegistry) + continue + } + mat, err := c.fetchAndWrite(ctx, ref, targetDir) + if err != nil { + res.skipped = append(res.skipped, skippedSkill{skillID: ref.SkillID, reason: err}) + continue + } + res.materialized = append(res.materialized, *mat) + } + return res, nil +} + +// resolveSelection turns a selection into a concrete list of skillRefs to fetch. +func (c *client) resolveSelection(ctx context.Context, sel selection) ([]skillRef, error) { + switch { + case len(sel.SkillRefs) > 0: + return sel.SkillRefs, nil + case sel.Query != "": + ids, err := c.retrieveSkillIDs(ctx, sel.Query, sel.TopK) + if err != nil { + return nil, err + } + return toRefs(ids), nil + case sel.All: + ids, err := c.listSkillIDs(ctx) + if err != nil { + return nil, err + } + return toRefs(ids), nil + default: + return nil, errors.New("empty selection: set All, SkillRefs, or Query") + } +} + +// fetchAndWrite fetches one skill's payload (latest or pinned) and unzips it. +func (c *client) fetchAndWrite(ctx context.Context, ref skillRef, targetDir string) (*MaterializedSkill, error) { + if ref.SkillID == "" { + return nil, errors.New("skill ref has empty SkillID") + } + payloadB64, revision, err := c.fetchPayload(ctx, ref) + if err != nil { + return nil, err + } + zipped, err := base64.StdEncoding.DecodeString(payloadB64) + if err != nil { + return nil, fmt.Errorf("decoding payload: %w", err) + } + skillDir := filepath.Join(targetDir, ref.SkillID) + // Replace any prior materialization of this skill so stale files don't linger. + if err := os.RemoveAll(skillDir); err != nil { + return nil, fmt.Errorf("clearing %q: %w", skillDir, err) + } + if err := safeUnzip(zipped, skillDir, c.caps); err != nil { + // Leave no partial dir behind on failure. + _ = os.RemoveAll(skillDir) + return nil, fmt.Errorf("unzipping: %w", err) + } + return &MaterializedSkill{SkillID: ref.SkillID, Revision: revision, Dir: skillDir}, nil +} + +// fetchPayload returns the base64 zippedFilesystem and the concrete revision id +// for a skill ref (GetSkill for latest, GetSkillRevision when pinned). +func (c *client) fetchPayload(ctx context.Context, ref skillRef) (payloadB64, revision string, err error) { + if ref.Revision != "" { + // GetSkillRevision: payload is nested under "skill". + var resp skillRevisionResponse + if err := c.getJSON(ctx, c.baseURL+"/"+ref.SkillID+"/revisions/"+ref.Revision, &resp); err != nil { + return "", "", err + } + if resp.Skill.ZippedFilesystem == "" { + return "", "", errors.New("GetSkillRevision: empty zippedFilesystem") + } + return resp.Skill.ZippedFilesystem, ref.Revision, nil + } + // GetSkill: payload at top level. + var resp skillResponse + if err := c.getJSON(ctx, c.baseURL+"/"+ref.SkillID, &resp); err != nil { + return "", "", err + } + if resp.ZippedFilesystem == "" { + return "", "", errors.New("GetSkill: empty zippedFilesystem") + } + return resp.ZippedFilesystem, resp.currentRevision(), nil +} + +// listSkillIDs pages through ListSkills and returns all skill ids. +func (c *client) listSkillIDs(ctx context.Context) ([]string, error) { + var ids []string + pageToken := "" + for { + u := c.baseURL + if pageToken != "" { + u += "?pageToken=" + url.QueryEscape(pageToken) + } + var resp listSkillsResponse + if err := c.getJSON(ctx, u, &resp); err != nil { + return nil, err + } + for _, s := range resp.Skills { + if id := s.id(); id != "" { + ids = append(ids, id) + } + } + if resp.NextPageToken == "" { + return ids, nil + } + pageToken = resp.NextPageToken + } +} + +// retrieveSkillIDs runs semantic search (skills:retrieve) and returns skill ids. +func (c *client) retrieveSkillIDs(ctx context.Context, query string, topK int) ([]string, error) { + u := c.baseURL + ":retrieve?query=" + url.QueryEscape(query) + if topK > 0 { + u += fmt.Sprintf("&topK=%d", topK) + } + var resp retrieveSkillsResponse + if err := c.getJSON(ctx, u, &resp); err != nil { + return nil, err + } + var ids []string + for _, r := range resp.RetrievedSkills { + if id := lastSegment(r.SkillName); id != "" { + ids = append(ids, id) + } + } + return ids, nil +} + +// getJSON issues an authenticated GET and decodes the JSON body into out. +func (c *client) getJSON(ctx context.Context, rawURL string, out any) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) + if err != nil { + return err + } + tok, err := c.ts.Token() + if err != nil { + return fmt.Errorf("obtaining token: %w", err) + } + tok.SetAuthHeader(req) + req.Header.Set("Content-Type", "application/json") + + resp, err := c.http.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("HTTP %d: %s", resp.StatusCode, truncate(string(body), 512)) + } + if err := json.Unmarshal(body, out); err != nil { + return fmt.Errorf("decoding response: %w", err) + } + return nil +} + +// --- Wire types (AIP-standard shapes; see the design doc appendix). --- + +type listSkillsResponse struct { + Skills []skillResponse `json:"skills"` + NextPageToken string `json:"nextPageToken"` +} + +// retrieveSkillsResponse is a skills:retrieve result. Each hit carries the skill +// resource name in a flat "skillName" field (not a nested skill object). +type retrieveSkillsResponse struct { + RetrievedSkills []struct { + SkillName string `json:"skillName"` // projects/.../skills/{id} + } `json:"retrievedSkills"` +} + +// skillResponse is a Skill resource. GetSkill returns the payload at top level. +type skillResponse struct { + Name string `json:"name"` // projects/.../skills/{id} + DisplayName string `json:"displayName"` + Description string `json:"description"` + State string `json:"state"` + DefaultRevision string `json:"defaultRevision"` + ZippedFilesystem string `json:"zippedFilesystem"` +} + +// skillRevisionResponse is a GetSkillRevision result; the payload is nested +// under "skill". +type skillRevisionResponse struct { + Name string `json:"name"` // projects/.../skills/{id}/revisions/{rev} + Skill skillResponse `json:"skill"` +} + +// id extracts the trailing skill id from a resource name. +func (s skillResponse) id() string { return lastSegment(s.Name) } + +// currentRevision reports the skill's default revision id, if the server provided +// one (best-effort; may be empty for "latest"). +func (s skillResponse) currentRevision() string { return lastSegment(s.DefaultRevision) } + +func lastSegment(resourceName string) string { + if resourceName == "" { + return "" + } + parts := strings.Split(resourceName, "/") + return parts[len(parts)-1] +} + +func toRefs(ids []string) []skillRef { + refs := make([]skillRef, 0, len(ids)) + for _, id := range ids { + refs = append(refs, skillRef{SkillID: id}) + } + return refs +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "...(truncated)" +} diff --git a/internal/skills/geminienterprise/client_test.go b/internal/skills/geminienterprise/client_test.go new file mode 100644 index 00000000..39285b5c --- /dev/null +++ b/internal/skills/geminienterprise/client_test.go @@ -0,0 +1,414 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package geminienterprise + +import ( + "archive/zip" + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "golang.org/x/oauth2" +) + +// zipFixtureB64 builds an in-memory zip from name->content and returns the base64 +// string the registry would return in `zippedFilesystem`. +func zipFixtureB64(t *testing.T, files map[string]string) string { + t.Helper() + var buf bytes.Buffer + zw := zip.NewWriter(&buf) + for name, content := range files { + w, err := zw.Create(name) + if err != nil { + t.Fatalf("zip create %q: %v", name, err) + } + if _, err := w.Write([]byte(content)); err != nil { + t.Fatalf("zip write %q: %v", name, err) + } + } + if err := zw.Close(); err != nil { + t.Fatalf("zip close: %v", err) + } + return base64.StdEncoding.EncodeToString(buf.Bytes()) +} + +// rawZipFixture returns the raw zip bytes (for direct safeUnzip tests). +func rawZipFixture(t *testing.T, files map[string]string) []byte { + t.Helper() + b, err := base64.StdEncoding.DecodeString(zipFixtureB64(t, files)) + if err != nil { + t.Fatalf("decode fixture: %v", err) + } + return b +} + +// newFakeClient returns a client pointed at an httptest.Server whose handler is +// provided by the caller, with a no-op token source. +func newFakeClient(t *testing.T, handler http.HandlerFunc) *client { + t.Helper() + srv := httptest.NewServer(handler) + t.Cleanup(srv.Close) + c, err := newClient(clientOptions{ + Project: "test-project", + Location: "us-central1", + Endpoint: srv.URL, + TokenSource: oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "fake"}), + }) + if err != nil { + t.Fatalf("newClient: %v", err) + } + return c +} + +func TestFetchByID_Latest(t *testing.T) { + payload := zipFixtureB64(t, map[string]string{ + "SKILL.md": "---\nname: emoji\n---\n# do stuff", + "scripts/x.sh": "echo hi", + }) + c := newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + // GetSkill: .../skills/emoji -> zippedFilesystem at top level. + if !strings.HasSuffix(r.URL.Path, "/skills/emoji") { + http.Error(w, "unexpected path "+r.URL.Path, http.StatusNotFound) + return + } + writeJSON(t, w, skillResponse{ + Name: "projects/p/locations/us-central1/skills/emoji", + DefaultRevision: "projects/p/locations/us-central1/skills/emoji/revisions/rev-7", + ZippedFilesystem: payload, + }) + }) + + dir := t.TempDir() + res, err := c.fetch(context.Background(), + selection{SkillRefs: []skillRef{{SkillID: "emoji"}}}, dir, newClaimSet(), 0) + if err != nil { + t.Fatalf("fetch: %v", err) + } + if len(res.skipped) != 0 { + t.Fatalf("unexpected skipped: %+v", res.skipped) + } + if len(res.materialized) != 1 { + t.Fatalf("materialized = %d, want 1", len(res.materialized)) + } + got := res.materialized[0] + if got.SkillID != "emoji" || got.Revision != "rev-7" { + t.Errorf("got %+v, want SkillID=emoji Revision=rev-7", got) + } + assertFile(t, filepath.Join(dir, "emoji", "SKILL.md"), "---\nname: emoji\n---\n# do stuff") + assertFile(t, filepath.Join(dir, "emoji", "scripts", "x.sh"), "echo hi") +} + +func TestFetchByID_PinnedRevision(t *testing.T) { + payload := zipFixtureB64(t, map[string]string{"SKILL.md": "pinned"}) + c := newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + // GetSkillRevision: .../skills/emoji/revisions/rev-3 -> nested under "skill". + if !strings.HasSuffix(r.URL.Path, "/skills/emoji/revisions/rev-3") { + http.Error(w, "unexpected path "+r.URL.Path, http.StatusNotFound) + return + } + writeJSON(t, w, skillRevisionResponse{ + Name: "projects/p/locations/us-central1/skills/emoji/revisions/rev-3", + Skill: skillResponse{ZippedFilesystem: payload}, + }) + }) + + dir := t.TempDir() + res, err := c.fetch(context.Background(), + selection{SkillRefs: []skillRef{{SkillID: "emoji", Revision: "rev-3"}}}, dir, newClaimSet(), 0) + if err != nil { + t.Fatalf("fetch: %v", err) + } + if len(res.materialized) != 1 || res.materialized[0].Revision != "rev-3" { + t.Fatalf("got %+v, want one skill pinned at rev-3", res.materialized) + } + assertFile(t, filepath.Join(dir, "emoji", "SKILL.md"), "pinned") +} + +func TestFetchAll_Paginated(t *testing.T) { + pa := zipFixtureB64(t, map[string]string{"SKILL.md": "a"}) + pb := zipFixtureB64(t, map[string]string{"SKILL.md": "b"}) + c := newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/skills") && r.URL.Query().Get("pageToken") == "": + writeJSON(t, w, listSkillsResponse{ + Skills: []skillResponse{{Name: ".../skills/a"}}, + NextPageToken: "TOK", + }) + case strings.HasSuffix(r.URL.Path, "/skills") && r.URL.Query().Get("pageToken") == "TOK": + writeJSON(t, w, listSkillsResponse{Skills: []skillResponse{{Name: ".../skills/b"}}}) + case strings.HasSuffix(r.URL.Path, "/skills/a"): + writeJSON(t, w, skillResponse{Name: ".../skills/a", ZippedFilesystem: pa}) + case strings.HasSuffix(r.URL.Path, "/skills/b"): + writeJSON(t, w, skillResponse{Name: ".../skills/b", ZippedFilesystem: pb}) + default: + http.Error(w, "unexpected "+r.URL.String(), http.StatusNotFound) + } + }) + + dir := t.TempDir() + res, err := c.fetch(context.Background(), selection{All: true}, dir, newClaimSet(), 0) + if err != nil { + t.Fatalf("fetch: %v", err) + } + if len(res.materialized) != 2 { + t.Fatalf("materialized = %d, want 2 (paged)", len(res.materialized)) + } + assertFile(t, filepath.Join(dir, "a", "SKILL.md"), "a") + assertFile(t, filepath.Join(dir, "b", "SKILL.md"), "b") +} + +func TestFetchByQuery(t *testing.T) { + payload := zipFixtureB64(t, map[string]string{"SKILL.md": "found"}) + c := newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.Contains(r.URL.Path, "/skills:retrieve"): + if got := r.URL.Query().Get("query"); got != "emoji stuff" { + t.Errorf("query = %q, want %q", got, "emoji stuff") + } + writeJSON(t, w, retrieveSkillsResponse{ + RetrievedSkills: []struct { + SkillName string `json:"skillName"` + }{{SkillName: ".../skills/emoji"}}, + }) + case strings.HasSuffix(r.URL.Path, "/skills/emoji"): + writeJSON(t, w, skillResponse{Name: ".../skills/emoji", ZippedFilesystem: payload}) + default: + http.Error(w, "unexpected "+r.URL.String(), http.StatusNotFound) + } + }) + + dir := t.TempDir() + res, err := c.fetch(context.Background(), + selection{Query: "emoji stuff", TopK: 3}, dir, newClaimSet(), 0) + if err != nil { + t.Fatalf("fetch: %v", err) + } + if len(res.materialized) != 1 { + t.Fatalf("materialized = %d, want 1", len(res.materialized)) + } + assertFile(t, filepath.Join(dir, "emoji", "SKILL.md"), "found") +} + +func TestFetchFailSafe_OneBadSkillSkipped(t *testing.T) { + good := zipFixtureB64(t, map[string]string{"SKILL.md": "ok"}) + c := newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/skills/good"): + writeJSON(t, w, skillResponse{Name: ".../skills/good", ZippedFilesystem: good}) + case strings.HasSuffix(r.URL.Path, "/skills/bad"): + http.Error(w, "boom", http.StatusInternalServerError) + default: + http.Error(w, "unexpected", http.StatusNotFound) + } + }) + + dir := t.TempDir() + res, err := c.fetch(context.Background(), + selection{SkillRefs: []skillRef{{SkillID: "good"}, {SkillID: "bad"}}}, dir, newClaimSet(), 0) + if err != nil { + t.Fatalf("fetch returned whole-op error, want fail-safe: %v", err) + } + if len(res.materialized) != 1 || res.materialized[0].SkillID != "good" { + t.Errorf("materialized = %+v, want only 'good'", res.materialized) + } + if len(res.skipped) != 1 || res.skipped[0].skillID != "bad" { + t.Errorf("skipped = %+v, want only 'bad'", res.skipped) + } + // The bad skill left no directory behind. + if _, err := os.Stat(filepath.Join(dir, "bad")); !os.IsNotExist(err) { + t.Errorf("expected no dir for 'bad', stat err = %v", err) + } +} + +func TestFetch_FirstWinsAcrossSharedClaimSet(t *testing.T) { + // Two fetches (simulating two registries) into the SAME dir with a SHARED + // claim set. Both offer skill id "dup"; only the first should be written. + payloadA := zipFixtureB64(t, map[string]string{"SKILL.md": "from-A"}) + payloadB := zipFixtureB64(t, map[string]string{"SKILL.md": "from-B"}) + + makeClient := func(payload string) *client { + return newFakeClient(t, func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/skills/dup") { + http.Error(w, "unexpected "+r.URL.Path, http.StatusNotFound) + return + } + writeJSON(t, w, skillResponse{Name: ".../skills/dup", ZippedFilesystem: payload}) + }) + } + + dir := t.TempDir() + claimed := newClaimSet() + sel := selection{SkillRefs: []skillRef{{SkillID: "dup"}}} + + // Registry 0 wins. + res0, err := makeClient(payloadA).fetch(context.Background(), sel, dir, claimed, 0) + if err != nil { + t.Fatalf("fetch #0: %v", err) + } + if len(res0.materialized) != 1 { + t.Fatalf("registry 0 materialized %d, want 1", len(res0.materialized)) + } + + // Registry 1 offers the same id into the same dir -> skipped (not written). + res1, err := makeClient(payloadB).fetch(context.Background(), sel, dir, claimed, 1) + if err != nil { + t.Fatalf("fetch #1: %v", err) + } + if len(res1.materialized) != 0 { + t.Errorf("registry 1 materialized %d, want 0 (first-wins)", len(res1.materialized)) + } + // The content on disk must be the FIRST writer's. + assertFile(t, filepath.Join(dir, "dup", "SKILL.md"), "from-A") +} + +func TestNewClientValidation(t *testing.T) { + if _, err := newClient(clientOptions{Location: "us-central1"}); err == nil { + t.Error("expected error for missing Project") + } + if _, err := newClient(clientOptions{Project: "p"}); err == nil { + t.Error("expected error for missing Location") + } +} + +// --- safeUnzip unit tests --- + +func TestSafeUnzip_Basic(t *testing.T) { + z := rawZipFixture(t, map[string]string{ + "SKILL.md": "hi", + "scripts/tool.sh": "run", + "references/a.txt": "ref", + }) + dir := t.TempDir() + if err := safeUnzip(z, filepath.Join(dir, "s"), unzipCaps{}.withDefaults()); err != nil { + t.Fatalf("safeUnzip: %v", err) + } + assertFile(t, filepath.Join(dir, "s", "SKILL.md"), "hi") + assertFile(t, filepath.Join(dir, "s", "scripts", "tool.sh"), "run") +} + +func TestSafeUnzip_PreservesExecutableMode(t *testing.T) { + // Build a zip with an executable script (0755) and a normal file (0644), + // setting per-entry modes via CreateHeader. + var buf bytes.Buffer + zw := zip.NewWriter(&buf) + writeEntry := func(name, content string, mode os.FileMode) { + hdr := &zip.FileHeader{Name: name, Method: zip.Deflate} + hdr.SetMode(mode) + w, err := zw.CreateHeader(hdr) + if err != nil { + t.Fatalf("CreateHeader %q: %v", name, err) + } + if _, err := w.Write([]byte(content)); err != nil { + t.Fatalf("write %q: %v", name, err) + } + } + writeEntry("scripts/tool.sh", "#!/bin/sh\necho hi\n", 0o755) + writeEntry("SKILL.md", "hi", 0o644) + if err := zw.Close(); err != nil { + t.Fatalf("zip close: %v", err) + } + + dir := t.TempDir() + if err := safeUnzip(buf.Bytes(), filepath.Join(dir, "s"), unzipCaps{}.withDefaults()); err != nil { + t.Fatalf("safeUnzip: %v", err) + } + + scriptInfo, err := os.Stat(filepath.Join(dir, "s", "scripts", "tool.sh")) + if err != nil { + t.Fatal(err) + } + if scriptInfo.Mode().Perm()&0o100 == 0 { + t.Errorf("script mode = %v, want owner-execute bit set (0755 preserved)", scriptInfo.Mode().Perm()) + } + mdInfo, err := os.Stat(filepath.Join(dir, "s", "SKILL.md")) + if err != nil { + t.Fatal(err) + } + if mdInfo.Mode().Perm()&0o111 != 0 { + t.Errorf("SKILL.md mode = %v, want no execute bits (0644)", mdInfo.Mode().Perm()) + } +} + +func TestSafeUnzip_ZipSlipRejected(t *testing.T) { + for _, bad := range []string{"../evil.txt", "a/../../evil.txt", "/etc/evil"} { + z := rawZipFixture(t, map[string]string{bad: "x"}) + dir := t.TempDir() + err := safeUnzip(z, filepath.Join(dir, "s"), unzipCaps{}.withDefaults()) + if err == nil { + t.Errorf("entry %q: expected rejection, got nil", bad) + continue + } + if _, statErr := os.Stat(filepath.Join(dir, "evil.txt")); !os.IsNotExist(statErr) { + t.Errorf("entry %q: file escaped destination", bad) + } + } +} + +func TestSafeUnzip_TooManyFiles(t *testing.T) { + files := map[string]string{} + for i := 0; i < 5; i++ { + files[fmt.Sprintf("f%d.txt", i)] = "x" + } + z := rawZipFixture(t, files) + err := safeUnzip(z, t.TempDir(), unzipCaps{MaxFiles: 3}.withDefaults()) + if err == nil || !strings.Contains(err.Error(), "exceeds cap") { + t.Fatalf("expected MaxFiles cap error, got %v", err) + } +} + +func TestSafeUnzip_TooLarge(t *testing.T) { + z := rawZipFixture(t, map[string]string{"big.txt": strings.Repeat("A", 1000)}) + err := safeUnzip(z, t.TempDir(), unzipCaps{MaxTotalUnzippedBytes: 100}.withDefaults()) + if err == nil || !strings.Contains(err.Error(), "size") { + t.Fatalf("expected size cap error, got %v", err) + } +} + +func TestSafeUnzip_TooDeep(t *testing.T) { + z := rawZipFixture(t, map[string]string{"a/b/c/d/e/f/g/h/i/j/deep.txt": "x"}) + err := safeUnzip(z, t.TempDir(), unzipCaps{MaxDepth: 3}.withDefaults()) + if err == nil || !strings.Contains(err.Error(), "depth") { + t.Fatalf("expected depth cap error, got %v", err) + } +} + +// --- helpers --- + +func writeJSON(t *testing.T, w http.ResponseWriter, v any) { + t.Helper() + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(v); err != nil { + t.Fatalf("encode response: %v", err) + } +} + +func assertFile(t *testing.T, path, want string) { + t.Helper() + got, err := os.ReadFile(path) + if err != nil { + t.Fatalf("reading %q: %v", path, err) + } + if string(got) != want { + t.Errorf("%q = %q, want %q", path, got, want) + } +} diff --git a/internal/skills/geminienterprise/materialize.go b/internal/skills/geminienterprise/materialize.go new file mode 100644 index 00000000..9c2f9421 --- /dev/null +++ b/internal/skills/geminienterprise/materialize.go @@ -0,0 +1,173 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package geminienterprise turns a harness's skills configuration into on-disk +// skill folders. It sources agentskills.io skills from the Gemini Enterprise +// Skill Registry (a managed, versioned catalog exposed over the Vertex AI +// v1beta1 REST API) and writes each skill to //. +// +// It is harness-agnostic: it only writes files and reports what it wrote (see +// Result). It knows nothing about specific harnesses, SKILLS_DIR, or discovery +// pointers — callers decide how a given harness is told where its skills are. +// +// Scope: read-only. This package never creates, updates, or deletes registry +// skills (authoring is out of scope). +package geminienterprise + +import ( + "context" + "log" + "os" + "strings" + + "github.com/google/ax/internal/config" +) + +// Environment fallbacks for project/location (the registry target_dir is a +// required config field, so it has no env fallback). +const ( + envCloudProject = "GOOGLE_CLOUD_PROJECT" + envCloudLocation = "GOOGLE_CLOUD_LOCATION" + + defaultRegistryLocation = "us-central1" +) + +// Result reports what Materialize wrote: one Written entry per registry that +// produced skills, grouping the skills with the directory they landed in. +type Result struct { + Written []Written +} + +// Written groups the skills materialized from one registry with the directory +// they were written into (each skill at //). +type Written struct { + Dir string + Skills []MaterializedSkill +} + +// Empty reports whether nothing was materialized. +func (r Result) Empty() bool { return len(r.Written) == 0 } + +// MaterializedSkill records one skill written to disk. +type MaterializedSkill struct { + SkillID string + Revision string + Dir string // path of the written skill folder (//) +} + +// Materialize materializes every enabled registry in sc into its configured +// target_dir (each skill at //) and reports what it wrote. +// +// target_dir is a required, validated config field (see config.SkillsConfig), so +// this does not fall back to any env var or the working directory. It is +// fail-safe: disabled registries are skipped, and any registry error is logged +// and swallowed so a skill problem never blocks harness creation. When the same +// skill id would be written into the same dir more than once (within one +// registry's selection, or across registries sharing a dir), the FIRST writer +// wins and later duplicates are skipped with a warning. Substrate/pod +// materialization is a separate, later path; this wires the local flow only. +func Materialize(ctx context.Context, sc config.SkillsConfig) Result { + var res Result + // claimed tracks (target_dir, skill-id) pairs already written across all + // registries so the FIRST writer of an id into a dir wins; later duplicates + // (within one registry's selection, or across registries sharing a dir) are + // skipped with a warning instead of silently overwriting. + claimed := newClaimSet() + // TODO: fetches are sequential (both here and per-skill in client.fetch); + // consider bounded-concurrent download/unzip to speed up materialization. + // TODO: bound overall materialization with a deadline (only per-call HTTP + // timeouts exist today); e.g. a default ~120s that users can override. + for i := range sc.Registries { + rc := sc.Registries[i] + if !rc.Enabled { + continue + } + project := firstNonEmpty(rc.Project, os.Getenv(envCloudProject)) + if project == "" { + log.Printf("skills: registries[%d] enabled but no project (config or %s); skipping", i, envCloudProject) + continue + } + location := firstNonEmpty(rc.Location, os.Getenv(envCloudLocation), defaultRegistryLocation) + + c, err := newClient(clientOptions{Project: project, Location: location}) + if err != nil { + log.Printf("skills: registries[%d] client init failed: %v; skipping", i, err) + continue + } + out, err := c.fetch(ctx, selectionFromConfig(rc), rc.TargetDir, claimed, i) + if err != nil { + log.Printf("skills: registries[%d] fetch failed: %v; continuing", i, err) + continue + } + for _, s := range out.skipped { + log.Printf("skills: registries[%d] skipped %q: %v", i, s.skillID, s.reason) + } + if len(out.materialized) == 0 { + continue + } + log.Printf("skills: registries[%d] materialized %d skill(s) into %s (skipped %d)", + i, len(out.materialized), rc.TargetDir, len(out.skipped)) + res.Written = append(res.Written, Written{Dir: rc.TargetDir, Skills: out.materialized}) + } + return res +} + +// claimSet tracks which (dir, skill-id) pairs have already been materialized, so +// the first writer of a given id into a given dir wins. +type claimSet struct { + seen map[string]int // key -> registry index that first claimed it +} + +func newClaimSet() *claimSet { return &claimSet{seen: map[string]int{}} } + +func claimKey(dir, id string) string { return dir + "\x00" + id } + +// claim records (dir, id) as owned by registry regIdx and returns true if this +// is the first claim. If already claimed, it returns false and the index of the +// registry that won. +func (c *claimSet) claim(dir, id string, regIdx int) (won bool, byRegistry int) { + key := claimKey(dir, id) + if prev, ok := c.seen[key]; ok { + return false, prev + } + c.seen[key] = regIdx + return true, regIdx +} + +// selectionFromConfig maps a SkillsRegistryConfig to a selection: +// - explicit Skills list => by-id (with optional revision pin) +// - else Query => by-query +// - else => all +func selectionFromConfig(rc config.SkillsRegistryConfig) selection { + if len(rc.Skills) > 0 { + refs := make([]skillRef, 0, len(rc.Skills)) + for _, s := range rc.Skills { + refs = append(refs, skillRef{SkillID: s.ID, Revision: s.Revision}) + } + return selection{SkillRefs: refs} + } + if rc.Query != nil && strings.TrimSpace(rc.Query.Text) != "" { + return selection{Query: rc.Query.Text, TopK: rc.Query.TopK} + } + return selection{All: true} +} + +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if v != "" { + return v + } + } + return "" +} diff --git a/internal/skills/geminienterprise/materialize_test.go b/internal/skills/geminienterprise/materialize_test.go new file mode 100644 index 00000000..2779a377 --- /dev/null +++ b/internal/skills/geminienterprise/materialize_test.go @@ -0,0 +1,100 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package geminienterprise + +import ( + "context" + "testing" + + "github.com/google/ax/internal/config" +) + +func TestSelectionFromConfig(t *testing.T) { + t.Run("explicit skills take precedence", func(t *testing.T) { + rc := config.SkillsRegistryConfig{ + Skills: []config.SkillRefConfig{{ID: "emoji"}, {ID: "lowercase", Revision: "rev-3"}}, + Query: &config.SkillsQueryConfig{Text: "ignored"}, + } + sel := selectionFromConfig(rc) + if sel.All || sel.Query != "" { + t.Fatalf("got %+v, want by-id", sel) + } + if len(sel.SkillRefs) != 2 || + sel.SkillRefs[0].SkillID != "emoji" || sel.SkillRefs[0].Revision != "" || + sel.SkillRefs[1].SkillID != "lowercase" || sel.SkillRefs[1].Revision != "rev-3" { + t.Fatalf("refs = %+v", sel.SkillRefs) + } + }) + + t.Run("query when no explicit skills", func(t *testing.T) { + sel := selectionFromConfig(config.SkillsRegistryConfig{ + Query: &config.SkillsQueryConfig{Text: "find gcp", TopK: 5}, + }) + if sel.Query != "find gcp" || sel.TopK != 5 { + t.Fatalf("got %+v, want query=find gcp topK=5", sel) + } + }) + + t.Run("empty query text falls through to all", func(t *testing.T) { + sel := selectionFromConfig(config.SkillsRegistryConfig{Query: &config.SkillsQueryConfig{Text: " "}}) + if !sel.All { + t.Fatalf("got %+v, want All (blank query text)", sel) + } + }) + + t.Run("all when nothing set", func(t *testing.T) { + if sel := selectionFromConfig(config.SkillsRegistryConfig{}); !sel.All { + t.Fatalf("got %+v, want All", sel) + } + }) +} + +func TestClaimSet_FirstWins(t *testing.T) { + cs := newClaimSet() + if won, by := cs.claim("/dir", "s1", 0); !won || by != 0 { + t.Fatalf("first claim = (%v,%d), want (true,0)", won, by) + } + // Same (dir,id) again -> loses, reports the first registry (0). + if won, by := cs.claim("/dir", "s1", 2); won || by != 0 { + t.Fatalf("dup claim = (%v,%d), want (false,0)", won, by) + } + // Same id, different dir -> independent, wins. + if won, _ := cs.claim("/other", "s1", 1); !won { + t.Fatal("same id in different dir should win") + } + // Different id, same dir -> wins. + if won, _ := cs.claim("/dir", "s2", 1); !won { + t.Fatal("different id in same dir should win") + } +} + +func TestMaterialize_DisabledIsEmpty(t *testing.T) { + // Disabled config => no-op, empty result, no error/panic. + if res := Materialize(context.Background(), config.SkillsConfig{}); !res.Empty() { + t.Errorf("disabled Materialize = %+v, want empty", res) + } +} + +func TestMaterialize_EnabledNoProjectIsEmpty(t *testing.T) { + // Enabled with a target_dir but no project (config empty, GOOGLE_CLOUD_PROJECT + // unset) => fail-safe empty result (no panic, no materialization). + t.Setenv(envCloudProject, "") + sc := config.SkillsConfig{Registries: []config.SkillsRegistryConfig{ + {Enabled: true, TargetDir: t.TempDir()}, + }} + if res := Materialize(context.Background(), sc); !res.Empty() { + t.Errorf("no-project Materialize = %+v, want empty", res) + } +} diff --git a/internal/skills/geminienterprise/unzip.go b/internal/skills/geminienterprise/unzip.go new file mode 100644 index 00000000..df1046b9 --- /dev/null +++ b/internal/skills/geminienterprise/unzip.go @@ -0,0 +1,202 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package geminienterprise + +import ( + "archive/zip" + "bytes" + "fmt" + "io" + "os" + "path/filepath" + "strings" +) + +// unzipCaps bounds a defensive local unzip of an untrusted skill payload. Even +// though the registry validates payloads server-side, the local unzip must +// independently guard against Zip-Slip (entries with "..", absolute paths, or +// symlinks) and resource exhaustion. +type unzipCaps struct { + // MaxFiles caps the number of entries per skill archive. + MaxFiles int + // MaxTotalUnzippedBytes caps the total unzipped size per skill archive. + MaxTotalUnzippedBytes int64 + // MaxDepth caps directory nesting depth within a skill archive. + MaxDepth int +} + +// The registry's documented payload caps, used as defaults. +const ( + defaultMaxFiles = 10_000 + defaultMaxTotalBytes = 500 << 20 // 500 MiB + defaultMaxDepth = 8 +) + +// withDefaults fills any zero-valued cap with the registry's documented limit. +func (c unzipCaps) withDefaults() unzipCaps { + if c.MaxFiles <= 0 { + c.MaxFiles = defaultMaxFiles + } + if c.MaxTotalUnzippedBytes <= 0 { + c.MaxTotalUnzippedBytes = defaultMaxTotalBytes + } + if c.MaxDepth <= 0 { + c.MaxDepth = defaultMaxDepth + } + return c +} + +// safeUnzip extracts a zip archive into destDir, defensively rejecting unsafe +// entries and enforcing caps. destDir must not exist yet or must be safe to +// write into; callers clear it beforehand. +// +// Guards: +// - Zip-Slip: every entry must resolve to a path inside destDir. +// - Absolute paths and paths containing ".." are rejected. +// - Symlinks (and any non-regular, non-dir mode) are rejected. +// - MaxFiles / MaxDepth / MaxTotalUnzippedBytes are enforced; the running +// total is checked while copying so a lying uncompressed-size can't be used +// to blow past the cap. +func safeUnzip(archive []byte, destDir string, caps unzipCaps) error { + zr, err := zip.NewReader(bytes.NewReader(archive), int64(len(archive))) + if err != nil { + return fmt.Errorf("opening zip: %w", err) + } + if len(zr.File) > caps.MaxFiles { + return fmt.Errorf("archive has %d entries, exceeds cap %d", len(zr.File), caps.MaxFiles) + } + + // Resolve destDir to an absolute, clean base for containment checks. + absDest, err := filepath.Abs(destDir) + if err != nil { + return fmt.Errorf("resolving dest: %w", err) + } + if err := os.MkdirAll(absDest, 0o755); err != nil { + return fmt.Errorf("creating dest: %w", err) + } + + var total int64 + for _, f := range zr.File { + if err := validateEntryName(f.Name, caps.MaxDepth); err != nil { + return err + } + // Reject symlinks and any special modes; only dirs and regular files. + mode := f.Mode() + if mode&os.ModeSymlink != 0 { + return fmt.Errorf("entry %q is a symlink (rejected)", f.Name) + } + if !mode.IsDir() && !mode.IsRegular() { + return fmt.Errorf("entry %q has unsupported mode %v (rejected)", f.Name, mode) + } + + target := filepath.Join(absDest, filepath.FromSlash(f.Name)) + // Containment: target must be within absDest. + if !withinBase(absDest, target) { + return fmt.Errorf("entry %q escapes destination (zip-slip)", f.Name) + } + + if f.FileInfo().IsDir() { + if err := os.MkdirAll(target, 0o755); err != nil { + return fmt.Errorf("creating dir %q: %w", f.Name, err) + } + continue + } + + if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil { + return fmt.Errorf("creating parent of %q: %w", f.Name, err) + } + written, err := writeCappedFile(f, target, caps.MaxTotalUnzippedBytes-total) + if err != nil { + return err + } + total += written + if total > caps.MaxTotalUnzippedBytes { + return fmt.Errorf("archive exceeds unzipped size cap %d bytes", caps.MaxTotalUnzippedBytes) + } + } + return nil +} + +// validateEntryName rejects absolute paths, "..", and over-deep nesting. +func validateEntryName(name string, maxDepth int) error { + if name == "" { + return fmt.Errorf("empty entry name") + } + // Reject Windows-style and POSIX absolute paths. + if filepath.IsAbs(name) || strings.HasPrefix(name, "/") || strings.HasPrefix(name, `\`) { + return fmt.Errorf("entry %q is an absolute path (rejected)", name) + } + slashed := strings.ReplaceAll(name, `\`, "/") + for _, seg := range strings.Split(slashed, "/") { + if seg == ".." { + return fmt.Errorf("entry %q contains '..' (rejected)", name) + } + } + depth := 0 + for _, seg := range strings.Split(strings.Trim(slashed, "/"), "/") { + if seg != "" && seg != "." { + depth++ + } + } + if depth > maxDepth { + return fmt.Errorf("entry %q nesting depth %d exceeds cap %d", name, depth, maxDepth) + } + return nil +} + +// withinBase reports whether target is inside base (after cleaning). +func withinBase(base, target string) bool { + rel, err := filepath.Rel(base, target) + if err != nil { + return false + } + return rel != ".." && !strings.HasPrefix(rel, ".."+string(os.PathSeparator)) +} + +// writeCappedFile copies one zip entry to disk, refusing to write more than +// remaining bytes (so a mismatched declared size can't overrun the cap). +func writeCappedFile(f *zip.File, target string, remaining int64) (int64, error) { + if remaining < 0 { + return 0, fmt.Errorf("unzipped size cap exceeded before %q", f.Name) + } + rc, err := f.Open() + if err != nil { + return 0, fmt.Errorf("opening entry %q: %w", f.Name, err) + } + defer rc.Close() + + // Preserve the archive entry's permission bits so executable skill scripts + // (e.g. under scripts/) stay executable after extraction. Chmod after create + // so the mode sticks regardless of the process umask. + perm := f.Mode().Perm() + out, err := os.OpenFile(target, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm) + if err != nil { + return 0, fmt.Errorf("creating file %q: %w", f.Name, err) + } + defer out.Close() + if err := out.Chmod(perm); err != nil { + return 0, fmt.Errorf("setting mode on %q: %w", f.Name, err) + } + + // Limit the copy to remaining+1 so we can detect overrun deterministically. + n, err := io.Copy(out, io.LimitReader(rc, remaining+1)) + if err != nil { + return n, fmt.Errorf("writing %q: %w", f.Name, err) + } + if n > remaining { + return n, fmt.Errorf("entry %q exceeds remaining unzipped size budget", f.Name) + } + return n, nil +} diff --git a/internal/telemetry/telemetry.go b/internal/telemetry/telemetry.go new file mode 100644 index 00000000..ac0da9ea --- /dev/null +++ b/internal/telemetry/telemetry.go @@ -0,0 +1,106 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package telemetry + +import ( + "context" + "fmt" + "os" + + "cloud.google.com/go/compute/metadata" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc" + "go.opentelemetry.io/otel/propagation" + "go.opentelemetry.io/otel/sdk/resource" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + semconv "go.opentelemetry.io/otel/semconv/v1.26.0" + oauth2google "golang.org/x/oauth2/google" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/google" +) + +const defaultEndpoint = "telemetry.googleapis.com" + +// WithGCPCredentials returns an option that configures the OTLP exporter to use Google Cloud credentials. +func WithGCPCredentials() otlptracegrpc.Option { + bundle := google.NewDefaultCredentials() + return otlptracegrpc.WithDialOption( + grpc.WithTransportCredentials(bundle.TransportCredentials()), + grpc.WithPerRPCCredentials(bundle.PerRPCCredentials()), + ) +} + +// SetTraceProvider initializes the OpenTelemetry SDK. +// It returns a shutdown function that should be called when the application exits. +func SetTraceProvider(ctx context.Context, service string, opts ...otlptracegrpc.Option) (func(context.Context) error, error) { + // 1. Set global propagator. This is crucial for context propagation over gRPC/HTTP. + otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator( + propagation.TraceContext{}, + propagation.Baggage{}, + )) + + // 2. Create OTLP Exporter. + exporter, err := otlptracegrpc.New(ctx, opts...) + if err != nil { + return nil, fmt.Errorf("failed to create OTLP trace exporter: %w", err) + } + + // 3. Define Resource. + attrs := []attribute.KeyValue{ + semconv.ServiceNameKey.String(service), + } + if projectID := detectGCPProjectID(ctx); projectID != "" { + attrs = append(attrs, attribute.String("gcp.project_id", projectID)) + } + + res, err := resource.New(ctx, + resource.WithAttributes(attrs...), + resource.WithProcess(), + resource.WithTelemetrySDK(), + ) + if err != nil { + return nil, fmt.Errorf("failed to create resource: %w", err) + } + + // 4. Create TracerProvider. + bsp := sdktrace.NewBatchSpanProcessor(exporter) + tp := sdktrace.NewTracerProvider( + sdktrace.WithSampler(sdktrace.AlwaysSample()), + sdktrace.WithResource(res), + sdktrace.WithSpanProcessor(bsp), + ) + + otel.SetTracerProvider(tp) + + return func(shutdownCtx context.Context) error { + return tp.Shutdown(shutdownCtx) + }, nil +} + +func detectGCPProjectID(ctx context.Context) string { + if proj := os.Getenv("GOOGLE_CLOUD_PROJECT"); proj != "" { + return proj + } + if metadata.OnGCEWithContext(ctx) { + if proj, err := metadata.ProjectIDWithContext(ctx); err == nil { + return proj + } + } + if creds, err := oauth2google.FindDefaultCredentials(ctx); err == nil && creds.ProjectID != "" { + return creds.ProjectID + } + return "" +} diff --git a/manifests/README.md b/manifests/README.md new file mode 100644 index 00000000..fe3678de --- /dev/null +++ b/manifests/README.md @@ -0,0 +1,193 @@ +# AX Harness Deployment on Kubernetes + +> [!WARNING] +> +> This path is experimental and incomplete: the manifests, scripts, and +> runtime behavior will change and may break without notice. + +This directory contains Kubernetes manifests and configurations to deploy +and verify the AX on Kubernetes using Agent Substrate. + +The target Kubernetes cluster is assumed to have +[Agent Substrate](https://github.com/agent-substrate/substrate) installed. + +--- + +## 🚀 Deploying to Agent Substrate + +### 1. Build and Deploy + +> [!NOTE] +> Do not manually edit `manifests/ax-deployment.yaml`. The installation script automatically injects your `${GEMINI_API_KEY}`, `${AX_SNAPSHOTS_BUCKET}`, and the built `${AX_IMAGE}`, `${AX_ANTIGRAVITY_IMAGE}`, and `${ATEOM_IMAGE}` references during deployment. + +The installation script builds the required images and applies the resolved +manifests to your cluster: + +- the **ax** image (the Go `ax` binary only), built from the `ax` + target of `cmd/ax/Dockerfile`; used by the ax-server and the Go interactions + harness, +- the **ax-antigravity** image (the Antigravity Python + sidecar: SDK, localharness, agent), built from the `antigravity` target of + `cmd/ax/Dockerfile`; used by the antigravity harness actor, +- the **ateom-gvisor** worker image, built with `ko` from the `go.mod` pinned + substrate module. + +#### Build prerequisites + +The ax-antigravity image bundles the antigravity SDK, installed from PyPI at +build time. Both images target the cluster's **linux/amd64** +nodes and are built with `--platform linux/amd64`. + +You also need a container engine to build and push the ax images. The script +auto-detects one (preferring a **running** docker, then podman); force a choice +with `CONTAINER_ENGINE=docker` or `CONTAINER_ENGINE=podman`: + +- **Docker** — Docker Desktop (macOS; cross-builds linux/amd64 via emulation) or + Docker Engine (Linux; native). +- **Podman** — on macOS, start a machine first with `podman machine init && + podman machine start` (cross-builds linux/amd64 via emulation); on Linux it + runs natively (podman/buildah >= 4.0). + +#### Registry authentication + +`GOOGLE_CLOUD_PROJECT` sets `AX_IMAGE_REPO=gcr.io/$GOOGLE_CLOUD_PROJECT`. The deploy pushes three +images — the **ax** and **ax-antigravity** images (via your container engine) and +the **ateom** image (via `ko`) — and all authenticate through the gcloud +credential helper: + +```bash +gcloud auth login # authenticate gcloud +gcloud auth configure-docker # set up the gcr.io credential helper +``` + +#### Deploy + +The event log is stored in Postgres. By default ax-server connects to an +**existing** Postgres that you provide via the `AX_EVENTLOG_DSN` env var (bring your own database). Pass `--deploy-postgres` to also +create a **bundled** Postgres in-cluster instead (for testing). + +```bash +export GOOGLE_CLOUD_PROJECT="ax-substrate" # Your GCP project ID +export GEMINI_API_KEY="your-api-key" +export AX_SNAPSHOTS_BUCKET="snapshot-substrate-test-$GOOGLE_CLOUD_PROJECT" + +# Connect to your existing Postgres: +export AX_EVENTLOG_DSN="postgres://user:pass@host:5432/db?sslmode=require" +./hack/install-ax.sh --deploy-ax-server + +# Or deploy a bundled Postgres for testing: +./hack/install-ax.sh --deploy-ax-server --deploy-postgres +``` + +The bundled Postgres uses an auto-generated password. To get its DSN: + +```bash +kubectl get secret ax-eventlog-postgres -n ax -o go-template='{{.data.dsn | base64decode}}' +``` + +#### Vertex AI access for the Antigravity Interactions harness + +> [!NOTE] +> You may skip this if you only use the default harness. + +The Antigravity **Interactions** harness (`ax harness antigravity-interactions`, +ActorTemplate `ax-harness-interactions-template`) calls the Vertex AI GenAI API +and authenticates with the actor's Google Cloud credentials — unlike the default +Antigravity harness, which uses `GEMINI_API_KEY`. + +The worker pods (WorkerPool `ax-harness-workerpool`, namespace `ax`) have no GSA +annotation, so with Workload Identity the actor authenticates **directly as the +Kubernetes ServiceAccount principal** `ax/default`. Grant that principal +`roles/aiplatform.user` on the project the harness reads from `GOOGLE_CLOUD_PROJECT`, or +Vertex calls fail with `PermissionDenied (403)`. IAM changes can take a minute or two to propagate. + +```bash +PROJECT_NUMBER="$(gcloud projects describe "${GOOGLE_CLOUD_PROJECT}" --format='value(projectNumber)')" + +gcloud projects add-iam-policy-binding "${GOOGLE_CLOUD_PROJECT}" \ + --role=roles/aiplatform.user \ + --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${GOOGLE_CLOUD_PROJECT}.svc.id.goog/subject/ns/ax/sa/default" \ + --condition=None +``` + +### 2. Port-Forward Services + +```bash +# Port-forward the ax-server ReplicaSet +kubectl port-forward -n ax rs/ax-server 8494:8494 +``` + +### 3. Test End-to-End + +Run an execution targeting the port-forwarded server. + +```bash +ax exec --server=localhost:8494 --input="hello, who are you?" +``` + +The server should respond with something like: +```text +Conversation: fb344a18-3720-4c4f-8a6e-2ce34db975b3 + +⏺ hello, who are you? + +I am a helpful assistant. How can I help you today? +``` +*The request is served by the antigravity harness actor running on Substrate.* + +## 🧹 How to Uninstall + +To remove the AX server and its components, run: + +```bash +./hack/install-ax.sh --delete-ax-server +``` + +> [!NOTE] +> The event-log database is preserved by default. If you want to +> delete everything including the data, after the command above, be careful and +> run: +> +> ```bash +> kubectl delete namespace ax +> ``` + +--- + +## 🛠️ Inspection & Diagnostics + +Use the **`kubectl ate`** CLI tool to inspect the live states of +active actors and allocated standby worker pool instances: + +```bash +kubectl ate get actors -a ax + +kubectl ate get workers +``` + +List the pods running in the `ax` namespace: + +```bash +# Add `-o wide` to see node/IP assignments, or `-w` to watch status changes. +kubectl get pods -n ax +``` + +## Substrate compatibility + +AX pins [Agent Substrate](https://github.com/agent-substrate/substrate) in +`go.mod`, and the **ateom** worker image is built from that pinned version. The +cluster's substrate **CRDs and control plane** must be compatible with the +manifest AX applies. + +When installing substrate, keep three things aligned: the ax `go.mod` pin = your +local substrate checkout = the cluster's installed substrate. + +```bash +# Get AX's pinned substrate commit: +commit=$(go list -m -f '{{.Version}}' github.com/agent-substrate/substrate | sed 's/.*-//') +echo "$commit" # e.g. fe93d160a1df + +# Check it out on a normal branch in your substrate clone (avoids a detached HEAD): +git -C fetch origin +git -C switch -C ax-pinned "$commit" +``` diff --git a/manifests/ax-deployment.yaml b/manifests/ax-deployment.yaml new file mode 100644 index 00000000..fda42868 --- /dev/null +++ b/manifests/ax-deployment.yaml @@ -0,0 +1,154 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# TODO(jbd): This yaml will eventually become the ax-deployment.yaml.tmpl. + +# --------------------------------------------------------------------------- +# Namespaces +# --------------------------------------------------------------------------- +# Namespace "ax" is reserved for built-in harnesses. Users who bring their own +# harness images should add their own namespaces. +apiVersion: v1 +kind: Namespace +metadata: + name: ax +--- +# --------------------------------------------------------------------------- +# Built-in harnesses +# --------------------------------------------------------------------------- +apiVersion: ate.dev/v1alpha1 +kind: WorkerPool +metadata: + name: ax-harness-workerpool + namespace: ax + labels: + workload: ax-harness +spec: + replicas: 5 + ateomImage: ${ATEOM_IMAGE} +--- +apiVersion: ate.dev/v1alpha1 +kind: ActorTemplate +metadata: + name: ax-harness-antigravity-template + namespace: ax +spec: + workerSelector: + matchLabels: + workload: ax-harness + pauseImage: "gcr.io/gke-release/pause@sha256:bcbd57ba5653580ec647b16d8163cdd1112df3609129b01f912a8032e48265da" + containers: + - name: "axantigravity" + image: ${AX_ANTIGRAVITY_IMAGE} + # Substrate ignores the image CMD/ENV/WORKDIR, so the harness command and + # environment are specified here. `ax harness antigravity` forks the + # Antigravity Python sidecar, which serves on port 80 because substrate + # DNATs workerPodIP:80. + command: ["/ax-app/ax", "harness", "antigravity", + "--host", "0.0.0.0", "--port", "80"] + env: + - name: GEMINI_API_KEY + value: "${GEMINI_API_KEY}" + - name: PYTHONPATH + value: "/ax-app:/ax-app/python" + - name: PYTHONUNBUFFERED + value: "1" + - name: SKILLS_DIR + value: "/ax/skills" + - name: AX_HARNESS_WORKDIR + value: "/workspace" + readyz: + httpGet: + path: /readyz + port: 8081 + snapshotsConfig: + location: gs://${AX_SNAPSHOTS_BUCKET}/axantigravity/ +--- +apiVersion: ate.dev/v1alpha1 +kind: ActorTemplate +metadata: + name: ax-harness-interactions-template + namespace: ax +spec: + workerSelector: + matchLabels: + workload: ax-harness + pauseImage: "gcr.io/gke-release/pause@sha256:bcbd57ba5653580ec647b16d8163cdd1112df3609129b01f912a8032e48265da" + containers: + - name: "axinteractions" + image: ${AX_IMAGE} + # Substrate ignores the image CMD/ENV/WORKDIR, so the harness command and + # environment are specified here. `ax harness antigravity-interactions` runs + # the Go Interactions harness on port 80 because substrate DNATs workerPodIP:80. + command: ["/ax-app/ax", "harness", "antigravity-interactions", + "--host", "0.0.0.0", "--port", "80"] + env: + - name: GOOGLE_CLOUD_PROJECT + value: "${GOOGLE_CLOUD_PROJECT}" + - name: AX_HARNESS_WORKDIR + value: "/workspace" + # Content of manifests/ax.yaml + - name: AX_CONFIG_CONTENT + value: "${AX_CONFIG_CONTENT}" + readyz: + httpGet: + path: /readyz + port: 8081 + snapshotsConfig: + location: gs://${AX_SNAPSHOTS_BUCKET}/axinteractions/ +--- + +# --------------------------------------------------------------------------- +# AX Server +# --------------------------------------------------------------------------- +apiVersion: apps/v1 +kind: ReplicaSet +metadata: + name: ax-server + namespace: ax + labels: + app: ax-server +spec: + replicas: 3 + selector: + matchLabels: + app: ax-server + template: + metadata: + labels: + app: ax-server + spec: + containers: + - name: ax-server + image: ${AX_IMAGE} + command: ["/ax-app/ax", "serve", "--config", "/etc/ax/ax.yaml"] + ports: + - containerPort: 8494 + env: + - name: GEMINI_API_KEY + value: "${GEMINI_API_KEY}" + - name: AX_SUBSTRATE + value: "1" + - name: AX_EVENTLOG_DSN + valueFrom: + secretKeyRef: + name: ax-eventlog-postgres + key: dsn + volumeMounts: + - name: ax-config + mountPath: /etc/ax + volumes: + - name: ax-config + configMap: + name: ax-server-config diff --git a/manifests/ax-postgres.yaml b/manifests/ax-postgres.yaml new file mode 100644 index 00000000..3009e2e2 --- /dev/null +++ b/manifests/ax-postgres.yaml @@ -0,0 +1,87 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# --------------------------------------------------------------------------- +# Event log storage (PostgreSQL) +# +# Applied by hack/install-ax.sh only when --deploy-postgres is passed. By default +# ax-server connects to an existing Postgres via AX_EVENTLOG_DSN and these +# resources are not deployed. The ax-eventlog-postgres Secret (holding the +# password used below and the DSN) is created by install-ax.sh. +# --------------------------------------------------------------------------- +apiVersion: v1 +kind: Service +metadata: + name: ax-eventlog-postgres + namespace: ax + labels: + app: ax-eventlog-postgres +spec: + selector: + app: ax-eventlog-postgres + ports: + - port: 5432 + targetPort: 5432 +--- +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: ax-eventlog-postgres + namespace: ax + labels: + app: ax-eventlog-postgres +spec: + serviceName: ax-eventlog-postgres + replicas: 1 + selector: + matchLabels: + app: ax-eventlog-postgres + template: + metadata: + labels: + app: ax-eventlog-postgres + spec: + containers: + - name: postgres + image: postgres:16 + ports: + - containerPort: 5432 + env: + - name: POSTGRES_DB + value: axeventlog + - name: POSTGRES_USER + value: axuser + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: ax-eventlog-postgres + key: password + - name: PGDATA + value: /var/lib/postgresql/data/pgdata + volumeMounts: + - name: data + mountPath: /var/lib/postgresql/data + readinessProbe: + exec: + command: ["pg_isready", "-U", "axuser", "-d", "axeventlog"] + initialDelaySeconds: 5 + periodSeconds: 5 + volumeClaimTemplates: + - metadata: + name: data + spec: + accessModes: ["ReadWriteOnce"] + resources: + requests: + storage: 10Gi diff --git a/manifests/ax.yaml b/manifests/ax.yaml new file mode 100644 index 00000000..8c243fef --- /dev/null +++ b/manifests/ax.yaml @@ -0,0 +1,30 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Deploy configuration for AX on Agent Substrate. install-ax.sh renders this +# single source into two places: +# - the ax-server-config ConfigMap (mounted at /etc/ax/ax.yaml), read by the +# controller (`ax serve`); and +# - the AX_CONFIG_CONTENT env (base64 of this file) on each ActorTemplate, +# so a substrate actor can read its config. +version: v1alpha + +eventlog: + postgres: + dsn: ${AX_EVENTLOG_DSN} + +harnesses: + antigravity: + default: true + antigravity-interactions: {} diff --git a/proto/ax.pb.go b/proto/ax.pb.go new file mode 100644 index 00000000..94b1ffe2 --- /dev/null +++ b/proto/ax.pb.go @@ -0,0 +1,1112 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.36.11 +// protoc v5.28.2 +// source: proto/ax.proto + +package proto + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + structpb "google.golang.org/protobuf/types/known/structpb" + reflect "reflect" + sync "sync" + unsafe "unsafe" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +// State represents the state of an execution +type State int32 + +const ( + State_STATE_UNSPECIFIED State = 0 // Unspecified state + State_STATE_PENDING State = 1 + State_STATE_FAILED State = 2 + State_STATE_COMPLETED State = 3 + State_STATE_CANCELED State = 4 +) + +// Enum value maps for State. +var ( + State_name = map[int32]string{ + 0: "STATE_UNSPECIFIED", + 1: "STATE_PENDING", + 2: "STATE_FAILED", + 3: "STATE_COMPLETED", + 4: "STATE_CANCELED", + } + State_value = map[string]int32{ + "STATE_UNSPECIFIED": 0, + "STATE_PENDING": 1, + "STATE_FAILED": 2, + "STATE_COMPLETED": 3, + "STATE_CANCELED": 4, + } +) + +func (x State) Enum() *State { + p := new(State) + *p = x + return p +} + +func (x State) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (State) Descriptor() protoreflect.EnumDescriptor { + return file_proto_ax_proto_enumTypes[0].Descriptor() +} + +func (State) Type() protoreflect.EnumType { + return &file_proto_ax_proto_enumTypes[0] +} + +func (x State) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use State.Descriptor instead. +func (State) EnumDescriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{0} +} + +// CancelReason explains why an in-flight execution is canceled. +type CancelReason int32 + +const ( + CancelReason_CANCEL_REASON_UNSPECIFIED CancelReason = 0 + CancelReason_CANCEL_REASON_USER_REQUESTED CancelReason = 1 + CancelReason_CANCEL_REASON_TIMEOUT CancelReason = 2 + CancelReason_CANCEL_REASON_INTERNAL_ERROR CancelReason = 3 +) + +// Enum value maps for CancelReason. +var ( + CancelReason_name = map[int32]string{ + 0: "CANCEL_REASON_UNSPECIFIED", + 1: "CANCEL_REASON_USER_REQUESTED", + 2: "CANCEL_REASON_TIMEOUT", + 3: "CANCEL_REASON_INTERNAL_ERROR", + } + CancelReason_value = map[string]int32{ + "CANCEL_REASON_UNSPECIFIED": 0, + "CANCEL_REASON_USER_REQUESTED": 1, + "CANCEL_REASON_TIMEOUT": 2, + "CANCEL_REASON_INTERNAL_ERROR": 3, + } +) + +func (x CancelReason) Enum() *CancelReason { + p := new(CancelReason) + *p = x + return p +} + +func (x CancelReason) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (CancelReason) Descriptor() protoreflect.EnumDescriptor { + return file_proto_ax_proto_enumTypes[1].Descriptor() +} + +func (CancelReason) Type() protoreflect.EnumType { + return &file_proto_ax_proto_enumTypes[1] +} + +func (x CancelReason) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use CancelReason.Descriptor instead. +func (CancelReason) EnumDescriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{1} +} + +// Message is a message in the history. +type Message struct { + state protoimpl.MessageState `protogen:"open.v1"` + Role string `protobuf:"bytes,1,opt,name=role,proto3" json:"role,omitempty"` // user, assistant, or model + Content *Content `protobuf:"bytes,2,opt,name=content,proto3" json:"content,omitempty"` // content of the message + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Message) Reset() { + *x = Message{} + mi := &file_proto_ax_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Message) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Message) ProtoMessage() {} + +func (x *Message) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[0] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use Message.ProtoReflect.Descriptor instead. +func (*Message) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{0} +} + +func (x *Message) GetRole() string { + if x != nil { + return x.Role + } + return "" +} + +func (x *Message) GetContent() *Content { + if x != nil { + return x.Content + } + return nil +} + +// A conversation is the historical session that consist of +// a number of execution. A conversation cannot be continued +// before the last execution is completed or failed. +type ConversationEvent struct { + state protoimpl.MessageState `protogen:"open.v1"` + ConversationId string `protobuf:"bytes,1,opt,name=conversation_id,json=conversationId,proto3" json:"conversation_id,omitempty"` + Seq int32 `protobuf:"varint,2,opt,name=seq,proto3" json:"seq,omitempty"` + ExecId string `protobuf:"bytes,3,opt,name=exec_id,json=execId,proto3" json:"exec_id,omitempty"` + HarnessId string `protobuf:"bytes,4,opt,name=harness_id,json=harnessId,proto3" json:"harness_id,omitempty"` + HarnessConfig *structpb.Struct `protobuf:"bytes,5,opt,name=harness_config,json=harnessConfig,proto3" json:"harness_config,omitempty"` + Messages []*Message `protobuf:"bytes,6,rep,name=messages,proto3" json:"messages,omitempty"` + State State `protobuf:"varint,7,opt,name=state,proto3,enum=ax.State" json:"state,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ConversationEvent) Reset() { + *x = ConversationEvent{} + mi := &file_proto_ax_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ConversationEvent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ConversationEvent) ProtoMessage() {} + +func (x *ConversationEvent) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ConversationEvent.ProtoReflect.Descriptor instead. +func (*ConversationEvent) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{1} +} + +func (x *ConversationEvent) GetConversationId() string { + if x != nil { + return x.ConversationId + } + return "" +} + +func (x *ConversationEvent) GetSeq() int32 { + if x != nil { + return x.Seq + } + return 0 +} + +func (x *ConversationEvent) GetExecId() string { + if x != nil { + return x.ExecId + } + return "" +} + +func (x *ConversationEvent) GetHarnessId() string { + if x != nil { + return x.HarnessId + } + return "" +} + +func (x *ConversationEvent) GetHarnessConfig() *structpb.Struct { + if x != nil { + return x.HarnessConfig + } + return nil +} + +func (x *ConversationEvent) GetMessages() []*Message { + if x != nil { + return x.Messages + } + return nil +} + +func (x *ConversationEvent) GetState() State { + if x != nil { + return x.State + } + return State_STATE_UNSPECIFIED +} + +type HarnessStart struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Per-execution harness configuration. + HarnessConfig []byte `protobuf:"bytes,1,opt,name=harness_config,json=harnessConfig,proto3" json:"harness_config,omitempty"` + Messages []*Message `protobuf:"bytes,2,rep,name=messages,proto3" json:"messages,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *HarnessStart) Reset() { + *x = HarnessStart{} + mi := &file_proto_ax_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HarnessStart) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HarnessStart) ProtoMessage() {} + +func (x *HarnessStart) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[2] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HarnessStart.ProtoReflect.Descriptor instead. +func (*HarnessStart) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{2} +} + +func (x *HarnessStart) GetHarnessConfig() []byte { + if x != nil { + return x.HarnessConfig + } + return nil +} + +func (x *HarnessStart) GetMessages() []*Message { + if x != nil { + return x.Messages + } + return nil +} + +// HarnessCancel aborts the in-flight harness execution. +type HarnessCancel struct { + state protoimpl.MessageState `protogen:"open.v1"` + Reason CancelReason `protobuf:"varint,1,opt,name=reason,proto3,enum=ax.CancelReason" json:"reason,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *HarnessCancel) Reset() { + *x = HarnessCancel{} + mi := &file_proto_ax_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HarnessCancel) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HarnessCancel) ProtoMessage() {} + +func (x *HarnessCancel) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[3] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HarnessCancel.ProtoReflect.Descriptor instead. +func (*HarnessCancel) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{3} +} + +func (x *HarnessCancel) GetReason() CancelReason { + if x != nil { + return x.Reason + } + return CancelReason_CANCEL_REASON_UNSPECIFIED +} + +type HarnessRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + ConversationId string `protobuf:"bytes,1,opt,name=conversation_id,json=conversationId,proto3" json:"conversation_id,omitempty"` + HarnessId string `protobuf:"bytes,2,opt,name=harness_id,json=harnessId,proto3" json:"harness_id,omitempty"` + // Types that are valid to be assigned to Type: + // + // *HarnessRequest_Start + // *HarnessRequest_Cancel + Type isHarnessRequest_Type `protobuf_oneof:"type"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *HarnessRequest) Reset() { + *x = HarnessRequest{} + mi := &file_proto_ax_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HarnessRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HarnessRequest) ProtoMessage() {} + +func (x *HarnessRequest) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[4] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HarnessRequest.ProtoReflect.Descriptor instead. +func (*HarnessRequest) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{4} +} + +func (x *HarnessRequest) GetConversationId() string { + if x != nil { + return x.ConversationId + } + return "" +} + +func (x *HarnessRequest) GetHarnessId() string { + if x != nil { + return x.HarnessId + } + return "" +} + +func (x *HarnessRequest) GetType() isHarnessRequest_Type { + if x != nil { + return x.Type + } + return nil +} + +func (x *HarnessRequest) GetStart() *HarnessStart { + if x != nil { + if x, ok := x.Type.(*HarnessRequest_Start); ok { + return x.Start + } + } + return nil +} + +func (x *HarnessRequest) GetCancel() *HarnessCancel { + if x != nil { + if x, ok := x.Type.(*HarnessRequest_Cancel); ok { + return x.Cancel + } + } + return nil +} + +type isHarnessRequest_Type interface { + isHarnessRequest_Type() +} + +type HarnessRequest_Start struct { + Start *HarnessStart `protobuf:"bytes,3,opt,name=start,proto3,oneof"` +} + +type HarnessRequest_Cancel struct { + Cancel *HarnessCancel `protobuf:"bytes,4,opt,name=cancel,proto3,oneof"` +} + +func (*HarnessRequest_Start) isHarnessRequest_Type() {} + +func (*HarnessRequest_Cancel) isHarnessRequest_Type() {} + +type HarnessOutputs struct { + state protoimpl.MessageState `protogen:"open.v1"` + Messages []*Message `protobuf:"bytes,1,rep,name=messages,proto3" json:"messages,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *HarnessOutputs) Reset() { + *x = HarnessOutputs{} + mi := &file_proto_ax_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HarnessOutputs) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HarnessOutputs) ProtoMessage() {} + +func (x *HarnessOutputs) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HarnessOutputs.ProtoReflect.Descriptor instead. +func (*HarnessOutputs) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{5} +} + +func (x *HarnessOutputs) GetMessages() []*Message { + if x != nil { + return x.Messages + } + return nil +} + +type Error struct { + state protoimpl.MessageState `protogen:"open.v1"` + // A status code, corresponds to google.golang.org/grpc/codes. + Code int32 `protobuf:"varint,1,opt,name=code,proto3" json:"code,omitempty"` + // A description of the error. + Description string `protobuf:"bytes,2,opt,name=description,proto3" json:"description,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Error) Reset() { + *x = Error{} + mi := &file_proto_ax_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Error) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Error) ProtoMessage() {} + +func (x *Error) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[6] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use Error.ProtoReflect.Descriptor instead. +func (*Error) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{6} +} + +func (x *Error) GetCode() int32 { + if x != nil { + return x.Code + } + return 0 +} + +func (x *Error) GetDescription() string { + if x != nil { + return x.Description + } + return "" +} + +type HarnessEnd struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Terminal state for the harness execution. + State State `protobuf:"varint,1,opt,name=state,proto3,enum=ax.State" json:"state,omitempty"` + // Optional error details. + Error *Error `protobuf:"bytes,2,opt,name=error,proto3" json:"error,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *HarnessEnd) Reset() { + *x = HarnessEnd{} + mi := &file_proto_ax_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HarnessEnd) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HarnessEnd) ProtoMessage() {} + +func (x *HarnessEnd) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[7] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HarnessEnd.ProtoReflect.Descriptor instead. +func (*HarnessEnd) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{7} +} + +func (x *HarnessEnd) GetState() State { + if x != nil { + return x.State + } + return State_STATE_UNSPECIFIED +} + +func (x *HarnessEnd) GetError() *Error { + if x != nil { + return x.Error + } + return nil +} + +type HarnessResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + ConversationId string `protobuf:"bytes,1,opt,name=conversation_id,json=conversationId,proto3" json:"conversation_id,omitempty"` + // Types that are valid to be assigned to Type: + // + // *HarnessResponse_Outputs + // *HarnessResponse_End + Type isHarnessResponse_Type `protobuf_oneof:"type"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *HarnessResponse) Reset() { + *x = HarnessResponse{} + mi := &file_proto_ax_proto_msgTypes[8] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HarnessResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HarnessResponse) ProtoMessage() {} + +func (x *HarnessResponse) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[8] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HarnessResponse.ProtoReflect.Descriptor instead. +func (*HarnessResponse) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{8} +} + +func (x *HarnessResponse) GetConversationId() string { + if x != nil { + return x.ConversationId + } + return "" +} + +func (x *HarnessResponse) GetType() isHarnessResponse_Type { + if x != nil { + return x.Type + } + return nil +} + +func (x *HarnessResponse) GetOutputs() *HarnessOutputs { + if x != nil { + if x, ok := x.Type.(*HarnessResponse_Outputs); ok { + return x.Outputs + } + } + return nil +} + +func (x *HarnessResponse) GetEnd() *HarnessEnd { + if x != nil { + if x, ok := x.Type.(*HarnessResponse_End); ok { + return x.End + } + } + return nil +} + +type isHarnessResponse_Type interface { + isHarnessResponse_Type() +} + +type HarnessResponse_Outputs struct { + Outputs *HarnessOutputs `protobuf:"bytes,2,opt,name=outputs,proto3,oneof"` +} + +type HarnessResponse_End struct { + End *HarnessEnd `protobuf:"bytes,3,opt,name=end,proto3,oneof"` +} + +func (*HarnessResponse_Outputs) isHarnessResponse_Type() {} + +func (*HarnessResponse_End) isHarnessResponse_Type() {} + +// ExecRequest for executing. +type ExecRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + ConversationId string `protobuf:"bytes,1,opt,name=conversation_id,json=conversationId,proto3" json:"conversation_id,omitempty"` // Unique conversation identifier + Inputs []*Message `protobuf:"bytes,2,rep,name=inputs,proto3" json:"inputs,omitempty"` // New inputs + LastSeq int32 `protobuf:"varint,3,opt,name=last_seq,json=lastSeq,proto3" json:"last_seq,omitempty"` // Last sequence number seen by the client + HarnessId string `protobuf:"bytes,4,opt,name=harness_id,json=harnessId,proto3" json:"harness_id,omitempty"` // Harness ID, empty selects the default harness + HarnessConfig []byte `protobuf:"bytes,5,opt,name=harness_config,json=harnessConfig,proto3" json:"harness_config,omitempty"` // Per-request harness configuration (opaque JSON), if any + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ExecRequest) Reset() { + *x = ExecRequest{} + mi := &file_proto_ax_proto_msgTypes[9] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ExecRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ExecRequest) ProtoMessage() {} + +func (x *ExecRequest) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[9] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ExecRequest.ProtoReflect.Descriptor instead. +func (*ExecRequest) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{9} +} + +func (x *ExecRequest) GetConversationId() string { + if x != nil { + return x.ConversationId + } + return "" +} + +func (x *ExecRequest) GetInputs() []*Message { + if x != nil { + return x.Inputs + } + return nil +} + +func (x *ExecRequest) GetLastSeq() int32 { + if x != nil { + return x.LastSeq + } + return 0 +} + +func (x *ExecRequest) GetHarnessId() string { + if x != nil { + return x.HarnessId + } + return "" +} + +func (x *ExecRequest) GetHarnessConfig() []byte { + if x != nil { + return x.HarnessConfig + } + return nil +} + +// ExecResponse contains the result of an execution. +type ExecResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Outputs []*Message `protobuf:"bytes,1,rep,name=outputs,proto3" json:"outputs,omitempty"` // Output content + Seq int32 `protobuf:"varint,2,opt,name=seq,proto3" json:"seq,omitempty"` // Seq of the outputs + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ExecResponse) Reset() { + *x = ExecResponse{} + mi := &file_proto_ax_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ExecResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ExecResponse) ProtoMessage() {} + +func (x *ExecResponse) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[10] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ExecResponse.ProtoReflect.Descriptor instead. +func (*ExecResponse) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{10} +} + +func (x *ExecResponse) GetOutputs() []*Message { + if x != nil { + return x.Outputs + } + return nil +} + +func (x *ExecResponse) GetSeq() int32 { + if x != nil { + return x.Seq + } + return 0 +} + +type DeleteConversationRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + ConversationId string `protobuf:"bytes,1,opt,name=conversation_id,json=conversationId,proto3" json:"conversation_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *DeleteConversationRequest) Reset() { + *x = DeleteConversationRequest{} + mi := &file_proto_ax_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *DeleteConversationRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeleteConversationRequest) ProtoMessage() {} + +func (x *DeleteConversationRequest) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[11] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeleteConversationRequest.ProtoReflect.Descriptor instead. +func (*DeleteConversationRequest) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{11} +} + +func (x *DeleteConversationRequest) GetConversationId() string { + if x != nil { + return x.ConversationId + } + return "" +} + +type DeleteConversationResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *DeleteConversationResponse) Reset() { + *x = DeleteConversationResponse{} + mi := &file_proto_ax_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *DeleteConversationResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeleteConversationResponse) ProtoMessage() {} + +func (x *DeleteConversationResponse) ProtoReflect() protoreflect.Message { + mi := &file_proto_ax_proto_msgTypes[12] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeleteConversationResponse.ProtoReflect.Descriptor instead. +func (*DeleteConversationResponse) Descriptor() ([]byte, []int) { + return file_proto_ax_proto_rawDescGZIP(), []int{12} +} + +var File_proto_ax_proto protoreflect.FileDescriptor + +const file_proto_ax_proto_rawDesc = "" + + "\n" + + "\x0eproto/ax.proto\x12\x02ax\x1a\x1cgoogle/protobuf/struct.proto\x1a\x13proto/content.proto\"D\n" + + "\aMessage\x12\x12\n" + + "\x04role\x18\x01 \x01(\tR\x04role\x12%\n" + + "\acontent\x18\x02 \x01(\v2\v.ax.ContentR\acontent\"\x90\x02\n" + + "\x11ConversationEvent\x12'\n" + + "\x0fconversation_id\x18\x01 \x01(\tR\x0econversationId\x12\x10\n" + + "\x03seq\x18\x02 \x01(\x05R\x03seq\x12\x17\n" + + "\aexec_id\x18\x03 \x01(\tR\x06execId\x12\x1d\n" + + "\n" + + "harness_id\x18\x04 \x01(\tR\tharnessId\x12>\n" + + "\x0eharness_config\x18\x05 \x01(\v2\x17.google.protobuf.StructR\rharnessConfig\x12'\n" + + "\bmessages\x18\x06 \x03(\v2\v.ax.MessageR\bmessages\x12\x1f\n" + + "\x05state\x18\a \x01(\x0e2\t.ax.StateR\x05state\"^\n" + + "\fHarnessStart\x12%\n" + + "\x0eharness_config\x18\x01 \x01(\fR\rharnessConfig\x12'\n" + + "\bmessages\x18\x02 \x03(\v2\v.ax.MessageR\bmessages\"9\n" + + "\rHarnessCancel\x12(\n" + + "\x06reason\x18\x01 \x01(\x0e2\x10.ax.CancelReasonR\x06reason\"\xb7\x01\n" + + "\x0eHarnessRequest\x12'\n" + + "\x0fconversation_id\x18\x01 \x01(\tR\x0econversationId\x12\x1d\n" + + "\n" + + "harness_id\x18\x02 \x01(\tR\tharnessId\x12(\n" + + "\x05start\x18\x03 \x01(\v2\x10.ax.HarnessStartH\x00R\x05start\x12+\n" + + "\x06cancel\x18\x04 \x01(\v2\x11.ax.HarnessCancelH\x00R\x06cancelB\x06\n" + + "\x04type\"9\n" + + "\x0eHarnessOutputs\x12'\n" + + "\bmessages\x18\x01 \x03(\v2\v.ax.MessageR\bmessages\"=\n" + + "\x05Error\x12\x12\n" + + "\x04code\x18\x01 \x01(\x05R\x04code\x12 \n" + + "\vdescription\x18\x02 \x01(\tR\vdescription\"N\n" + + "\n" + + "HarnessEnd\x12\x1f\n" + + "\x05state\x18\x01 \x01(\x0e2\t.ax.StateR\x05state\x12\x1f\n" + + "\x05error\x18\x02 \x01(\v2\t.ax.ErrorR\x05error\"\x96\x01\n" + + "\x0fHarnessResponse\x12'\n" + + "\x0fconversation_id\x18\x01 \x01(\tR\x0econversationId\x12.\n" + + "\aoutputs\x18\x02 \x01(\v2\x12.ax.HarnessOutputsH\x00R\aoutputs\x12\"\n" + + "\x03end\x18\x03 \x01(\v2\x0e.ax.HarnessEndH\x00R\x03endB\x06\n" + + "\x04type\"\xbc\x01\n" + + "\vExecRequest\x12'\n" + + "\x0fconversation_id\x18\x01 \x01(\tR\x0econversationId\x12#\n" + + "\x06inputs\x18\x02 \x03(\v2\v.ax.MessageR\x06inputs\x12\x19\n" + + "\blast_seq\x18\x03 \x01(\x05R\alastSeq\x12\x1d\n" + + "\n" + + "harness_id\x18\x04 \x01(\tR\tharnessId\x12%\n" + + "\x0eharness_config\x18\x05 \x01(\fR\rharnessConfig\"G\n" + + "\fExecResponse\x12%\n" + + "\aoutputs\x18\x01 \x03(\v2\v.ax.MessageR\aoutputs\x12\x10\n" + + "\x03seq\x18\x02 \x01(\x05R\x03seq\"D\n" + + "\x19DeleteConversationRequest\x12'\n" + + "\x0fconversation_id\x18\x01 \x01(\tR\x0econversationId\"\x1c\n" + + "\x1aDeleteConversationResponse*l\n" + + "\x05State\x12\x15\n" + + "\x11STATE_UNSPECIFIED\x10\x00\x12\x11\n" + + "\rSTATE_PENDING\x10\x01\x12\x10\n" + + "\fSTATE_FAILED\x10\x02\x12\x13\n" + + "\x0fSTATE_COMPLETED\x10\x03\x12\x12\n" + + "\x0eSTATE_CANCELED\x10\x04*\x8c\x01\n" + + "\fCancelReason\x12\x1d\n" + + "\x19CANCEL_REASON_UNSPECIFIED\x10\x00\x12 \n" + + "\x1cCANCEL_REASON_USER_REQUESTED\x10\x01\x12\x19\n" + + "\x15CANCEL_REASON_TIMEOUT\x10\x02\x12 \n" + + "\x1cCANCEL_REASON_INTERNAL_ERROR\x10\x032H\n" + + "\x0eHarnessService\x126\n" + + "\aConnect\x12\x12.ax.HarnessRequest\x1a\x13.ax.HarnessResponse(\x010\x012?\n" + + "\x10ExecutionService\x12+\n" + + "\x04Exec\x12\x0f.ax.ExecRequest\x1a\x10.ax.ExecResponse0\x012j\n" + + "\x13ConversationService\x12S\n" + + "\x12DeleteConversation\x12\x1d.ax.DeleteConversationRequest\x1a\x1e.ax.DeleteConversationResponseB\x1cZ\x1agithub.com/google/ax/protob\x06proto3" + +var ( + file_proto_ax_proto_rawDescOnce sync.Once + file_proto_ax_proto_rawDescData []byte +) + +func file_proto_ax_proto_rawDescGZIP() []byte { + file_proto_ax_proto_rawDescOnce.Do(func() { + file_proto_ax_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_proto_ax_proto_rawDesc), len(file_proto_ax_proto_rawDesc))) + }) + return file_proto_ax_proto_rawDescData +} + +var file_proto_ax_proto_enumTypes = make([]protoimpl.EnumInfo, 2) +var file_proto_ax_proto_msgTypes = make([]protoimpl.MessageInfo, 13) +var file_proto_ax_proto_goTypes = []any{ + (State)(0), // 0: ax.State + (CancelReason)(0), // 1: ax.CancelReason + (*Message)(nil), // 2: ax.Message + (*ConversationEvent)(nil), // 3: ax.ConversationEvent + (*HarnessStart)(nil), // 4: ax.HarnessStart + (*HarnessCancel)(nil), // 5: ax.HarnessCancel + (*HarnessRequest)(nil), // 6: ax.HarnessRequest + (*HarnessOutputs)(nil), // 7: ax.HarnessOutputs + (*Error)(nil), // 8: ax.Error + (*HarnessEnd)(nil), // 9: ax.HarnessEnd + (*HarnessResponse)(nil), // 10: ax.HarnessResponse + (*ExecRequest)(nil), // 11: ax.ExecRequest + (*ExecResponse)(nil), // 12: ax.ExecResponse + (*DeleteConversationRequest)(nil), // 13: ax.DeleteConversationRequest + (*DeleteConversationResponse)(nil), // 14: ax.DeleteConversationResponse + (*Content)(nil), // 15: ax.Content + (*structpb.Struct)(nil), // 16: google.protobuf.Struct +} +var file_proto_ax_proto_depIdxs = []int32{ + 15, // 0: ax.Message.content:type_name -> ax.Content + 16, // 1: ax.ConversationEvent.harness_config:type_name -> google.protobuf.Struct + 2, // 2: ax.ConversationEvent.messages:type_name -> ax.Message + 0, // 3: ax.ConversationEvent.state:type_name -> ax.State + 2, // 4: ax.HarnessStart.messages:type_name -> ax.Message + 1, // 5: ax.HarnessCancel.reason:type_name -> ax.CancelReason + 4, // 6: ax.HarnessRequest.start:type_name -> ax.HarnessStart + 5, // 7: ax.HarnessRequest.cancel:type_name -> ax.HarnessCancel + 2, // 8: ax.HarnessOutputs.messages:type_name -> ax.Message + 0, // 9: ax.HarnessEnd.state:type_name -> ax.State + 8, // 10: ax.HarnessEnd.error:type_name -> ax.Error + 7, // 11: ax.HarnessResponse.outputs:type_name -> ax.HarnessOutputs + 9, // 12: ax.HarnessResponse.end:type_name -> ax.HarnessEnd + 2, // 13: ax.ExecRequest.inputs:type_name -> ax.Message + 2, // 14: ax.ExecResponse.outputs:type_name -> ax.Message + 6, // 15: ax.HarnessService.Connect:input_type -> ax.HarnessRequest + 11, // 16: ax.ExecutionService.Exec:input_type -> ax.ExecRequest + 13, // 17: ax.ConversationService.DeleteConversation:input_type -> ax.DeleteConversationRequest + 10, // 18: ax.HarnessService.Connect:output_type -> ax.HarnessResponse + 12, // 19: ax.ExecutionService.Exec:output_type -> ax.ExecResponse + 14, // 20: ax.ConversationService.DeleteConversation:output_type -> ax.DeleteConversationResponse + 18, // [18:21] is the sub-list for method output_type + 15, // [15:18] is the sub-list for method input_type + 15, // [15:15] is the sub-list for extension type_name + 15, // [15:15] is the sub-list for extension extendee + 0, // [0:15] is the sub-list for field type_name +} + +func init() { file_proto_ax_proto_init() } +func file_proto_ax_proto_init() { + if File_proto_ax_proto != nil { + return + } + file_proto_content_proto_init() + file_proto_ax_proto_msgTypes[4].OneofWrappers = []any{ + (*HarnessRequest_Start)(nil), + (*HarnessRequest_Cancel)(nil), + } + file_proto_ax_proto_msgTypes[8].OneofWrappers = []any{ + (*HarnessResponse_Outputs)(nil), + (*HarnessResponse_End)(nil), + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_proto_ax_proto_rawDesc), len(file_proto_ax_proto_rawDesc)), + NumEnums: 2, + NumMessages: 13, + NumExtensions: 0, + NumServices: 3, + }, + GoTypes: file_proto_ax_proto_goTypes, + DependencyIndexes: file_proto_ax_proto_depIdxs, + EnumInfos: file_proto_ax_proto_enumTypes, + MessageInfos: file_proto_ax_proto_msgTypes, + }.Build() + File_proto_ax_proto = out.File + file_proto_ax_proto_goTypes = nil + file_proto_ax_proto_depIdxs = nil +} diff --git a/proto/ax.proto b/proto/ax.proto new file mode 100644 index 00000000..675fb885 --- /dev/null +++ b/proto/ax.proto @@ -0,0 +1,149 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +syntax = "proto3"; + +package ax; + +import "google/protobuf/struct.proto"; +import "proto/content.proto"; + +option go_package = "github.com/google/ax/proto"; + +// WARNING: This file is in active development and +// significant changes can be made with breaking changes. + +// Message is a message in the history. +message Message { + string role = 1; // user, assistant, or model + Content content = 2; // content of the message +} + +// A conversation is the historical session that consist of +// a number of execution. A conversation cannot be continued +// before the last execution is completed or failed. +message ConversationEvent { + string conversation_id = 1; + int32 seq = 2; + string exec_id = 3; + string harness_id = 4; + google.protobuf.Struct harness_config = 5; + repeated Message messages = 6; + State state = 7; +} + +message HarnessStart { + // Per-execution harness configuration. + bytes harness_config = 1; + repeated Message messages = 2; +} + +// HarnessCancel aborts the in-flight harness execution. +message HarnessCancel { + CancelReason reason = 1; +} + +message HarnessRequest { + string conversation_id = 1; + string harness_id = 2; + oneof type { + HarnessStart start = 3; + HarnessCancel cancel = 4; + } +} + +message HarnessOutputs { + repeated Message messages = 1; +} + +message Error { + // A status code, corresponds to google.golang.org/grpc/codes. + int32 code = 1; + // A description of the error. + string description = 2; +} + +message HarnessEnd { + // Terminal state for the harness execution. + State state = 1; + // Optional error details. + Error error = 2; +} + +message HarnessResponse { + string conversation_id = 1; + oneof type { + HarnessOutputs outputs = 2; + HarnessEnd end = 3; + } +} + +service HarnessService { + // Connect drives one harness execution. The client sends + // HarnessRequest{start} (and may send one HarnessRequest{cancel} + // mid-stream), and the server streams zero or more HarnessResponse{outputs} + // frames terminated by exactly one HarnessResponse{end}. + rpc Connect(stream HarnessRequest) returns (stream HarnessResponse); +} + +// State represents the state of an execution +enum State { + STATE_UNSPECIFIED = 0; // Unspecified state + STATE_PENDING = 1; + STATE_FAILED = 2; + STATE_COMPLETED = 3; + STATE_CANCELED = 4; +} + +// CancelReason explains why an in-flight execution is canceled. +enum CancelReason { + CANCEL_REASON_UNSPECIFIED = 0; + CANCEL_REASON_USER_REQUESTED = 1; + CANCEL_REASON_TIMEOUT = 2; + CANCEL_REASON_INTERNAL_ERROR = 3; +} + +// ExecRequest for executing. +message ExecRequest { + string conversation_id = 1; // Unique conversation identifier + repeated Message inputs = 2; // New inputs + int32 last_seq = 3; // Last sequence number seen by the client + + string harness_id = 4; // Harness ID, empty selects the default harness + bytes harness_config = 5; // Per-request harness configuration (opaque JSON), if any +} + +// ExecResponse contains the result of an execution. +message ExecResponse { + repeated Message outputs = 1; // Output content + int32 seq = 2; // Seq of the outputs +} + +service ExecutionService { + // Exec executes an agentic task or resumes an existing one with streaming responses + // If the conversation_id already exists, it will be resumed. + rpc Exec(ExecRequest) returns (stream ExecResponse); +} + +message DeleteConversationRequest { + string conversation_id = 1; +} + +message DeleteConversationResponse {} + +service ConversationService { + // Deletes conversational events and all event log resources + // for its children executions. + rpc DeleteConversation(DeleteConversationRequest) returns (DeleteConversationResponse); +} diff --git a/proto/ax_grpc.pb.go b/proto/ax_grpc.pb.go new file mode 100644 index 00000000..dc1e0b81 --- /dev/null +++ b/proto/ax_grpc.pb.go @@ -0,0 +1,352 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Code generated by protoc-gen-go-grpc. DO NOT EDIT. +// versions: +// - protoc-gen-go-grpc v1.5.1 +// - protoc v5.28.2 +// source: proto/ax.proto + +package proto + +import ( + context "context" + grpc "google.golang.org/grpc" + codes "google.golang.org/grpc/codes" + status "google.golang.org/grpc/status" +) + +// This is a compile-time assertion to ensure that this generated file +// is compatible with the grpc package it is being compiled against. +// Requires gRPC-Go v1.64.0 or later. +const _ = grpc.SupportPackageIsVersion9 + +const ( + HarnessService_Connect_FullMethodName = "/ax.HarnessService/Connect" +) + +// HarnessServiceClient is the client API for HarnessService service. +// +// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. +type HarnessServiceClient interface { + // Connect drives one harness execution. The client sends + // HarnessRequest{start} (and may send one HarnessRequest{cancel} + // mid-stream), and the server streams zero or more HarnessResponse{outputs} + // frames terminated by exactly one HarnessResponse{end}. + Connect(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[HarnessRequest, HarnessResponse], error) +} + +type harnessServiceClient struct { + cc grpc.ClientConnInterface +} + +func NewHarnessServiceClient(cc grpc.ClientConnInterface) HarnessServiceClient { + return &harnessServiceClient{cc} +} + +func (c *harnessServiceClient) Connect(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[HarnessRequest, HarnessResponse], error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + stream, err := c.cc.NewStream(ctx, &HarnessService_ServiceDesc.Streams[0], HarnessService_Connect_FullMethodName, cOpts...) + if err != nil { + return nil, err + } + x := &grpc.GenericClientStream[HarnessRequest, HarnessResponse]{ClientStream: stream} + return x, nil +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type HarnessService_ConnectClient = grpc.BidiStreamingClient[HarnessRequest, HarnessResponse] + +// HarnessServiceServer is the server API for HarnessService service. +// All implementations must embed UnimplementedHarnessServiceServer +// for forward compatibility. +type HarnessServiceServer interface { + // Connect drives one harness execution. The client sends + // HarnessRequest{start} (and may send one HarnessRequest{cancel} + // mid-stream), and the server streams zero or more HarnessResponse{outputs} + // frames terminated by exactly one HarnessResponse{end}. + Connect(grpc.BidiStreamingServer[HarnessRequest, HarnessResponse]) error + mustEmbedUnimplementedHarnessServiceServer() +} + +// UnimplementedHarnessServiceServer must be embedded to have +// forward compatible implementations. +// +// NOTE: this should be embedded by value instead of pointer to avoid a nil +// pointer dereference when methods are called. +type UnimplementedHarnessServiceServer struct{} + +func (UnimplementedHarnessServiceServer) Connect(grpc.BidiStreamingServer[HarnessRequest, HarnessResponse]) error { + return status.Errorf(codes.Unimplemented, "method Connect not implemented") +} +func (UnimplementedHarnessServiceServer) mustEmbedUnimplementedHarnessServiceServer() {} +func (UnimplementedHarnessServiceServer) testEmbeddedByValue() {} + +// UnsafeHarnessServiceServer may be embedded to opt out of forward compatibility for this service. +// Use of this interface is not recommended, as added methods to HarnessServiceServer will +// result in compilation errors. +type UnsafeHarnessServiceServer interface { + mustEmbedUnimplementedHarnessServiceServer() +} + +func RegisterHarnessServiceServer(s grpc.ServiceRegistrar, srv HarnessServiceServer) { + // If the following call pancis, it indicates UnimplementedHarnessServiceServer was + // embedded by pointer and is nil. This will cause panics if an + // unimplemented method is ever invoked, so we test this at initialization + // time to prevent it from happening at runtime later due to I/O. + if t, ok := srv.(interface{ testEmbeddedByValue() }); ok { + t.testEmbeddedByValue() + } + s.RegisterService(&HarnessService_ServiceDesc, srv) +} + +func _HarnessService_Connect_Handler(srv interface{}, stream grpc.ServerStream) error { + return srv.(HarnessServiceServer).Connect(&grpc.GenericServerStream[HarnessRequest, HarnessResponse]{ServerStream: stream}) +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type HarnessService_ConnectServer = grpc.BidiStreamingServer[HarnessRequest, HarnessResponse] + +// HarnessService_ServiceDesc is the grpc.ServiceDesc for HarnessService service. +// It's only intended for direct use with grpc.RegisterService, +// and not to be introspected or modified (even as a copy) +var HarnessService_ServiceDesc = grpc.ServiceDesc{ + ServiceName: "ax.HarnessService", + HandlerType: (*HarnessServiceServer)(nil), + Methods: []grpc.MethodDesc{}, + Streams: []grpc.StreamDesc{ + { + StreamName: "Connect", + Handler: _HarnessService_Connect_Handler, + ServerStreams: true, + ClientStreams: true, + }, + }, + Metadata: "proto/ax.proto", +} + +const ( + ExecutionService_Exec_FullMethodName = "/ax.ExecutionService/Exec" +) + +// ExecutionServiceClient is the client API for ExecutionService service. +// +// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. +type ExecutionServiceClient interface { + // Exec executes an agentic task or resumes an existing one with streaming responses + // If the conversation_id already exists, it will be resumed. + Exec(ctx context.Context, in *ExecRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ExecResponse], error) +} + +type executionServiceClient struct { + cc grpc.ClientConnInterface +} + +func NewExecutionServiceClient(cc grpc.ClientConnInterface) ExecutionServiceClient { + return &executionServiceClient{cc} +} + +func (c *executionServiceClient) Exec(ctx context.Context, in *ExecRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ExecResponse], error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + stream, err := c.cc.NewStream(ctx, &ExecutionService_ServiceDesc.Streams[0], ExecutionService_Exec_FullMethodName, cOpts...) + if err != nil { + return nil, err + } + x := &grpc.GenericClientStream[ExecRequest, ExecResponse]{ClientStream: stream} + if err := x.ClientStream.SendMsg(in); err != nil { + return nil, err + } + if err := x.ClientStream.CloseSend(); err != nil { + return nil, err + } + return x, nil +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type ExecutionService_ExecClient = grpc.ServerStreamingClient[ExecResponse] + +// ExecutionServiceServer is the server API for ExecutionService service. +// All implementations must embed UnimplementedExecutionServiceServer +// for forward compatibility. +type ExecutionServiceServer interface { + // Exec executes an agentic task or resumes an existing one with streaming responses + // If the conversation_id already exists, it will be resumed. + Exec(*ExecRequest, grpc.ServerStreamingServer[ExecResponse]) error + mustEmbedUnimplementedExecutionServiceServer() +} + +// UnimplementedExecutionServiceServer must be embedded to have +// forward compatible implementations. +// +// NOTE: this should be embedded by value instead of pointer to avoid a nil +// pointer dereference when methods are called. +type UnimplementedExecutionServiceServer struct{} + +func (UnimplementedExecutionServiceServer) Exec(*ExecRequest, grpc.ServerStreamingServer[ExecResponse]) error { + return status.Errorf(codes.Unimplemented, "method Exec not implemented") +} +func (UnimplementedExecutionServiceServer) mustEmbedUnimplementedExecutionServiceServer() {} +func (UnimplementedExecutionServiceServer) testEmbeddedByValue() {} + +// UnsafeExecutionServiceServer may be embedded to opt out of forward compatibility for this service. +// Use of this interface is not recommended, as added methods to ExecutionServiceServer will +// result in compilation errors. +type UnsafeExecutionServiceServer interface { + mustEmbedUnimplementedExecutionServiceServer() +} + +func RegisterExecutionServiceServer(s grpc.ServiceRegistrar, srv ExecutionServiceServer) { + // If the following call pancis, it indicates UnimplementedExecutionServiceServer was + // embedded by pointer and is nil. This will cause panics if an + // unimplemented method is ever invoked, so we test this at initialization + // time to prevent it from happening at runtime later due to I/O. + if t, ok := srv.(interface{ testEmbeddedByValue() }); ok { + t.testEmbeddedByValue() + } + s.RegisterService(&ExecutionService_ServiceDesc, srv) +} + +func _ExecutionService_Exec_Handler(srv interface{}, stream grpc.ServerStream) error { + m := new(ExecRequest) + if err := stream.RecvMsg(m); err != nil { + return err + } + return srv.(ExecutionServiceServer).Exec(m, &grpc.GenericServerStream[ExecRequest, ExecResponse]{ServerStream: stream}) +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type ExecutionService_ExecServer = grpc.ServerStreamingServer[ExecResponse] + +// ExecutionService_ServiceDesc is the grpc.ServiceDesc for ExecutionService service. +// It's only intended for direct use with grpc.RegisterService, +// and not to be introspected or modified (even as a copy) +var ExecutionService_ServiceDesc = grpc.ServiceDesc{ + ServiceName: "ax.ExecutionService", + HandlerType: (*ExecutionServiceServer)(nil), + Methods: []grpc.MethodDesc{}, + Streams: []grpc.StreamDesc{ + { + StreamName: "Exec", + Handler: _ExecutionService_Exec_Handler, + ServerStreams: true, + }, + }, + Metadata: "proto/ax.proto", +} + +const ( + ConversationService_DeleteConversation_FullMethodName = "/ax.ConversationService/DeleteConversation" +) + +// ConversationServiceClient is the client API for ConversationService service. +// +// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. +type ConversationServiceClient interface { + // Deletes conversational events and all event log resources + // for its children executions. + DeleteConversation(ctx context.Context, in *DeleteConversationRequest, opts ...grpc.CallOption) (*DeleteConversationResponse, error) +} + +type conversationServiceClient struct { + cc grpc.ClientConnInterface +} + +func NewConversationServiceClient(cc grpc.ClientConnInterface) ConversationServiceClient { + return &conversationServiceClient{cc} +} + +func (c *conversationServiceClient) DeleteConversation(ctx context.Context, in *DeleteConversationRequest, opts ...grpc.CallOption) (*DeleteConversationResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(DeleteConversationResponse) + err := c.cc.Invoke(ctx, ConversationService_DeleteConversation_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +// ConversationServiceServer is the server API for ConversationService service. +// All implementations must embed UnimplementedConversationServiceServer +// for forward compatibility. +type ConversationServiceServer interface { + // Deletes conversational events and all event log resources + // for its children executions. + DeleteConversation(context.Context, *DeleteConversationRequest) (*DeleteConversationResponse, error) + mustEmbedUnimplementedConversationServiceServer() +} + +// UnimplementedConversationServiceServer must be embedded to have +// forward compatible implementations. +// +// NOTE: this should be embedded by value instead of pointer to avoid a nil +// pointer dereference when methods are called. +type UnimplementedConversationServiceServer struct{} + +func (UnimplementedConversationServiceServer) DeleteConversation(context.Context, *DeleteConversationRequest) (*DeleteConversationResponse, error) { + return nil, status.Errorf(codes.Unimplemented, "method DeleteConversation not implemented") +} +func (UnimplementedConversationServiceServer) mustEmbedUnimplementedConversationServiceServer() {} +func (UnimplementedConversationServiceServer) testEmbeddedByValue() {} + +// UnsafeConversationServiceServer may be embedded to opt out of forward compatibility for this service. +// Use of this interface is not recommended, as added methods to ConversationServiceServer will +// result in compilation errors. +type UnsafeConversationServiceServer interface { + mustEmbedUnimplementedConversationServiceServer() +} + +func RegisterConversationServiceServer(s grpc.ServiceRegistrar, srv ConversationServiceServer) { + // If the following call pancis, it indicates UnimplementedConversationServiceServer was + // embedded by pointer and is nil. This will cause panics if an + // unimplemented method is ever invoked, so we test this at initialization + // time to prevent it from happening at runtime later due to I/O. + if t, ok := srv.(interface{ testEmbeddedByValue() }); ok { + t.testEmbeddedByValue() + } + s.RegisterService(&ConversationService_ServiceDesc, srv) +} + +func _ConversationService_DeleteConversation_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(DeleteConversationRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(ConversationServiceServer).DeleteConversation(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: ConversationService_DeleteConversation_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(ConversationServiceServer).DeleteConversation(ctx, req.(*DeleteConversationRequest)) + } + return interceptor(ctx, in, info, handler) +} + +// ConversationService_ServiceDesc is the grpc.ServiceDesc for ConversationService service. +// It's only intended for direct use with grpc.RegisterService, +// and not to be introspected or modified (even as a copy) +var ConversationService_ServiceDesc = grpc.ServiceDesc{ + ServiceName: "ax.ConversationService", + HandlerType: (*ConversationServiceServer)(nil), + Methods: []grpc.MethodDesc{ + { + MethodName: "DeleteConversation", + Handler: _ConversationService_DeleteConversation_Handler, + }, + }, + Streams: []grpc.StreamDesc{}, + Metadata: "proto/ax.proto", +} diff --git a/proto/content.pb.go b/proto/content.pb.go new file mode 100644 index 00000000..1cae4f60 --- /dev/null +++ b/proto/content.pb.go @@ -0,0 +1,1868 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.36.11 +// protoc v5.28.2 +// source: proto/content.proto + +package proto + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + structpb "google.golang.org/protobuf/types/known/structpb" + reflect "reflect" + sync "sync" + unsafe "unsafe" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +// Resolution for input media (images/video). +type MediaResolution int32 + +const ( + MediaResolution_MEDIA_RESOLUTION_UNSPECIFIED MediaResolution = 0 + MediaResolution_LOW MediaResolution = 1 + MediaResolution_MEDIUM MediaResolution = 2 + MediaResolution_HIGH MediaResolution = 3 + MediaResolution_ULTRA_HIGH MediaResolution = 4 +) + +// Enum value maps for MediaResolution. +var ( + MediaResolution_name = map[int32]string{ + 0: "MEDIA_RESOLUTION_UNSPECIFIED", + 1: "LOW", + 2: "MEDIUM", + 3: "HIGH", + 4: "ULTRA_HIGH", + } + MediaResolution_value = map[string]int32{ + "MEDIA_RESOLUTION_UNSPECIFIED": 0, + "LOW": 1, + "MEDIUM": 2, + "HIGH": 3, + "ULTRA_HIGH": 4, + } +) + +func (x MediaResolution) Enum() *MediaResolution { + p := new(MediaResolution) + *p = x + return p +} + +func (x MediaResolution) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (MediaResolution) Descriptor() protoreflect.EnumDescriptor { + return file_proto_content_proto_enumTypes[0].Descriptor() +} + +func (MediaResolution) Type() protoreflect.EnumType { + return &file_proto_content_proto_enumTypes[0] +} + +func (x MediaResolution) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use MediaResolution.Descriptor instead. +func (MediaResolution) EnumDescriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{0} +} + +type ImageContent_MimeType int32 + +const ( + ImageContent_TYPE_UNSPECIFIED ImageContent_MimeType = 0 + ImageContent_TYPE_PNG ImageContent_MimeType = 1 // image/png + ImageContent_TYPE_JPEG ImageContent_MimeType = 2 // image/jpeg + ImageContent_TYPE_WEBP ImageContent_MimeType = 3 // image/webp + ImageContent_TYPE_HEIC ImageContent_MimeType = 4 // image/heic + ImageContent_TYPE_HEIF ImageContent_MimeType = 5 // image/heif + ImageContent_TYPE_GIF ImageContent_MimeType = 7 // image/gif + ImageContent_TYPE_BMP ImageContent_MimeType = 8 // image/bmp + ImageContent_TYPE_TIFF ImageContent_MimeType = 9 // image/tiff +) + +// Enum value maps for ImageContent_MimeType. +var ( + ImageContent_MimeType_name = map[int32]string{ + 0: "TYPE_UNSPECIFIED", + 1: "TYPE_PNG", + 2: "TYPE_JPEG", + 3: "TYPE_WEBP", + 4: "TYPE_HEIC", + 5: "TYPE_HEIF", + 7: "TYPE_GIF", + 8: "TYPE_BMP", + 9: "TYPE_TIFF", + } + ImageContent_MimeType_value = map[string]int32{ + "TYPE_UNSPECIFIED": 0, + "TYPE_PNG": 1, + "TYPE_JPEG": 2, + "TYPE_WEBP": 3, + "TYPE_HEIC": 4, + "TYPE_HEIF": 5, + "TYPE_GIF": 7, + "TYPE_BMP": 8, + "TYPE_TIFF": 9, + } +) + +func (x ImageContent_MimeType) Enum() *ImageContent_MimeType { + p := new(ImageContent_MimeType) + *p = x + return p +} + +func (x ImageContent_MimeType) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (ImageContent_MimeType) Descriptor() protoreflect.EnumDescriptor { + return file_proto_content_proto_enumTypes[1].Descriptor() +} + +func (ImageContent_MimeType) Type() protoreflect.EnumType { + return &file_proto_content_proto_enumTypes[1] +} + +func (x ImageContent_MimeType) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use ImageContent_MimeType.Descriptor instead. +func (ImageContent_MimeType) EnumDescriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{10, 0} +} + +type AudioContent_MimeType int32 + +const ( + AudioContent_TYPE_UNSPECIFIED AudioContent_MimeType = 0 + AudioContent_TYPE_WAV AudioContent_MimeType = 1 // audio/wav + AudioContent_TYPE_MP3 AudioContent_MimeType = 2 // audio/mp3 + AudioContent_TYPE_AIFF AudioContent_MimeType = 3 // audio/aiff + AudioContent_TYPE_AAC AudioContent_MimeType = 4 // audio/aac + AudioContent_TYPE_OGG AudioContent_MimeType = 5 // audio/ogg + AudioContent_TYPE_FLAC AudioContent_MimeType = 6 // audio/flac + AudioContent_TYPE_MPEG AudioContent_MimeType = 7 // audio/mpeg + AudioContent_TYPE_M4A AudioContent_MimeType = 8 // audio/m4a + AudioContent_TYPE_L16 AudioContent_MimeType = 9 // audio/l16 + AudioContent_TYPE_S16LE AudioContent_MimeType = 10 // audio/s16le + AudioContent_TYPE_OPUS AudioContent_MimeType = 11 // audio/opus + AudioContent_TYPE_ALAW AudioContent_MimeType = 12 // audio/alaw + AudioContent_TYPE_MULAW AudioContent_MimeType = 13 // audio/mulaw +) + +// Enum value maps for AudioContent_MimeType. +var ( + AudioContent_MimeType_name = map[int32]string{ + 0: "TYPE_UNSPECIFIED", + 1: "TYPE_WAV", + 2: "TYPE_MP3", + 3: "TYPE_AIFF", + 4: "TYPE_AAC", + 5: "TYPE_OGG", + 6: "TYPE_FLAC", + 7: "TYPE_MPEG", + 8: "TYPE_M4A", + 9: "TYPE_L16", + 10: "TYPE_S16LE", + 11: "TYPE_OPUS", + 12: "TYPE_ALAW", + 13: "TYPE_MULAW", + } + AudioContent_MimeType_value = map[string]int32{ + "TYPE_UNSPECIFIED": 0, + "TYPE_WAV": 1, + "TYPE_MP3": 2, + "TYPE_AIFF": 3, + "TYPE_AAC": 4, + "TYPE_OGG": 5, + "TYPE_FLAC": 6, + "TYPE_MPEG": 7, + "TYPE_M4A": 8, + "TYPE_L16": 9, + "TYPE_S16LE": 10, + "TYPE_OPUS": 11, + "TYPE_ALAW": 12, + "TYPE_MULAW": 13, + } +) + +func (x AudioContent_MimeType) Enum() *AudioContent_MimeType { + p := new(AudioContent_MimeType) + *p = x + return p +} + +func (x AudioContent_MimeType) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (AudioContent_MimeType) Descriptor() protoreflect.EnumDescriptor { + return file_proto_content_proto_enumTypes[2].Descriptor() +} + +func (AudioContent_MimeType) Type() protoreflect.EnumType { + return &file_proto_content_proto_enumTypes[2] +} + +func (x AudioContent_MimeType) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use AudioContent_MimeType.Descriptor instead. +func (AudioContent_MimeType) EnumDescriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{11, 0} +} + +type DocumentContent_MimeType int32 + +const ( + DocumentContent_TYPE_UNSPECIFIED DocumentContent_MimeType = 0 + DocumentContent_TYPE_PDF DocumentContent_MimeType = 1 // application/pdf + DocumentContent_TYPE_JSON DocumentContent_MimeType = 2 // application/json + DocumentContent_TYPE_PYTHON DocumentContent_MimeType = 3 // text/x-python +) + +// Enum value maps for DocumentContent_MimeType. +var ( + DocumentContent_MimeType_name = map[int32]string{ + 0: "TYPE_UNSPECIFIED", + 1: "TYPE_PDF", + 2: "TYPE_JSON", + 3: "TYPE_PYTHON", + } + DocumentContent_MimeType_value = map[string]int32{ + "TYPE_UNSPECIFIED": 0, + "TYPE_PDF": 1, + "TYPE_JSON": 2, + "TYPE_PYTHON": 3, + } +) + +func (x DocumentContent_MimeType) Enum() *DocumentContent_MimeType { + p := new(DocumentContent_MimeType) + *p = x + return p +} + +func (x DocumentContent_MimeType) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (DocumentContent_MimeType) Descriptor() protoreflect.EnumDescriptor { + return file_proto_content_proto_enumTypes[3].Descriptor() +} + +func (DocumentContent_MimeType) Type() protoreflect.EnumType { + return &file_proto_content_proto_enumTypes[3] +} + +func (x DocumentContent_MimeType) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use DocumentContent_MimeType.Descriptor instead. +func (DocumentContent_MimeType) EnumDescriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{12, 0} +} + +type VideoContent_MimeType int32 + +const ( + VideoContent_TYPE_UNSPECIFIED VideoContent_MimeType = 0 + VideoContent_TYPE_MP4 VideoContent_MimeType = 1 // video/mp4 + VideoContent_TYPE_MPEG VideoContent_MimeType = 2 // video/mpeg + VideoContent_TYPE_MPG VideoContent_MimeType = 3 // video/mpg + VideoContent_TYPE_MOV VideoContent_MimeType = 4 // video/mov + VideoContent_TYPE_AVI VideoContent_MimeType = 5 // video/avi + VideoContent_TYPE_X_FLV VideoContent_MimeType = 6 // video/x-flv + VideoContent_TYPE_WEBM VideoContent_MimeType = 7 // video/webm + VideoContent_TYPE_WMV VideoContent_MimeType = 8 // video/wmv + VideoContent_TYPE_3GPP VideoContent_MimeType = 9 // video/3gpp +) + +// Enum value maps for VideoContent_MimeType. +var ( + VideoContent_MimeType_name = map[int32]string{ + 0: "TYPE_UNSPECIFIED", + 1: "TYPE_MP4", + 2: "TYPE_MPEG", + 3: "TYPE_MPG", + 4: "TYPE_MOV", + 5: "TYPE_AVI", + 6: "TYPE_X_FLV", + 7: "TYPE_WEBM", + 8: "TYPE_WMV", + 9: "TYPE_3GPP", + } + VideoContent_MimeType_value = map[string]int32{ + "TYPE_UNSPECIFIED": 0, + "TYPE_MP4": 1, + "TYPE_MPEG": 2, + "TYPE_MPG": 3, + "TYPE_MOV": 4, + "TYPE_AVI": 5, + "TYPE_X_FLV": 6, + "TYPE_WEBM": 7, + "TYPE_WMV": 8, + "TYPE_3GPP": 9, + } +) + +func (x VideoContent_MimeType) Enum() *VideoContent_MimeType { + p := new(VideoContent_MimeType) + *p = x + return p +} + +func (x VideoContent_MimeType) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (VideoContent_MimeType) Descriptor() protoreflect.EnumDescriptor { + return file_proto_content_proto_enumTypes[4].Descriptor() +} + +func (VideoContent_MimeType) Type() protoreflect.EnumType { + return &file_proto_content_proto_enumTypes[4] +} + +func (x VideoContent_MimeType) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use VideoContent_MimeType.Descriptor instead. +func (VideoContent_MimeType) EnumDescriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{13, 0} +} + +// TextContent represents a text content. +type TextContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Text string `protobuf:"bytes,3,opt,name=text,proto3" json:"text,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *TextContent) Reset() { + *x = TextContent{} + mi := &file_proto_content_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *TextContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TextContent) ProtoMessage() {} + +func (x *TextContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[0] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use TextContent.ProtoReflect.Descriptor instead. +func (*TextContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{0} +} + +func (x *TextContent) GetText() string { + if x != nil { + return x.Text + } + return "" +} + +type ApprovalDecision struct { + state protoimpl.MessageState `protogen:"open.v1"` + Approved bool `protobuf:"varint,1,opt,name=approved,proto3" json:"approved,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ApprovalDecision) Reset() { + *x = ApprovalDecision{} + mi := &file_proto_content_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ApprovalDecision) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ApprovalDecision) ProtoMessage() {} + +func (x *ApprovalDecision) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ApprovalDecision.ProtoReflect.Descriptor instead. +func (*ApprovalDecision) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{1} +} + +func (x *ApprovalDecision) GetApproved() bool { + if x != nil { + return x.Approved + } + return false +} + +type DeclineDecision struct { + state protoimpl.MessageState `protogen:"open.v1"` + Declined bool `protobuf:"varint,1,opt,name=declined,proto3" json:"declined,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *DeclineDecision) Reset() { + *x = DeclineDecision{} + mi := &file_proto_content_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *DeclineDecision) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeclineDecision) ProtoMessage() {} + +func (x *DeclineDecision) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[2] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeclineDecision.ProtoReflect.Descriptor instead. +func (*DeclineDecision) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{2} +} + +func (x *DeclineDecision) GetDeclined() bool { + if x != nil { + return x.Declined + } + return false +} + +type ConfirmationContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Id string `protobuf:"bytes,3,opt,name=id,proto3" json:"id,omitempty"` + Question string `protobuf:"bytes,4,opt,name=question,proto3" json:"question,omitempty"` + // Types that are valid to be assigned to Decision: + // + // *ConfirmationContent_Approval + // *ConfirmationContent_Decline + Decision isConfirmationContent_Decision `protobuf_oneof:"decision"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ConfirmationContent) Reset() { + *x = ConfirmationContent{} + mi := &file_proto_content_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ConfirmationContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ConfirmationContent) ProtoMessage() {} + +func (x *ConfirmationContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[3] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ConfirmationContent.ProtoReflect.Descriptor instead. +func (*ConfirmationContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{3} +} + +func (x *ConfirmationContent) GetId() string { + if x != nil { + return x.Id + } + return "" +} + +func (x *ConfirmationContent) GetQuestion() string { + if x != nil { + return x.Question + } + return "" +} + +func (x *ConfirmationContent) GetDecision() isConfirmationContent_Decision { + if x != nil { + return x.Decision + } + return nil +} + +func (x *ConfirmationContent) GetApproval() *ApprovalDecision { + if x != nil { + if x, ok := x.Decision.(*ConfirmationContent_Approval); ok { + return x.Approval + } + } + return nil +} + +func (x *ConfirmationContent) GetDecline() *DeclineDecision { + if x != nil { + if x, ok := x.Decision.(*ConfirmationContent_Decline); ok { + return x.Decline + } + } + return nil +} + +type isConfirmationContent_Decision interface { + isConfirmationContent_Decision() +} + +type ConfirmationContent_Approval struct { + Approval *ApprovalDecision `protobuf:"bytes,5,opt,name=approval,proto3,oneof"` +} + +type ConfirmationContent_Decline struct { + Decline *DeclineDecision `protobuf:"bytes,6,opt,name=decline,proto3,oneof"` +} + +func (*ConfirmationContent_Approval) isConfirmationContent_Decision() {} + +func (*ConfirmationContent_Decline) isConfirmationContent_Decision() {} + +type ThoughtSummaryContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Types that are valid to be assigned to Type: + // + // *ThoughtSummaryContent_Text + Type isThoughtSummaryContent_Type `protobuf_oneof:"type"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ThoughtSummaryContent) Reset() { + *x = ThoughtSummaryContent{} + mi := &file_proto_content_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ThoughtSummaryContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ThoughtSummaryContent) ProtoMessage() {} + +func (x *ThoughtSummaryContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[4] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ThoughtSummaryContent.ProtoReflect.Descriptor instead. +func (*ThoughtSummaryContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{4} +} + +func (x *ThoughtSummaryContent) GetType() isThoughtSummaryContent_Type { + if x != nil { + return x.Type + } + return nil +} + +func (x *ThoughtSummaryContent) GetText() *TextContent { + if x != nil { + if x, ok := x.Type.(*ThoughtSummaryContent_Text); ok { + return x.Text + } + } + return nil +} + +type isThoughtSummaryContent_Type interface { + isThoughtSummaryContent_Type() +} + +type ThoughtSummaryContent_Text struct { + Text *TextContent `protobuf:"bytes,1,opt,name=text,proto3,oneof"` +} + +func (*ThoughtSummaryContent_Text) isThoughtSummaryContent_Type() {} + +type ThoughtContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Signature []byte `protobuf:"bytes,7,opt,name=signature,proto3" json:"signature,omitempty"` + Summary []*ThoughtSummaryContent `protobuf:"bytes,9,rep,name=summary,proto3" json:"summary,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ThoughtContent) Reset() { + *x = ThoughtContent{} + mi := &file_proto_content_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ThoughtContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ThoughtContent) ProtoMessage() {} + +func (x *ThoughtContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ThoughtContent.ProtoReflect.Descriptor instead. +func (*ThoughtContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{5} +} + +func (x *ThoughtContent) GetSignature() []byte { + if x != nil { + return x.Signature + } + return nil +} + +func (x *ThoughtContent) GetSummary() []*ThoughtSummaryContent { + if x != nil { + return x.Summary + } + return nil +} + +type ToolCallContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` // A unique ID for this specific tool call. + Signature []byte `protobuf:"bytes,9,opt,name=signature,proto3" json:"signature,omitempty"` // A signature hash for backend validation. + // Types that are valid to be assigned to Type: + // + // *ToolCallContent_FunctionCall + Type isToolCallContent_Type `protobuf_oneof:"type"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ToolCallContent) Reset() { + *x = ToolCallContent{} + mi := &file_proto_content_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ToolCallContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ToolCallContent) ProtoMessage() {} + +func (x *ToolCallContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[6] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ToolCallContent.ProtoReflect.Descriptor instead. +func (*ToolCallContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{6} +} + +func (x *ToolCallContent) GetId() string { + if x != nil { + return x.Id + } + return "" +} + +func (x *ToolCallContent) GetSignature() []byte { + if x != nil { + return x.Signature + } + return nil +} + +func (x *ToolCallContent) GetType() isToolCallContent_Type { + if x != nil { + return x.Type + } + return nil +} + +func (x *ToolCallContent) GetFunctionCall() *FunctionCallContent { + if x != nil { + if x, ok := x.Type.(*ToolCallContent_FunctionCall); ok { + return x.FunctionCall + } + } + return nil +} + +type isToolCallContent_Type interface { + isToolCallContent_Type() +} + +type ToolCallContent_FunctionCall struct { + FunctionCall *FunctionCallContent `protobuf:"bytes,2,opt,name=function_call,json=functionCall,proto3,oneof"` +} + +func (*ToolCallContent_FunctionCall) isToolCallContent_Type() {} + +type ToolResultContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + CallId string `protobuf:"bytes,8,opt,name=call_id,json=callId,proto3" json:"call_id,omitempty"` // ID to match the ID from the function call block. + Signature []byte `protobuf:"bytes,9,opt,name=signature,proto3" json:"signature,omitempty"` // A signature hash for backend validation. + // Types that are valid to be assigned to Type: + // + // *ToolResultContent_FunctionResult + Type isToolResultContent_Type `protobuf_oneof:"type"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ToolResultContent) Reset() { + *x = ToolResultContent{} + mi := &file_proto_content_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ToolResultContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ToolResultContent) ProtoMessage() {} + +func (x *ToolResultContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[7] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ToolResultContent.ProtoReflect.Descriptor instead. +func (*ToolResultContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{7} +} + +func (x *ToolResultContent) GetCallId() string { + if x != nil { + return x.CallId + } + return "" +} + +func (x *ToolResultContent) GetSignature() []byte { + if x != nil { + return x.Signature + } + return nil +} + +func (x *ToolResultContent) GetType() isToolResultContent_Type { + if x != nil { + return x.Type + } + return nil +} + +func (x *ToolResultContent) GetFunctionResult() *FunctionResultContent { + if x != nil { + if x, ok := x.Type.(*ToolResultContent_FunctionResult); ok { + return x.FunctionResult + } + } + return nil +} + +type isToolResultContent_Type interface { + isToolResultContent_Type() +} + +type ToolResultContent_FunctionResult struct { + FunctionResult *FunctionResultContent `protobuf:"bytes,2,opt,name=function_result,json=functionResult,proto3,oneof"` +} + +func (*ToolResultContent_FunctionResult) isToolResultContent_Type() {} + +type FunctionCallContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Name string `protobuf:"bytes,3,opt,name=name,proto3" json:"name,omitempty"` + Arguments *structpb.Struct `protobuf:"bytes,4,opt,name=arguments,proto3" json:"arguments,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FunctionCallContent) Reset() { + *x = FunctionCallContent{} + mi := &file_proto_content_proto_msgTypes[8] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FunctionCallContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FunctionCallContent) ProtoMessage() {} + +func (x *FunctionCallContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[8] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FunctionCallContent.ProtoReflect.Descriptor instead. +func (*FunctionCallContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{8} +} + +func (x *FunctionCallContent) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +func (x *FunctionCallContent) GetArguments() *structpb.Struct { + if x != nil { + return x.Arguments + } + return nil +} + +type FunctionResultContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Name string `protobuf:"bytes,8,opt,name=name,proto3" json:"name,omitempty"` + // Types that are valid to be assigned to Result: + // + // *FunctionResultContent_Response + Result isFunctionResultContent_Result `protobuf_oneof:"result"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FunctionResultContent) Reset() { + *x = FunctionResultContent{} + mi := &file_proto_content_proto_msgTypes[9] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FunctionResultContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FunctionResultContent) ProtoMessage() {} + +func (x *FunctionResultContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[9] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FunctionResultContent.ProtoReflect.Descriptor instead. +func (*FunctionResultContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{9} +} + +func (x *FunctionResultContent) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +func (x *FunctionResultContent) GetResult() isFunctionResultContent_Result { + if x != nil { + return x.Result + } + return nil +} + +func (x *FunctionResultContent) GetResponse() *structpb.Struct { + if x != nil { + if x, ok := x.Result.(*FunctionResultContent_Response); ok { + return x.Response + } + } + return nil +} + +type isFunctionResultContent_Result interface { + isFunctionResultContent_Result() +} + +type FunctionResultContent_Response struct { + Response *structpb.Struct `protobuf:"bytes,3,opt,name=response,proto3,oneof"` +} + +func (*FunctionResultContent_Response) isFunctionResultContent_Result() {} + +// An image content block. +type ImageContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + MimeType ImageContent_MimeType `protobuf:"varint,1,opt,name=mime_type,json=mimeType,proto3,enum=ax.ImageContent_MimeType" json:"mime_type,omitempty"` + // Types that are valid to be assigned to DataOrUri: + // + // *ImageContent_Data + // *ImageContent_Uri + DataOrUri isImageContent_DataOrUri `protobuf_oneof:"data_or_uri"` + Resolution MediaResolution `protobuf:"varint,5,opt,name=resolution,proto3,enum=ax.MediaResolution" json:"resolution,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ImageContent) Reset() { + *x = ImageContent{} + mi := &file_proto_content_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ImageContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ImageContent) ProtoMessage() {} + +func (x *ImageContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[10] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ImageContent.ProtoReflect.Descriptor instead. +func (*ImageContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{10} +} + +func (x *ImageContent) GetMimeType() ImageContent_MimeType { + if x != nil { + return x.MimeType + } + return ImageContent_TYPE_UNSPECIFIED +} + +func (x *ImageContent) GetDataOrUri() isImageContent_DataOrUri { + if x != nil { + return x.DataOrUri + } + return nil +} + +func (x *ImageContent) GetData() []byte { + if x != nil { + if x, ok := x.DataOrUri.(*ImageContent_Data); ok { + return x.Data + } + } + return nil +} + +func (x *ImageContent) GetUri() string { + if x != nil { + if x, ok := x.DataOrUri.(*ImageContent_Uri); ok { + return x.Uri + } + } + return "" +} + +func (x *ImageContent) GetResolution() MediaResolution { + if x != nil { + return x.Resolution + } + return MediaResolution_MEDIA_RESOLUTION_UNSPECIFIED +} + +type isImageContent_DataOrUri interface { + isImageContent_DataOrUri() +} + +type ImageContent_Data struct { + Data []byte `protobuf:"bytes,2,opt,name=data,proto3,oneof"` +} + +type ImageContent_Uri struct { + Uri string `protobuf:"bytes,6,opt,name=uri,proto3,oneof"` +} + +func (*ImageContent_Data) isImageContent_DataOrUri() {} + +func (*ImageContent_Uri) isImageContent_DataOrUri() {} + +// An audio content block. +type AudioContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + MimeType AudioContent_MimeType `protobuf:"varint,1,opt,name=mime_type,json=mimeType,proto3,enum=ax.AudioContent_MimeType" json:"mime_type,omitempty"` + // Types that are valid to be assigned to DataOrUri: + // + // *AudioContent_Data + // *AudioContent_Uri + DataOrUri isAudioContent_DataOrUri `protobuf_oneof:"data_or_uri"` + Channels int32 `protobuf:"varint,7,opt,name=channels,proto3" json:"channels,omitempty"` + SampleRate int32 `protobuf:"varint,8,opt,name=sample_rate,json=sampleRate,proto3" json:"sample_rate,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *AudioContent) Reset() { + *x = AudioContent{} + mi := &file_proto_content_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *AudioContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*AudioContent) ProtoMessage() {} + +func (x *AudioContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[11] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use AudioContent.ProtoReflect.Descriptor instead. +func (*AudioContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{11} +} + +func (x *AudioContent) GetMimeType() AudioContent_MimeType { + if x != nil { + return x.MimeType + } + return AudioContent_TYPE_UNSPECIFIED +} + +func (x *AudioContent) GetDataOrUri() isAudioContent_DataOrUri { + if x != nil { + return x.DataOrUri + } + return nil +} + +func (x *AudioContent) GetData() []byte { + if x != nil { + if x, ok := x.DataOrUri.(*AudioContent_Data); ok { + return x.Data + } + } + return nil +} + +func (x *AudioContent) GetUri() string { + if x != nil { + if x, ok := x.DataOrUri.(*AudioContent_Uri); ok { + return x.Uri + } + } + return "" +} + +func (x *AudioContent) GetChannels() int32 { + if x != nil { + return x.Channels + } + return 0 +} + +func (x *AudioContent) GetSampleRate() int32 { + if x != nil { + return x.SampleRate + } + return 0 +} + +type isAudioContent_DataOrUri interface { + isAudioContent_DataOrUri() +} + +type AudioContent_Data struct { + Data []byte `protobuf:"bytes,2,opt,name=data,proto3,oneof"` +} + +type AudioContent_Uri struct { + Uri string `protobuf:"bytes,5,opt,name=uri,proto3,oneof"` +} + +func (*AudioContent_Data) isAudioContent_DataOrUri() {} + +func (*AudioContent_Uri) isAudioContent_DataOrUri() {} + +// A document content block. +type DocumentContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + MimeType DocumentContent_MimeType `protobuf:"varint,1,opt,name=mime_type,json=mimeType,proto3,enum=ax.DocumentContent_MimeType" json:"mime_type,omitempty"` + // Types that are valid to be assigned to DataOrUri: + // + // *DocumentContent_Data + // *DocumentContent_Uri + DataOrUri isDocumentContent_DataOrUri `protobuf_oneof:"data_or_uri"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *DocumentContent) Reset() { + *x = DocumentContent{} + mi := &file_proto_content_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *DocumentContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DocumentContent) ProtoMessage() {} + +func (x *DocumentContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[12] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DocumentContent.ProtoReflect.Descriptor instead. +func (*DocumentContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{12} +} + +func (x *DocumentContent) GetMimeType() DocumentContent_MimeType { + if x != nil { + return x.MimeType + } + return DocumentContent_TYPE_UNSPECIFIED +} + +func (x *DocumentContent) GetDataOrUri() isDocumentContent_DataOrUri { + if x != nil { + return x.DataOrUri + } + return nil +} + +func (x *DocumentContent) GetData() []byte { + if x != nil { + if x, ok := x.DataOrUri.(*DocumentContent_Data); ok { + return x.Data + } + } + return nil +} + +func (x *DocumentContent) GetUri() string { + if x != nil { + if x, ok := x.DataOrUri.(*DocumentContent_Uri); ok { + return x.Uri + } + } + return "" +} + +type isDocumentContent_DataOrUri interface { + isDocumentContent_DataOrUri() +} + +type DocumentContent_Data struct { + Data []byte `protobuf:"bytes,2,opt,name=data,proto3,oneof"` +} + +type DocumentContent_Uri struct { + Uri string `protobuf:"bytes,5,opt,name=uri,proto3,oneof"` +} + +func (*DocumentContent_Data) isDocumentContent_DataOrUri() {} + +func (*DocumentContent_Uri) isDocumentContent_DataOrUri() {} + +// A video content block. +type VideoContent struct { + state protoimpl.MessageState `protogen:"open.v1"` + MimeType VideoContent_MimeType `protobuf:"varint,1,opt,name=mime_type,json=mimeType,proto3,enum=ax.VideoContent_MimeType" json:"mime_type,omitempty"` + // Types that are valid to be assigned to DataOrUri: + // + // *VideoContent_Data + // *VideoContent_Uri + DataOrUri isVideoContent_DataOrUri `protobuf_oneof:"data_or_uri"` + Resolution MediaResolution `protobuf:"varint,5,opt,name=resolution,proto3,enum=ax.MediaResolution" json:"resolution,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *VideoContent) Reset() { + *x = VideoContent{} + mi := &file_proto_content_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *VideoContent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*VideoContent) ProtoMessage() {} + +func (x *VideoContent) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[13] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use VideoContent.ProtoReflect.Descriptor instead. +func (*VideoContent) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{13} +} + +func (x *VideoContent) GetMimeType() VideoContent_MimeType { + if x != nil { + return x.MimeType + } + return VideoContent_TYPE_UNSPECIFIED +} + +func (x *VideoContent) GetDataOrUri() isVideoContent_DataOrUri { + if x != nil { + return x.DataOrUri + } + return nil +} + +func (x *VideoContent) GetData() []byte { + if x != nil { + if x, ok := x.DataOrUri.(*VideoContent_Data); ok { + return x.Data + } + } + return nil +} + +func (x *VideoContent) GetUri() string { + if x != nil { + if x, ok := x.DataOrUri.(*VideoContent_Uri); ok { + return x.Uri + } + } + return "" +} + +func (x *VideoContent) GetResolution() MediaResolution { + if x != nil { + return x.Resolution + } + return MediaResolution_MEDIA_RESOLUTION_UNSPECIFIED +} + +type isVideoContent_DataOrUri interface { + isVideoContent_DataOrUri() +} + +type VideoContent_Data struct { + Data []byte `protobuf:"bytes,2,opt,name=data,proto3,oneof"` +} + +type VideoContent_Uri struct { + Uri string `protobuf:"bytes,6,opt,name=uri,proto3,oneof"` +} + +func (*VideoContent_Data) isVideoContent_DataOrUri() {} + +func (*VideoContent_Uri) isVideoContent_DataOrUri() {} + +// Content represents a content input or output. +type Content struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Types that are valid to be assigned to Type: + // + // *Content_Thought + // *Content_Text + // *Content_Image + // *Content_Audio + // *Content_Document + // *Content_Video + // *Content_Confirmation + // *Content_ToolCall + // *Content_ToolResult + Type isContent_Type `protobuf_oneof:"type"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Content) Reset() { + *x = Content{} + mi := &file_proto_content_proto_msgTypes[14] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Content) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Content) ProtoMessage() {} + +func (x *Content) ProtoReflect() protoreflect.Message { + mi := &file_proto_content_proto_msgTypes[14] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use Content.ProtoReflect.Descriptor instead. +func (*Content) Descriptor() ([]byte, []int) { + return file_proto_content_proto_rawDescGZIP(), []int{14} +} + +func (x *Content) GetType() isContent_Type { + if x != nil { + return x.Type + } + return nil +} + +func (x *Content) GetThought() *ThoughtContent { + if x != nil { + if x, ok := x.Type.(*Content_Thought); ok { + return x.Thought + } + } + return nil +} + +func (x *Content) GetText() *TextContent { + if x != nil { + if x, ok := x.Type.(*Content_Text); ok { + return x.Text + } + } + return nil +} + +func (x *Content) GetImage() *ImageContent { + if x != nil { + if x, ok := x.Type.(*Content_Image); ok { + return x.Image + } + } + return nil +} + +func (x *Content) GetAudio() *AudioContent { + if x != nil { + if x, ok := x.Type.(*Content_Audio); ok { + return x.Audio + } + } + return nil +} + +func (x *Content) GetDocument() *DocumentContent { + if x != nil { + if x, ok := x.Type.(*Content_Document); ok { + return x.Document + } + } + return nil +} + +func (x *Content) GetVideo() *VideoContent { + if x != nil { + if x, ok := x.Type.(*Content_Video); ok { + return x.Video + } + } + return nil +} + +func (x *Content) GetConfirmation() *ConfirmationContent { + if x != nil { + if x, ok := x.Type.(*Content_Confirmation); ok { + return x.Confirmation + } + } + return nil +} + +func (x *Content) GetToolCall() *ToolCallContent { + if x != nil { + if x, ok := x.Type.(*Content_ToolCall); ok { + return x.ToolCall + } + } + return nil +} + +func (x *Content) GetToolResult() *ToolResultContent { + if x != nil { + if x, ok := x.Type.(*Content_ToolResult); ok { + return x.ToolResult + } + } + return nil +} + +type isContent_Type interface { + isContent_Type() +} + +type Content_Thought struct { + Thought *ThoughtContent `protobuf:"bytes,5,opt,name=thought,proto3,oneof"` +} + +type Content_Text struct { + Text *TextContent `protobuf:"bytes,10,opt,name=text,proto3,oneof"` +} + +type Content_Image struct { + Image *ImageContent `protobuf:"bytes,11,opt,name=image,proto3,oneof"` +} + +type Content_Audio struct { + Audio *AudioContent `protobuf:"bytes,12,opt,name=audio,proto3,oneof"` +} + +type Content_Document struct { + Document *DocumentContent `protobuf:"bytes,13,opt,name=document,proto3,oneof"` +} + +type Content_Video struct { + Video *VideoContent `protobuf:"bytes,14,opt,name=video,proto3,oneof"` +} + +type Content_Confirmation struct { + Confirmation *ConfirmationContent `protobuf:"bytes,26,opt,name=confirmation,proto3,oneof"` // TODO(jbd): Remove out of the Content. +} + +type Content_ToolCall struct { + ToolCall *ToolCallContent `protobuf:"bytes,24,opt,name=tool_call,json=toolCall,proto3,oneof"` +} + +type Content_ToolResult struct { + ToolResult *ToolResultContent `protobuf:"bytes,25,opt,name=tool_result,json=toolResult,proto3,oneof"` +} + +func (*Content_Thought) isContent_Type() {} + +func (*Content_Text) isContent_Type() {} + +func (*Content_Image) isContent_Type() {} + +func (*Content_Audio) isContent_Type() {} + +func (*Content_Document) isContent_Type() {} + +func (*Content_Video) isContent_Type() {} + +func (*Content_Confirmation) isContent_Type() {} + +func (*Content_ToolCall) isContent_Type() {} + +func (*Content_ToolResult) isContent_Type() {} + +var File_proto_content_proto protoreflect.FileDescriptor + +const file_proto_content_proto_rawDesc = "" + + "\n" + + "\x13proto/content.proto\x12\x02ax\x1a\x1cgoogle/protobuf/struct.proto\"3\n" + + "\vTextContent\x12\x12\n" + + "\x04text\x18\x03 \x01(\tR\x04textJ\x04\b\x01\x10\x02J\x04\b\x02\x10\x03R\x04type\".\n" + + "\x10ApprovalDecision\x12\x1a\n" + + "\bapproved\x18\x01 \x01(\bR\bapproved\";\n" + + "\x0fDeclineDecision\x12\x1a\n" + + "\bdeclined\x18\x01 \x01(\bR\bdeclinedJ\x04\b\x02\x10\x03R\x06reason\"\xc4\x01\n" + + "\x13ConfirmationContent\x12\x0e\n" + + "\x02id\x18\x03 \x01(\tR\x02id\x12\x1a\n" + + "\bquestion\x18\x04 \x01(\tR\bquestion\x122\n" + + "\bapproval\x18\x05 \x01(\v2\x14.ax.ApprovalDecisionH\x00R\bapproval\x12/\n" + + "\adecline\x18\x06 \x01(\v2\x13.ax.DeclineDecisionH\x00R\adeclineB\n" + + "\n" + + "\bdecisionJ\x04\b\x01\x10\x02J\x04\b\x02\x10\x03R\x04type\"L\n" + + "\x15ThoughtSummaryContent\x12%\n" + + "\x04text\x18\x01 \x01(\v2\x0f.ax.TextContentH\x00R\x04textB\x06\n" + + "\x04typeJ\x04\b\x02\x10\x03\"\x8d\x01\n" + + "\x0eThoughtContent\x12\x1c\n" + + "\tsignature\x18\a \x01(\fR\tsignature\x123\n" + + "\asummary\x18\t \x03(\v2\x19.ax.ThoughtSummaryContentR\asummaryJ\x04\b\x01\x10\x02J\x04\b\x02\x10\x03J\x04\b\x03\x10\x04J\x04\b\x04\x10\x05J\x04\b\x05\x10\x06J\x04\b\b\x10\tR\x04type\"\x87\x01\n" + + "\x0fToolCallContent\x12\x0e\n" + + "\x02id\x18\x01 \x01(\tR\x02id\x12\x1c\n" + + "\tsignature\x18\t \x01(\fR\tsignature\x12>\n" + + "\rfunction_call\x18\x02 \x01(\v2\x17.ax.FunctionCallContentH\x00R\ffunctionCallB\x06\n" + + "\x04type\"\x9e\x01\n" + + "\x11ToolResultContent\x12\x17\n" + + "\acall_id\x18\b \x01(\tR\x06callId\x12\x1c\n" + + "\tsignature\x18\t \x01(\fR\tsignature\x12D\n" + + "\x0ffunction_result\x18\x02 \x01(\v2\x19.ax.FunctionResultContentH\x00R\x0efunctionResultB\x06\n" + + "\x04typeJ\x04\b\x01\x10\x02\"l\n" + + "\x13FunctionCallContent\x12\x12\n" + + "\x04name\x18\x03 \x01(\tR\x04name\x125\n" + + "\targuments\x18\x04 \x01(\v2\x17.google.protobuf.StructR\targumentsJ\x04\b\x01\x10\x02J\x04\b\x02\x10\x03\"~\n" + + "\x15FunctionResultContent\x12\x12\n" + + "\x04name\x18\b \x01(\tR\x04name\x125\n" + + "\bresponse\x18\x03 \x01(\v2\x17.google.protobuf.StructH\x00R\bresponseB\b\n" + + "\x06resultJ\x04\b\x01\x10\x02J\x04\b\x04\x10\x05R\x04type\"\xde\x02\n" + + "\fImageContent\x126\n" + + "\tmime_type\x18\x01 \x01(\x0e2\x19.ax.ImageContent.MimeTypeR\bmimeType\x12\x14\n" + + "\x04data\x18\x02 \x01(\fH\x00R\x04data\x12\x12\n" + + "\x03uri\x18\x06 \x01(\tH\x00R\x03uri\x123\n" + + "\n" + + "resolution\x18\x05 \x01(\x0e2\x13.ax.MediaResolutionR\n" + + "resolution\"\x9b\x01\n" + + "\bMimeType\x12\x14\n" + + "\x10TYPE_UNSPECIFIED\x10\x00\x12\f\n" + + "\bTYPE_PNG\x10\x01\x12\r\n" + + "\tTYPE_JPEG\x10\x02\x12\r\n" + + "\tTYPE_WEBP\x10\x03\x12\r\n" + + "\tTYPE_HEIC\x10\x04\x12\r\n" + + "\tTYPE_HEIF\x10\x05\x12\f\n" + + "\bTYPE_GIF\x10\a\x12\f\n" + + "\bTYPE_BMP\x10\b\x12\r\n" + + "\tTYPE_TIFF\x10\t\"\x04\b\x06\x10\x06B\r\n" + + "\vdata_or_uriJ\x04\b\x03\x10\x04R\x04type\"\xb6\x03\n" + + "\fAudioContent\x126\n" + + "\tmime_type\x18\x01 \x01(\x0e2\x19.ax.AudioContent.MimeTypeR\bmimeType\x12\x14\n" + + "\x04data\x18\x02 \x01(\fH\x00R\x04data\x12\x12\n" + + "\x03uri\x18\x05 \x01(\tH\x00R\x03uri\x12\x1a\n" + + "\bchannels\x18\a \x01(\x05R\bchannels\x12\x1f\n" + + "\vsample_rate\x18\b \x01(\x05R\n" + + "sampleRate\"\xdf\x01\n" + + "\bMimeType\x12\x14\n" + + "\x10TYPE_UNSPECIFIED\x10\x00\x12\f\n" + + "\bTYPE_WAV\x10\x01\x12\f\n" + + "\bTYPE_MP3\x10\x02\x12\r\n" + + "\tTYPE_AIFF\x10\x03\x12\f\n" + + "\bTYPE_AAC\x10\x04\x12\f\n" + + "\bTYPE_OGG\x10\x05\x12\r\n" + + "\tTYPE_FLAC\x10\x06\x12\r\n" + + "\tTYPE_MPEG\x10\a\x12\f\n" + + "\bTYPE_M4A\x10\b\x12\f\n" + + "\bTYPE_L16\x10\t\x12\x0e\n" + + "\n" + + "TYPE_S16LE\x10\n" + + "\x12\r\n" + + "\tTYPE_OPUS\x10\v\x12\r\n" + + "\tTYPE_ALAW\x10\f\x12\x0e\n" + + "\n" + + "TYPE_MULAW\x10\rB\r\n" + + "\vdata_or_uriJ\x04\b\x03\x10\x04J\x04\b\x06\x10\aR\x04typeR\x04rate\"\xe1\x01\n" + + "\x0fDocumentContent\x129\n" + + "\tmime_type\x18\x01 \x01(\x0e2\x1c.ax.DocumentContent.MimeTypeR\bmimeType\x12\x14\n" + + "\x04data\x18\x02 \x01(\fH\x00R\x04data\x12\x12\n" + + "\x03uri\x18\x05 \x01(\tH\x00R\x03uri\"N\n" + + "\bMimeType\x12\x14\n" + + "\x10TYPE_UNSPECIFIED\x10\x00\x12\f\n" + + "\bTYPE_PDF\x10\x01\x12\r\n" + + "\tTYPE_JSON\x10\x02\x12\x0f\n" + + "\vTYPE_PYTHON\x10\x03B\r\n" + + "\vdata_or_uriJ\x04\b\x03\x10\x04R\x04type\"\xe6\x02\n" + + "\fVideoContent\x126\n" + + "\tmime_type\x18\x01 \x01(\x0e2\x19.ax.VideoContent.MimeTypeR\bmimeType\x12\x14\n" + + "\x04data\x18\x02 \x01(\fH\x00R\x04data\x12\x12\n" + + "\x03uri\x18\x06 \x01(\tH\x00R\x03uri\x123\n" + + "\n" + + "resolution\x18\x05 \x01(\x0e2\x13.ax.MediaResolutionR\n" + + "resolution\"\xa3\x01\n" + + "\bMimeType\x12\x14\n" + + "\x10TYPE_UNSPECIFIED\x10\x00\x12\f\n" + + "\bTYPE_MP4\x10\x01\x12\r\n" + + "\tTYPE_MPEG\x10\x02\x12\f\n" + + "\bTYPE_MPG\x10\x03\x12\f\n" + + "\bTYPE_MOV\x10\x04\x12\f\n" + + "\bTYPE_AVI\x10\x05\x12\x0e\n" + + "\n" + + "TYPE_X_FLV\x10\x06\x12\r\n" + + "\tTYPE_WEBM\x10\a\x12\f\n" + + "\bTYPE_WMV\x10\b\x12\r\n" + + "\tTYPE_3GPP\x10\tB\r\n" + + "\vdata_or_uriJ\x04\b\x03\x10\x04R\x04type\"\xd8\x03\n" + + "\aContent\x12.\n" + + "\athought\x18\x05 \x01(\v2\x12.ax.ThoughtContentH\x00R\athought\x12%\n" + + "\x04text\x18\n" + + " \x01(\v2\x0f.ax.TextContentH\x00R\x04text\x12(\n" + + "\x05image\x18\v \x01(\v2\x10.ax.ImageContentH\x00R\x05image\x12(\n" + + "\x05audio\x18\f \x01(\v2\x10.ax.AudioContentH\x00R\x05audio\x121\n" + + "\bdocument\x18\r \x01(\v2\x13.ax.DocumentContentH\x00R\bdocument\x12(\n" + + "\x05video\x18\x0e \x01(\v2\x10.ax.VideoContentH\x00R\x05video\x12=\n" + + "\fconfirmation\x18\x1a \x01(\v2\x17.ax.ConfirmationContentH\x00R\fconfirmation\x122\n" + + "\ttool_call\x18\x18 \x01(\v2\x13.ax.ToolCallContentH\x00R\btoolCall\x128\n" + + "\vtool_result\x18\x19 \x01(\v2\x15.ax.ToolResultContentH\x00R\n" + + "toolResultB\x06\n" + + "\x04typeJ\x04\b\x01\x10\x05J\x04\b\x06\x10\n" + + "J\x04\b\x0f\x10\x18*b\n" + + "\x0fMediaResolution\x12 \n" + + "\x1cMEDIA_RESOLUTION_UNSPECIFIED\x10\x00\x12\a\n" + + "\x03LOW\x10\x01\x12\n" + + "\n" + + "\x06MEDIUM\x10\x02\x12\b\n" + + "\x04HIGH\x10\x03\x12\x0e\n" + + "\n" + + "ULTRA_HIGH\x10\x04B\x1cZ\x1agithub.com/google/ax/protob\x06proto3" + +var ( + file_proto_content_proto_rawDescOnce sync.Once + file_proto_content_proto_rawDescData []byte +) + +func file_proto_content_proto_rawDescGZIP() []byte { + file_proto_content_proto_rawDescOnce.Do(func() { + file_proto_content_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_proto_content_proto_rawDesc), len(file_proto_content_proto_rawDesc))) + }) + return file_proto_content_proto_rawDescData +} + +var file_proto_content_proto_enumTypes = make([]protoimpl.EnumInfo, 5) +var file_proto_content_proto_msgTypes = make([]protoimpl.MessageInfo, 15) +var file_proto_content_proto_goTypes = []any{ + (MediaResolution)(0), // 0: ax.MediaResolution + (ImageContent_MimeType)(0), // 1: ax.ImageContent.MimeType + (AudioContent_MimeType)(0), // 2: ax.AudioContent.MimeType + (DocumentContent_MimeType)(0), // 3: ax.DocumentContent.MimeType + (VideoContent_MimeType)(0), // 4: ax.VideoContent.MimeType + (*TextContent)(nil), // 5: ax.TextContent + (*ApprovalDecision)(nil), // 6: ax.ApprovalDecision + (*DeclineDecision)(nil), // 7: ax.DeclineDecision + (*ConfirmationContent)(nil), // 8: ax.ConfirmationContent + (*ThoughtSummaryContent)(nil), // 9: ax.ThoughtSummaryContent + (*ThoughtContent)(nil), // 10: ax.ThoughtContent + (*ToolCallContent)(nil), // 11: ax.ToolCallContent + (*ToolResultContent)(nil), // 12: ax.ToolResultContent + (*FunctionCallContent)(nil), // 13: ax.FunctionCallContent + (*FunctionResultContent)(nil), // 14: ax.FunctionResultContent + (*ImageContent)(nil), // 15: ax.ImageContent + (*AudioContent)(nil), // 16: ax.AudioContent + (*DocumentContent)(nil), // 17: ax.DocumentContent + (*VideoContent)(nil), // 18: ax.VideoContent + (*Content)(nil), // 19: ax.Content + (*structpb.Struct)(nil), // 20: google.protobuf.Struct +} +var file_proto_content_proto_depIdxs = []int32{ + 6, // 0: ax.ConfirmationContent.approval:type_name -> ax.ApprovalDecision + 7, // 1: ax.ConfirmationContent.decline:type_name -> ax.DeclineDecision + 5, // 2: ax.ThoughtSummaryContent.text:type_name -> ax.TextContent + 9, // 3: ax.ThoughtContent.summary:type_name -> ax.ThoughtSummaryContent + 13, // 4: ax.ToolCallContent.function_call:type_name -> ax.FunctionCallContent + 14, // 5: ax.ToolResultContent.function_result:type_name -> ax.FunctionResultContent + 20, // 6: ax.FunctionCallContent.arguments:type_name -> google.protobuf.Struct + 20, // 7: ax.FunctionResultContent.response:type_name -> google.protobuf.Struct + 1, // 8: ax.ImageContent.mime_type:type_name -> ax.ImageContent.MimeType + 0, // 9: ax.ImageContent.resolution:type_name -> ax.MediaResolution + 2, // 10: ax.AudioContent.mime_type:type_name -> ax.AudioContent.MimeType + 3, // 11: ax.DocumentContent.mime_type:type_name -> ax.DocumentContent.MimeType + 4, // 12: ax.VideoContent.mime_type:type_name -> ax.VideoContent.MimeType + 0, // 13: ax.VideoContent.resolution:type_name -> ax.MediaResolution + 10, // 14: ax.Content.thought:type_name -> ax.ThoughtContent + 5, // 15: ax.Content.text:type_name -> ax.TextContent + 15, // 16: ax.Content.image:type_name -> ax.ImageContent + 16, // 17: ax.Content.audio:type_name -> ax.AudioContent + 17, // 18: ax.Content.document:type_name -> ax.DocumentContent + 18, // 19: ax.Content.video:type_name -> ax.VideoContent + 8, // 20: ax.Content.confirmation:type_name -> ax.ConfirmationContent + 11, // 21: ax.Content.tool_call:type_name -> ax.ToolCallContent + 12, // 22: ax.Content.tool_result:type_name -> ax.ToolResultContent + 23, // [23:23] is the sub-list for method output_type + 23, // [23:23] is the sub-list for method input_type + 23, // [23:23] is the sub-list for extension type_name + 23, // [23:23] is the sub-list for extension extendee + 0, // [0:23] is the sub-list for field type_name +} + +func init() { file_proto_content_proto_init() } +func file_proto_content_proto_init() { + if File_proto_content_proto != nil { + return + } + file_proto_content_proto_msgTypes[3].OneofWrappers = []any{ + (*ConfirmationContent_Approval)(nil), + (*ConfirmationContent_Decline)(nil), + } + file_proto_content_proto_msgTypes[4].OneofWrappers = []any{ + (*ThoughtSummaryContent_Text)(nil), + } + file_proto_content_proto_msgTypes[6].OneofWrappers = []any{ + (*ToolCallContent_FunctionCall)(nil), + } + file_proto_content_proto_msgTypes[7].OneofWrappers = []any{ + (*ToolResultContent_FunctionResult)(nil), + } + file_proto_content_proto_msgTypes[9].OneofWrappers = []any{ + (*FunctionResultContent_Response)(nil), + } + file_proto_content_proto_msgTypes[10].OneofWrappers = []any{ + (*ImageContent_Data)(nil), + (*ImageContent_Uri)(nil), + } + file_proto_content_proto_msgTypes[11].OneofWrappers = []any{ + (*AudioContent_Data)(nil), + (*AudioContent_Uri)(nil), + } + file_proto_content_proto_msgTypes[12].OneofWrappers = []any{ + (*DocumentContent_Data)(nil), + (*DocumentContent_Uri)(nil), + } + file_proto_content_proto_msgTypes[13].OneofWrappers = []any{ + (*VideoContent_Data)(nil), + (*VideoContent_Uri)(nil), + } + file_proto_content_proto_msgTypes[14].OneofWrappers = []any{ + (*Content_Thought)(nil), + (*Content_Text)(nil), + (*Content_Image)(nil), + (*Content_Audio)(nil), + (*Content_Document)(nil), + (*Content_Video)(nil), + (*Content_Confirmation)(nil), + (*Content_ToolCall)(nil), + (*Content_ToolResult)(nil), + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_proto_content_proto_rawDesc), len(file_proto_content_proto_rawDesc)), + NumEnums: 5, + NumMessages: 15, + NumExtensions: 0, + NumServices: 0, + }, + GoTypes: file_proto_content_proto_goTypes, + DependencyIndexes: file_proto_content_proto_depIdxs, + EnumInfos: file_proto_content_proto_enumTypes, + MessageInfos: file_proto_content_proto_msgTypes, + }.Build() + File_proto_content_proto = out.File + file_proto_content_proto_goTypes = nil + file_proto_content_proto_depIdxs = nil +} diff --git a/proto/content.proto b/proto/content.proto new file mode 100644 index 00000000..20cf98f1 --- /dev/null +++ b/proto/content.proto @@ -0,0 +1,225 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +syntax = "proto3"; + +package ax; + +import "google/protobuf/struct.proto"; + +option go_package = "github.com/google/ax/proto"; + +// TextContent represents a text content. +message TextContent { + reserved 1, 2; + reserved "type"; + string text = 3; +} + +message ApprovalDecision { + bool approved = 1; +} + +message DeclineDecision { + reserved 2; + reserved "reason"; + bool declined = 1; +} + +message ConfirmationContent { + reserved 1, 2; + reserved "type"; + string id = 3; + string question = 4; + oneof decision { + ApprovalDecision approval = 5; + DeclineDecision decline = 6; + } +} + +message ThoughtSummaryContent { + reserved 2; + + oneof type { + TextContent text = 1; + } +} + +message ThoughtContent { + reserved 1, 2, 3, 4, 5, 8; + reserved "type"; + + bytes signature = 7; + repeated ThoughtSummaryContent summary = 9; +} + +message ToolCallContent { + string id = 1; // A unique ID for this specific tool call. + bytes signature = 9; // A signature hash for backend validation. + oneof type { + FunctionCallContent function_call = 2; + } +} + +message ToolResultContent { + reserved 1; + + string call_id = 8; // ID to match the ID from the function call block. + bytes signature = 9; // A signature hash for backend validation. + + oneof type { + FunctionResultContent function_result = 2; + } +} + +message FunctionCallContent { + reserved 1, 2; + + string name = 3; + google.protobuf.Struct arguments = 4; +} + +message FunctionResultContent { + reserved 1, 4; + reserved "type"; + + string name = 8; + oneof result { + google.protobuf.Struct response = 3; + } +} + +// Resolution for input media (images/video). +enum MediaResolution { + MEDIA_RESOLUTION_UNSPECIFIED = 0; + LOW = 1; + MEDIUM = 2; + HIGH = 3; + ULTRA_HIGH = 4; +} + +// An image content block. +message ImageContent { + reserved 3; + reserved "type"; + + enum MimeType { + reserved 6; + TYPE_UNSPECIFIED = 0; + TYPE_PNG = 1; // image/png + TYPE_JPEG = 2; // image/jpeg + TYPE_WEBP = 3; // image/webp + TYPE_HEIC = 4; // image/heic + TYPE_HEIF = 5; // image/heif + TYPE_GIF = 7; // image/gif + TYPE_BMP = 8; // image/bmp + TYPE_TIFF = 9; // image/tiff + } + MimeType mime_type = 1; + oneof data_or_uri { + bytes data = 2; + string uri = 6; + } + MediaResolution resolution = 5; +} + +// An audio content block. +message AudioContent { + reserved 3, 6; + reserved "type", "rate"; + + enum MimeType { + TYPE_UNSPECIFIED = 0; + TYPE_WAV = 1; // audio/wav + TYPE_MP3 = 2; // audio/mp3 + TYPE_AIFF = 3; // audio/aiff + TYPE_AAC = 4; // audio/aac + TYPE_OGG = 5; // audio/ogg + TYPE_FLAC = 6; // audio/flac + TYPE_MPEG = 7; // audio/mpeg + TYPE_M4A = 8; // audio/m4a + TYPE_L16 = 9; // audio/l16 + TYPE_S16LE = 10; // audio/s16le + TYPE_OPUS = 11; // audio/opus + TYPE_ALAW = 12; // audio/alaw + TYPE_MULAW = 13; // audio/mulaw + } + MimeType mime_type = 1; + oneof data_or_uri { + bytes data = 2; + string uri = 5; + } + int32 channels = 7; + int32 sample_rate = 8; +} + +// A document content block. +message DocumentContent { + reserved 3; + reserved "type"; + + enum MimeType { + TYPE_UNSPECIFIED = 0; + TYPE_PDF = 1; // application/pdf + TYPE_JSON = 2; // application/json + TYPE_PYTHON = 3; // text/x-python + } + MimeType mime_type = 1; + oneof data_or_uri { + bytes data = 2; + string uri = 5; + } +} + +// A video content block. +message VideoContent { + reserved 3; + reserved "type"; + + enum MimeType { + TYPE_UNSPECIFIED = 0; + TYPE_MP4 = 1; // video/mp4 + TYPE_MPEG = 2; // video/mpeg + TYPE_MPG = 3; // video/mpg + TYPE_MOV = 4; // video/mov + TYPE_AVI = 5; // video/avi + TYPE_X_FLV = 6; // video/x-flv + TYPE_WEBM = 7; // video/webm + TYPE_WMV = 8; // video/wmv + TYPE_3GPP = 9; // video/3gpp + } + MimeType mime_type = 1; + oneof data_or_uri { + bytes data = 2; + string uri = 6; + } + MediaResolution resolution = 5; +} + +// Content represents a content input or output. +message Content { + reserved 1 to 4, 6 to 9, 15 to 23; + + oneof type { + ThoughtContent thought = 5; + TextContent text = 10; + ImageContent image = 11; + AudioContent audio = 12; + DocumentContent document = 13; + VideoContent video = 14; + ConfirmationContent confirmation = 26; // TODO(jbd): Remove out of the Content. + ToolCallContent tool_call = 24; + ToolResultContent tool_result = 25; + } +} diff --git a/proto/gar.pb.go b/proto/gar.pb.go deleted file mode 100644 index cd325fd0..00000000 --- a/proto/gar.pb.go +++ /dev/null @@ -1,958 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Code generated by protoc-gen-go. DO NOT EDIT. -// versions: -// protoc-gen-go v1.35.1 -// protoc v5.28.2 -// source: proto/gar.proto - -package proto - -import ( - protoreflect "google.golang.org/protobuf/reflect/protoreflect" - protoimpl "google.golang.org/protobuf/runtime/protoimpl" - timestamppb "google.golang.org/protobuf/types/known/timestamppb" - reflect "reflect" - sync "sync" -) - -const ( - // Verify that this generated code is sufficiently up-to-date. - _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) - // Verify that runtime/protoimpl is sufficiently up-to-date. - _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) -) - -// State represents the state of a session -type State int32 - -const ( - State_STATE_UNSPECIFIED State = 0 // Unspecified state - State_STATE_RUNNING State = 1 // Session is currently running - State_STATE_FAILED State = 2 // Session failed - State_STATE_COMPLETED State = 3 // Completed, can accept new triggers -) - -// Enum value maps for State. -var ( - State_name = map[int32]string{ - 0: "STATE_UNSPECIFIED", - 1: "STATE_RUNNING", - 2: "STATE_FAILED", - 3: "STATE_COMPLETED", - } - State_value = map[string]int32{ - "STATE_UNSPECIFIED": 0, - "STATE_RUNNING": 1, - "STATE_FAILED": 2, - "STATE_COMPLETED": 3, - } -) - -func (x State) Enum() *State { - p := new(State) - *p = x - return p -} - -func (x State) String() string { - return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) -} - -func (State) Descriptor() protoreflect.EnumDescriptor { - return file_proto_gar_proto_enumTypes[0].Descriptor() -} - -func (State) Type() protoreflect.EnumType { - return &file_proto_gar_proto_enumTypes[0] -} - -func (x State) Number() protoreflect.EnumNumber { - return protoreflect.EnumNumber(x) -} - -// Deprecated: Use State.Descriptor instead. -func (State) EnumDescriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{0} -} - -// Content represents a message with role, type, mimetype, and data fields -type Content struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Role string `protobuf:"bytes,1,opt,name=role,proto3" json:"role,omitempty"` // The role of the content (e.g., "user", "assistant", "system") - Type string `protobuf:"bytes,2,opt,name=type,proto3" json:"type,omitempty"` // The type of content (e.g., "text", "image") - Mimetype string `protobuf:"bytes,3,opt,name=mimetype,proto3" json:"mimetype,omitempty"` // MIME type of the data (e.g., "text/plain", "image/png") - Data string `protobuf:"bytes,4,opt,name=data,proto3" json:"data,omitempty"` // The actual content data - CheckpointId string `protobuf:"bytes,5,opt,name=checkpoint_id,json=checkpointId,proto3" json:"checkpoint_id,omitempty"` // Optional: Checkpoint ID for this content -} - -func (x *Content) Reset() { - *x = Content{} - mi := &file_proto_gar_proto_msgTypes[0] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *Content) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*Content) ProtoMessage() {} - -func (x *Content) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[0] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use Content.ProtoReflect.Descriptor instead. -func (*Content) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{0} -} - -func (x *Content) GetRole() string { - if x != nil { - return x.Role - } - return "" -} - -func (x *Content) GetType() string { - if x != nil { - return x.Type - } - return "" -} - -func (x *Content) GetMimetype() string { - if x != nil { - return x.Mimetype - } - return "" -} - -func (x *Content) GetData() string { - if x != nil { - return x.Data - } - return "" -} - -func (x *Content) GetCheckpointId() string { - if x != nil { - return x.CheckpointId - } - return "" -} - -// HealthCheckRequest for agent health checks -type HealthCheckRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields -} - -func (x *HealthCheckRequest) Reset() { - *x = HealthCheckRequest{} - mi := &file_proto_gar_proto_msgTypes[1] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *HealthCheckRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*HealthCheckRequest) ProtoMessage() {} - -func (x *HealthCheckRequest) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[1] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use HealthCheckRequest.ProtoReflect.Descriptor instead. -func (*HealthCheckRequest) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{1} -} - -// HealthCheckResponse contains agent health status -type HealthCheckResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Healthy bool `protobuf:"varint,1,opt,name=healthy,proto3" json:"healthy,omitempty"` - Message string `protobuf:"bytes,2,opt,name=message,proto3" json:"message,omitempty"` -} - -func (x *HealthCheckResponse) Reset() { - *x = HealthCheckResponse{} - mi := &file_proto_gar_proto_msgTypes[2] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *HealthCheckResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*HealthCheckResponse) ProtoMessage() {} - -func (x *HealthCheckResponse) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[2] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use HealthCheckResponse.ProtoReflect.Descriptor instead. -func (*HealthCheckResponse) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{2} -} - -func (x *HealthCheckResponse) GetHealthy() bool { - if x != nil { - return x.Healthy - } - return false -} - -func (x *HealthCheckResponse) GetMessage() string { - if x != nil { - return x.Message - } - return "" -} - -// TriggerSessionRequest for triggering a new session -type TriggerSessionRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` // Unique session identifier - CheckpointId string `protobuf:"bytes,2,opt,name=checkpoint_id,json=checkpointId,proto3" json:"checkpoint_id,omitempty"` // Optional: Resume from specific checkpoint UUID (empty for latest) - Inputs []*Content `protobuf:"bytes,3,rep,name=inputs,proto3" json:"inputs,omitempty"` // Input content to process -} - -func (x *TriggerSessionRequest) Reset() { - *x = TriggerSessionRequest{} - mi := &file_proto_gar_proto_msgTypes[3] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *TriggerSessionRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*TriggerSessionRequest) ProtoMessage() {} - -func (x *TriggerSessionRequest) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[3] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use TriggerSessionRequest.ProtoReflect.Descriptor instead. -func (*TriggerSessionRequest) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{3} -} - -func (x *TriggerSessionRequest) GetSessionId() string { - if x != nil { - return x.SessionId - } - return "" -} - -func (x *TriggerSessionRequest) GetCheckpointId() string { - if x != nil { - return x.CheckpointId - } - return "" -} - -func (x *TriggerSessionRequest) GetInputs() []*Content { - if x != nil { - return x.Inputs - } - return nil -} - -// TriggerSessionResponse contains the result of triggering a session -type TriggerSessionResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` // Session ID (in case it was generated or modified by server) - State State `protobuf:"varint,2,opt,name=state,proto3,enum=proto.State" json:"state,omitempty"` // Session state - Output *Content `protobuf:"bytes,3,opt,name=output,proto3" json:"output,omitempty"` // Output content -} - -func (x *TriggerSessionResponse) Reset() { - *x = TriggerSessionResponse{} - mi := &file_proto_gar_proto_msgTypes[4] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *TriggerSessionResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*TriggerSessionResponse) ProtoMessage() {} - -func (x *TriggerSessionResponse) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[4] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use TriggerSessionResponse.ProtoReflect.Descriptor instead. -func (*TriggerSessionResponse) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{4} -} - -func (x *TriggerSessionResponse) GetSessionId() string { - if x != nil { - return x.SessionId - } - return "" -} - -func (x *TriggerSessionResponse) GetState() State { - if x != nil { - return x.State - } - return State_STATE_UNSPECIFIED -} - -func (x *TriggerSessionResponse) GetOutput() *Content { - if x != nil { - return x.Output - } - return nil -} - -// GetSessionRequest for retrieving session details -type GetSessionRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` -} - -func (x *GetSessionRequest) Reset() { - *x = GetSessionRequest{} - mi := &file_proto_gar_proto_msgTypes[5] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *GetSessionRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*GetSessionRequest) ProtoMessage() {} - -func (x *GetSessionRequest) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[5] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use GetSessionRequest.ProtoReflect.Descriptor instead. -func (*GetSessionRequest) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{5} -} - -func (x *GetSessionRequest) GetSessionId() string { - if x != nil { - return x.SessionId - } - return "" -} - -// SessionInfo contains session details -type SessionInfo struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - State State `protobuf:"varint,1,opt,name=state,proto3,enum=proto.State" json:"state,omitempty"` - ActiveAgents []string `protobuf:"bytes,2,rep,name=active_agents,json=activeAgents,proto3" json:"active_agents,omitempty"` - CheckpointCount int32 `protobuf:"varint,3,opt,name=checkpoint_count,json=checkpointCount,proto3" json:"checkpoint_count,omitempty"` - CreatedAt *timestamppb.Timestamp `protobuf:"bytes,4,opt,name=created_at,json=createdAt,proto3" json:"created_at,omitempty"` // Session creation time - UpdatedAt *timestamppb.Timestamp `protobuf:"bytes,5,opt,name=updated_at,json=updatedAt,proto3" json:"updated_at,omitempty"` // Last update time -} - -func (x *SessionInfo) Reset() { - *x = SessionInfo{} - mi := &file_proto_gar_proto_msgTypes[6] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *SessionInfo) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*SessionInfo) ProtoMessage() {} - -func (x *SessionInfo) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[6] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use SessionInfo.ProtoReflect.Descriptor instead. -func (*SessionInfo) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{6} -} - -func (x *SessionInfo) GetState() State { - if x != nil { - return x.State - } - return State_STATE_UNSPECIFIED -} - -func (x *SessionInfo) GetActiveAgents() []string { - if x != nil { - return x.ActiveAgents - } - return nil -} - -func (x *SessionInfo) GetCheckpointCount() int32 { - if x != nil { - return x.CheckpointCount - } - return 0 -} - -func (x *SessionInfo) GetCreatedAt() *timestamppb.Timestamp { - if x != nil { - return x.CreatedAt - } - return nil -} - -func (x *SessionInfo) GetUpdatedAt() *timestamppb.Timestamp { - if x != nil { - return x.UpdatedAt - } - return nil -} - -// GetSessionResponse contains session information -type GetSessionResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Session *SessionInfo `protobuf:"bytes,1,opt,name=session,proto3" json:"session,omitempty"` -} - -func (x *GetSessionResponse) Reset() { - *x = GetSessionResponse{} - mi := &file_proto_gar_proto_msgTypes[7] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *GetSessionResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*GetSessionResponse) ProtoMessage() {} - -func (x *GetSessionResponse) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[7] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use GetSessionResponse.ProtoReflect.Descriptor instead. -func (*GetSessionResponse) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{7} -} - -func (x *GetSessionResponse) GetSession() *SessionInfo { - if x != nil { - return x.Session - } - return nil -} - -// RegisterAgentRequest for registering an agent -type RegisterAgentRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - AgentId string `protobuf:"bytes,1,opt,name=agent_id,json=agentId,proto3" json:"agent_id,omitempty"` - Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` // Human-readable name for the agent - Description string `protobuf:"bytes,3,opt,name=description,proto3" json:"description,omitempty"` // Description of agent capabilities - Address string `protobuf:"bytes,4,opt,name=address,proto3" json:"address,omitempty"` // gRPC address for remote agents - Metadata map[string]string `protobuf:"bytes,5,rep,name=metadata,proto3" json:"metadata,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"` -} - -func (x *RegisterAgentRequest) Reset() { - *x = RegisterAgentRequest{} - mi := &file_proto_gar_proto_msgTypes[8] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *RegisterAgentRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*RegisterAgentRequest) ProtoMessage() {} - -func (x *RegisterAgentRequest) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[8] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use RegisterAgentRequest.ProtoReflect.Descriptor instead. -func (*RegisterAgentRequest) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{8} -} - -func (x *RegisterAgentRequest) GetAgentId() string { - if x != nil { - return x.AgentId - } - return "" -} - -func (x *RegisterAgentRequest) GetName() string { - if x != nil { - return x.Name - } - return "" -} - -func (x *RegisterAgentRequest) GetDescription() string { - if x != nil { - return x.Description - } - return "" -} - -func (x *RegisterAgentRequest) GetAddress() string { - if x != nil { - return x.Address - } - return "" -} - -func (x *RegisterAgentRequest) GetMetadata() map[string]string { - if x != nil { - return x.Metadata - } - return nil -} - -// RegisterAgentResponse contains registration result -type RegisterAgentResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields -} - -func (x *RegisterAgentResponse) Reset() { - *x = RegisterAgentResponse{} - mi := &file_proto_gar_proto_msgTypes[9] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *RegisterAgentResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*RegisterAgentResponse) ProtoMessage() {} - -func (x *RegisterAgentResponse) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[9] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use RegisterAgentResponse.ProtoReflect.Descriptor instead. -func (*RegisterAgentResponse) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{9} -} - -// UnregisterAgentRequest for unregistering an agent -type UnregisterAgentRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - AgentId string `protobuf:"bytes,1,opt,name=agent_id,json=agentId,proto3" json:"agent_id,omitempty"` -} - -func (x *UnregisterAgentRequest) Reset() { - *x = UnregisterAgentRequest{} - mi := &file_proto_gar_proto_msgTypes[10] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *UnregisterAgentRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*UnregisterAgentRequest) ProtoMessage() {} - -func (x *UnregisterAgentRequest) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[10] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use UnregisterAgentRequest.ProtoReflect.Descriptor instead. -func (*UnregisterAgentRequest) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{10} -} - -func (x *UnregisterAgentRequest) GetAgentId() string { - if x != nil { - return x.AgentId - } - return "" -} - -// UnregisterAgentResponse contains unregistration result -type UnregisterAgentResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields -} - -func (x *UnregisterAgentResponse) Reset() { - *x = UnregisterAgentResponse{} - mi := &file_proto_gar_proto_msgTypes[11] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) -} - -func (x *UnregisterAgentResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*UnregisterAgentResponse) ProtoMessage() {} - -func (x *UnregisterAgentResponse) ProtoReflect() protoreflect.Message { - mi := &file_proto_gar_proto_msgTypes[11] - if x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use UnregisterAgentResponse.ProtoReflect.Descriptor instead. -func (*UnregisterAgentResponse) Descriptor() ([]byte, []int) { - return file_proto_gar_proto_rawDescGZIP(), []int{11} -} - -var File_proto_gar_proto protoreflect.FileDescriptor - -var file_proto_gar_proto_rawDesc = []byte{ - 0x0a, 0x0f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x67, 0x61, 0x72, 0x2e, 0x70, 0x72, 0x6f, 0x74, - 0x6f, 0x12, 0x05, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x1a, 0x1f, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, - 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, - 0x61, 0x6d, 0x70, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x86, 0x01, 0x0a, 0x07, 0x43, 0x6f, - 0x6e, 0x74, 0x65, 0x6e, 0x74, 0x12, 0x12, 0x0a, 0x04, 0x72, 0x6f, 0x6c, 0x65, 0x18, 0x01, 0x20, - 0x01, 0x28, 0x09, 0x52, 0x04, 0x72, 0x6f, 0x6c, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, - 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x12, 0x1a, 0x0a, - 0x08, 0x6d, 0x69, 0x6d, 0x65, 0x74, 0x79, 0x70, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, - 0x08, 0x6d, 0x69, 0x6d, 0x65, 0x74, 0x79, 0x70, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x64, 0x61, 0x74, - 0x61, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x64, 0x61, 0x74, 0x61, 0x12, 0x23, 0x0a, - 0x0d, 0x63, 0x68, 0x65, 0x63, 0x6b, 0x70, 0x6f, 0x69, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x05, - 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, 0x63, 0x68, 0x65, 0x63, 0x6b, 0x70, 0x6f, 0x69, 0x6e, 0x74, - 0x49, 0x64, 0x22, 0x14, 0x0a, 0x12, 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x43, 0x68, 0x65, 0x63, - 0x6b, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, 0x49, 0x0a, 0x13, 0x48, 0x65, 0x61, 0x6c, - 0x74, 0x68, 0x43, 0x68, 0x65, 0x63, 0x6b, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, - 0x18, 0x0a, 0x07, 0x68, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, - 0x52, 0x07, 0x68, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x79, 0x12, 0x18, 0x0a, 0x07, 0x6d, 0x65, 0x73, - 0x73, 0x61, 0x67, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, - 0x61, 0x67, 0x65, 0x22, 0x83, 0x01, 0x0a, 0x15, 0x54, 0x72, 0x69, 0x67, 0x67, 0x65, 0x72, 0x53, - 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1d, 0x0a, - 0x0a, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, - 0x09, 0x52, 0x09, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x23, 0x0a, 0x0d, - 0x63, 0x68, 0x65, 0x63, 0x6b, 0x70, 0x6f, 0x69, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, - 0x01, 0x28, 0x09, 0x52, 0x0c, 0x63, 0x68, 0x65, 0x63, 0x6b, 0x70, 0x6f, 0x69, 0x6e, 0x74, 0x49, - 0x64, 0x12, 0x26, 0x0a, 0x06, 0x69, 0x6e, 0x70, 0x75, 0x74, 0x73, 0x18, 0x03, 0x20, 0x03, 0x28, - 0x0b, 0x32, 0x0e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x43, 0x6f, 0x6e, 0x74, 0x65, 0x6e, - 0x74, 0x52, 0x06, 0x69, 0x6e, 0x70, 0x75, 0x74, 0x73, 0x22, 0x83, 0x01, 0x0a, 0x16, 0x54, 0x72, - 0x69, 0x67, 0x67, 0x65, 0x72, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, - 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x1d, 0x0a, 0x0a, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x5f, - 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, - 0x6e, 0x49, 0x64, 0x12, 0x22, 0x0a, 0x05, 0x73, 0x74, 0x61, 0x74, 0x65, 0x18, 0x02, 0x20, 0x01, - 0x28, 0x0e, 0x32, 0x0c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x65, - 0x52, 0x05, 0x73, 0x74, 0x61, 0x74, 0x65, 0x12, 0x26, 0x0a, 0x06, 0x6f, 0x75, 0x74, 0x70, 0x75, - 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, - 0x43, 0x6f, 0x6e, 0x74, 0x65, 0x6e, 0x74, 0x52, 0x06, 0x6f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x22, - 0x32, 0x0a, 0x11, 0x47, 0x65, 0x74, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, - 0x75, 0x65, 0x73, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x5f, - 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, - 0x6e, 0x49, 0x64, 0x22, 0xf7, 0x01, 0x0a, 0x0b, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x49, - 0x6e, 0x66, 0x6f, 0x12, 0x22, 0x0a, 0x05, 0x73, 0x74, 0x61, 0x74, 0x65, 0x18, 0x01, 0x20, 0x01, - 0x28, 0x0e, 0x32, 0x0c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x65, - 0x52, 0x05, 0x73, 0x74, 0x61, 0x74, 0x65, 0x12, 0x23, 0x0a, 0x0d, 0x61, 0x63, 0x74, 0x69, 0x76, - 0x65, 0x5f, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x73, 0x18, 0x02, 0x20, 0x03, 0x28, 0x09, 0x52, 0x0c, - 0x61, 0x63, 0x74, 0x69, 0x76, 0x65, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x73, 0x12, 0x29, 0x0a, 0x10, - 0x63, 0x68, 0x65, 0x63, 0x6b, 0x70, 0x6f, 0x69, 0x6e, 0x74, 0x5f, 0x63, 0x6f, 0x75, 0x6e, 0x74, - 0x18, 0x03, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0f, 0x63, 0x68, 0x65, 0x63, 0x6b, 0x70, 0x6f, 0x69, - 0x6e, 0x74, 0x43, 0x6f, 0x75, 0x6e, 0x74, 0x12, 0x39, 0x0a, 0x0a, 0x63, 0x72, 0x65, 0x61, 0x74, - 0x65, 0x64, 0x5f, 0x61, 0x74, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, - 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, - 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x09, 0x63, 0x72, 0x65, 0x61, 0x74, 0x65, 0x64, - 0x41, 0x74, 0x12, 0x39, 0x0a, 0x0a, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x64, 0x5f, 0x61, 0x74, - 0x18, 0x05, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, - 0x6d, 0x70, 0x52, 0x09, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x64, 0x41, 0x74, 0x22, 0x42, 0x0a, - 0x12, 0x47, 0x65, 0x74, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, - 0x6e, 0x73, 0x65, 0x12, 0x2c, 0x0a, 0x07, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x01, - 0x20, 0x01, 0x28, 0x0b, 0x32, 0x12, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x53, 0x65, 0x73, - 0x73, 0x69, 0x6f, 0x6e, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x07, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, - 0x6e, 0x22, 0x85, 0x02, 0x0a, 0x14, 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x41, 0x67, - 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x19, 0x0a, 0x08, 0x61, 0x67, - 0x65, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x61, 0x67, - 0x65, 0x6e, 0x74, 0x49, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, - 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x20, 0x0a, 0x0b, 0x64, 0x65, 0x73, - 0x63, 0x72, 0x69, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0b, - 0x64, 0x65, 0x73, 0x63, 0x72, 0x69, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x18, 0x0a, 0x07, 0x61, - 0x64, 0x64, 0x72, 0x65, 0x73, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x61, 0x64, - 0x64, 0x72, 0x65, 0x73, 0x73, 0x12, 0x45, 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, - 0x61, 0x18, 0x05, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x29, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, - 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, - 0x75, 0x65, 0x73, 0x74, 0x2e, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x45, 0x6e, 0x74, - 0x72, 0x79, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x1a, 0x3b, 0x0a, 0x0d, - 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, - 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, - 0x14, 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, - 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, 0x22, 0x17, 0x0a, 0x15, 0x52, 0x65, 0x67, - 0x69, 0x73, 0x74, 0x65, 0x72, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, - 0x73, 0x65, 0x22, 0x33, 0x0a, 0x16, 0x55, 0x6e, 0x72, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, - 0x41, 0x67, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x19, 0x0a, 0x08, - 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, - 0x61, 0x67, 0x65, 0x6e, 0x74, 0x49, 0x64, 0x22, 0x19, 0x0a, 0x17, 0x55, 0x6e, 0x72, 0x65, 0x67, - 0x69, 0x73, 0x74, 0x65, 0x72, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, - 0x73, 0x65, 0x2a, 0x58, 0x0a, 0x05, 0x53, 0x74, 0x61, 0x74, 0x65, 0x12, 0x15, 0x0a, 0x11, 0x53, - 0x54, 0x41, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x53, 0x50, 0x45, 0x43, 0x49, 0x46, 0x49, 0x45, 0x44, - 0x10, 0x00, 0x12, 0x11, 0x0a, 0x0d, 0x53, 0x54, 0x41, 0x54, 0x45, 0x5f, 0x52, 0x55, 0x4e, 0x4e, - 0x49, 0x4e, 0x47, 0x10, 0x01, 0x12, 0x10, 0x0a, 0x0c, 0x53, 0x54, 0x41, 0x54, 0x45, 0x5f, 0x46, - 0x41, 0x49, 0x4c, 0x45, 0x44, 0x10, 0x02, 0x12, 0x13, 0x0a, 0x0f, 0x53, 0x54, 0x41, 0x54, 0x45, - 0x5f, 0x43, 0x4f, 0x4d, 0x50, 0x4c, 0x45, 0x54, 0x45, 0x44, 0x10, 0x03, 0x32, 0x83, 0x01, 0x0a, - 0x0c, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x2d, 0x0a, - 0x07, 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x12, 0x0e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, - 0x2e, 0x43, 0x6f, 0x6e, 0x74, 0x65, 0x6e, 0x74, 0x1a, 0x0e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, - 0x2e, 0x43, 0x6f, 0x6e, 0x74, 0x65, 0x6e, 0x74, 0x28, 0x01, 0x30, 0x01, 0x12, 0x44, 0x0a, 0x0b, - 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x43, 0x68, 0x65, 0x63, 0x6b, 0x12, 0x19, 0x2e, 0x70, 0x72, - 0x6f, 0x74, 0x6f, 0x2e, 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x43, 0x68, 0x65, 0x63, 0x6b, 0x52, - 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x48, - 0x65, 0x61, 0x6c, 0x74, 0x68, 0x43, 0x68, 0x65, 0x63, 0x6b, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, - 0x73, 0x65, 0x32, 0xbe, 0x02, 0x0a, 0x0a, 0x47, 0x41, 0x52, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, - 0x65, 0x12, 0x4f, 0x0a, 0x0e, 0x54, 0x72, 0x69, 0x67, 0x67, 0x65, 0x72, 0x53, 0x65, 0x73, 0x73, - 0x69, 0x6f, 0x6e, 0x12, 0x1c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x54, 0x72, 0x69, 0x67, - 0x67, 0x65, 0x72, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, - 0x74, 0x1a, 0x1d, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x54, 0x72, 0x69, 0x67, 0x67, 0x65, - 0x72, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, - 0x30, 0x01, 0x12, 0x41, 0x0a, 0x0a, 0x47, 0x65, 0x74, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, - 0x12, 0x18, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x53, 0x65, 0x73, 0x73, - 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x19, 0x2e, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, - 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x4a, 0x0a, 0x0d, 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, - 0x72, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x12, 0x1b, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, - 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, - 0x65, 0x73, 0x74, 0x1a, 0x1c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x67, 0x69, - 0x73, 0x74, 0x65, 0x72, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, - 0x65, 0x12, 0x50, 0x0a, 0x0f, 0x55, 0x6e, 0x72, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x41, - 0x67, 0x65, 0x6e, 0x74, 0x12, 0x1d, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x55, 0x6e, 0x72, - 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, - 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x55, 0x6e, 0x72, 0x65, - 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, - 0x6e, 0x73, 0x65, 0x42, 0x1d, 0x5a, 0x1b, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, - 0x6d, 0x2f, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x67, 0x61, 0x72, 0x2f, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, -} - -var ( - file_proto_gar_proto_rawDescOnce sync.Once - file_proto_gar_proto_rawDescData = file_proto_gar_proto_rawDesc -) - -func file_proto_gar_proto_rawDescGZIP() []byte { - file_proto_gar_proto_rawDescOnce.Do(func() { - file_proto_gar_proto_rawDescData = protoimpl.X.CompressGZIP(file_proto_gar_proto_rawDescData) - }) - return file_proto_gar_proto_rawDescData -} - -var file_proto_gar_proto_enumTypes = make([]protoimpl.EnumInfo, 1) -var file_proto_gar_proto_msgTypes = make([]protoimpl.MessageInfo, 13) -var file_proto_gar_proto_goTypes = []any{ - (State)(0), // 0: proto.State - (*Content)(nil), // 1: proto.Content - (*HealthCheckRequest)(nil), // 2: proto.HealthCheckRequest - (*HealthCheckResponse)(nil), // 3: proto.HealthCheckResponse - (*TriggerSessionRequest)(nil), // 4: proto.TriggerSessionRequest - (*TriggerSessionResponse)(nil), // 5: proto.TriggerSessionResponse - (*GetSessionRequest)(nil), // 6: proto.GetSessionRequest - (*SessionInfo)(nil), // 7: proto.SessionInfo - (*GetSessionResponse)(nil), // 8: proto.GetSessionResponse - (*RegisterAgentRequest)(nil), // 9: proto.RegisterAgentRequest - (*RegisterAgentResponse)(nil), // 10: proto.RegisterAgentResponse - (*UnregisterAgentRequest)(nil), // 11: proto.UnregisterAgentRequest - (*UnregisterAgentResponse)(nil), // 12: proto.UnregisterAgentResponse - nil, // 13: proto.RegisterAgentRequest.MetadataEntry - (*timestamppb.Timestamp)(nil), // 14: google.protobuf.Timestamp -} -var file_proto_gar_proto_depIdxs = []int32{ - 1, // 0: proto.TriggerSessionRequest.inputs:type_name -> proto.Content - 0, // 1: proto.TriggerSessionResponse.state:type_name -> proto.State - 1, // 2: proto.TriggerSessionResponse.output:type_name -> proto.Content - 0, // 3: proto.SessionInfo.state:type_name -> proto.State - 14, // 4: proto.SessionInfo.created_at:type_name -> google.protobuf.Timestamp - 14, // 5: proto.SessionInfo.updated_at:type_name -> google.protobuf.Timestamp - 7, // 6: proto.GetSessionResponse.session:type_name -> proto.SessionInfo - 13, // 7: proto.RegisterAgentRequest.metadata:type_name -> proto.RegisterAgentRequest.MetadataEntry - 1, // 8: proto.AgentService.Process:input_type -> proto.Content - 2, // 9: proto.AgentService.HealthCheck:input_type -> proto.HealthCheckRequest - 4, // 10: proto.GARService.TriggerSession:input_type -> proto.TriggerSessionRequest - 6, // 11: proto.GARService.GetSession:input_type -> proto.GetSessionRequest - 9, // 12: proto.GARService.RegisterAgent:input_type -> proto.RegisterAgentRequest - 11, // 13: proto.GARService.UnregisterAgent:input_type -> proto.UnregisterAgentRequest - 1, // 14: proto.AgentService.Process:output_type -> proto.Content - 3, // 15: proto.AgentService.HealthCheck:output_type -> proto.HealthCheckResponse - 5, // 16: proto.GARService.TriggerSession:output_type -> proto.TriggerSessionResponse - 8, // 17: proto.GARService.GetSession:output_type -> proto.GetSessionResponse - 10, // 18: proto.GARService.RegisterAgent:output_type -> proto.RegisterAgentResponse - 12, // 19: proto.GARService.UnregisterAgent:output_type -> proto.UnregisterAgentResponse - 14, // [14:20] is the sub-list for method output_type - 8, // [8:14] is the sub-list for method input_type - 8, // [8:8] is the sub-list for extension type_name - 8, // [8:8] is the sub-list for extension extendee - 0, // [0:8] is the sub-list for field type_name -} - -func init() { file_proto_gar_proto_init() } -func file_proto_gar_proto_init() { - if File_proto_gar_proto != nil { - return - } - type x struct{} - out := protoimpl.TypeBuilder{ - File: protoimpl.DescBuilder{ - GoPackagePath: reflect.TypeOf(x{}).PkgPath(), - RawDescriptor: file_proto_gar_proto_rawDesc, - NumEnums: 1, - NumMessages: 13, - NumExtensions: 0, - NumServices: 2, - }, - GoTypes: file_proto_gar_proto_goTypes, - DependencyIndexes: file_proto_gar_proto_depIdxs, - EnumInfos: file_proto_gar_proto_enumTypes, - MessageInfos: file_proto_gar_proto_msgTypes, - }.Build() - File_proto_gar_proto = out.File - file_proto_gar_proto_rawDesc = nil - file_proto_gar_proto_goTypes = nil - file_proto_gar_proto_depIdxs = nil -} diff --git a/proto/gar.proto b/proto/gar.proto deleted file mode 100644 index 0823f82c..00000000 --- a/proto/gar.proto +++ /dev/null @@ -1,132 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -syntax = "proto3"; - -package proto; - -import "google/protobuf/timestamp.proto"; - -option go_package = "github.com/google/gar/proto"; - -// Content represents a message with role, type, mimetype, and data fields -message Content { - string role = 1; // The role of the content (e.g., "user", "assistant", "system") - string type = 2; // The type of content (e.g., "text", "image") - string mimetype = 3; // MIME type of the data (e.g., "text/plain", "image/png") - string data = 4; // The actual content data - string checkpoint_id = 5; // Optional: Checkpoint ID for this content - - // TODO: Replace Content with Interactions Content. -} - -// HealthCheckRequest for agent health checks -message HealthCheckRequest {} - -// HealthCheckResponse contains agent health status -message HealthCheckResponse { - bool healthy = 1; - string message = 2; -} - -// AgentService defines the gRPC service for agent communication -service AgentService { - // Process handles bidirectional streaming of content between controller and agent - rpc Process(stream Content) returns (stream Content); - - // HealthCheck checks if the agent is healthy and responsive - rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse); -} - -// State represents the state of a session -enum State { - STATE_UNSPECIFIED = 0; // Unspecified state - STATE_RUNNING = 1; // Session is currently running - STATE_FAILED = 2; // Session failed - STATE_COMPLETED = 3; // Completed, can accept new triggers - - // TODO(jbd): Add STATE_COMPLETED and only allow triggers if - // the state is STATE_UNSPECIFIED or STATE_COMPLETED. -} - -// TriggerSessionRequest for triggering a new session -message TriggerSessionRequest { - string session_id = 1; // Unique session identifier - string checkpoint_id = 2; // Optional: Resume from specific checkpoint UUID (empty for latest) - repeated Content inputs = 3; // Input content to process -} - -// TriggerSessionResponse contains the result of triggering a session -message TriggerSessionResponse { - string session_id = 1; // Session ID (in case it was generated or modified by server) - State state = 2; // Session state - Content output = 3; // Output content -} - -// GetSessionRequest for retrieving session details -message GetSessionRequest { - string session_id = 1; -} - -// SessionInfo contains session details -message SessionInfo { - State state = 1; - repeated string active_agents = 2; - int32 checkpoint_count = 3; - google.protobuf.Timestamp created_at = 4; // Session creation time - google.protobuf.Timestamp updated_at = 5; // Last update time -} - -// GetSessionResponse contains session information -message GetSessionResponse { - SessionInfo session = 1; -} - -// RegisterAgentRequest for registering an agent -message RegisterAgentRequest { - string agent_id = 1; - string name = 2; // Human-readable name for the agent - string description = 3; // Description of agent capabilities - string address = 4; // gRPC address for remote agents - map metadata = 5; -} - -// RegisterAgentResponse contains registration result -message RegisterAgentResponse { -} - -// UnregisterAgentRequest for unregistering an agent -message UnregisterAgentRequest { - string agent_id = 1; -} - -// UnregisterAgentResponse contains unregistration result -message UnregisterAgentResponse { -} - -// GARService defines the gRPC service for GAR operations -service GARService { - // TriggerSession triggers a new agentic loop session or resumes an existing one with streaming responses - // If the session_id already exists, it will be resumed from the last checkpoint - rpc TriggerSession(TriggerSessionRequest) returns (stream TriggerSessionResponse); - - // GetSession retrieves session details - rpc GetSession(GetSessionRequest) returns (GetSessionResponse); - - // RegisterAgent registers a new agent with the controller - rpc RegisterAgent(RegisterAgentRequest) returns (RegisterAgentResponse); - - // UnregisterAgent removes an agent from the controller - rpc UnregisterAgent(UnregisterAgentRequest) returns (UnregisterAgentResponse); -} diff --git a/proto/gar_grpc.pb.go b/proto/gar_grpc.pb.go deleted file mode 100644 index e5e4fbad..00000000 --- a/proto/gar_grpc.pb.go +++ /dev/null @@ -1,411 +0,0 @@ -// Copyright 2026 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Code generated by protoc-gen-go-grpc. DO NOT EDIT. -// versions: -// - protoc-gen-go-grpc v1.5.1 -// - protoc v5.28.2 -// source: proto/gar.proto - -package proto - -import ( - context "context" - - grpc "google.golang.org/grpc" - codes "google.golang.org/grpc/codes" - status "google.golang.org/grpc/status" -) - -// This is a compile-time assertion to ensure that this generated file -// is compatible with the grpc package it is being compiled against. -// Requires gRPC-Go v1.64.0 or later. -const _ = grpc.SupportPackageIsVersion9 - -const ( - AgentService_Process_FullMethodName = "/proto.AgentService/Process" - AgentService_HealthCheck_FullMethodName = "/proto.AgentService/HealthCheck" -) - -// AgentServiceClient is the client API for AgentService service. -// -// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. -// -// AgentService defines the gRPC service for agent communication -type AgentServiceClient interface { - // Process handles bidirectional streaming of content between controller and agent - Process(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[Content, Content], error) - // HealthCheck checks if the agent is healthy and responsive - HealthCheck(ctx context.Context, in *HealthCheckRequest, opts ...grpc.CallOption) (*HealthCheckResponse, error) -} - -type agentServiceClient struct { - cc grpc.ClientConnInterface -} - -func NewAgentServiceClient(cc grpc.ClientConnInterface) AgentServiceClient { - return &agentServiceClient{cc} -} - -func (c *agentServiceClient) Process(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[Content, Content], error) { - cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) - stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[0], AgentService_Process_FullMethodName, cOpts...) - if err != nil { - return nil, err - } - x := &grpc.GenericClientStream[Content, Content]{ClientStream: stream} - return x, nil -} - -// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. -type AgentService_ProcessClient = grpc.BidiStreamingClient[Content, Content] - -func (c *agentServiceClient) HealthCheck(ctx context.Context, in *HealthCheckRequest, opts ...grpc.CallOption) (*HealthCheckResponse, error) { - cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) - out := new(HealthCheckResponse) - err := c.cc.Invoke(ctx, AgentService_HealthCheck_FullMethodName, in, out, cOpts...) - if err != nil { - return nil, err - } - return out, nil -} - -// AgentServiceServer is the server API for AgentService service. -// All implementations must embed UnimplementedAgentServiceServer -// for forward compatibility. -// -// AgentService defines the gRPC service for agent communication -type AgentServiceServer interface { - // Process handles bidirectional streaming of content between controller and agent - Process(grpc.BidiStreamingServer[Content, Content]) error - // HealthCheck checks if the agent is healthy and responsive - HealthCheck(context.Context, *HealthCheckRequest) (*HealthCheckResponse, error) - mustEmbedUnimplementedAgentServiceServer() -} - -// UnimplementedAgentServiceServer must be embedded to have -// forward compatible implementations. -// -// NOTE: this should be embedded by value instead of pointer to avoid a nil -// pointer dereference when methods are called. -type UnimplementedAgentServiceServer struct{} - -func (UnimplementedAgentServiceServer) Process(grpc.BidiStreamingServer[Content, Content]) error { - return status.Errorf(codes.Unimplemented, "method Process not implemented") -} -func (UnimplementedAgentServiceServer) HealthCheck(context.Context, *HealthCheckRequest) (*HealthCheckResponse, error) { - return nil, status.Errorf(codes.Unimplemented, "method HealthCheck not implemented") -} -func (UnimplementedAgentServiceServer) mustEmbedUnimplementedAgentServiceServer() {} -func (UnimplementedAgentServiceServer) testEmbeddedByValue() {} - -// UnsafeAgentServiceServer may be embedded to opt out of forward compatibility for this service. -// Use of this interface is not recommended, as added methods to AgentServiceServer will -// result in compilation errors. -type UnsafeAgentServiceServer interface { - mustEmbedUnimplementedAgentServiceServer() -} - -func RegisterAgentServiceServer(s grpc.ServiceRegistrar, srv AgentServiceServer) { - // If the following call pancis, it indicates UnimplementedAgentServiceServer was - // embedded by pointer and is nil. This will cause panics if an - // unimplemented method is ever invoked, so we test this at initialization - // time to prevent it from happening at runtime later due to I/O. - if t, ok := srv.(interface{ testEmbeddedByValue() }); ok { - t.testEmbeddedByValue() - } - s.RegisterService(&AgentService_ServiceDesc, srv) -} - -func _AgentService_Process_Handler(srv interface{}, stream grpc.ServerStream) error { - return srv.(AgentServiceServer).Process(&grpc.GenericServerStream[Content, Content]{ServerStream: stream}) -} - -// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. -type AgentService_ProcessServer = grpc.BidiStreamingServer[Content, Content] - -func _AgentService_HealthCheck_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { - in := new(HealthCheckRequest) - if err := dec(in); err != nil { - return nil, err - } - if interceptor == nil { - return srv.(AgentServiceServer).HealthCheck(ctx, in) - } - info := &grpc.UnaryServerInfo{ - Server: srv, - FullMethod: AgentService_HealthCheck_FullMethodName, - } - handler := func(ctx context.Context, req interface{}) (interface{}, error) { - return srv.(AgentServiceServer).HealthCheck(ctx, req.(*HealthCheckRequest)) - } - return interceptor(ctx, in, info, handler) -} - -// AgentService_ServiceDesc is the grpc.ServiceDesc for AgentService service. -// It's only intended for direct use with grpc.RegisterService, -// and not to be introspected or modified (even as a copy) -var AgentService_ServiceDesc = grpc.ServiceDesc{ - ServiceName: "proto.AgentService", - HandlerType: (*AgentServiceServer)(nil), - Methods: []grpc.MethodDesc{ - { - MethodName: "HealthCheck", - Handler: _AgentService_HealthCheck_Handler, - }, - }, - Streams: []grpc.StreamDesc{ - { - StreamName: "Process", - Handler: _AgentService_Process_Handler, - ServerStreams: true, - ClientStreams: true, - }, - }, - Metadata: "proto/gar.proto", -} - -const ( - GARService_TriggerSession_FullMethodName = "/proto.GARService/TriggerSession" - GARService_GetSession_FullMethodName = "/proto.GARService/GetSession" - GARService_RegisterAgent_FullMethodName = "/proto.GARService/RegisterAgent" - GARService_UnregisterAgent_FullMethodName = "/proto.GARService/UnregisterAgent" -) - -// GARServiceClient is the client API for GARService service. -// -// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. -// -// GARService defines the gRPC service for GAR operations -type GARServiceClient interface { - // TriggerSession triggers a new agentic loop session or resumes an existing one with streaming responses - // If the session_id already exists, it will be resumed from the last checkpoint - TriggerSession(ctx context.Context, in *TriggerSessionRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[TriggerSessionResponse], error) - // GetSession retrieves session details - GetSession(ctx context.Context, in *GetSessionRequest, opts ...grpc.CallOption) (*GetSessionResponse, error) - // RegisterAgent registers a new agent with the controller - RegisterAgent(ctx context.Context, in *RegisterAgentRequest, opts ...grpc.CallOption) (*RegisterAgentResponse, error) - // UnregisterAgent removes an agent from the controller - UnregisterAgent(ctx context.Context, in *UnregisterAgentRequest, opts ...grpc.CallOption) (*UnregisterAgentResponse, error) -} - -type gARServiceClient struct { - cc grpc.ClientConnInterface -} - -func NewGARServiceClient(cc grpc.ClientConnInterface) GARServiceClient { - return &gARServiceClient{cc} -} - -func (c *gARServiceClient) TriggerSession(ctx context.Context, in *TriggerSessionRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[TriggerSessionResponse], error) { - cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) - stream, err := c.cc.NewStream(ctx, &GARService_ServiceDesc.Streams[0], GARService_TriggerSession_FullMethodName, cOpts...) - if err != nil { - return nil, err - } - x := &grpc.GenericClientStream[TriggerSessionRequest, TriggerSessionResponse]{ClientStream: stream} - if err := x.ClientStream.SendMsg(in); err != nil { - return nil, err - } - if err := x.ClientStream.CloseSend(); err != nil { - return nil, err - } - return x, nil -} - -// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. -type GARService_TriggerSessionClient = grpc.ServerStreamingClient[TriggerSessionResponse] - -func (c *gARServiceClient) GetSession(ctx context.Context, in *GetSessionRequest, opts ...grpc.CallOption) (*GetSessionResponse, error) { - cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) - out := new(GetSessionResponse) - err := c.cc.Invoke(ctx, GARService_GetSession_FullMethodName, in, out, cOpts...) - if err != nil { - return nil, err - } - return out, nil -} - -func (c *gARServiceClient) RegisterAgent(ctx context.Context, in *RegisterAgentRequest, opts ...grpc.CallOption) (*RegisterAgentResponse, error) { - cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) - out := new(RegisterAgentResponse) - err := c.cc.Invoke(ctx, GARService_RegisterAgent_FullMethodName, in, out, cOpts...) - if err != nil { - return nil, err - } - return out, nil -} - -func (c *gARServiceClient) UnregisterAgent(ctx context.Context, in *UnregisterAgentRequest, opts ...grpc.CallOption) (*UnregisterAgentResponse, error) { - cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) - out := new(UnregisterAgentResponse) - err := c.cc.Invoke(ctx, GARService_UnregisterAgent_FullMethodName, in, out, cOpts...) - if err != nil { - return nil, err - } - return out, nil -} - -// GARServiceServer is the server API for GARService service. -// All implementations must embed UnimplementedGARServiceServer -// for forward compatibility. -// -// GARService defines the gRPC service for GAR operations -type GARServiceServer interface { - // TriggerSession triggers a new agentic loop session or resumes an existing one with streaming responses - // If the session_id already exists, it will be resumed from the last checkpoint - TriggerSession(*TriggerSessionRequest, grpc.ServerStreamingServer[TriggerSessionResponse]) error - // GetSession retrieves session details - GetSession(context.Context, *GetSessionRequest) (*GetSessionResponse, error) - // RegisterAgent registers a new agent with the controller - RegisterAgent(context.Context, *RegisterAgentRequest) (*RegisterAgentResponse, error) - // UnregisterAgent removes an agent from the controller - UnregisterAgent(context.Context, *UnregisterAgentRequest) (*UnregisterAgentResponse, error) - mustEmbedUnimplementedGARServiceServer() -} - -// UnimplementedGARServiceServer must be embedded to have -// forward compatible implementations. -// -// NOTE: this should be embedded by value instead of pointer to avoid a nil -// pointer dereference when methods are called. -type UnimplementedGARServiceServer struct{} - -func (UnimplementedGARServiceServer) TriggerSession(*TriggerSessionRequest, grpc.ServerStreamingServer[TriggerSessionResponse]) error { - return status.Errorf(codes.Unimplemented, "method TriggerSession not implemented") -} -func (UnimplementedGARServiceServer) GetSession(context.Context, *GetSessionRequest) (*GetSessionResponse, error) { - return nil, status.Errorf(codes.Unimplemented, "method GetSession not implemented") -} -func (UnimplementedGARServiceServer) RegisterAgent(context.Context, *RegisterAgentRequest) (*RegisterAgentResponse, error) { - return nil, status.Errorf(codes.Unimplemented, "method RegisterAgent not implemented") -} -func (UnimplementedGARServiceServer) UnregisterAgent(context.Context, *UnregisterAgentRequest) (*UnregisterAgentResponse, error) { - return nil, status.Errorf(codes.Unimplemented, "method UnregisterAgent not implemented") -} -func (UnimplementedGARServiceServer) mustEmbedUnimplementedGARServiceServer() {} -func (UnimplementedGARServiceServer) testEmbeddedByValue() {} - -// UnsafeGARServiceServer may be embedded to opt out of forward compatibility for this service. -// Use of this interface is not recommended, as added methods to GARServiceServer will -// result in compilation errors. -type UnsafeGARServiceServer interface { - mustEmbedUnimplementedGARServiceServer() -} - -func RegisterGARServiceServer(s grpc.ServiceRegistrar, srv GARServiceServer) { - // If the following call pancis, it indicates UnimplementedGARServiceServer was - // embedded by pointer and is nil. This will cause panics if an - // unimplemented method is ever invoked, so we test this at initialization - // time to prevent it from happening at runtime later due to I/O. - if t, ok := srv.(interface{ testEmbeddedByValue() }); ok { - t.testEmbeddedByValue() - } - s.RegisterService(&GARService_ServiceDesc, srv) -} - -func _GARService_TriggerSession_Handler(srv interface{}, stream grpc.ServerStream) error { - m := new(TriggerSessionRequest) - if err := stream.RecvMsg(m); err != nil { - return err - } - return srv.(GARServiceServer).TriggerSession(m, &grpc.GenericServerStream[TriggerSessionRequest, TriggerSessionResponse]{ServerStream: stream}) -} - -// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. -type GARService_TriggerSessionServer = grpc.ServerStreamingServer[TriggerSessionResponse] - -func _GARService_GetSession_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { - in := new(GetSessionRequest) - if err := dec(in); err != nil { - return nil, err - } - if interceptor == nil { - return srv.(GARServiceServer).GetSession(ctx, in) - } - info := &grpc.UnaryServerInfo{ - Server: srv, - FullMethod: GARService_GetSession_FullMethodName, - } - handler := func(ctx context.Context, req interface{}) (interface{}, error) { - return srv.(GARServiceServer).GetSession(ctx, req.(*GetSessionRequest)) - } - return interceptor(ctx, in, info, handler) -} - -func _GARService_RegisterAgent_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { - in := new(RegisterAgentRequest) - if err := dec(in); err != nil { - return nil, err - } - if interceptor == nil { - return srv.(GARServiceServer).RegisterAgent(ctx, in) - } - info := &grpc.UnaryServerInfo{ - Server: srv, - FullMethod: GARService_RegisterAgent_FullMethodName, - } - handler := func(ctx context.Context, req interface{}) (interface{}, error) { - return srv.(GARServiceServer).RegisterAgent(ctx, req.(*RegisterAgentRequest)) - } - return interceptor(ctx, in, info, handler) -} - -func _GARService_UnregisterAgent_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { - in := new(UnregisterAgentRequest) - if err := dec(in); err != nil { - return nil, err - } - if interceptor == nil { - return srv.(GARServiceServer).UnregisterAgent(ctx, in) - } - info := &grpc.UnaryServerInfo{ - Server: srv, - FullMethod: GARService_UnregisterAgent_FullMethodName, - } - handler := func(ctx context.Context, req interface{}) (interface{}, error) { - return srv.(GARServiceServer).UnregisterAgent(ctx, req.(*UnregisterAgentRequest)) - } - return interceptor(ctx, in, info, handler) -} - -// GARService_ServiceDesc is the grpc.ServiceDesc for GARService service. -// It's only intended for direct use with grpc.RegisterService, -// and not to be introspected or modified (even as a copy) -var GARService_ServiceDesc = grpc.ServiceDesc{ - ServiceName: "proto.GARService", - HandlerType: (*GARServiceServer)(nil), - Methods: []grpc.MethodDesc{ - { - MethodName: "GetSession", - Handler: _GARService_GetSession_Handler, - }, - { - MethodName: "RegisterAgent", - Handler: _GARService_RegisterAgent_Handler, - }, - { - MethodName: "UnregisterAgent", - Handler: _GARService_UnregisterAgent_Handler, - }, - }, - Streams: []grpc.StreamDesc{ - { - StreamName: "TriggerSession", - Handler: _GARService_TriggerSession_Handler, - ServerStreams: true, - }, - }, - Metadata: "proto/gar.proto", -} diff --git a/python/antigravity/__init__.py b/python/antigravity/__init__.py new file mode 100644 index 00000000..58d482ea --- /dev/null +++ b/python/antigravity/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/python/antigravity/config.json b/python/antigravity/config.json new file mode 100644 index 00000000..9f78bbb5 --- /dev/null +++ b/python/antigravity/config.json @@ -0,0 +1,3 @@ +{ + "model": "gemini-3.1-pro-preview" +} \ No newline at end of file diff --git a/python/antigravity/harness_server.py b/python/antigravity/harness_server.py new file mode 100644 index 00000000..e8b53907 --- /dev/null +++ b/python/antigravity/harness_server.py @@ -0,0 +1,525 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# NOTE ON ARCHITECTURE: +# This gRPC server implements the AX HarnessService protocol. It embeds the +# Antigravity weather agent logic directly, serving it over production gRPC. + +import argparse + +import asyncio +import json +import logging +import os +import pathlib +import re +import sys +from typing import TypedDict +import grpc +from grpc_health.v1 import health, health_pb2, health_pb2_grpc +from google.protobuf.struct_pb2 import Struct + + +from python.proto import ax_pb2 +from python.proto import ax_pb2_grpc +from python.proto import content_pb2 +from google.antigravity import Agent, AgentConfig, LocalAgentConfig +from google.antigravity.types import Text, Thought, ToolCall + +# Fields that come from outside harness_config and must not be set through it: +# - conversation_id: taken from the runtime request (request.conversation_id). +# - save_dir: derived at the server level from the configured state_dir. +# 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. +_NON_HARNESS_CONFIG_FIELDS = frozenset({"conversation_id", "save_dir"}) + + +class HarnessConfigError(ValueError): + """Raised when request harness_config is not a valid overlay.""" + + +class ConversationIdError(ValueError): + """Raised when a request's conversation_id is unusable as a save_dir name.""" + + +def _validate_conversation_id(conversation_id: str) -> None: + """Guards conversation_id for safe use as a save_dir path component. + + Rejects empty ids and path separators / "." / ".." so a request cannot + escape state_dir. The id-format contract (length, charset) is left to the + Antigravity harness (forwarded cascade_id) to avoid the layers drifting. + """ + if not conversation_id: + raise ConversationIdError("conversation_id must be set") + if "/" in conversation_id or "\\" in conversation_id: + raise ConversationIdError( + "conversation_id must not contain a path separator, got " + f"{conversation_id!r}" + ) + if conversation_id in (".", ".."): + raise ConversationIdError( + f"conversation_id must not be a path component, got {conversation_id!r}" + ) + + +class VertexKwargs(TypedDict, total=False): + """Typed subset of LocalAgentConfig kwargs needed to enable Vertex AI. + + `total=False` so {} is a valid value (returned when env does not request + Vertex). When populated, all three keys are present. + """ + + vertex: bool + project: str + location: str + + +def _env_use_vertex() -> bool: + """True if env requests the Vertex AI backend (vs. AI Studio API key).""" + return os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "").lower() in ( + "true", + "1", + ) or os.environ.get("GOOGLE_GENAI_USE_ENTERPRISE", "").lower() in ("true", "1") + + +def _vertex_kwargs_from_env() -> VertexKwargs: + """Returns LocalAgentConfig kwargs from GOOGLE_CLOUD_{PROJECT,LOCATION} env. + + Temporary override until AGY supports reading these env vars natively. + Returns {} when env does not request Vertex (caller's programmatic config + stands as-is). When env requests Vertex, returns VertexKwargs populated + for passing to LocalAgentConfig.__init__ so AGY's @model_validator picks + them up. + + Raises ValueError if env requests Vertex but project/location are missing. + + TODO: remove once AGY reads these env vars natively. + """ + if not _env_use_vertex(): + return {} + + project = os.environ.get("GOOGLE_CLOUD_PROJECT", "") + location = os.environ.get("GOOGLE_CLOUD_LOCATION", "") + + missing = [ + name + for name, value in ( + ("project (set GOOGLE_CLOUD_PROJECT)", project), + ("location (set GOOGLE_CLOUD_LOCATION)", location), + ) + if not value + ] + if missing: + raise ValueError( + "Vertex AI backend requested but missing required config: " + + ", ".join(missing) + ) + + print(f"Vertex AI backend configured: project={project} location={location}") + return {"vertex": True, "project": project, "location": location} + + +def _build_default_config() -> LocalAgentConfig: + """Builds the default agent config the sidecar serves on startup. + + Vertex configuration is read from env via `_vertex_kwargs_from_env`. + + TODO(#194): per-request `harness_config` will override fields of this + default on a per-conversation basis. Until then, every conversation uses + this config. + """ + return LocalAgentConfig( + system_instructions="You are a helpful agent.", + **_vertex_kwargs_from_env(), + ) + + +def _has_credentials(config: AgentConfig | None) -> bool: + """Checks if Gemini credentials are set per AGY's accepted sources. + + Mirrors AGY's own validation. AGY accepts exactly these sources: + 1. GEMINI_API_KEY environment variable (read directly by AGY). + 2. config.api_key set programmatically (AI Studio path). + 3. config.vertex=True + config.{project,location} set (Vertex path). + 4. config.vertex=True + config.api_key set (Vertex Express Mode; + covered by case 2). + + Anything else (e.g. vertex=True without project/location) is rejected + by AGY at request time, so we reject it here at startup too. + """ + # Check env - AGY reads GEMINI_API_KEY directly from os.environ. + if os.environ.get("GEMINI_API_KEY"): + return True + + # Check passed in config + if config: + if getattr(config, "api_key", None): + return True + if getattr(config, "vertex", False): + # Vertex requires project + location, unless an api_key (Express + # Mode) is set - but Express Mode would have returned True above. + if getattr(config, "project", None) and getattr(config, "location", None): + return True + + return False + + +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 + # its own save_dir so at most one .db lives there; the SDK conv_id is its stem. + dbs = list(pathlib.Path(save_dir).glob("*.db")) + return dbs[0].stem if dbs else None + + +def _reject_disallowed_fields(overrides: dict[str, object]) -> None: + """Best-effort validation of a request harness_config overlay's keys. + + Rejects fields managed outside harness_config and unknown top-level + fields (typos that LocalAgentConfig's extra="ignore" would otherwise + silently drop). + Top-level only: nested-key and value/type validation is delegated to the + SDK's own LocalAgentConfig validation when the config is constructed. + """ + managed = sorted(set(overrides) & _NON_HARNESS_CONFIG_FIELDS) + if managed: + raise HarnessConfigError( + f"field(s) managed outside harness_config cannot be set: {', '.join(managed)}" + ) + unknown = sorted(set(overrides) - set(LocalAgentConfig.model_fields)) + if unknown: + raise HarnessConfigError(f"unknown config field(s): {', '.join(unknown)}") + + +class AntigravityHarnessServiceServicer(ax_pb2_grpc.HarnessServiceServicer): + """Implements the ax.HarnessService protocol over gRPC.""" + + 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, 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") + _reject_disallowed_fields(overrides) + else: + overrides = {} + + # Persistence values managed outside harness_config 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 + # open across turns until the client half-closes. + async for request in request_iterator: + if request.WhichOneof("type") != "start": + continue # cancel frames not handled yet + + async for response in self._run_turn(request): + yield response + + async def _run_turn(self, request): + print(f"[gRPC] Connect turn requested. conv_id={request.conversation_id}") + + # Guard conversation_id for safe use as a save_dir path component + # below. The id-format contract (length, charset) is owned by the + # Antigravity harness via the forwarded cascade_id. + try: + _validate_conversation_id(request.conversation_id) + except ConversationIdError 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 conversation_id: {exc}", + ), + ), + ) + return + + # 1. Retrieve messages. + ax_messages = request.start.messages + if not ax_messages: + latest_query_text = "" + else: + latest_message = ax_messages[-1] + + if latest_message.content.WhichOneof("type") != "text": + 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="Latest message must contain text content", + ), + ), + ) + return + latest_query_text = latest_message.content.text.text + + if not self._default_config: + yield ax_pb2.HarnessResponse( + conversation_id=request.conversation_id, + end=ax_pb2.HarnessEnd( + state=ax_pb2.STATE_FAILED, + error=ax_pb2.Error( + code=9, # FAILED_PRECONDITION + description="Agent config is not loaded on the server", + ), + ), + ) + return + try: + 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 + + print(f"[gRPC] Running chat query: {latest_query_text}") + response = await conversation.chat(latest_query_text) + + # To avoid streaming individual tokens inside TextContent messages (which is not + # supported by the Interactions proto/TextContent specifications), we buffer + # contiguous blocks of text and thought chunks, yielding them only when the + # contiguous block ends or a different chunk type is received. + text_chunks = [] + thought_chunks = [] + + def flush_text(): + if not text_chunks: + return None + msg = ax_pb2.Message( + role="assistant", + content=content_pb2.Content( + text=content_pb2.TextContent(text="".join(text_chunks)) + ), + ) + text_chunks.clear() + return ax_pb2.HarnessResponse( + conversation_id=request.conversation_id, + outputs=ax_pb2.HarnessOutputs(messages=[msg]), + ) + + def flush_thought(): + if not thought_chunks: + return None + + # Normalize Gemini's thinking output: collapse runs of 3+ newlines + # (emitted between reasoning sections) to exactly one blank line + # and strip trailing whitespace. Then append exactly one trailing + # newline so downstream displays render a blank line between the + # thinking block and whatever follows (next thought, tool call, + # or answer text). Without the trailing newline the display's + # transition logic emits only one newline, gluing blocks together. + raw_text = "".join(thought_chunks) + clean_text = re.sub(r"\n{3,}", "\n\n", raw_text).rstrip() + "\n" + + summary = [ + content_pb2.ThoughtSummaryContent( + text=content_pb2.TextContent(text=clean_text) + ) + ] + thought_chunks.clear() + msg = ax_pb2.Message( + role="model", + content=content_pb2.Content( + thought=content_pb2.ThoughtContent(summary=summary) + ), + ) + return ax_pb2.HarnessResponse( + conversation_id=request.conversation_id, + outputs=ax_pb2.HarnessOutputs(messages=[msg]), + ) + + async for chunk in response.chunks: + if isinstance(chunk, Text): + if resp := flush_thought(): + yield resp + text_chunks.append(chunk.text) + elif isinstance(chunk, Thought): + if resp := flush_text(): + yield resp + thought_chunks.append(chunk.text) + elif isinstance(chunk, ToolCall): + # Flush all pending text/thought buffers before dispatching the tool call + if resp := flush_text(): + yield resp + if resp := flush_thought(): + yield resp + + struct_args = Struct() + struct_args.update(chunk.args) + + func_call = content_pb2.FunctionCallContent( + name=str(chunk.name), arguments=struct_args + ) + msg = ax_pb2.Message( + role="model", + content=content_pb2.Content( + tool_call=content_pb2.ToolCallContent( + id=chunk.id or "", function_call=func_call + ) + ), + ) + yield ax_pb2.HarnessResponse( + conversation_id=request.conversation_id, + outputs=ax_pb2.HarnessOutputs(messages=[msg]), + ) + + # Flush any remaining text/thought buffers after the generator loop ends + if resp := flush_text(): + yield resp + if resp := flush_thought(): + yield resp + + # Yield completion end frame + yield ax_pb2.HarnessResponse( + conversation_id=request.conversation_id, + end=ax_pb2.HarnessEnd(state=ax_pb2.STATE_COMPLETED), + ) + 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( + conversation_id=request.conversation_id, + end=ax_pb2.HarnessEnd( + state=ax_pb2.STATE_FAILED, + error=ax_pb2.Error( + code=13, # INTERNAL + description=f"Agent execution terminated due to error. ({str(e)})", + ), + ), + ) + return + + +async def _serve( + host: str, port: int, default_config: AgentConfig, state_dir: pathlib.Path +): + server = grpc.aio.server() + servicer = AntigravityHarnessServiceServicer(default_config, state_dir) + ax_pb2_grpc.add_HarnessServiceServicer_to_server(servicer, server) + + # Serve the standard gRPC health protocol. + health_servicer = health.aio.HealthServicer() + health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server) + await health_servicer.set("", health_pb2.HealthCheckResponse.SERVING) + + listen_addr = f"{host}:{port}" + server.add_insecure_port(listen_addr) + print(f"Starting gRPC harness server on {listen_addr}...") + await server.start() + await server.wait_for_termination() + + +def _enhance_config_from_env(config) -> None: + skills_dir = os.environ.get("SKILLS_DIR") + if skills_dir and os.path.isdir(skills_dir): + print(f"Adding preinstalled skills directory to agent config: {skills_dir}") + if not hasattr(config, "skills_paths") or config.skills_paths is None: + config.skills_paths = [] + config.skills_paths = list(config.skills_paths) + if skills_dir not in config.skills_paths: + config.skills_paths.append(skills_dir) + + +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: + default_config = _build_default_config() + _enhance_config_from_env(default_config) + if not _has_credentials(default_config): + raise ValueError( + "No Gemini credentials configured. Set GEMINI_API_KEY " + "(AI Studio) or GOOGLE_GENAI_USE_VERTEXAI=True + " + "GOOGLE_CLOUD_{PROJECT,LOCATION} (Vertex AI)." + ) + except ValueError as e: + # Single startup-config exit point. + print(f"ERROR: {e}", file=sys.stderr) + sys.exit(1) + + 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 new file mode 100644 index 00000000..3d06d98f --- /dev/null +++ b/python/antigravity/harness_server_test.py @@ -0,0 +1,737 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# 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 ConversationIdError +from python.antigravity.harness_server import HarnessConfigError +from python.antigravity.harness_server import _validate_conversation_id +from google.antigravity import LocalAgentConfig + +@pytest.fixture +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, tmp_path): + async def _run(): + # 1. Start temporary local gRPC server on random open port + server = grpc.aio.server() + 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() + + # 2. Connect async stub channel + addr = f"localhost:{port}" + async with grpc.aio.insecure_channel(addr) as channel: + stub = ax_pb2_grpc.HarnessServiceStub(channel) + + # Mock the underlying Antigravity SDK class calls + class MockConversation: + def __init__(self): + self._steps = [] + async def chat(self, text): + class MockResponse: + def __init__(self): + self.chunks = self._chunk_generator() + async def _chunk_generator(self): + from google.antigravity.types import Text, Thought + yield Thought(text="Thinking details", step_index=0) + yield Text(text="Hello human", step_index=0) + return MockResponse() + + class MockAgent: + def __init__(self, config): + self.conversation = MockConversation() + async def __aenter__(self): + return self + async def __aexit__(self, exc_type, exc, tb): + pass + + monkeypatch.setattr("python.antigravity.harness_server.Agent", MockAgent) + + # 3. Construct and fire a HarnessRequest{start} over the bidi stream + start_payload = ax_pb2.HarnessStart( + messages=[ + ax_pb2.Message(role="user", content=content_pb2.Content(text=content_pb2.TextContent(text="Hi"))) + ] + ) + req = ax_pb2.HarnessRequest( + conversation_id="conv-test", + harness_id="antigravity", + start=start_payload + ) + + async def request_iter(): + yield req + + responses = [] + async for resp in stub.Connect(request_iter()): + responses.append(resp) + + # 4. Assert outputs are correctly mapped and completed + assert len(responses) == 3 # Thought + Text + End + assert responses[0].outputs.messages[0].content.thought.summary[0].text.text == "Thinking details\n" + assert responses[1].outputs.messages[0].content.text.text == "Hello human" + assert responses[2].WhichOneof('type') == 'end' + assert responses[2].end.state == ax_pb2.STATE_COMPLETED + + await server.stop(0) + + asyncio.run(_run()) + + +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 the + configured state_dir. Same AX conv_id -> same save_dir + (SDK-native resume).""" + + async def _run(): + server = grpc.aio.server() + 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() + + addr = f"localhost:{port}" + async with grpc.aio.insecure_channel(addr) as channel: + stub = ax_pb2_grpc.HarnessServiceStub(channel) + + class MockConversation: + async def chat(self, text): + class MockResponse: + def __init__(self): + self.chunks = self._chunk_generator() + async def _chunk_generator(self): + from google.antigravity.types import Text + yield Text(text="Response", step_index=0) + return MockResponse() + + agent_instances = [] + class MockAgent: + def __init__(self, config): + self.config = config + self.conversation = MockConversation() + agent_instances.append(self) + async def __aenter__(self): + return self + async def __aexit__(self, exc_type, exc, tb): + pass + + monkeypatch.setattr("python.antigravity.harness_server.Agent", MockAgent) + + async def fire(conv_id): + req = ax_pb2.HarnessRequest( + conversation_id=conv_id, + harness_id="antigravity", + start=ax_pb2.HarnessStart( + messages=[ax_pb2.Message(role="user", + content=content_pb2.Content(text=content_pb2.TextContent(text="Hi")))] + ), + ) + async def req_iter(): + yield req + async for _ in stub.Connect(req_iter()): + pass + + await fire("conv-1") + await fire("conv-1") + await fire("conv-2") + + assert len(agent_instances) == 3 + 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] == 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] + + await server.stop(0) + + asyncio.run(_run()) + + +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, 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) + port = server.add_insecure_port("localhost:0") + await server.start() + try: + async with grpc.aio.insecure_channel(f"localhost:{port}") as channel: + stub = health_pb2_grpc.HealthStub(channel) + resp = await stub.Check(health_pb2.HealthCheckRequest(service="")) + assert resp.status == health_pb2.HealthCheckResponse.SERVING + finally: + await server.stop(0) + + asyncio.run(_run()) + + +def test_has_credentials_missing(monkeypatch): + """Returns False when neither env nor config provides credentials.""" + from python.antigravity.harness_server import _has_credentials + + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + monkeypatch.delenv("GOOGLE_GENAI_USE_ENTERPRISE", raising=False) + + cfg = LocalAgentConfig(system_instructions="test") + assert _has_credentials(cfg) is False + + +def test_has_credentials_vertex_requires_project_and_location(monkeypatch): + """vertex=True alone is not enough; AGY requires project+location too.""" + from python.antigravity.harness_server import _has_credentials + + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + + cfg_vertex_only = LocalAgentConfig(system_instructions="test", vertex=True) + assert _has_credentials(cfg_vertex_only) is False + + cfg_vertex_proj_only = LocalAgentConfig(system_instructions="test", vertex=True, project="p") + assert _has_credentials(cfg_vertex_proj_only) is False + + cfg_vertex_loc_only = LocalAgentConfig(system_instructions="test", vertex=True, location="us-central1") + assert _has_credentials(cfg_vertex_loc_only) is False + + cfg_vertex_full = LocalAgentConfig(system_instructions="test", vertex=True, project="p", location="us-central1") + assert _has_credentials(cfg_vertex_full) is True + + +def test_has_credentials_vertex_express_mode(monkeypatch): + """vertex=True + api_key (Express Mode) is accepted even without project/location.""" + from python.antigravity.harness_server import _has_credentials + + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + + cfg = LocalAgentConfig(system_instructions="test", vertex=True, api_key="express-key") + assert _has_credentials(cfg) is True + + +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) + monkeypatch.delenv("GOOGLE_GENAI_USE_ENTERPRISE", raising=False) + + # Config with API key programmatically set + cfg = LocalAgentConfig(system_instructions="Test instructions", api_key="mock-config-api-key") + + async def _run(): + server = grpc.aio.server() + servicer = AntigravityHarnessServiceServicer(cfg, tmp_path) + ax_pb2_grpc.add_HarnessServiceServicer_to_server(servicer, server) + port = server.add_insecure_port("localhost:0") + await server.start() + + addr = f"localhost:{port}" + async with grpc.aio.insecure_channel(addr) as channel: + stub = ax_pb2_grpc.HarnessServiceStub(channel) + + # Mock Agent so we can test programmatic config logic passes + class MockConversation: + def __init__(self): + self._steps = [] + async def chat(self, text): + class MockResponse: + def __init__(self): + self.chunks = self._chunk_generator() + async def _chunk_generator(self): + from google.antigravity.types import Text + yield Text(text="Passed check", step_index=0) + return MockResponse() + + class MockAgent: + def __init__(self, config): + self.conversation = MockConversation() + async def __aenter__(self): + return self + async def __aexit__(self, exc_type, exc, tb): + pass + monkeypatch.setattr("python.antigravity.harness_server.Agent", MockAgent) + + start_payload = ax_pb2.HarnessStart( + messages=[ + ax_pb2.Message(role="user", content=content_pb2.Content(text=content_pb2.TextContent(text="Hi"))) + ] + ) + req = ax_pb2.HarnessRequest( + conversation_id="conv-test-prog", + harness_id="antigravity", + start=start_payload + ) + + async def request_iter(): + yield req + + responses = [] + async for resp in stub.Connect(request_iter()): + responses.append(resp) + + assert len(responses) == 2 # Text + End + assert responses[0].outputs.messages[0].content.text.text == "Passed check" + assert responses[1].end.state == ax_pb2.STATE_COMPLETED + + await server.stop(0) + + asyncio.run(_run()) + + +def test_enhance_config_from_env(monkeypatch, tmp_path): + from python.antigravity.harness_server import _enhance_config_from_env + from google.antigravity import LocalAgentConfig + import os + + # Create a mock skills dir + skills_dir = tmp_path / "skills" + skills_dir.mkdir() + + cfg = LocalAgentConfig(system_instructions="test") + + # Test: Using SKILLS_DIR env var + monkeypatch.setenv("SKILLS_DIR", str(skills_dir)) + _enhance_config_from_env(cfg) + assert str(skills_dir) in cfg.skills_paths + + +def test_grpc_connect_buffering(mock_config, monkeypatch, tmp_path): + async def _run(): + server = grpc.aio.server() + 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() + + addr = f"localhost:{port}" + async with grpc.aio.insecure_channel(addr) as channel: + stub = ax_pb2_grpc.HarnessServiceStub(channel) + + class MockConversation: + def __init__(self): + self._steps = [] + async def chat(self, text): + class MockResponse: + def __init__(self): + self.chunks = self._chunk_generator() + async def _chunk_generator(self): + from google.antigravity.types import Text, Thought, ToolCall + yield Thought(text="Think1", step_index=0) + yield Thought(text=" Think2", step_index=0) + yield ToolCall(name="tool1", args={}, id="call1") + yield Text(text="Hello", step_index=0) + yield Text(text=" human", step_index=0) + return MockResponse() + + class MockAgent: + def __init__(self, config): + self.conversation = MockConversation() + async def __aenter__(self): + return self + async def __aexit__(self, exc_type, exc, tb): + pass + monkeypatch.setattr("python.antigravity.harness_server.Agent", MockAgent) + + start_payload = ax_pb2.HarnessStart( + messages=[ + ax_pb2.Message(role="user", content=content_pb2.Content(text=content_pb2.TextContent(text="Hi"))) + ] + ) + req = ax_pb2.HarnessRequest( + conversation_id="conv-test-buffer", + harness_id="antigravity", + start=start_payload + ) + + async def request_iter(): + yield req + + responses = [] + async for resp in stub.Connect(request_iter()): + responses.append(resp) + + # Responses should be: + # 1. Thought ("Think1 Think2") - flushed when ToolCall is encountered + # 2. ToolCall ("tool1") - processed immediately + # 3. Text ("Hello human") - flushed at the end + # 4. End frame + assert len(responses) == 4 + + # Assert 1st response: Thought summary text is "Think1 Think2" + assert responses[0].outputs.messages[0].content.WhichOneof('type') == 'thought' + assert responses[0].outputs.messages[0].content.thought.summary[0].text.text == "Think1 Think2\n" + + # Assert 2nd response: ToolCall name is "tool1" + assert responses[1].outputs.messages[0].content.WhichOneof('type') == 'tool_call' + assert responses[1].outputs.messages[0].content.tool_call.function_call.name == "tool1" + + # Assert 3rd response: Text content is "Hello human" + assert responses[2].outputs.messages[0].content.WhichOneof('type') == 'text' + assert responses[2].outputs.messages[0].content.text.text == "Hello human" + + # Assert 4th response: Completion end frame + assert responses[3].WhichOneof('type') == 'end' + assert responses[3].end.state == ax_pb2.STATE_COMPLETED + + await server.stop(0) + + asyncio.run(_run()) + +def test_vertex_kwargs_from_env_returns_kwargs(monkeypatch): + """GOOGLE_GENAI_USE_VERTEXAI + GOOGLE_CLOUD_{PROJECT,LOCATION} -> kwargs dict.""" + from python.antigravity.harness_server import _vertex_kwargs_from_env + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "True") + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "env-project") + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-east1") + assert _vertex_kwargs_from_env() == { + "vertex": True, + "project": "env-project", + "location": "us-east1", + } + + +def test_vertex_kwargs_from_env_no_op_without_vertex(monkeypatch): + """Without GOOGLE_GENAI_USE_VERTEXAI, returns empty dict.""" + from python.antigravity.harness_server import _vertex_kwargs_from_env + monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + monkeypatch.delenv("GOOGLE_GENAI_USE_ENTERPRISE", raising=False) + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "should-be-ignored") + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "should-be-ignored") + assert _vertex_kwargs_from_env() == {} + + +def test_vertex_kwargs_from_env_raises_when_project_missing(monkeypatch): + """Vertex requested with no project -> ValueError naming the env var.""" + from python.antigravity.harness_server import _vertex_kwargs_from_env + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "True") + monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False) + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-east1") + with pytest.raises(ValueError, match="GOOGLE_CLOUD_PROJECT"): + _vertex_kwargs_from_env() + + +def test_vertex_kwargs_from_env_raises_when_location_missing(monkeypatch): + """Vertex requested with no location -> ValueError naming the env var.""" + from python.antigravity.harness_server import _vertex_kwargs_from_env + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "True") + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "env-project") + monkeypatch.delenv("GOOGLE_CLOUD_LOCATION", raising=False) + with pytest.raises(ValueError, match="GOOGLE_CLOUD_LOCATION"): + _vertex_kwargs_from_env() + + +def test_vertex_kwargs_from_env_enterprise_alias(monkeypatch): + """GOOGLE_GENAI_USE_ENTERPRISE is an alias for VERTEXAI.""" + from python.antigravity.harness_server import _vertex_kwargs_from_env + monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", "true") + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "ent-project") + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-central1") + assert _vertex_kwargs_from_env() == { + "vertex": True, + "project": "ent-project", + "location": "us-central1", + } + + +def test_build_default_config_picks_up_vertex_env(monkeypatch): + """End-to-end: env vars flow through _build_default_config into LocalAgentConfig.""" + from python.antigravity.harness_server import _build_default_config + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "True") + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "env-project") + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-east1") + cfg = _build_default_config() + assert cfg.vertex is True + assert cfg.project == "env-project" + assert cfg.location == "us-east1" + + +def test_servicer_requires_default_config(): + """Constructor takes a default config; passing nothing is a TypeError.""" + with pytest.raises(TypeError, match="default_config"): + AntigravityHarnessServiceServicer() + + +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, tmp_path) + servicer._default_config = None + server = grpc.aio.server() + ax_pb2_grpc.add_HarnessServiceServicer_to_server(servicer, server) + port = server.add_insecure_port("localhost:0") + await server.start() + try: + async with grpc.aio.insecure_channel(f"localhost:{port}") as channel: + stub = ax_pb2_grpc.HarnessServiceStub(channel) + req = ax_pb2.HarnessRequest( + conversation_id="conv-guard", + harness_id="antigravity", + start=ax_pb2.HarnessStart(messages=[ + ax_pb2.Message(role="user", + content=content_pb2.Content(text=content_pb2.TextContent(text="Hi"))), + ]), + ) + async def request_iter(): + yield req + responses = [] + async for resp in stub.Connect(request_iter()): + responses.append(resp) + assert len(responses) == 1 + assert responses[0].end.state == ax_pb2.STATE_FAILED + assert responses[0].end.error.code == 9 + assert "Agent config is not loaded" in responses[0].end.error.description + finally: + 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(), "managed outside harness_config"), + (json.dumps({"conversation_id": "other"}).encode(), "managed outside harness_config"), + ( + json.dumps({"capabilities": {"enabled_tools": ["not-a-tool"]}}).encode(), + "validation error", + ), + (json.dumps({"system_instruction": "typo"}).encode(), "unknown config field"), + (json.dumps({"model": "m", "frobnicate": True}).encode(), "unknown config field"), +]) +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()) + + +def test_harness_config_unknown_field_names_are_reported(mock_config, tmp_path): + # The error lists the offending field(s), sorted, so a typo is actionable; + # any unknown field rejects the whole overlay (no silent drop) and valid + # fields in the same request are not flagged. + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + raw = json.dumps({"zzz_bad": 1, "aaa_bad": 2, "system_instructions": "ok"}).encode() + with pytest.raises(HarnessConfigError) as excinfo: + servicer._build_config_for("conv-1", raw) + msg = str(excinfo.value) + assert "unknown config field(s): aaa_bad, zzz_bad" in msg + assert "system_instructions" not in msg + + + +@pytest.mark.parametrize("conv_id", [ + "conv-1", # short ids are fine here; the harness owns the format contract + "conv-test", + "11111111-2222-3333-4444-555555555555", + "a", # single char, still a safe dir name +]) +def test_validate_conversation_id_accepts_path_safe(conv_id): + # Should not raise: these are all usable as a save_dir path component. + _validate_conversation_id(conv_id) + + +@pytest.mark.parametrize(("conv_id", "error"), [ + ("", "must be set"), + ("..", "path component"), + (".", "path component"), + ("../escape", "path separator"), + ("a/b", "path separator"), + ("nested/../conv", "path separator"), + ("back\\slash", "path separator"), +]) +def test_validate_conversation_id_rejects_unsafe(conv_id, error): + with pytest.raises(ConversationIdError, match=error): + _validate_conversation_id(conv_id) + + +def test_run_turn_unsafe_conversation_id_maps_to_invalid_argument(mock_config, tmp_path): + # An unsafe conversation_id is rejected at the boundary: the turn yields a + # single STATE_FAILED end frame with INVALID_ARGUMENT, before any agent runs. + async def _run(): + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + req = ax_pb2.HarnessRequest( + conversation_id="../escape", + harness_id="antigravity", + start=ax_pb2.HarnessStart( + 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 conversation_id" in responses[0].end.error.description + + asyncio.run(_run()) + + +def test_run_turn_rejects_conversation_id_before_creating_save_dir(mock_config, tmp_path): + # Validation must happen before conversation_id is used as a storage + # directory name, so a rejected id leaves no directory behind under (or + # outside) state_dir. + async def _run(): + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + req = ax_pb2.HarnessRequest( + conversation_id="../escape", + harness_id="antigravity", + start=ax_pb2.HarnessStart( + 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 list(tmp_path.iterdir()) == [] + + asyncio.run(_run()) + + +def test_run_turn_empty_messages_resumes_with_empty_prompt(mock_config, monkeypatch, tmp_path): + """Antigravity SDK to resume with empty input.""" + + captured = {} + + class MockConversation: + async def chat(self, text): + captured["prompt"] = text + class MockResponse: + def __init__(self): + self.chunks = self._chunk_generator() + async def _chunk_generator(self): + from google.antigravity.types import Text + yield Text(text="continued essay", step_index=0) + return MockResponse() + + class MockAgent: + def __init__(self, config): + self.conversation = MockConversation() + async def __aenter__(self): + return self + async def __aexit__(self, exc_type, exc, tb): + pass + + monkeypatch.setattr("python.antigravity.harness_server.Agent", MockAgent) + + async def _run(): + servicer = AntigravityHarnessServiceServicer(mock_config, tmp_path) + req = ax_pb2.HarnessRequest( + conversation_id="conv-1", + harness_id="antigravity", + start=ax_pb2.HarnessStart(messages=[]), # empty = resume + ) + responses = [r async for r in servicer._run_turn(req)] + + assert captured["prompt"] == "" + assert all(r.end.state != ax_pb2.STATE_FAILED for r in responses if r.HasField("end")) + assert any(r.HasField("end") and r.end.state == ax_pb2.STATE_COMPLETED for r in responses) + + asyncio.run(_run()) diff --git a/python/antigravity/requirements.txt b/python/antigravity/requirements.txt new file mode 100644 index 00000000..e65fc40a --- /dev/null +++ b/python/antigravity/requirements.txt @@ -0,0 +1,4 @@ +# Runtime dependencies for the antigravity HarnessService server +# (python/antigravity/harness_server.py). +google-antigravity==0.1.7 +grpcio-health-checking==1.81.0 diff --git a/python/example_agent.py b/python/example_agent.py deleted file mode 100644 index fee97ef9..00000000 --- a/python/example_agent.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -#!/usr/bin/env python3 -""" -Example Python agent using the GAR framework. - -This demonstrates a simple agent that uppercases input text. -""" - -from gar import Agent -import proto.gar_pb2 as pb2 - - -def process(session_id, inputs): - """Process incoming content list and yield responses""" - for content in inputs: - yield pb2.Content( - role="assistant", - type="text", - mimetype="text/plain", - data=f"Python processed (session {session_id}): {content.data.upper()}" - ) - - -def health_check(): - """Health check function that always returns healthy""" - return True, "OK", {} - - -if __name__ == "__main__": - agent = Agent( - process_func=process, - health_check_func=health_check - ) - agent.serve(port=50051) diff --git a/python/gar.py b/python/gar.py deleted file mode 100644 index 51273d36..00000000 --- a/python/gar.py +++ /dev/null @@ -1,136 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -GAR Agent Framework for Python - -A simple framework for building Python agents that work with the GAR orchestrator. -""" - -import grpc -from concurrent import futures -from typing import Callable, Iterator, Optional, Dict - - -class Agent: - """ - Agent provides a simple framework for building Python agents. - - Usage: - def process(session_id, inputs): - for content in inputs: - yield Content(role="assistant", type="text", - mimetype="text/plain", data=f"Processed: {content.data}") - - agent = Agent(process_func=process) - agent.serve(port=50051) - """ - - def __init__( - self, - process_func: Callable, - health_check_func: Optional[Callable] = None - ): - """ - Initialize an agent. - - Args: - process_func: Function that takes (session_id: str, inputs: list) and yields Content responses - health_check_func: Optional function that returns (healthy: bool, message: str, metadata: dict) - """ - self.process_func = process_func - self.health_check_func = health_check_func - - def _create_servicer(self, pb2, pb2_grpc): - """Create the gRPC servicer implementation.""" - agent = self - - class AgentServicer(pb2_grpc.AgentServiceServicer): - def Process(self, request_iterator, context): - # Extract session_id from gRPC metadata - metadata = dict(context.invocation_metadata()) - session_id = metadata.get('session-id', '') - - # Collect all content into a list - inputs = list(request_iterator) - - # Process the list of content with session_id - for response in agent.process_func(session_id, inputs): - if response: - yield response - - def HealthCheck(self, request, context): - if agent.health_check_func: - healthy, message, metadata = agent.health_check_func() - else: - healthy, message, metadata = True, "Agent is healthy", {} - - return pb2.HealthCheckResponse( - healthy=healthy, - message=message, - metadata=metadata or {} - ) - - return AgentServicer() - - def serve(self, port: int = 50051, max_workers: int = 10): - """ - Start the gRPC server. - - Args: - port: Port to listen on (default: 50051) - max_workers: Maximum number of worker threads (default: 10) - """ - # Import proto files (assuming they've been generated) - try: - import proto.gar_pb2 as pb2 - import proto.gar_pb2_grpc as pb2_grpc - except ImportError: - raise ImportError( - "Proto files not found. Generate them first:\n" - " python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. proto/gar.proto" - ) - - server = grpc.server(futures.ThreadPoolExecutor(max_workers=max_workers)) - pb2_grpc.add_AgentServiceServicer_to_server( - self._create_servicer(pb2, pb2_grpc), - server - ) - server.add_insecure_port(f'[::]:{port}') - server.start() - print(f"Agent listening on port {port}") - - try: - server.wait_for_termination() - except KeyboardInterrupt: - print("\nShutting down agent...") - server.stop(grace=5) - - -# Convenience function for quick agent creation -def create_agent( - process_func: Callable, - health_check_func: Optional[Callable] = None, - port: int = 50051 -): - """ - Create and start an agent in one call. - - Args: - process_func: Function that takes (session_id: str, inputs: list) and yields Content responses - health_check_func: Optional function that returns (healthy: bool, message: str, metadata: dict) - port: Port to listen on (default: 50051) - """ - agent = Agent(process_func, health_check_func) - agent.serve(port=port) diff --git a/python/proto/__init__.py b/python/proto/__init__.py new file mode 100644 index 00000000..ce346f19 --- /dev/null +++ b/python/proto/__init__.py @@ -0,0 +1,21 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generated proto files.""" + +import sys +import os + +# Add this directory to sys.path to allow generated files to find each other. +sys.path.insert(0, os.path.dirname(__file__)) diff --git a/python/proto/ax_pb2.py b/python/proto/ax_pb2.py new file mode 100644 index 00000000..d1471030 --- /dev/null +++ b/python/proto/ax_pb2.py @@ -0,0 +1,77 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: proto/ax.proto +# Protobuf Python Version: 4.25.1 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2 +from proto import content_pb2 as proto_dot_content__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0eproto/ax.proto\x12\x02\x61x\x1a\x1cgoogle/protobuf/struct.proto\x1a\x13proto/content.proto\"5\n\x07Message\x12\x0c\n\x04role\x18\x01 \x01(\t\x12\x1c\n\x07\x63ontent\x18\x02 \x01(\x0b\x32\x0b.ax.Content\"\xc8\x01\n\x11\x43onversationEvent\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x0b\n\x03seq\x18\x02 \x01(\x05\x12\x0f\n\x07\x65xec_id\x18\x03 \x01(\t\x12\x12\n\nharness_id\x18\x04 \x01(\t\x12/\n\x0eharness_config\x18\x05 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x1d\n\x08messages\x18\x06 \x03(\x0b\x32\x0b.ax.Message\x12\x18\n\x05state\x18\x07 \x01(\x0e\x32\t.ax.State\"E\n\x0cHarnessStart\x12\x16\n\x0eharness_config\x18\x01 \x01(\x0c\x12\x1d\n\x08messages\x18\x02 \x03(\x0b\x32\x0b.ax.Message\"1\n\rHarnessCancel\x12 \n\x06reason\x18\x01 \x01(\x0e\x32\x10.ax.CancelReason\"\x8d\x01\n\x0eHarnessRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x12\n\nharness_id\x18\x02 \x01(\t\x12!\n\x05start\x18\x03 \x01(\x0b\x32\x10.ax.HarnessStartH\x00\x12#\n\x06\x63\x61ncel\x18\x04 \x01(\x0b\x32\x11.ax.HarnessCancelH\x00\x42\x06\n\x04type\"/\n\x0eHarnessOutputs\x12\x1d\n\x08messages\x18\x01 \x03(\x0b\x32\x0b.ax.Message\"*\n\x05\x45rror\x12\x0c\n\x04\x63ode\x18\x01 \x01(\x05\x12\x13\n\x0b\x64\x65scription\x18\x02 \x01(\t\"@\n\nHarnessEnd\x12\x18\n\x05state\x18\x01 \x01(\x0e\x32\t.ax.State\x12\x18\n\x05\x65rror\x18\x02 \x01(\x0b\x32\t.ax.Error\"x\n\x0fHarnessResponse\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12%\n\x07outputs\x18\x02 \x01(\x0b\x32\x12.ax.HarnessOutputsH\x00\x12\x1d\n\x03\x65nd\x18\x03 \x01(\x0b\x32\x0e.ax.HarnessEndH\x00\x42\x06\n\x04type\"\x81\x01\n\x0b\x45xecRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x1b\n\x06inputs\x18\x02 \x03(\x0b\x32\x0b.ax.Message\x12\x10\n\x08last_seq\x18\x03 \x01(\x05\x12\x12\n\nharness_id\x18\x04 \x01(\t\x12\x16\n\x0eharness_config\x18\x05 \x01(\x0c\"9\n\x0c\x45xecResponse\x12\x1c\n\x07outputs\x18\x01 \x03(\x0b\x32\x0b.ax.Message\x12\x0b\n\x03seq\x18\x02 \x01(\x05\"4\n\x19\x44\x65leteConversationRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\"\x1c\n\x1a\x44\x65leteConversationResponse*l\n\x05State\x12\x15\n\x11STATE_UNSPECIFIED\x10\x00\x12\x11\n\rSTATE_PENDING\x10\x01\x12\x10\n\x0cSTATE_FAILED\x10\x02\x12\x13\n\x0fSTATE_COMPLETED\x10\x03\x12\x12\n\x0eSTATE_CANCELED\x10\x04*\x8c\x01\n\x0c\x43\x61ncelReason\x12\x1d\n\x19\x43\x41NCEL_REASON_UNSPECIFIED\x10\x00\x12 \n\x1c\x43\x41NCEL_REASON_USER_REQUESTED\x10\x01\x12\x19\n\x15\x43\x41NCEL_REASON_TIMEOUT\x10\x02\x12 \n\x1c\x43\x41NCEL_REASON_INTERNAL_ERROR\x10\x03\x32H\n\x0eHarnessService\x12\x36\n\x07\x43onnect\x12\x12.ax.HarnessRequest\x1a\x13.ax.HarnessResponse(\x01\x30\x01\x32?\n\x10\x45xecutionService\x12+\n\x04\x45xec\x12\x0f.ax.ExecRequest\x1a\x10.ax.ExecResponse0\x01\x32j\n\x13\x43onversationService\x12S\n\x12\x44\x65leteConversation\x12\x1d.ax.DeleteConversationRequest\x1a\x1e.ax.DeleteConversationResponseB\x1cZ\x1agithub.com/google/ax/protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'proto.ax_pb2', _globals) +if _descriptor._USE_C_DESCRIPTORS == False: + _globals['DESCRIPTOR']._options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\032github.com/google/ax/proto' + _globals['_STATE']._serialized_start=1153 + _globals['_STATE']._serialized_end=1261 + _globals['_CANCELREASON']._serialized_start=1264 + _globals['_CANCELREASON']._serialized_end=1404 + _globals['_MESSAGE']._serialized_start=73 + _globals['_MESSAGE']._serialized_end=126 + _globals['_CONVERSATIONEVENT']._serialized_start=129 + _globals['_CONVERSATIONEVENT']._serialized_end=329 + _globals['_HARNESSSTART']._serialized_start=331 + _globals['_HARNESSSTART']._serialized_end=400 + _globals['_HARNESSCANCEL']._serialized_start=402 + _globals['_HARNESSCANCEL']._serialized_end=451 + _globals['_HARNESSREQUEST']._serialized_start=454 + _globals['_HARNESSREQUEST']._serialized_end=595 + _globals['_HARNESSOUTPUTS']._serialized_start=597 + _globals['_HARNESSOUTPUTS']._serialized_end=644 + _globals['_ERROR']._serialized_start=646 + _globals['_ERROR']._serialized_end=688 + _globals['_HARNESSEND']._serialized_start=690 + _globals['_HARNESSEND']._serialized_end=754 + _globals['_HARNESSRESPONSE']._serialized_start=756 + _globals['_HARNESSRESPONSE']._serialized_end=876 + _globals['_EXECREQUEST']._serialized_start=879 + _globals['_EXECREQUEST']._serialized_end=1008 + _globals['_EXECRESPONSE']._serialized_start=1010 + _globals['_EXECRESPONSE']._serialized_end=1067 + _globals['_DELETECONVERSATIONREQUEST']._serialized_start=1069 + _globals['_DELETECONVERSATIONREQUEST']._serialized_end=1121 + _globals['_DELETECONVERSATIONRESPONSE']._serialized_start=1123 + _globals['_DELETECONVERSATIONRESPONSE']._serialized_end=1151 + _globals['_HARNESSSERVICE']._serialized_start=1406 + _globals['_HARNESSSERVICE']._serialized_end=1478 + _globals['_EXECUTIONSERVICE']._serialized_start=1480 + _globals['_EXECUTIONSERVICE']._serialized_end=1543 + _globals['_CONVERSATIONSERVICE']._serialized_start=1545 + _globals['_CONVERSATIONSERVICE']._serialized_end=1651 +# @@protoc_insertion_point(module_scope) diff --git a/python/proto/ax_pb2_grpc.py b/python/proto/ax_pb2_grpc.py new file mode 100644 index 00000000..b35d9534 --- /dev/null +++ b/python/proto/ax_pb2_grpc.py @@ -0,0 +1,210 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + +from proto import ax_pb2 as proto_dot_ax__pb2 + + +class HarnessServiceStub(object): + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.Connect = channel.stream_stream( + '/ax.HarnessService/Connect', + request_serializer=proto_dot_ax__pb2.HarnessRequest.SerializeToString, + response_deserializer=proto_dot_ax__pb2.HarnessResponse.FromString, + ) + + +class HarnessServiceServicer(object): + """Missing associated documentation comment in .proto file.""" + + def Connect(self, request_iterator, context): + """Connect drives one harness execution. The client sends + HarnessRequest{start} (and may send one HarnessRequest{cancel} + mid-stream), and the server streams zero or more HarnessResponse{outputs} + frames terminated by exactly one HarnessResponse{end}. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_HarnessServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'Connect': grpc.stream_stream_rpc_method_handler( + servicer.Connect, + request_deserializer=proto_dot_ax__pb2.HarnessRequest.FromString, + response_serializer=proto_dot_ax__pb2.HarnessResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'ax.HarnessService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + + + # This class is part of an EXPERIMENTAL API. +class HarnessService(object): + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def Connect(request_iterator, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.stream_stream(request_iterator, target, '/ax.HarnessService/Connect', + proto_dot_ax__pb2.HarnessRequest.SerializeToString, + proto_dot_ax__pb2.HarnessResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + +class ExecutionServiceStub(object): + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.Exec = channel.unary_stream( + '/ax.ExecutionService/Exec', + request_serializer=proto_dot_ax__pb2.ExecRequest.SerializeToString, + response_deserializer=proto_dot_ax__pb2.ExecResponse.FromString, + ) + + +class ExecutionServiceServicer(object): + """Missing associated documentation comment in .proto file.""" + + def Exec(self, request, context): + """Exec executes an agentic task or resumes an existing one with streaming responses + If the conversation_id already exists, it will be resumed. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_ExecutionServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'Exec': grpc.unary_stream_rpc_method_handler( + servicer.Exec, + request_deserializer=proto_dot_ax__pb2.ExecRequest.FromString, + response_serializer=proto_dot_ax__pb2.ExecResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'ax.ExecutionService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + + + # This class is part of an EXPERIMENTAL API. +class ExecutionService(object): + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def Exec(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream(request, target, '/ax.ExecutionService/Exec', + proto_dot_ax__pb2.ExecRequest.SerializeToString, + proto_dot_ax__pb2.ExecResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + +class ConversationServiceStub(object): + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.DeleteConversation = channel.unary_unary( + '/ax.ConversationService/DeleteConversation', + request_serializer=proto_dot_ax__pb2.DeleteConversationRequest.SerializeToString, + response_deserializer=proto_dot_ax__pb2.DeleteConversationResponse.FromString, + ) + + +class ConversationServiceServicer(object): + """Missing associated documentation comment in .proto file.""" + + def DeleteConversation(self, request, context): + """Deletes conversational events and all event log resources + for its children executions. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_ConversationServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'DeleteConversation': grpc.unary_unary_rpc_method_handler( + servicer.DeleteConversation, + request_deserializer=proto_dot_ax__pb2.DeleteConversationRequest.FromString, + response_serializer=proto_dot_ax__pb2.DeleteConversationResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'ax.ConversationService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + + + # This class is part of an EXPERIMENTAL API. +class ConversationService(object): + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def DeleteConversation(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/ax.ConversationService/DeleteConversation', + proto_dot_ax__pb2.DeleteConversationRequest.SerializeToString, + proto_dot_ax__pb2.DeleteConversationResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) diff --git a/python/proto/content_pb2.py b/python/proto/content_pb2.py new file mode 100644 index 00000000..c20ef7bf --- /dev/null +++ b/python/proto/content_pb2.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: proto/content.proto +# Protobuf Python Version: 4.25.1 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x13proto/content.proto\x12\x02\x61x\x1a\x1cgoogle/protobuf/struct.proto\"-\n\x0bTextContent\x12\x0c\n\x04text\x18\x03 \x01(\tJ\x04\x08\x01\x10\x02J\x04\x08\x02\x10\x03R\x04type\"$\n\x10\x41pprovalDecision\x12\x10\n\x08\x61pproved\x18\x01 \x01(\x08\"1\n\x0f\x44\x65\x63lineDecision\x12\x10\n\x08\x64\x65\x63lined\x18\x01 \x01(\x08J\x04\x08\x02\x10\x03R\x06reason\"\xa3\x01\n\x13\x43onfirmationContent\x12\n\n\x02id\x18\x03 \x01(\t\x12\x10\n\x08question\x18\x04 \x01(\t\x12(\n\x08\x61pproval\x18\x05 \x01(\x0b\x32\x14.ax.ApprovalDecisionH\x00\x12&\n\x07\x64\x65\x63line\x18\x06 \x01(\x0b\x32\x13.ax.DeclineDecisionH\x00\x42\n\n\x08\x64\x65\x63isionJ\x04\x08\x01\x10\x02J\x04\x08\x02\x10\x03R\x04type\"F\n\x15ThoughtSummaryContent\x12\x1f\n\x04text\x18\x01 \x01(\x0b\x32\x0f.ax.TextContentH\x00\x42\x06\n\x04typeJ\x04\x08\x02\x10\x03\"y\n\x0eThoughtContent\x12\x11\n\tsignature\x18\x07 \x01(\x0c\x12*\n\x07summary\x18\t \x03(\x0b\x32\x19.ax.ThoughtSummaryContentJ\x04\x08\x01\x10\x02J\x04\x08\x02\x10\x03J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05J\x04\x08\x05\x10\x06J\x04\x08\x08\x10\tR\x04type\"j\n\x0fToolCallContent\x12\n\n\x02id\x18\x01 \x01(\t\x12\x11\n\tsignature\x18\t \x01(\x0c\x12\x30\n\rfunction_call\x18\x02 \x01(\x0b\x32\x17.ax.FunctionCallContentH\x00\x42\x06\n\x04type\"{\n\x11ToolResultContent\x12\x0f\n\x07\x63\x61ll_id\x18\x08 \x01(\t\x12\x11\n\tsignature\x18\t \x01(\x0c\x12\x34\n\x0f\x66unction_result\x18\x02 \x01(\x0b\x32\x19.ax.FunctionResultContentH\x00\x42\x06\n\x04typeJ\x04\x08\x01\x10\x02\"[\n\x13\x46unctionCallContent\x12\x0c\n\x04name\x18\x03 \x01(\t\x12*\n\targuments\x18\x04 \x01(\x0b\x32\x17.google.protobuf.StructJ\x04\x08\x01\x10\x02J\x04\x08\x02\x10\x03\"n\n\x15\x46unctionResultContent\x12\x0c\n\x04name\x18\x08 \x01(\t\x12+\n\x08response\x18\x03 \x01(\x0b\x32\x17.google.protobuf.StructH\x00\x42\x08\n\x06resultJ\x04\x08\x01\x10\x02J\x04\x08\x04\x10\x05R\x04type\"\xbd\x02\n\x0cImageContent\x12,\n\tmime_type\x18\x01 \x01(\x0e\x32\x19.ax.ImageContent.MimeType\x12\x0e\n\x04\x64\x61ta\x18\x02 \x01(\x0cH\x00\x12\r\n\x03uri\x18\x06 \x01(\tH\x00\x12\'\n\nresolution\x18\x05 \x01(\x0e\x32\x13.ax.MediaResolution\"\x9b\x01\n\x08MimeType\x12\x14\n\x10TYPE_UNSPECIFIED\x10\x00\x12\x0c\n\x08TYPE_PNG\x10\x01\x12\r\n\tTYPE_JPEG\x10\x02\x12\r\n\tTYPE_WEBP\x10\x03\x12\r\n\tTYPE_HEIC\x10\x04\x12\r\n\tTYPE_HEIF\x10\x05\x12\x0c\n\x08TYPE_GIF\x10\x07\x12\x0c\n\x08TYPE_BMP\x10\x08\x12\r\n\tTYPE_TIFF\x10\t\"\x04\x08\x06\x10\x06\x42\r\n\x0b\x64\x61ta_or_uriJ\x04\x08\x03\x10\x04R\x04type\"\x8b\x03\n\x0c\x41udioContent\x12,\n\tmime_type\x18\x01 \x01(\x0e\x32\x19.ax.AudioContent.MimeType\x12\x0e\n\x04\x64\x61ta\x18\x02 \x01(\x0cH\x00\x12\r\n\x03uri\x18\x05 \x01(\tH\x00\x12\x10\n\x08\x63hannels\x18\x07 \x01(\x05\x12\x13\n\x0bsample_rate\x18\x08 \x01(\x05\"\xdf\x01\n\x08MimeType\x12\x14\n\x10TYPE_UNSPECIFIED\x10\x00\x12\x0c\n\x08TYPE_WAV\x10\x01\x12\x0c\n\x08TYPE_MP3\x10\x02\x12\r\n\tTYPE_AIFF\x10\x03\x12\x0c\n\x08TYPE_AAC\x10\x04\x12\x0c\n\x08TYPE_OGG\x10\x05\x12\r\n\tTYPE_FLAC\x10\x06\x12\r\n\tTYPE_MPEG\x10\x07\x12\x0c\n\x08TYPE_M4A\x10\x08\x12\x0c\n\x08TYPE_L16\x10\t\x12\x0e\n\nTYPE_S16LE\x10\n\x12\r\n\tTYPE_OPUS\x10\x0b\x12\r\n\tTYPE_ALAW\x10\x0c\x12\x0e\n\nTYPE_MULAW\x10\rB\r\n\x0b\x64\x61ta_or_uriJ\x04\x08\x03\x10\x04J\x04\x08\x06\x10\x07R\x04typeR\x04rate\"\xcc\x01\n\x0f\x44ocumentContent\x12/\n\tmime_type\x18\x01 \x01(\x0e\x32\x1c.ax.DocumentContent.MimeType\x12\x0e\n\x04\x64\x61ta\x18\x02 \x01(\x0cH\x00\x12\r\n\x03uri\x18\x05 \x01(\tH\x00\"N\n\x08MimeType\x12\x14\n\x10TYPE_UNSPECIFIED\x10\x00\x12\x0c\n\x08TYPE_PDF\x10\x01\x12\r\n\tTYPE_JSON\x10\x02\x12\x0f\n\x0bTYPE_PYTHON\x10\x03\x42\r\n\x0b\x64\x61ta_or_uriJ\x04\x08\x03\x10\x04R\x04type\"\xc5\x02\n\x0cVideoContent\x12,\n\tmime_type\x18\x01 \x01(\x0e\x32\x19.ax.VideoContent.MimeType\x12\x0e\n\x04\x64\x61ta\x18\x02 \x01(\x0cH\x00\x12\r\n\x03uri\x18\x06 \x01(\tH\x00\x12\'\n\nresolution\x18\x05 \x01(\x0e\x32\x13.ax.MediaResolution\"\xa3\x01\n\x08MimeType\x12\x14\n\x10TYPE_UNSPECIFIED\x10\x00\x12\x0c\n\x08TYPE_MP4\x10\x01\x12\r\n\tTYPE_MPEG\x10\x02\x12\x0c\n\x08TYPE_MPG\x10\x03\x12\x0c\n\x08TYPE_MOV\x10\x04\x12\x0c\n\x08TYPE_AVI\x10\x05\x12\x0e\n\nTYPE_X_FLV\x10\x06\x12\r\n\tTYPE_WEBM\x10\x07\x12\x0c\n\x08TYPE_WMV\x10\x08\x12\r\n\tTYPE_3GPP\x10\tB\r\n\x0b\x64\x61ta_or_uriJ\x04\x08\x03\x10\x04R\x04type\"\x86\x03\n\x07\x43ontent\x12%\n\x07thought\x18\x05 \x01(\x0b\x32\x12.ax.ThoughtContentH\x00\x12\x1f\n\x04text\x18\n \x01(\x0b\x32\x0f.ax.TextContentH\x00\x12!\n\x05image\x18\x0b \x01(\x0b\x32\x10.ax.ImageContentH\x00\x12!\n\x05\x61udio\x18\x0c \x01(\x0b\x32\x10.ax.AudioContentH\x00\x12\'\n\x08\x64ocument\x18\r \x01(\x0b\x32\x13.ax.DocumentContentH\x00\x12!\n\x05video\x18\x0e \x01(\x0b\x32\x10.ax.VideoContentH\x00\x12/\n\x0c\x63onfirmation\x18\x1a \x01(\x0b\x32\x17.ax.ConfirmationContentH\x00\x12(\n\ttool_call\x18\x18 \x01(\x0b\x32\x13.ax.ToolCallContentH\x00\x12,\n\x0btool_result\x18\x19 \x01(\x0b\x32\x15.ax.ToolResultContentH\x00\x42\x06\n\x04typeJ\x04\x08\x01\x10\x05J\x04\x08\x06\x10\nJ\x04\x08\x0f\x10\x18*b\n\x0fMediaResolution\x12 \n\x1cMEDIA_RESOLUTION_UNSPECIFIED\x10\x00\x12\x07\n\x03LOW\x10\x01\x12\n\n\x06MEDIUM\x10\x02\x12\x08\n\x04HIGH\x10\x03\x12\x0e\n\nULTRA_HIGH\x10\x04\x42\x1cZ\x1agithub.com/google/ax/protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'proto.content_pb2', _globals) +if _descriptor._USE_C_DESCRIPTORS == False: + _globals['DESCRIPTOR']._options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\032github.com/google/ax/proto' + _globals['_MEDIARESOLUTION']._serialized_start=2638 + _globals['_MEDIARESOLUTION']._serialized_end=2736 + _globals['_TEXTCONTENT']._serialized_start=57 + _globals['_TEXTCONTENT']._serialized_end=102 + _globals['_APPROVALDECISION']._serialized_start=104 + _globals['_APPROVALDECISION']._serialized_end=140 + _globals['_DECLINEDECISION']._serialized_start=142 + _globals['_DECLINEDECISION']._serialized_end=191 + _globals['_CONFIRMATIONCONTENT']._serialized_start=194 + _globals['_CONFIRMATIONCONTENT']._serialized_end=357 + _globals['_THOUGHTSUMMARYCONTENT']._serialized_start=359 + _globals['_THOUGHTSUMMARYCONTENT']._serialized_end=429 + _globals['_THOUGHTCONTENT']._serialized_start=431 + _globals['_THOUGHTCONTENT']._serialized_end=552 + _globals['_TOOLCALLCONTENT']._serialized_start=554 + _globals['_TOOLCALLCONTENT']._serialized_end=660 + _globals['_TOOLRESULTCONTENT']._serialized_start=662 + _globals['_TOOLRESULTCONTENT']._serialized_end=785 + _globals['_FUNCTIONCALLCONTENT']._serialized_start=787 + _globals['_FUNCTIONCALLCONTENT']._serialized_end=878 + _globals['_FUNCTIONRESULTCONTENT']._serialized_start=880 + _globals['_FUNCTIONRESULTCONTENT']._serialized_end=990 + _globals['_IMAGECONTENT']._serialized_start=993 + _globals['_IMAGECONTENT']._serialized_end=1310 + _globals['_IMAGECONTENT_MIMETYPE']._serialized_start=1128 + _globals['_IMAGECONTENT_MIMETYPE']._serialized_end=1283 + _globals['_AUDIOCONTENT']._serialized_start=1313 + _globals['_AUDIOCONTENT']._serialized_end=1708 + _globals['_AUDIOCONTENT_MIMETYPE']._serialized_start=1446 + _globals['_AUDIOCONTENT_MIMETYPE']._serialized_end=1669 + _globals['_DOCUMENTCONTENT']._serialized_start=1711 + _globals['_DOCUMENTCONTENT']._serialized_end=1915 + _globals['_DOCUMENTCONTENT_MIMETYPE']._serialized_start=1810 + _globals['_DOCUMENTCONTENT_MIMETYPE']._serialized_end=1888 + _globals['_VIDEOCONTENT']._serialized_start=1918 + _globals['_VIDEOCONTENT']._serialized_end=2243 + _globals['_VIDEOCONTENT_MIMETYPE']._serialized_start=2053 + _globals['_VIDEOCONTENT_MIMETYPE']._serialized_end=2216 + _globals['_CONTENT']._serialized_start=2246 + _globals['_CONTENT']._serialized_end=2636 +# @@protoc_insertion_point(module_scope) diff --git a/python/proto/content_pb2_grpc.py b/python/proto/content_pb2_grpc.py new file mode 100644 index 00000000..fe9b2e77 --- /dev/null +++ b/python/proto/content_pb2_grpc.py @@ -0,0 +1,18 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + diff --git a/python/python.go b/python/python.go new file mode 100644 index 00000000..faa502e3 --- /dev/null +++ b/python/python.go @@ -0,0 +1,24 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package python embeds the Python directory assets (such as antigravity and proto modules) +// for use by the sidecar and harnesses. +package python + +import "embed" + +// FS contains the embedded python/ directory assets. +// +//go:embed antigravity proto antigravity/_*.py proto/_*.py +var FS embed.FS