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) + } +} diff --git a/contract_test/contract_test.go b/contract_test/contract_test.go index 7d2d748..65cfe41 100644 --- a/contract_test/contract_test.go +++ b/contract_test/contract_test.go @@ -191,6 +191,8 @@ var operations = []operation{ {operationID: "get-task-builder-batch-task-responses", call: func(c *client.Client) { c.GetAITaskBuilderResponses("batch-id") }}, {operationID: "get-task-builder-batch-report", skip: "OUTOFSCOPE: no CLI command for batch report; GetAITaskBuilderTasks calls /tasks which is not in the spec"}, {operationID: "duplicate-task-builder-batch", skip: "OUTOFSCOPE: no CLI command for duplicating a batch"}, + {operationID: "sync-task-builder-batch", call: func(c *client.Client) { c.SyncAITaskBuilderBatch("batch-id") }}, + {operationID: "get-batch-sync-status", call: func(c *client.Client) { c.GetAITaskBuilderBatchSyncStatus("batch-id", "sync-id") }}, {operationID: "request-batch-export", call: func(c *client.Client) { c.InitiateBatchExport("batch-id") }}, {operationID: "get-batch-export-status", call: func(c *client.Client) { c.GetBatchExportStatus("batch-id", "export-id") }}, @@ -198,12 +200,15 @@ var operations = []operation{ {operationID: "create-task-builder-dataset", call: func(c *client.Client) { c.CreateAITaskBuilderDataset("ws-id", client.CreateAITaskBuilderDatasetPayload{Name: "t"}) }}, + {operationID: "update-task-builder-dataset", skip: "OUTOFSCOPE: no CLI command for updating a dataset"}, + {operationID: "append-dataset-datapoints", skip: "OUTOFSCOPE: no CLI command for appending datapoints to a dataset"}, {operationID: "get-dataset-upload-url", call: func(c *client.Client) { c.GetAITaskBuilderDatasetUploadURL("ds-id", "data.jsonl") }}, {operationID: "get-task-builder-dataset", skip: "OUTOFSCOPE: no CLI command for retrieving a specific dataset by ID"}, {operationID: "get-task-builder-dataset-status", call: func(c *client.Client) { c.GetAITaskBuilderDatasetStatus("ds-id") }}, {operationID: "get-dataset-import-status", call: func(c *client.Client) { c.GetAITaskBuilderDatasetImportStatus("ds-id", "import-id") }}, + {operationID: "get-schema-migration-status", skip: "OUTOFSCOPE: no CLI command for schema migration status"}, // AI Task Builder — Instructions {operationID: "get-task-builder-instructions", skip: "OUTOFSCOPE: no CLI command for getting task builder instructions"}, @@ -306,10 +311,6 @@ var operations = []operation{ c.UpdateCredentialPool("pool-id", "user,pass\nuser2,pass2") }}, - // Costs - {operationID: "get-organisation-cost", skip: "OUTOFSCOPE: no CLI command for organisation cost"}, - {operationID: "get-workspace-cost", skip: "OUTOFSCOPE: no CLI command for workspace cost"}, - // Reward Recommendations {operationID: "calculate-reward-recommendations", skip: "OUTOFSCOPE: reward recommendations not needed in CLI"},