From 2977b3416bf9988aa254194cb927a8154ae4f935 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Mon, 12 Oct 2026 11:31:36 +0000 Subject: [PATCH 1/4] use agent enum for agent-shim events --- experimental/ssh/internal/client/agentshim.go | 13 +++++++++++- .../ssh/internal/client/agentshim_test.go | 13 ++++++++++-- libs/telemetry/protos/ssh_agent_shim.go | 21 ++++++++++++++++--- libs/telemetry/protos/ssh_agent_shim_test.go | 2 +- 4 files changed, 42 insertions(+), 7 deletions(-) diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index 1c0fd2dbe5c..a6bb843f39c 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -238,13 +238,24 @@ func logAgentShimEvent(ctx context.Context, agentName string, outcome agentShimO // separated from logAgentShimEvent so the field mapping can be unit tested. func buildAgentShimEvent(agentName string, outcome agentShimOutcome, setupDuration time.Duration) *protos.SshAgentShimEvent { return &protos.SshAgentShimEvent{ - AgentName: agentName, + Agent: agentNameToAgentEnum(agentName), IsSuccess: outcome.err == nil, ErrorCategory: outcome.category(), SetupDurationMs: setupDuration.Milliseconds(), } } +func agentNameToAgentEnum(agentName string) protos.SshAgentShimAgent { + switch agentName { + case "claude": + return protos.SshAgentShimAgentClaudeCode + case "codex": + return protos.SshAgentShimAgentCodex + default: + return protos.SshAgentShimAgentUnspecified + } +} + // uploadAgentShimTelemetry flushes buffered telemetry before the shim execs into the agent. // A successful launch replaces this process (see execProcess), so the end-of-command upload // in cmd/root.Execute never runs; this rebuilds the ExecutionContext that upload would have diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index 7583180a598..d8d99c243d2 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -565,10 +565,19 @@ func TestProbeModelServicesPageError(t *testing.T) { } func TestBuildAgentShimEvent(t *testing.T) { + for _, agent := range supportedAgents { + // This test will fail when you add a new supported agent without adding a new + // `SshAgentShimAgent` value in `libs/telemetry/protos/ssh_agent_shim.go` + t.Run(fmt.Sprintf("maps %s to an enum for telemetry", agent.name), func(t *testing.T) { + got := buildAgentShimEvent(agent.name, agentShimOutcome{}, 0) + assert.NotEqual(t, got.Agent, protos.SshAgentShimAgentUnspecified) + }) + } + t.Run("successful launch reports success and no category", func(t *testing.T) { got := buildAgentShimEvent("claude", agentShimOutcome{}, 1500*time.Millisecond) assert.Equal(t, &protos.SshAgentShimEvent{ - AgentName: "claude", + Agent: protos.SshAgentShimAgentClaudeCode, IsSuccess: true, ErrorCategory: protos.SshAgentShimErrorCategoryUnspecified, SetupDurationMs: 1500, @@ -580,7 +589,7 @@ func TestBuildAgentShimEvent(t *testing.T) { err: categorize(protos.SshAgentShimErrorCategoryGatewayAuthFailed, errors.New("boom")), }, 200*time.Millisecond) assert.Equal(t, &protos.SshAgentShimEvent{ - AgentName: "codex", + Agent: protos.SshAgentShimAgentCodex, IsSuccess: false, ErrorCategory: protos.SshAgentShimErrorCategoryGatewayAuthFailed, SetupDurationMs: 200, diff --git a/libs/telemetry/protos/ssh_agent_shim.go b/libs/telemetry/protos/ssh_agent_shim.go index ec52fab0a96..126ab9b128d 100644 --- a/libs/telemetry/protos/ssh_agent_shim.go +++ b/libs/telemetry/protos/ssh_agent_shim.go @@ -81,6 +81,21 @@ const ( SshAgentShimErrorCategoryUnknown SshAgentShimErrorCategory = "UNKNOWN" ) +// The agent for an `ssh agent-shim ` command invocation. Values are in +// screaming snake case to match what JSON->proto serialization expects. +type SshAgentShimAgent string + +const ( + // Unknown agent was specified. + SshAgentShimAgentUnspecified SshAgentShimAgent = "TYPE_UNSPECIFIED" + + // Claude Code CLI. + SshAgentShimAgentClaudeCode SshAgentShimAgent = "CLAUDE_CODE" + + // Codex CLI. + SshAgentShimAgentCodex SshAgentShimAgent = "CODEX" +) + // SshAgentShimEvent is emitted when a user runs `ssh agent-shim ` on the remote // driver to launch a coding agent against the Unity AI Gateway. // @@ -89,9 +104,9 @@ const ( // uploaded from the shim just before the exec. A failed launch returns normally and is // uploaded by cmd/root as usual. type SshAgentShimEvent struct { - // Name of the agent the user asked to launch (e.g. "claude", "codex"). Drawn from a - // fixed set of CLI-defined identifiers, never user-authored text. - AgentName string `json:"agent_name"` + // The agent the user asked to launch, mapped to a string matching the proto enum + // (e.g. "CLAUDE_CODE", "CODEX", etc). + Agent SshAgentShimAgent `json:"agent"` // Whether the agent was launched. False means setup failed before the agent ran. // Populated on every event, so no omitempty: a genuine false must stay distinguishable diff --git a/libs/telemetry/protos/ssh_agent_shim_test.go b/libs/telemetry/protos/ssh_agent_shim_test.go index 0f8d1a31c53..15f5fb1025a 100644 --- a/libs/telemetry/protos/ssh_agent_shim_test.go +++ b/libs/telemetry/protos/ssh_agent_shim_test.go @@ -21,7 +21,7 @@ func TestSshAgentShimEventEncodesFailureExplicitly(t *testing.T) { assert.Equal(t, false, got["is_success"], "is_success must be sent as false, not omitted") // The remaining fields must appear in the payload at all: omitempty would drop the // zero value and make a failure indistinguishable from an unreported one. - for _, field := range []string{"agent_name", "error_category", "setup_duration_ms"} { + for _, field := range []string{"agent", "error_category", "setup_duration_ms"} { assert.Contains(t, got, field, "%s must be sent, not omitted", field) } } From 584355c39e0bf268f5b6d80d5ab81618f48873f2 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Mon, 12 Oct 2026 12:22:08 +0000 Subject: [PATCH 2/4] fix lint --- experimental/ssh/internal/client/agentshim_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index d8d99c243d2..a4ebe011878 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -570,7 +570,7 @@ func TestBuildAgentShimEvent(t *testing.T) { // `SshAgentShimAgent` value in `libs/telemetry/protos/ssh_agent_shim.go` t.Run(fmt.Sprintf("maps %s to an enum for telemetry", agent.name), func(t *testing.T) { got := buildAgentShimEvent(agent.name, agentShimOutcome{}, 0) - assert.NotEqual(t, got.Agent, protos.SshAgentShimAgentUnspecified) + assert.NotEqual(t, protos.SshAgentShimAgentUnspecified, got.Agent) }) } From dd62dac51b2d38690fd97cd963b72be9d23fbc03 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Tue, 13 Oct 2026 07:55:28 +0000 Subject: [PATCH 3/4] move telemetry enum into agentSpec --- experimental/ssh/internal/client/agentshim.go | 27 +++++++++---------- .../ssh/internal/client/agentshim_test.go | 9 ++++--- libs/telemetry/protos/ssh_agent_shim.go | 2 +- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index a6bb843f39c..d596a4e49d9 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -85,7 +85,8 @@ const ( ) type agentSpec struct { - name string + name string + telemetryName protos.SshAgentShimAgent // contextFlag passes the context file via this CLI flag (Claude's --append-system-prompt-file). contextFlag string // contextHomeFile writes the context to this $HOME instructions file (codex's AGENTS.md). @@ -94,8 +95,8 @@ type agentSpec struct { // Currently limited to the agents whose Unity Gateway CLI command supports --workspace var supportedAgents = []agentSpec{ - {name: "claude", contextFlag: "--append-system-prompt-file"}, - {name: "codex", contextHomeFile: ".codex/AGENTS.md"}, + {name: "claude", telemetryName: protos.SshAgentShimAgentClaudeCode, contextFlag: "--append-system-prompt-file"}, + {name: "codex", telemetryName: protos.SshAgentShimAgentCodex, contextHomeFile: ".codex/AGENTS.md"}, } // SupportedAgentNames lists the agents the ssh command registers a subcommand for. @@ -237,25 +238,21 @@ func logAgentShimEvent(ctx context.Context, agentName string, outcome agentShimO // buildAgentShimEvent maps the agent name and outcome onto the telemetry event. It is // separated from logAgentShimEvent so the field mapping can be unit tested. func buildAgentShimEvent(agentName string, outcome agentShimOutcome, setupDuration time.Duration) *protos.SshAgentShimEvent { + // lookup agent by name so this function is the single decision point for + // telemetry and is easily testable + telemetryName := protos.SshAgentShimAgentUnspecified + agent, ok := agentByName(agentName) + if ok { + telemetryName = agent.telemetryName + } return &protos.SshAgentShimEvent{ - Agent: agentNameToAgentEnum(agentName), + Agent: telemetryName, IsSuccess: outcome.err == nil, ErrorCategory: outcome.category(), SetupDurationMs: setupDuration.Milliseconds(), } } -func agentNameToAgentEnum(agentName string) protos.SshAgentShimAgent { - switch agentName { - case "claude": - return protos.SshAgentShimAgentClaudeCode - case "codex": - return protos.SshAgentShimAgentCodex - default: - return protos.SshAgentShimAgentUnspecified - } -} - // uploadAgentShimTelemetry flushes buffered telemetry before the shim execs into the agent. // A successful launch replaces this process (see execProcess), so the end-of-command upload // in cmd/root.Execute never runs; this rebuilds the ExecutionContext that upload would have diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index a4ebe011878..ada5d3ea3f2 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -566,14 +566,17 @@ func TestProbeModelServicesPageError(t *testing.T) { func TestBuildAgentShimEvent(t *testing.T) { for _, agent := range supportedAgents { - // This test will fail when you add a new supported agent without adding a new - // `SshAgentShimAgent` value in `libs/telemetry/protos/ssh_agent_shim.go` - t.Run(fmt.Sprintf("maps %s to an enum for telemetry", agent.name), func(t *testing.T) { + t.Run(fmt.Sprintf("%s has an appropriate enum for telemetry", agent.name), func(t *testing.T) { got := buildAgentShimEvent(agent.name, agentShimOutcome{}, 0) assert.NotEqual(t, protos.SshAgentShimAgentUnspecified, got.Agent) }) } + t.Run("maps an unknown agent to unspecified", func(t *testing.T) { + got := buildAgentShimEvent("wibble", agentShimOutcome{}, 0) + assert.Equal(t, protos.SshAgentShimAgentUnspecified, got.Agent) + }) + t.Run("successful launch reports success and no category", func(t *testing.T) { got := buildAgentShimEvent("claude", agentShimOutcome{}, 1500*time.Millisecond) assert.Equal(t, &protos.SshAgentShimEvent{ diff --git a/libs/telemetry/protos/ssh_agent_shim.go b/libs/telemetry/protos/ssh_agent_shim.go index 126ab9b128d..0fa50b07925 100644 --- a/libs/telemetry/protos/ssh_agent_shim.go +++ b/libs/telemetry/protos/ssh_agent_shim.go @@ -82,7 +82,7 @@ const ( ) // The agent for an `ssh agent-shim ` command invocation. Values are in -// screaming snake case to match what JSON->proto serialization expects. +// screaming snake case to match what JSON->proto deserialization expects. type SshAgentShimAgent string const ( From 600e290aacf9443f11f40db85fdfb2d3f10ba4ad Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Tue, 13 Oct 2026 14:54:43 +0000 Subject: [PATCH 4/4] lint --- experimental/ssh/internal/client/agentshim_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index ada5d3ea3f2..283d4288fbe 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -566,7 +566,7 @@ func TestProbeModelServicesPageError(t *testing.T) { func TestBuildAgentShimEvent(t *testing.T) { for _, agent := range supportedAgents { - t.Run(fmt.Sprintf("%s has an appropriate enum for telemetry", agent.name), func(t *testing.T) { + t.Run(agent.name+" has an appropriate enum for telemetry", func(t *testing.T) { got := buildAgentShimEvent(agent.name, agentShimOutcome{}, 0) assert.NotEqual(t, protos.SshAgentShimAgentUnspecified, got.Agent) })