diff --git a/agentenv/agentenv.go b/agentenv/agentenv.go new file mode 100644 index 0000000..8f4439c --- /dev/null +++ b/agentenv/agentenv.go @@ -0,0 +1,54 @@ +// Package agentenv detects well-known environment variables set by AI +// coding agents (Claude Code, Antigravity, etc.) and resolves the driving +// agent's name so the CLI can advertise it in its User-Agent header, letting +// the Prolific API attribute requests to the agent that drove them. +package agentenv + +import "os" + +type agentSource struct { + env string // env var variable to check + name string // the agent / model slug, empty is treated as 'forward value' +} + +var sources = []agentSource{ + {"CLAUDE_CODE", "claude-code"}, + {"ANTIGRAVITY_AGENT", "antigravity"}, + {"AI_AGENT", ""}, + {"LLM_AGENT", ""}, +} + +// Detected returns the name of the AI agent driving the CLI, or "" if none is +// detected. Branded tools map to a fixed slug; generic vars forward their +// value verbatim, provided it contains no control characters or whitespace. +// Unset, empty, or malformed values yield "". +func Detected() string { + for _, s := range sources { + val := os.Getenv(s.env) + if val == "" { + continue + } + name := s.name + if name == "" { // generic var: forward its value + name = val + } + if !validHeaderValue(name) { + continue + } + return name // first usable match wins + } + return "" +} + +// validHeaderValue reports whether s is safe to embed as a single +// space-separated User-Agent token: no control characters, and no +// whitespace (which would split the token across multiple segments). +func validHeaderValue(s string) bool { + for i := 0; i < len(s); i++ { + b := s[i] + if b < 0x20 || b == 0x7f || b == ' ' { + return false + } + } + return true +} diff --git a/agentenv/agentenv_test.go b/agentenv/agentenv_test.go new file mode 100644 index 0000000..1658359 --- /dev/null +++ b/agentenv/agentenv_test.go @@ -0,0 +1,102 @@ +package agentenv + +import ( + "testing" +) + +func TestDetected(t *testing.T) { + // Every env var Detected looks at. Each subtest zeroes all of them first + // so it's isolated from whatever environment its running in (e.g. running + // inside an AI agent, or a local dev machine) has set. + knownVars := []string{"CLAUDE_CODE", "ANTIGRAVITY_AGENT", "AI_AGENT", "LLM_AGENT"} + + tests := []struct { + name string + env map[string]string + want string + }{ + { + name: "no agent detected", + env: map[string]string{}, + want: "", + }, + { + name: "claude code, branded slug regardless of value", + env: map[string]string{"CLAUDE_CODE": "1"}, + want: "claude-code", + }, + { + name: "antigravity, branded slug regardless of value", + env: map[string]string{"ANTIGRAVITY_AGENT": "ignored-value"}, + want: "antigravity", + }, + { + name: "generic AI_AGENT forwards its value", + env: map[string]string{"AI_AGENT": "cursor"}, + want: "cursor", + }, + { + // tricky scenario, decision for branded + name: "branded wins over generic (precedence)", + env: map[string]string{"CLAUDE_CODE": "1", "AI_AGENT": "cursor"}, + want: "claude-code", + }, + { + name: "empty string treated as unset", + env: map[string]string{"CLAUDE_CODE": ""}, + want: "", + }, + { + name: "invalid generic value skipped, falls through to next source", + env: map[string]string{"AI_AGENT": "bad\nvalue", "LLM_AGENT": "gemma4"}, + want: "gemma4", + }, + { + name: "valid non-ASCII generic value forwarded unchanged", + env: map[string]string{"AI_AGENT": "claudé-日本"}, + want: "claudé-日本", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for _, k := range knownVars { + t.Setenv(k, "") // isolate from ambient env vars + } + for k, v := range tt.env { + t.Setenv(k, v) + } + + if got := Detected(); got != tt.want { + t.Fatalf("Detected() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestValidHeaderValue(t *testing.T) { + tests := []struct { + name string + value string + want bool + }{ + {"plain ascii", "claude-code", true}, + {"empty string", "", true}, + {"tab is rejected", "a\tb", false}, + {"space is rejected", "a b", false}, + {"utf-8 multibyte", "日本語", true}, + {"high byte 0xFF", "\xff", true}, + {"carriage return", "a\rb", false}, + {"line feed", "a\nb", false}, + {"null byte", "a\x00b", false}, + {"del", "a\x7fb", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := validHeaderValue(tt.value); got != tt.want { + t.Fatalf("validHeaderValue(%q) = %v, want %v", tt.value, got, tt.want) + } + }) + } +} diff --git a/client/client.go b/client/client.go index a63c813..c2a55aa 100644 --- a/client/client.go +++ b/client/client.go @@ -16,8 +16,10 @@ import ( "os" "strings" + "github.com/prolific-oss/cli/agentenv" "github.com/prolific-oss/cli/config" "github.com/prolific-oss/cli/model" + "github.com/prolific-oss/cli/version" "github.com/spf13/viper" "golang.org/x/exp/slices" ) @@ -188,8 +190,13 @@ func (c *Client) Execute(method, url string, body any, response any) (*http.Resp return nil, err } + userAgent := "prolific-oss/cli/" + version.Get() + if agent := agentenv.Detected(); agent != "" { + userAgent += " agent/" + agent + } + request.Header.Set("Content-Type", "application/json") - request.Header.Set("User-Agent", "prolific-oss/cli") + request.Header.Set("User-Agent", userAgent) request.Header.Set("Authorization", fmt.Sprintf("Token %s", c.Token)) if c.Debug { diff --git a/client/client_test.go b/client/client_test.go index 048dfe9..52dbbb4 100644 --- a/client/client_test.go +++ b/client/client_test.go @@ -1,7 +1,11 @@ package client import ( + "net/http" + "net/http/httptest" "testing" + + "github.com/prolific-oss/cli/version" ) func TestFormatBatchErrorBody(t *testing.T) { @@ -51,3 +55,31 @@ func TestFormatBatchErrorBody(t *testing.T) { }) } } +func TestExecuteSetsAgentInUserAgent(t *testing.T) { + // Isolate from ambient agent env vars (this shell has AI_AGENT set). + for _, k := range []string{"CLAUDE_CODE", "ANTIGRAVITY_AGENT", "AI_AGENT", "LLM_AGENT"} { + t.Setenv(k, "") + } + t.Setenv("CLAUDE_CODE", "1") + + var gotUserAgent string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotUserAgent = r.Header.Get("User-Agent") + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + c := Client{ + Client: server.Client(), + BaseURL: server.URL, + Token: "test-token", + } + + if _, err := c.Execute(http.MethodGet, "/studies", nil, nil); err != nil { + t.Fatalf("Execute returned error: %v", err) + } + + if want := "prolific-oss/cli/" + version.Get() + " agent/claude-code"; gotUserAgent != want { + t.Fatalf("User-Agent = %q, want %q", gotUserAgent, want) + } +}