Skip to content

Commit 8a24de2

Browse files
authored
Send agent enum for agent-shim telemetry events (#7019)
## Changes Replace `agent_name` with `agent` and make the values match JSON->proto deserialization for enums (e.g. "CLAUDE_CODE" instead of "claude") ## Why The original PR to add this event didn't follow telemetry best practices, this PR fixes that ## Tests automated tests pass
1 parent 7d3ea75 commit 8a24de2

4 files changed

Lines changed: 45 additions & 10 deletions

File tree

‎experimental/ssh/internal/client/agentshim.go‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,8 @@ const (
8585
)
8686

8787
type agentSpec struct {
88-
name string
88+
name string
89+
telemetryName protos.SshAgentShimAgent
8990
// contextFlag passes the context file via this CLI flag (Claude's --append-system-prompt-file).
9091
contextFlag string
9192
// contextHomeFile writes the context to this $HOME instructions file (codex's AGENTS.md).
@@ -94,8 +95,8 @@ type agentSpec struct {
9495

9596
// Currently limited to the agents whose Unity Gateway CLI command supports --workspace
9697
var supportedAgents = []agentSpec{
97-
{name: "claude", contextFlag: "--append-system-prompt-file"},
98-
{name: "codex", contextHomeFile: ".codex/AGENTS.md"},
98+
{name: "claude", telemetryName: protos.SshAgentShimAgentClaudeCode, contextFlag: "--append-system-prompt-file"},
99+
{name: "codex", telemetryName: protos.SshAgentShimAgentCodex, contextHomeFile: ".codex/AGENTS.md"},
99100
}
100101

101102
// SupportedAgentNames lists the agents the ssh command registers a subcommand for.
@@ -237,8 +238,15 @@ func logAgentShimEvent(ctx context.Context, agentName string, outcome agentShimO
237238
// buildAgentShimEvent maps the agent name and outcome onto the telemetry event. It is
238239
// separated from logAgentShimEvent so the field mapping can be unit tested.
239240
func buildAgentShimEvent(agentName string, outcome agentShimOutcome, setupDuration time.Duration) *protos.SshAgentShimEvent {
241+
// lookup agent by name so this function is the single decision point for
242+
// telemetry and is easily testable
243+
telemetryName := protos.SshAgentShimAgentUnspecified
244+
agent, ok := agentByName(agentName)
245+
if ok {
246+
telemetryName = agent.telemetryName
247+
}
240248
return &protos.SshAgentShimEvent{
241-
AgentName: agentName,
249+
Agent: telemetryName,
242250
IsSuccess: outcome.err == nil,
243251
ErrorCategory: outcome.category(),
244252
SetupDurationMs: setupDuration.Milliseconds(),

‎experimental/ssh/internal/client/agentshim_test.go‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -565,10 +565,22 @@ func TestProbeModelServicesPageError(t *testing.T) {
565565
}
566566

567567
func TestBuildAgentShimEvent(t *testing.T) {
568+
for _, agent := range supportedAgents {
569+
t.Run(agent.name+" has an appropriate enum for telemetry", func(t *testing.T) {
570+
got := buildAgentShimEvent(agent.name, agentShimOutcome{}, 0)
571+
assert.NotEqual(t, protos.SshAgentShimAgentUnspecified, got.Agent)
572+
})
573+
}
574+
575+
t.Run("maps an unknown agent to unspecified", func(t *testing.T) {
576+
got := buildAgentShimEvent("wibble", agentShimOutcome{}, 0)
577+
assert.Equal(t, protos.SshAgentShimAgentUnspecified, got.Agent)
578+
})
579+
568580
t.Run("successful launch reports success and no category", func(t *testing.T) {
569581
got := buildAgentShimEvent("claude", agentShimOutcome{}, 1500*time.Millisecond)
570582
assert.Equal(t, &protos.SshAgentShimEvent{
571-
AgentName: "claude",
583+
Agent: protos.SshAgentShimAgentClaudeCode,
572584
IsSuccess: true,
573585
ErrorCategory: protos.SshAgentShimErrorCategoryUnspecified,
574586
SetupDurationMs: 1500,
@@ -580,7 +592,7 @@ func TestBuildAgentShimEvent(t *testing.T) {
580592
err: categorize(protos.SshAgentShimErrorCategoryGatewayAuthFailed, errors.New("boom")),
581593
}, 200*time.Millisecond)
582594
assert.Equal(t, &protos.SshAgentShimEvent{
583-
AgentName: "codex",
595+
Agent: protos.SshAgentShimAgentCodex,
584596
IsSuccess: false,
585597
ErrorCategory: protos.SshAgentShimErrorCategoryGatewayAuthFailed,
586598
SetupDurationMs: 200,

‎libs/telemetry/protos/ssh_agent_shim.go‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,21 @@ const (
8181
SshAgentShimErrorCategoryUnknown SshAgentShimErrorCategory = "UNKNOWN"
8282
)
8383

84+
// The agent for an `ssh agent-shim <agent>` command invocation. Values are in
85+
// screaming snake case to match what JSON->proto deserialization expects.
86+
type SshAgentShimAgent string
87+
88+
const (
89+
// Unknown agent was specified.
90+
SshAgentShimAgentUnspecified SshAgentShimAgent = "TYPE_UNSPECIFIED"
91+
92+
// Claude Code CLI.
93+
SshAgentShimAgentClaudeCode SshAgentShimAgent = "CLAUDE_CODE"
94+
95+
// Codex CLI.
96+
SshAgentShimAgentCodex SshAgentShimAgent = "CODEX"
97+
)
98+
8499
// SshAgentShimEvent is emitted when a user runs `ssh agent-shim <agent>` on the remote
85100
// driver to launch a coding agent against the Unity AI Gateway.
86101
//
@@ -89,9 +104,9 @@ const (
89104
// uploaded from the shim just before the exec. A failed launch returns normally and is
90105
// uploaded by cmd/root as usual.
91106
type SshAgentShimEvent struct {
92-
// Name of the agent the user asked to launch (e.g. "claude", "codex"). Drawn from a
93-
// fixed set of CLI-defined identifiers, never user-authored text.
94-
AgentName string `json:"agent_name"`
107+
// The agent the user asked to launch, mapped to a string matching the proto enum
108+
// (e.g. "CLAUDE_CODE", "CODEX", etc).
109+
Agent SshAgentShimAgent `json:"agent"`
95110

96111
// Whether the agent was launched. False means setup failed before the agent ran.
97112
// Populated on every event, so no omitempty: a genuine false must stay distinguishable

‎libs/telemetry/protos/ssh_agent_shim_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ func TestSshAgentShimEventEncodesFailureExplicitly(t *testing.T) {
2121
assert.Equal(t, false, got["is_success"], "is_success must be sent as false, not omitted")
2222
// The remaining fields must appear in the payload at all: omitempty would drop the
2323
// zero value and make a failure indistinguishable from an unreported one.
24-
for _, field := range []string{"agent_name", "error_category", "setup_duration_ms"} {
24+
for _, field := range []string{"agent", "error_category", "setup_duration_ms"} {
2525
assert.Contains(t, got, field, "%s must be sent, not omitted", field)
2626
}
2727
}

0 commit comments

Comments
 (0)