package main import ( "bufio" "encoding/json" "fmt" "io" "net/http" "net/http/httptest" "os" "path/filepath" "sync" "testing" "github.com/stretchr/testify/require" ) // testKey is a fake API key. const testKey = "sk-test-not-a-real-key" func TestClientSendsTheFixedRequestFields(t *testing.T) { server := newFakeResponses(t) client := newTestClient(t, server) _, err := client.Create(t.Context(), Request{ Model: "gpt-test", Effort: "high", Instructions: "You are a test.", Tools: []Tool{{ Type: "function", Name: "date", Parameters: map[string]any{"type": "object"}, }}, Input: []any{userMessage("\n\n", "hello")}, }) require.NoError(t, err) call := server.Call(0) require.Equal(t, "Bearer "+testKey, call.header.Get("Authorization")) require.Equal(t, "text/event-stream", call.header.Get("Accept")) require.Equal(t, "application/json", call.header.Get("Content-Type")) body := call.body require.Equal(t, "gpt-test", body["model"]) require.Equal(t, "You are a test.", body["instructions"]) require.Equal(t, false, body["store"]) require.Equal(t, true, body["stream"]) require.Equal(t, []any{"reasoning.encrypted_content"}, body["include"]) require.Equal(t, "agent-key", body["prompt_cache_key"]) require.Equal(t, map[string]any{"effort": "high"}, body["reasoning"]) require.NotContains(t, body, "max_output_tokens") require.NotContains(t, body, "metadata") require.Equal(t, []any{map[string]any{ "role": "user", "content": []any{ map[string]any{"type": "input_text", "text": "\n\n"}, map[string]any{"type": "input_text", "text": "hello"}, }, }}, body["input"]) require.Len(t, body["tools"], 1) } func TestClientReadsTheStream(t *testing.T) { server := newFakeResponses(t) server.reply = func(w http.ResponseWriter, call fakeCall) { writeEvents(w, map[string]any{"type": "response.created", "response": map[string]any{}}, map[string]any{"type": "response.output_item.done", "item": map[string]any{"type": "reasoning", "encrypted_content": "opaque"}}, map[string]any{"type": "response.output_text.delta", "delta": "Hel"}, map[string]any{"type": "response.output_item.done", "item": messageItem("Hello")}, map[string]any{"type": "response.output_item.done", "item": messageItem(" world")}, completedEvent(nil, 1200, 1024, 30), ) } response, err := newTestClient(t, server).Create(t.Context(), testRequest()) require.NoError(t, err) require.Equal(t, "Hello world", response.Text) require.Len(t, response.Output, 3) require.JSONEq(t, `{"type": "reasoning", "encrypted_content": "opaque"}`, string(response.Output[0])) require.Equal(t, Usage{ Input: 1200, Cached: 1024, Output: 30, }, response.Usage) } func TestClientTakesTheOutputOfTheCompletedResponse(t *testing.T) { server := newFakeResponses(t) server.reply = func(w http.ResponseWriter, call fakeCall) { writeEvents(w, completedEvent([]any{messageItem("only here")}, 10, 0, 2)) } response, err := newTestClient(t, server).Create(t.Context(), testRequest()) require.NoError(t, err) require.Equal(t, "only here", response.Text) } func TestClientRecordsEachCall(t *testing.T) { server := newFakeResponses(t) server.reply = func(w http.ResponseWriter, call fakeCall) { w.Header().Set("X-Codex-Primary-Used-Percent", "12") w.Header().Set("X-Codex-Plan-Type", "plus") w.Header().Set("X-Other", "not quota") w.Header().Set("X-Codex-Turn-State", "opaque state") if call.number == 1 { w.WriteHeader(http.StatusTooManyRequests) _, _ = io.WriteString(w, `{"detail": "slow down"}`) return } writeEvents(w, completedEvent([]any{messageItem("ok")}, 500, 400, 7)) } client := newTestClient(t, server) _, err := client.Create(t.Context(), testRequest()) require.NoError(t, err) _, err = client.Create(t.Context(), testRequest()) require.ErrorContains(t, err, "HTTP 429") require.ErrorContains(t, err, "slow down") records := readUsage(t, client.usage.path) require.Len(t, records, 2) quota := map[string]string{ "x-codex-primary-used-percent": "12", "x-codex-plan-type": "plus", } require.Equal(t, "gpt-test", records[0].Model) require.Equal(t, http.StatusOK, records[0].Status) require.Equal(t, 500, records[0].InputTokens) require.Equal(t, 400, records[0].CachedTokens) require.Equal(t, 7, records[0].OutputTokens) require.Equal(t, quota, records[0].Quota) require.Empty(t, records[0].Error) require.NotEmpty(t, records[0].Date) require.Equal(t, http.StatusTooManyRequests, records[1].Status) require.Equal(t, quota, records[1].Quota) require.Contains(t, records[1].Error, "HTTP 429") data, err := os.ReadFile(client.usage.path) require.NoError(t, err) require.NotContains(t, string(data), testKey) info, err := os.Stat(client.usage.path) require.NoError(t, err) require.Equal(t, os.FileMode(0o600), info.Mode().Perm()) } func TestClientFailsWhenTheResponseFails(t *testing.T) { testCases := []struct { name string events []any want string }{ { name: "failed response", events: []any{map[string]any{"type": "response.failed", "response": map[string]any{"error": map[string]any{"code": "server_error", "message": "busy"}}}}, want: "server_error: busy", }, { name: "incomplete response", events: []any{map[string]any{"type": "response.incomplete", "response": map[string]any{"incomplete_details": map[string]any{"reason": "content_filter"}}}}, want: "content_filter", }, { name: "error event", events: []any{map[string]any{"type": "error", "code": "rate_limit", "message": "too many"}}, want: "rate_limit: too many", }, { name: "stream ends early", events: []any{map[string]any{"type": "response.output_item.done", "item": messageItem("cut")}}, want: "ended before the response completed", }, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { server := newFakeResponses(t) server.reply = func(w http.ResponseWriter, call fakeCall) { writeEvents(w, tc.events...) } client := newTestClient(t, server) _, err := client.Create(t.Context(), testRequest()) require.ErrorContains(t, err, tc.want) records := readUsage(t, client.usage.path) require.Len(t, records, 1) require.Contains(t, records[0].Error, tc.want) }) } } // fakeResponses is an httptest Responses endpoint. It keeps each call, and // reply writes the response. The default reply is one short message. type fakeResponses struct { *httptest.Server mutex sync.Mutex calls []fakeCall reply func(http.ResponseWriter, fakeCall) } // fakeCall is one call to a fake endpoint, numbered from 0. type fakeCall struct { number int header http.Header raw []byte body map[string]any } func newFakeResponses(t *testing.T) *fakeResponses { t.Helper() f := &fakeResponses{ reply: func(w http.ResponseWriter, call fakeCall) { writeEvents(w, completedEvent([]any{messageItem("a line")}, 100, 0, 5)) }, } f.Server = httptest.NewServer(http.HandlerFunc(f.serve)) t.Cleanup(f.Close) return f } func (f *fakeResponses) serve(w http.ResponseWriter, r *http.Request) { raw, err := io.ReadAll(r.Body) if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } call := fakeCall{ header: r.Header.Clone(), raw: raw, } err = json.Unmarshal(raw, &call.body) if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } f.mutex.Lock() call.number = len(f.calls) f.calls = append(f.calls, call) reply := f.reply f.mutex.Unlock() reply(w, call) } // Call returns call number n. func (f *fakeResponses) Call(n int) fakeCall { f.mutex.Lock() defer f.mutex.Unlock() return f.calls[n] } // Calls returns the number of calls. func (f *fakeResponses) Calls() int { f.mutex.Lock() defer f.mutex.Unlock() return len(f.calls) } func newTestClient(t *testing.T, server *fakeResponses) *Client { t.Helper() return newClient(server.URL, apiKey(testKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl")) } func testRequest() Request { return Request{ Model: "gpt-test", Effort: "low", Input: []any{userMessage("hello")}, } } // writeEvents writes events as a server-sent event stream. func writeEvents(w http.ResponseWriter, events ...any) { w.Header().Set("Content-Type", "text/event-stream") for _, event := range events { data, err := json.Marshal(event) if err != nil { panic(err) } _, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data) if err != nil { panic(err) } } } func messageItem(text string) map[string]any { return map[string]any{ "type": "message", "role": "assistant", "content": []any{ map[string]any{"type": "output_text", "text": text}, }, } } func completedEvent(output []any, input, cached, out int) map[string]any { if output == nil { output = []any{} } return map[string]any{ "type": "response.completed", "response": map[string]any{ "output": output, "usage": map[string]any{ "input_tokens": input, "input_tokens_details": map[string]any{"cached_tokens": cached}, "output_tokens": out, }, }, } } func readUsage(t *testing.T, path string) []usageRecord { t.Helper() file, err := os.Open(path) require.NoError(t, err) defer func() { require.NoError(t, file.Close()) }() var records []usageRecord scanner := bufio.NewScanner(file) for scanner.Scan() { var record usageRecord require.NoError(t, json.Unmarshal(scanner.Bytes(), &record)) records = append(records, record) } require.NoError(t, scanner.Err()) return records }