Add ai detection to include the agent or model using the cli Closes #459

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
+}

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)
+			}
+		})
+	}
+}

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 {

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)
+	}
+}