From 2918e10b7ff79f7d1792f67bd7ff0c18718dc9ad Mon Sep 17 00:00:00 2001 From: SeanAlexanderHarris Date: Fri, 17 Jul 2026 01:16:04 +0100 Subject: [PATCH 1/3] feat(459): :sparkles: add ai detection to include the agent or model using the cli this should help with meta monitoring of agentic use of the CLI, assist with evals, potentially power optimisation for agent use --- agentenv/agentenv.go | 53 ++++++++++++++++++++ agentenv/agentenv_test.go | 101 ++++++++++++++++++++++++++++++++++++++ client/client.go | 9 +++- client/client_test.go | 32 ++++++++++++ 4 files changed, 194 insertions(+), 1 deletion(-) create mode 100644 agentenv/agentenv.go create mode 100644 agentenv/agentenv_test.go diff --git a/agentenv/agentenv.go b/agentenv/agentenv.go new file mode 100644 index 00000000..f202afba --- /dev/null +++ b/agentenv/agentenv.go @@ -0,0 +1,53 @@ +// 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 +// (sanitised) value. 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 "" +} + +func validHeaderValue(s string) bool { + for i := 0; i < len(s); i++ { + b := s[i] + if b == '\t' { + continue + } + if b < 0x20 || b == 0x7f { + return false + } + } + return true +} diff --git a/agentenv/agentenv_test.go b/agentenv/agentenv_test.go new file mode 100644 index 00000000..c6cf0d50 --- /dev/null +++ b/agentenv/agentenv_test.go @@ -0,0 +1,101 @@ +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 allowed", "a\tb", true}, + {"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 a63c8134..c2a55aa2 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 048dfe9b..0cf68da5 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", "GEMINI_CLI", "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) + } +} From 24ababd8a45971df8f4d78e197d492a1645fe018 Mon Sep 17 00:00:00 2001 From: SeanAlexanderHarris Date: Fri, 17 Jul 2026 02:13:35 +0100 Subject: [PATCH 2/3] fix(459): reject whitespace in agentenv header value validator Co-Authored-By: Claude Sonnet 5 --- agentenv/agentenv.go | 11 ++++++----- agentenv/agentenv_test.go | 3 ++- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/agentenv/agentenv.go b/agentenv/agentenv.go index f202afba..8f4439c5 100644 --- a/agentenv/agentenv.go +++ b/agentenv/agentenv.go @@ -20,7 +20,8 @@ var sources = []agentSource{ // 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 -// (sanitised) value. Unset, empty, or malformed values yield "". +// 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) @@ -39,13 +40,13 @@ func Detected() string { 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 == '\t' { - continue - } - if b < 0x20 || b == 0x7f { + if b < 0x20 || b == 0x7f || b == ' ' { return false } } diff --git a/agentenv/agentenv_test.go b/agentenv/agentenv_test.go index c6cf0d50..1658359d 100644 --- a/agentenv/agentenv_test.go +++ b/agentenv/agentenv_test.go @@ -82,7 +82,8 @@ func TestValidHeaderValue(t *testing.T) { }{ {"plain ascii", "claude-code", true}, {"empty string", "", true}, - {"tab is allowed", "a\tb", 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}, From 701deeddd965a13e7f2c3048b9efb5c0c1cf66af Mon Sep 17 00:00:00 2001 From: SeanAlexanderHarris Date: Fri, 17 Jul 2026 02:14:15 +0100 Subject: [PATCH 3/3] test(459): fix env-var isolation list in TestExecuteSetsAgentInUserAgent Co-Authored-By: Claude Sonnet 5 --- client/client_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/client/client_test.go b/client/client_test.go index 0cf68da5..52dbbb4e 100644 --- a/client/client_test.go +++ b/client/client_test.go @@ -57,7 +57,7 @@ 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", "GEMINI_CLI", "AI_AGENT", "LLM_AGENT"} { + for _, k := range []string{"CLAUDE_CODE", "ANTIGRAVITY_AGENT", "AI_AGENT", "LLM_AGENT"} { t.Setenv(k, "") } t.Setenv("CLAUDE_CODE", "1")