From 868ba2556f89022ab971aad38af21c2d0e3fe946 Mon Sep 17 00:00:00 2001 From: SeanAlexanderHarris Date: Fri, 17 Jul 2026 01:47:54 +0100 Subject: [PATCH 1/7] refactor(460): export agentenv header-value validator Co-Authored-By: Claude Sonnet 5 --- agentenv/agentenv.go | 6 +++--- agentenv/agentenv_test.go | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/agentenv/agentenv.go b/agentenv/agentenv.go index 8f4439c..19cd459 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,10 @@ func Detected() string { return "" } -// validHeaderValue reports whether s is safe to embed as a single +// 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 { +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) } }) } From 8dc673e7660bd95bc4b2f4b50d6775780e34ce88 Mon Sep 17 00:00:00 2001 From: SeanAlexanderHarris Date: Fri, 17 Jul 2026 01:49:59 +0100 Subject: [PATCH 2/7] feat(460): compose skill token into client User-Agent Co-Authored-By: Claude Sonnet 5 --- client/client.go | 33 +++++++++++++--- client/client_test.go | 77 ++++++++++++++++++++++++++++++++++++++ mock_client/mock_client.go | 14 +++++++ 3 files changed, 118 insertions(+), 6 deletions(-) diff --git a/client/client.go b/client/client.go index c2a55aa..5b2f53a 100644 --- a/client/client.go +++ b/client/client.go @@ -32,6 +32,10 @@ const DefaultRecordLimit = 200 // API represents what is allowed to be called on the Prolific client. type API interface { + // UserAgent returns the User-Agent string to send with requests that + // build their own *http.Request and bypass Execute. + UserAgent() string + GetMe() (*MeResponse, error) CreateStudy(model.CreateStudy) (*model.Study, error) @@ -152,6 +156,28 @@ type Client struct { BaseURL string Token string Debug bool + Skill string +} + +// 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 := "prolific-oss/cli/" + version.Get() + 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. @@ -190,13 +216,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..9be9aee 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", "GEMINI_CLI", "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/mock_client/mock_client.go b/mock_client/mock_client.go index 13a4e17..87a571f 100644 --- a/mock_client/mock_client.go +++ b/mock_client/mock_client.go @@ -1286,3 +1286,17 @@ func (mr *MockAPIMockRecorder) UpdateStudy(ID, study interface{}) *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateStudy", reflect.TypeOf((*MockAPI)(nil).UpdateStudy), ID, study) } + +// UserAgent mocks base method. +func (m *MockAPI) UserAgent() string { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UserAgent") + ret0, _ := ret[0].(string) + return ret0 +} + +// UserAgent indicates an expected call of UserAgent. +func (mr *MockAPIMockRecorder) UserAgent() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UserAgent", reflect.TypeOf((*MockAPI)(nil).UserAgent)) +} From 426221afa79ff1f57d26950854d3b382555c5338 Mon Sep 17 00:00:00 2001 From: SeanAlexanderHarris Date: Fri, 17 Jul 2026 01:50:49 +0100 Subject: [PATCH 3/7] feat(460): add --skill root flag Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 2 ++ cmd/root.go | 2 ++ cmd/root_test.go | 13 +++++++++++++ 3 files changed, 17 insertions(+) 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/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) + } +} From d98297ab0b8bb607027566f2bc29c3b46ad6f68f Mon Sep 17 00:00:00 2001 From: SeanAlexanderHarris Date: Fri, 17 Jul 2026 01:53:36 +0100 Subject: [PATCH 4/7] feat(460): send User-Agent on dataset upload requests Co-Authored-By: Claude Sonnet 5 --- cmd/aitaskbuilder/get_dataset_status_test.go | 9 ++++++++- cmd/aitaskbuilder/upload_dataset.go | 5 +++-- cmd/aitaskbuilder/upload_dataset_test.go | 6 ++++++ 3 files changed, 17 insertions(+), 3 deletions(-) diff --git a/cmd/aitaskbuilder/get_dataset_status_test.go b/cmd/aitaskbuilder/get_dataset_status_test.go index 991f91e..37eb807 100644 --- a/cmd/aitaskbuilder/get_dataset_status_test.go +++ b/cmd/aitaskbuilder/get_dataset_status_test.go @@ -16,10 +16,17 @@ import ( "github.com/prolific-oss/cli/model" ) +// testUserAgent is the default value returned by the UserAgent() mock set up +// in setupMockClient. It's a cross-cutting header sent with every request, +// so most tests don't assert on it directly and just need a harmless default. +const testUserAgent = "prolific-oss/cli/test agent/test-agent skill/test-skill" + func setupMockClient(t *testing.T) *mock_client.MockAPI { ctrl := gomock.NewController(t) t.Cleanup(func() { ctrl.Finish() }) - return mock_client.NewMockAPI(ctrl) + c := mock_client.NewMockAPI(ctrl) + c.EXPECT().UserAgent().Return(testUserAgent).AnyTimes() + return c } func TestNewGetDatasetStatusCommand(t *testing.T) { diff --git a/cmd/aitaskbuilder/upload_dataset.go b/cmd/aitaskbuilder/upload_dataset.go index 2ce801f..3a0d895 100644 --- a/cmd/aitaskbuilder/upload_dataset.go +++ b/cmd/aitaskbuilder/upload_dataset.go @@ -125,7 +125,7 @@ func uploadDatasetFile(client client.API, opts DatasetUploadOptions, w io.Writer fmt.Fprintf(w, "Uploading %s...\n", uploadRequest.DisplayName) // Upload file to the presigned URL - err = uploadFileToPresignedURL(uploadRequest.LocalPath, uploadResponse.UploadURL, uploadResponse.HTTPMethod, uploadResponse.ContentType) + err = uploadFileToPresignedURL(uploadRequest.LocalPath, uploadResponse.UploadURL, uploadResponse.HTTPMethod, uploadResponse.ContentType, client.UserAgent()) if err != nil { return fmt.Errorf("failed to upload file: %w", err) } @@ -211,7 +211,7 @@ func detectDatasetUploadFormat(fileName string) (model.DatasetImportFormat, erro } // uploadFileToPresignedURL uploads a file to a presigned URL using the API-provided method and content type. -func uploadFileToPresignedURL(filePath, uploadURL, method, contentType string) error { +func uploadFileToPresignedURL(filePath, uploadURL, method, contentType, userAgent string) error { if method == "" { return errors.New("upload response missing http method") } @@ -238,6 +238,7 @@ func uploadFileToPresignedURL(filePath, uploadURL, method, contentType string) e request.ContentLength = info.Size() request.Header.Set("Content-Type", contentType) + request.Header.Set("User-Agent", userAgent) response, err := http.DefaultClient.Do(request) if err != nil { diff --git a/cmd/aitaskbuilder/upload_dataset_test.go b/cmd/aitaskbuilder/upload_dataset_test.go index 929daf5..713e3ab 100644 --- a/cmd/aitaskbuilder/upload_dataset_test.go +++ b/cmd/aitaskbuilder/upload_dataset_test.go @@ -83,6 +83,7 @@ func TestDatasetUploadCommandUploadsDetectedCSV(t *testing.T) { var gotContentLength int64 var gotTransferEncoding []string var gotBody []byte + var gotUserAgent string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, err := io.ReadAll(r.Body) if err != nil { @@ -94,6 +95,7 @@ func TestDatasetUploadCommandUploadsDetectedCSV(t *testing.T) { gotContentLength = r.ContentLength gotTransferEncoding = r.TransferEncoding gotBody = body + gotUserAgent = r.Header.Get("User-Agent") w.WriteHeader(http.StatusOK) })) t.Cleanup(srv.Close) @@ -163,6 +165,10 @@ func TestDatasetUploadCommandUploadsDetectedCSV(t *testing.T) { t.Fatalf("expected uploaded file contents to match") } + if gotUserAgent != testUserAgent { + t.Fatalf("expected User-Agent %q, got %q", testUserAgent, gotUserAgent) + } + output := b.String() if !strings.Contains(output, "Import ID: import-123") { t.Fatalf("expected output to contain import ID, got %s", output) From f67683a22c57f13361859bf5bd1aabd861e9b192 Mon Sep 17 00:00:00 2001 From: SeanAlexanderHarris Date: Fri, 17 Jul 2026 01:55:35 +0100 Subject: [PATCH 5/7] feat(460): send User-Agent on collection export downloads Co-Authored-By: Claude Sonnet 5 --- cmd/collection/export.go | 11 ++++++----- cmd/collection/export_test.go | 21 ++++++++++++++++----- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/cmd/collection/export.go b/cmd/collection/export.go index cf134c7..467e708 100644 --- a/cmd/collection/export.go +++ b/cmd/collection/export.go @@ -107,7 +107,7 @@ func exportCollection(c client.API, opts ExportOptions, w io.Writer) error { // If already complete (cached result), download immediately. if initResult.Status == exportStatusComplete { - return downloadExport(initResult.URL, opts.Output, w) + return downloadExport(initResult.URL, opts.Output, c.UserAgent(), w) } if initResult.Status != exportStatusGenerating { @@ -134,7 +134,7 @@ func exportCollection(c client.API, opts ExportOptions, w io.Writer) error { switch pollResult.Status { case exportStatusComplete: - return downloadExport(pollResult.URL, opts.Output, w) + return downloadExport(pollResult.URL, opts.Output, c.UserAgent(), w) case exportStatusFailed: return fmt.Errorf("export generation failed for collection %s", collectionID) case exportStatusGenerating: @@ -145,16 +145,16 @@ func exportCollection(c client.API, opts ExportOptions, w io.Writer) error { } } -func downloadExport(url, outputPath string, w io.Writer) error { +func downloadExport(url, outputPath, userAgent string, w io.Writer) error { fmt.Fprintf(w, "\nExport ready. Downloading to %s...\n", outputPath) - if err := downloadFile(url, outputPath); err != nil { + if err := downloadFile(url, outputPath, userAgent); err != nil { return fmt.Errorf("error downloading export: %s", err.Error()) } fmt.Fprintf(w, "Export saved to %s\n", outputPath) return nil } -func downloadFile(rawURL, outputPath string) error { +func downloadFile(rawURL, outputPath, userAgent string) error { parsed, err := url.Parse(rawURL) if err != nil { return fmt.Errorf("invalid download URL: %w", err) @@ -170,6 +170,7 @@ func downloadFile(rawURL, outputPath string) error { if err != nil { return fmt.Errorf("failed to create download request: %w", err) } + req.Header.Set("User-Agent", userAgent) resp, err := downloadClient.Do(req) if err != nil { diff --git a/cmd/collection/export_test.go b/cmd/collection/export_test.go index 25fa54c..61b1a8f 100644 --- a/cmd/collection/export_test.go +++ b/cmd/collection/export_test.go @@ -19,18 +19,22 @@ import ( ) const testExportID = "export-job-uuid-123" +const testUserAgent = "prolific-oss/cli/test agent/test-agent skill/test-skill" // newZIPServer creates a TLS test server that returns the given bytes as a // download response. Tests must inject srv.Client() via SetDownloadClientForTesting // so that the HTTPS scheme check passes and the self-signed cert is trusted. -func newZIPServer(t *testing.T, content []byte) *httptest.Server { +// The returned pointer captures the User-Agent header of the last request. +func newZIPServer(t *testing.T, content []byte) (*httptest.Server, *string) { t.Helper() + var gotUserAgent string srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotUserAgent = r.Header.Get("User-Agent") w.WriteHeader(http.StatusOK) _, _ = w.Write(content) })) t.Cleanup(srv.Close) - return srv + return srv, &gotUserAgent } func TestNewExportCommand(t *testing.T) { @@ -68,13 +72,14 @@ func TestExportCommandRequiresCollectionID(t *testing.T) { // status "complete" immediately (server has a valid cached export). func TestExportCommandImmediateComplete(t *testing.T) { zipContent := []byte("PK\x03\x04fake zip content") - srv := newZIPServer(t, zipContent) + srv, gotUserAgent := newZIPServer(t, zipContent) defer collection.SetDownloadClientForTesting(srv.Client())() ctrl := gomock.NewController(t) defer ctrl.Finish() mockClient := mock_client.NewMockAPI(ctrl) + mockClient.EXPECT().UserAgent().Return(testUserAgent).Times(1) mockClient. EXPECT(). InitiateCollectionExport(gomock.Eq(testCollectionID)). @@ -112,6 +117,10 @@ func TestExportCommandImmediateComplete(t *testing.T) { if !strings.Contains(output, outputPath) { t.Errorf("expected output to mention file path %q, got: %s", outputPath, output) } + + if *gotUserAgent != testUserAgent { + t.Errorf("expected User-Agent %q, got %q", testUserAgent, *gotUserAgent) + } } // TestExportCommandPollingToComplete covers the normal async flow: @@ -120,13 +129,14 @@ func TestExportCommandPollingToComplete(t *testing.T) { defer collection.SetPollSleepForTesting(func(time.Duration) {})() zipContent := []byte("PK\x03\x04fake zip content") - srv := newZIPServer(t, zipContent) + srv, _ := newZIPServer(t, zipContent) defer collection.SetDownloadClientForTesting(srv.Client())() ctrl := gomock.NewController(t) defer ctrl.Finish() mockClient := mock_client.NewMockAPI(ctrl) + mockClient.EXPECT().UserAgent().Return(testUserAgent).Times(1) mockClient.EXPECT(). InitiateCollectionExport(gomock.Eq(testCollectionID)). Return(&client.CollectionExportResponse{ @@ -227,13 +237,14 @@ func TestExportCommandInitiateError(t *testing.T) { func TestExportCommandDefaultOutputPath(t *testing.T) { zipContent := []byte("PK\x03\x04fake zip content") - srv := newZIPServer(t, zipContent) + srv, _ := newZIPServer(t, zipContent) defer collection.SetDownloadClientForTesting(srv.Client())() ctrl := gomock.NewController(t) defer ctrl.Finish() mockClient := mock_client.NewMockAPI(ctrl) + mockClient.EXPECT().UserAgent().Return(testUserAgent).Times(1) mockClient.EXPECT(). InitiateCollectionExport(gomock.Eq(testCollectionID)). Return(&client.CollectionExportResponse{ From 3307ab25b27a6ce308997b797c5eb5641c1ec0b3 Mon Sep 17 00:00:00 2001 From: SeanAlexanderHarris Date: Fri, 17 Jul 2026 01:57:40 +0100 Subject: [PATCH 6/7] feat(460): send User-Agent on batch export downloads Co-Authored-By: Claude Sonnet 5 --- cmd/aitaskbuilder/batch_export.go | 11 ++++++----- cmd/aitaskbuilder/batch_export_test.go | 20 +++++++++++++++----- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/cmd/aitaskbuilder/batch_export.go b/cmd/aitaskbuilder/batch_export.go index 537a408..1156d41 100644 --- a/cmd/aitaskbuilder/batch_export.go +++ b/cmd/aitaskbuilder/batch_export.go @@ -102,7 +102,7 @@ func exportBatch(c client.API, opts BatchExportOptions, w io.Writer) error { // If already complete (cached result), download immediately. if initResult.Status == batchExportStatusComplete { - return batchDownloadExport(initResult.URL, opts.Output, w) + return batchDownloadExport(initResult.URL, opts.Output, c.UserAgent(), w) } if initResult.Status != batchExportStatusGenerating { @@ -129,7 +129,7 @@ func exportBatch(c client.API, opts BatchExportOptions, w io.Writer) error { switch pollResult.Status { case batchExportStatusComplete: - return batchDownloadExport(pollResult.URL, opts.Output, w) + return batchDownloadExport(pollResult.URL, opts.Output, c.UserAgent(), w) case batchExportStatusFailed: return fmt.Errorf("export generation failed for batch %s", batchID) case batchExportStatusGenerating: @@ -140,16 +140,16 @@ func exportBatch(c client.API, opts BatchExportOptions, w io.Writer) error { } } -func batchDownloadExport(rawURL, outputPath string, w io.Writer) error { +func batchDownloadExport(rawURL, outputPath, userAgent string, w io.Writer) error { fmt.Fprintf(w, "\nExport ready. Downloading to %s...\n", outputPath) - if err := batchDownloadFile(rawURL, outputPath); err != nil { + if err := batchDownloadFile(rawURL, outputPath, userAgent); err != nil { return fmt.Errorf("error downloading export: %s", err.Error()) } fmt.Fprintf(w, "Export saved to %s\n", outputPath) return nil } -func batchDownloadFile(rawURL, outputPath string) error { +func batchDownloadFile(rawURL, outputPath, userAgent string) error { parsed, err := url.Parse(rawURL) if err != nil { return fmt.Errorf("invalid download URL: %w", err) @@ -165,6 +165,7 @@ func batchDownloadFile(rawURL, outputPath string) error { if err != nil { return fmt.Errorf("failed to create download request: %w", err) } + req.Header.Set("User-Agent", userAgent) resp, err := batchExportDownloadClient.Do(req) if err != nil { diff --git a/cmd/aitaskbuilder/batch_export_test.go b/cmd/aitaskbuilder/batch_export_test.go index b331f64..63c5b9b 100644 --- a/cmd/aitaskbuilder/batch_export_test.go +++ b/cmd/aitaskbuilder/batch_export_test.go @@ -24,14 +24,17 @@ const testBatchID = "batch-id-123" // newBatchZIPServer creates a TLS test server that returns the given bytes as a // download response. Tests must inject srv.Client() via SetBatchExportDownloadClientForTesting // so that the HTTPS scheme check passes and the self-signed cert is trusted. -func newBatchZIPServer(t *testing.T, content []byte) *httptest.Server { +// The returned pointer captures the User-Agent header of the last request. +func newBatchZIPServer(t *testing.T, content []byte) (*httptest.Server, *string) { t.Helper() + var gotUserAgent string srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotUserAgent = r.Header.Get("User-Agent") w.WriteHeader(http.StatusOK) _, _ = w.Write(content) })) t.Cleanup(srv.Close) - return srv + return srv, &gotUserAgent } func TestNewBatchExportCommand(t *testing.T) { @@ -69,13 +72,14 @@ func TestBatchExportCommandRequiresBatchID(t *testing.T) { // status "complete" immediately (server has a valid cached export). func TestBatchExportCommandImmediateComplete(t *testing.T) { zipContent := []byte("PK\x03\x04fake zip content") - srv := newBatchZIPServer(t, zipContent) + srv, gotUserAgent := newBatchZIPServer(t, zipContent) defer aitaskbuilder.SetBatchExportDownloadClientForTesting(srv.Client())() ctrl := gomock.NewController(t) defer ctrl.Finish() mockClient := mock_client.NewMockAPI(ctrl) + mockClient.EXPECT().UserAgent().Return(testUserAgent).Times(1) mockClient. EXPECT(). InitiateBatchExport(gomock.Eq(testBatchID)). @@ -113,6 +117,10 @@ func TestBatchExportCommandImmediateComplete(t *testing.T) { if !strings.Contains(output, outputPath) { t.Errorf("expected output to mention file path %q, got: %s", outputPath, output) } + + if *gotUserAgent != testUserAgent { + t.Errorf("expected User-Agent %q, got %q", testUserAgent, *gotUserAgent) + } } // TestBatchExportCommandPollingToComplete covers the normal async flow: @@ -121,13 +129,14 @@ func TestBatchExportCommandPollingToComplete(t *testing.T) { defer aitaskbuilder.SetBatchExportPollSleepForTesting(func(time.Duration) {})() zipContent := []byte("PK\x03\x04fake zip content") - srv := newBatchZIPServer(t, zipContent) + srv, _ := newBatchZIPServer(t, zipContent) defer aitaskbuilder.SetBatchExportDownloadClientForTesting(srv.Client())() ctrl := gomock.NewController(t) defer ctrl.Finish() mockClient := mock_client.NewMockAPI(ctrl) + mockClient.EXPECT().UserAgent().Return(testUserAgent).Times(1) mockClient.EXPECT(). InitiateBatchExport(gomock.Eq(testBatchID)). Return(&client.BatchExportResponse{ @@ -228,13 +237,14 @@ func TestBatchExportCommandInitiateError(t *testing.T) { func TestBatchExportCommandDefaultOutputPath(t *testing.T) { zipContent := []byte("PK\x03\x04fake zip content") - srv := newBatchZIPServer(t, zipContent) + srv, _ := newBatchZIPServer(t, zipContent) defer aitaskbuilder.SetBatchExportDownloadClientForTesting(srv.Client())() ctrl := gomock.NewController(t) defer ctrl.Finish() mockClient := mock_client.NewMockAPI(ctrl) + mockClient.EXPECT().UserAgent().Return(testUserAgent).Times(1) mockClient.EXPECT(). InitiateBatchExport(gomock.Eq(testBatchID)). Return(&client.BatchExportResponse{ From dd74bd092c159f4b9f0c279ff99a0624f62c0172 Mon Sep 17 00:00:00 2001 From: SeanAlexanderHarris Date: Fri, 17 Jul 2026 02:17:49 +0100 Subject: [PATCH 7/7] test(460): fix env-var isolation list in TestExecuteSetsSkillInUserAgent 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 9be9aee..48b9057 100644 --- a/client/client_test.go +++ b/client/client_test.go @@ -105,7 +105,7 @@ func TestComposeUserAgent(t *testing.T) { func TestExecuteSetsSkillInUserAgent(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, "") }