diff --git a/CHANGELOG.md b/CHANGELOG.md index 223b28a..74d643e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## next +- Add `--skill` root flag to identify the AI skill/workflow invoking a command; folded into the `User-Agent` header sent with API requests + ## 1.0.3 diff --git a/agentenv/agentenv.go b/agentenv/agentenv.go index 8f4439c..f1f36c1 100644 --- a/agentenv/agentenv.go +++ b/agentenv/agentenv.go @@ -32,7 +32,7 @@ func Detected() string { if name == "" { // generic var: forward its value name = val } - if !validHeaderValue(name) { + if !ValidHeaderValue(name) { continue } return name // first usable match wins @@ -40,10 +40,11 @@ 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 { +// ValidHeaderValue reports whether s is safe to embed as a single +// space-separated User-Agent token: no ASCII control characters (0x00-0x1F), +// DEL (0x7F), or the ASCII space character. Other Unicode whitespace is not +// rejected but is not expected from the known env vars or --skill input. +func ValidHeaderValue(s string) bool { for i := 0; i < len(s); i++ { b := s[i] if b < 0x20 || b == 0x7f || b == ' ' { diff --git a/agentenv/agentenv_test.go b/agentenv/agentenv_test.go index 1658359..5077319 100644 --- a/agentenv/agentenv_test.go +++ b/agentenv/agentenv_test.go @@ -94,8 +94,8 @@ func TestValidHeaderValue(t *testing.T) { 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) + 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 5dbab5d..52a1298 100644 --- a/client/client.go +++ b/client/client.go @@ -15,6 +15,7 @@ import ( "net/url" "os" "strings" + "sync" "github.com/prolific-oss/cli/agentenv" "github.com/prolific-oss/cli/config" @@ -154,6 +155,35 @@ type Client struct { BaseURL string Token string Debug bool + Skill string +} + +// cliVersionPrefix is computed once since the CLI version can't change during +// the process lifetime, avoiding a repeated debug.ReadBuildInfo() call on +// every request in dev builds. +var cliVersionPrefix = sync.OnceValue(func() string { + return "prolific-oss/cli/" + version.Get() +}) + +// ComposeUserAgent builds the User-Agent string sent with every outgoing +// request: the CLI's own identity, followed by the detected driving AI +// agent (if any), followed by the invoking skill/workflow (if any and +// valid). Each token is independent and omitted when empty or invalid. +func ComposeUserAgent(skill string) string { + ua := cliVersionPrefix() + if agent := agentenv.Detected(); agent != "" { + ua += " agent/" + agent + } + if skill != "" && agentenv.ValidHeaderValue(skill) { + ua += " skill/" + skill + } + return ua +} + +// userAgent returns the composed User-Agent string for this client, folding +// in the configured Skill. +func (c *Client) userAgent() string { + return ComposeUserAgent(c.Skill) } // New will return a new Prolific client. @@ -192,13 +222,8 @@ 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", userAgent) + request.Header.Set("User-Agent", c.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 52dbbb4..48b9057 100644 --- a/client/client_test.go +++ b/client/client_test.go @@ -55,6 +55,83 @@ func TestFormatBatchErrorBody(t *testing.T) { }) } } +func TestComposeUserAgent(t *testing.T) { + knownVars := []string{"CLAUDE_CODE", "ANTIGRAVITY_AGENT", "AI_AGENT", "LLM_AGENT"} + + tests := []struct { + name string + skill string + agentEnv map[string]string + want string + }{ + { + name: "no skill, no agent", + skill: "", + want: "prolific-oss/cli/" + version.Get(), + }, + { + name: "skill only", + skill: "cli-command-create", + want: "prolific-oss/cli/" + version.Get() + " skill/cli-command-create", + }, + { + name: "agent and skill together", + skill: "cli-command-create", + agentEnv: map[string]string{"CLAUDE_CODE": "1"}, + want: "prolific-oss/cli/" + version.Get() + " agent/claude-code skill/cli-command-create", + }, + { + name: "invalid skill (control characters) is dropped", + skill: "bad\nvalue", + want: "prolific-oss/cli/" + version.Get(), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for _, k := range knownVars { + t.Setenv(k, "") // isolate from ambient agent env vars + } + for k, v := range tt.agentEnv { + t.Setenv(k, v) + } + + if got := ComposeUserAgent(tt.skill); got != tt.want { + t.Fatalf("ComposeUserAgent(%q) = %q, want %q", tt.skill, got, tt.want) + } + }) + } +} + +func TestExecuteSetsSkillInUserAgent(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, "") + } + + 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", + Skill: "cli-command-create", + } + + 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() + " skill/cli-command-create"; gotUserAgent != want { + t.Fatalf("User-Agent = %q, want %q", gotUserAgent, want) + } +} + 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"} { diff --git a/cmd/root.go b/cmd/root.go index 62e4319..7a8e490 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -69,6 +69,8 @@ func NewRootCommand() *cobra.Command { client := client.New() + cmd.PersistentFlags().StringVar(&client.Skill, "skill", "", "Optional identifier for the AI skill/workflow invoking this command; folded into the User-Agent header sent with API requests") + w := os.Stdout cmd.AddCommand( diff --git a/cmd/root_test.go b/cmd/root_test.go index f9adcbd..2384063 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -20,3 +20,16 @@ func TestNewGitHubCommand(t *testing.T) { t.Fatalf("expected use: %s; got %s", short, cmd.Short) } } + +func TestNewRootCommandRegistersSkillFlag(t *testing.T) { + root := cmd.NewRootCommand() + + flag := root.PersistentFlags().Lookup("skill") + if flag == nil { + t.Fatal("expected --skill persistent flag to be registered") + } + + if flag.DefValue != "" { + t.Fatalf("expected --skill default value to be empty, got %q", flag.DefValue) + } +} diff --git a/contract_test/contract_test.go b/contract_test/contract_test.go index f383dc3..8ff029b 100644 --- a/contract_test/contract_test.go +++ b/contract_test/contract_test.go @@ -208,7 +208,7 @@ var operations = []operation{ {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"}, + {operationID: "get-schema-migration-status", skip: "OUTOFSCOPE: no CLI command for dataset schema migration status"}, // AI Task Builder — Instructions {operationID: "get-task-builder-instructions", skip: "OUTOFSCOPE: no CLI command for getting task builder instructions"},