git.giacomolli.uk - onechat/blob - responses_test.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "encoding/json"
   6         "fmt"
   7         "io"
   8         "net/http"
   9         "net/http/httptest"
  10         "os"
  11         "path/filepath"
  12         "sync"
  13         "testing"
  14
  15         "github.com/stretchr/testify/require"
  16 )
  17
  18 // testKey is a fake API key.
  19 const testKey = "sk-test-not-a-real-key"
  20
  21 func TestClientSendsTheFixedRequestFields(t *testing.T) {
  22         server := newFakeResponses(t)
  23         client := newTestClient(t, server)
  24
  25         _, err := client.Create(t.Context(), Request{
  26                 Model:        "gpt-test",
  27                 Effort:       "high",
  28                 Instructions: "You are a test.",
  29                 Tools: []Tool{{
  30                         Type:       "function",
  31                         Name:       "date",
  32                         Parameters: map[string]any{"type": "object"},
  33                 }},
  34                 Input: []any{userMessage("<chat>\n</chat>\n", "hello")},
  35         })
  36         require.NoError(t, err)
  37
  38         call := server.Call(0)
  39         require.Equal(t, "Bearer "+testKey, call.header.Get("Authorization"))
  40         require.Equal(t, "text/event-stream", call.header.Get("Accept"))
  41         require.Equal(t, "application/json", call.header.Get("Content-Type"))
  42
  43         body := call.body
  44         require.Equal(t, "gpt-test", body["model"])
  45         require.Equal(t, "You are a test.", body["instructions"])
  46         require.Equal(t, false, body["store"])
  47         require.Equal(t, true, body["stream"])
  48         require.Equal(t, []any{"reasoning.encrypted_content"}, body["include"])
  49         require.Equal(t, "agent-key", body["prompt_cache_key"])
  50         require.Equal(t, map[string]any{"effort": "high"}, body["reasoning"])
  51         require.NotContains(t, body, "max_output_tokens")
  52         require.NotContains(t, body, "metadata")
  53         require.Equal(t, []any{map[string]any{
  54                 "role": "user",
  55                 "content": []any{
  56                         map[string]any{"type": "input_text", "text": "<chat>\n</chat>\n"},
  57                         map[string]any{"type": "input_text", "text": "hello"},
  58                 },
  59         }}, body["input"])
  60         require.Len(t, body["tools"], 1)
  61 }
  62
  63 func TestClientReadsTheStream(t *testing.T) {
  64         server := newFakeResponses(t)
  65         server.reply = func(w http.ResponseWriter, call fakeCall) {
  66                 writeEvents(w,
  67                         map[string]any{"type": "response.created", "response": map[string]any{}},
  68                         map[string]any{"type": "response.output_item.done", "item": map[string]any{"type": "reasoning", "encrypted_content": "opaque"}},
  69                         map[string]any{"type": "response.output_text.delta", "delta": "Hel"},
  70                         map[string]any{"type": "response.output_item.done", "item": messageItem("Hello")},
  71                         map[string]any{"type": "response.output_item.done", "item": messageItem(" world")},
  72                         completedEvent(nil, 1200, 1024, 30),
  73                 )
  74         }
  75
  76         response, err := newTestClient(t, server).Create(t.Context(), testRequest())
  77         require.NoError(t, err)
  78         require.Equal(t, "Hello world", response.Text)
  79         require.Len(t, response.Output, 3)
  80         require.JSONEq(t, `{"type": "reasoning", "encrypted_content": "opaque"}`, string(response.Output[0]))
  81         require.Equal(t, Usage{
  82                 Input:  1200,
  83                 Cached: 1024,
  84                 Output: 30,
  85         }, response.Usage)
  86 }
  87
  88 func TestClientTakesTheOutputOfTheCompletedResponse(t *testing.T) {
  89         server := newFakeResponses(t)
  90         server.reply = func(w http.ResponseWriter, call fakeCall) {
  91                 writeEvents(w, completedEvent([]any{messageItem("only here")}, 10, 0, 2))
  92         }
  93
  94         response, err := newTestClient(t, server).Create(t.Context(), testRequest())
  95         require.NoError(t, err)
  96         require.Equal(t, "only here", response.Text)
  97 }
  98
  99 func TestClientRecordsEachCall(t *testing.T) {
 100         server := newFakeResponses(t)
 101         server.reply = func(w http.ResponseWriter, call fakeCall) {
 102                 w.Header().Set("X-Codex-Primary-Used-Percent", "12")
 103                 w.Header().Set("X-Codex-Plan-Type", "plus")
 104                 w.Header().Set("X-Other", "not quota")
 105                 if call.number == 1 {
 106                         w.WriteHeader(http.StatusTooManyRequests)
 107                         _, _ = io.WriteString(w, `{"detail": "slow down"}`)
 108
 109                         return
 110                 }
 111
 112                 writeEvents(w, completedEvent([]any{messageItem("ok")}, 500, 400, 7))
 113         }
 114
 115         client := newTestClient(t, server)
 116         _, err := client.Create(t.Context(), testRequest())
 117         require.NoError(t, err)
 118
 119         _, err = client.Create(t.Context(), testRequest())
 120         require.ErrorContains(t, err, "HTTP 429")
 121         require.ErrorContains(t, err, "slow down")
 122
 123         records := readUsage(t, client.usage.path)
 124         require.Len(t, records, 2)
 125
 126         quota := map[string]string{
 127                 "x-codex-primary-used-percent": "12",
 128                 "x-codex-plan-type":            "plus",
 129         }
 130         require.Equal(t, "gpt-test", records[0].Model)
 131         require.Equal(t, http.StatusOK, records[0].Status)
 132         require.Equal(t, 500, records[0].InputTokens)
 133         require.Equal(t, 400, records[0].CachedTokens)
 134         require.Equal(t, 7, records[0].OutputTokens)
 135         require.Equal(t, quota, records[0].Quota)
 136         require.Empty(t, records[0].Error)
 137         require.NotEmpty(t, records[0].Date)
 138
 139         require.Equal(t, http.StatusTooManyRequests, records[1].Status)
 140         require.Equal(t, quota, records[1].Quota)
 141         require.Contains(t, records[1].Error, "HTTP 429")
 142
 143         data, err := os.ReadFile(client.usage.path)
 144         require.NoError(t, err)
 145         require.NotContains(t, string(data), testKey)
 146
 147         info, err := os.Stat(client.usage.path)
 148         require.NoError(t, err)
 149         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 150 }
 151
 152 func TestClientFailsWhenTheResponseFails(t *testing.T) {
 153         testCases := []struct {
 154                 name   string
 155                 events []any
 156                 want   string
 157         }{
 158                 {
 159                         name:   "failed response",
 160                         events: []any{map[string]any{"type": "response.failed", "response": map[string]any{"error": map[string]any{"code": "server_error", "message": "busy"}}}},
 161                         want:   "server_error: busy",
 162                 },
 163                 {
 164                         name:   "incomplete response",
 165                         events: []any{map[string]any{"type": "response.incomplete", "response": map[string]any{"incomplete_details": map[string]any{"reason": "content_filter"}}}},
 166                         want:   "content_filter",
 167                 },
 168                 {
 169                         name:   "error event",
 170                         events: []any{map[string]any{"type": "error", "code": "rate_limit", "message": "too many"}},
 171                         want:   "rate_limit: too many",
 172                 },
 173                 {
 174                         name:   "stream ends early",
 175                         events: []any{map[string]any{"type": "response.output_item.done", "item": messageItem("cut")}},
 176                         want:   "ended before the response completed",
 177                 },
 178         }
 179
 180         for _, tc := range testCases {
 181                 t.Run(tc.name, func(t *testing.T) {
 182                         server := newFakeResponses(t)
 183                         server.reply = func(w http.ResponseWriter, call fakeCall) {
 184                                 writeEvents(w, tc.events...)
 185                         }
 186
 187                         client := newTestClient(t, server)
 188                         _, err := client.Create(t.Context(), testRequest())
 189                         require.ErrorContains(t, err, tc.want)
 190
 191                         records := readUsage(t, client.usage.path)
 192                         require.Len(t, records, 1)
 193                         require.Contains(t, records[0].Error, tc.want)
 194                 })
 195         }
 196 }
 197
 198 // fakeResponses is an httptest Responses endpoint. It keeps each call, and
 199 // reply writes the response. The default reply is one short message.
 200 type fakeResponses struct {
 201         *httptest.Server
 202         mutex sync.Mutex
 203         calls []fakeCall
 204         reply func(http.ResponseWriter, fakeCall)
 205 }
 206
 207 // fakeCall is one call to a fake endpoint, numbered from 0.
 208 type fakeCall struct {
 209         number int
 210         header http.Header
 211         raw    []byte
 212         body   map[string]any
 213 }
 214
 215 func newFakeResponses(t *testing.T) *fakeResponses {
 216         t.Helper()
 217
 218         f := &fakeResponses{
 219                 reply: func(w http.ResponseWriter, call fakeCall) {
 220                         writeEvents(w, completedEvent([]any{messageItem("a line")}, 100, 0, 5))
 221                 },
 222         }
 223
 224         f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
 225         t.Cleanup(f.Close)
 226
 227         return f
 228 }
 229
 230 func (f *fakeResponses) serve(w http.ResponseWriter, r *http.Request) {
 231         raw, err := io.ReadAll(r.Body)
 232         if err != nil {
 233                 http.Error(w, err.Error(), http.StatusBadRequest)
 234                 return
 235         }
 236
 237         call := fakeCall{
 238                 header: r.Header.Clone(),
 239                 raw:    raw,
 240         }
 241
 242         err = json.Unmarshal(raw, &call.body)
 243         if err != nil {
 244                 http.Error(w, err.Error(), http.StatusBadRequest)
 245                 return
 246         }
 247
 248         f.mutex.Lock()
 249         call.number = len(f.calls)
 250         f.calls = append(f.calls, call)
 251         reply := f.reply
 252         f.mutex.Unlock()
 253
 254         reply(w, call)
 255 }
 256
 257 // Call returns call number n.
 258 func (f *fakeResponses) Call(n int) fakeCall {
 259         f.mutex.Lock()
 260         defer f.mutex.Unlock()
 261
 262         return f.calls[n]
 263 }
 264
 265 // Calls returns the number of calls.
 266 func (f *fakeResponses) Calls() int {
 267         f.mutex.Lock()
 268         defer f.mutex.Unlock()
 269
 270         return len(f.calls)
 271 }
 272
 273 func newTestClient(t *testing.T, server *fakeResponses) *Client {
 274         t.Helper()
 275
 276         return newClient(server.URL, apiKey(testKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
 277 }
 278
 279 func testRequest() Request {
 280         return Request{
 281                 Model:  "gpt-test",
 282                 Effort: "low",
 283                 Input:  []any{userMessage("hello")},
 284         }
 285 }
 286
 287 // writeEvents writes events as a server-sent event stream.
 288 func writeEvents(w http.ResponseWriter, events ...any) {
 289         w.Header().Set("Content-Type", "text/event-stream")
 290         for _, event := range events {
 291                 data, err := json.Marshal(event)
 292                 if err != nil {
 293                         panic(err)
 294                 }
 295
 296                 _, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data)
 297                 if err != nil {
 298                         panic(err)
 299                 }
 300         }
 301 }
 302
 303 func messageItem(text string) map[string]any {
 304         return map[string]any{
 305                 "type": "message",
 306                 "role": "assistant",
 307                 "content": []any{
 308                         map[string]any{"type": "output_text", "text": text},
 309                 },
 310         }
 311 }
 312
 313 func completedEvent(output []any, input, cached, out int) map[string]any {
 314         if output == nil {
 315                 output = []any{}
 316         }
 317
 318         return map[string]any{
 319                 "type": "response.completed",
 320                 "response": map[string]any{
 321                         "output": output,
 322                         "usage": map[string]any{
 323                                 "input_tokens":         input,
 324                                 "input_tokens_details": map[string]any{"cached_tokens": cached},
 325                                 "output_tokens":        out,
 326                         },
 327                 },
 328         }
 329 }
 330
 331 func readUsage(t *testing.T, path string) []usageRecord {
 332         t.Helper()
 333
 334         file, err := os.Open(path)
 335         require.NoError(t, err)
 336
 337         defer func() {
 338                 require.NoError(t, file.Close())
 339         }()
 340
 341         var records []usageRecord
 342         scanner := bufio.NewScanner(file)
 343         for scanner.Scan() {
 344                 var record usageRecord
 345                 require.NoError(t, json.Unmarshal(scanner.Bytes(), &record))
 346                 records = append(records, record)
 347         }
 348
 349         require.NoError(t, scanner.Err())
 350
 351         return records
 352 }