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                 w.Header().Set("X-Codex-Turn-State", "opaque state")
 106                 if call.number == 1 {
 107                         w.WriteHeader(http.StatusTooManyRequests)
 108                         _, _ = io.WriteString(w, `{"detail": "slow down"}`)
 109
 110                         return
 111                 }
 112
 113                 writeEvents(w, completedEvent([]any{messageItem("ok")}, 500, 400, 7))
 114         }
 115
 116         client := newTestClient(t, server)
 117         _, err := client.Create(t.Context(), testRequest())
 118         require.NoError(t, err)
 119
 120         _, err = client.Create(t.Context(), testRequest())
 121         require.ErrorContains(t, err, "HTTP 429")
 122         require.ErrorContains(t, err, "slow down")
 123
 124         records := readUsage(t, client.usage.path)
 125         require.Len(t, records, 2)
 126
 127         quota := map[string]string{
 128                 "x-codex-primary-used-percent": "12",
 129                 "x-codex-plan-type":            "plus",
 130         }
 131         require.Equal(t, "gpt-test", records[0].Model)
 132         require.Equal(t, http.StatusOK, records[0].Status)
 133         require.Equal(t, 500, records[0].InputTokens)
 134         require.Equal(t, 400, records[0].CachedTokens)
 135         require.Equal(t, 7, records[0].OutputTokens)
 136         require.Equal(t, quota, records[0].Quota)
 137         require.Empty(t, records[0].Error)
 138         require.NotEmpty(t, records[0].Date)
 139
 140         require.Equal(t, http.StatusTooManyRequests, records[1].Status)
 141         require.Equal(t, quota, records[1].Quota)
 142         require.Contains(t, records[1].Error, "HTTP 429")
 143
 144         data, err := os.ReadFile(client.usage.path)
 145         require.NoError(t, err)
 146         require.NotContains(t, string(data), testKey)
 147
 148         info, err := os.Stat(client.usage.path)
 149         require.NoError(t, err)
 150         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 151 }
 152
 153 func TestClientFailsWhenTheResponseFails(t *testing.T) {
 154         testCases := []struct {
 155                 name   string
 156                 events []any
 157                 want   string
 158         }{
 159                 {
 160                         name:   "failed response",
 161                         events: []any{map[string]any{"type": "response.failed", "response": map[string]any{"error": map[string]any{"code": "server_error", "message": "busy"}}}},
 162                         want:   "server_error: busy",
 163                 },
 164                 {
 165                         name:   "incomplete response",
 166                         events: []any{map[string]any{"type": "response.incomplete", "response": map[string]any{"incomplete_details": map[string]any{"reason": "content_filter"}}}},
 167                         want:   "content_filter",
 168                 },
 169                 {
 170                         name:   "error event",
 171                         events: []any{map[string]any{"type": "error", "code": "rate_limit", "message": "too many"}},
 172                         want:   "rate_limit: too many",
 173                 },
 174                 {
 175                         name:   "stream ends early",
 176                         events: []any{map[string]any{"type": "response.output_item.done", "item": messageItem("cut")}},
 177                         want:   "ended before the response completed",
 178                 },
 179         }
 180
 181         for _, tc := range testCases {
 182                 t.Run(tc.name, func(t *testing.T) {
 183                         server := newFakeResponses(t)
 184                         server.reply = func(w http.ResponseWriter, call fakeCall) {
 185                                 writeEvents(w, tc.events...)
 186                         }
 187
 188                         client := newTestClient(t, server)
 189                         _, err := client.Create(t.Context(), testRequest())
 190                         require.ErrorContains(t, err, tc.want)
 191
 192                         records := readUsage(t, client.usage.path)
 193                         require.Len(t, records, 1)
 194                         require.Contains(t, records[0].Error, tc.want)
 195                 })
 196         }
 197 }
 198
 199 // fakeResponses is an httptest Responses endpoint. It keeps each call, and
 200 // reply writes the response. The default reply is one short message.
 201 type fakeResponses struct {
 202         *httptest.Server
 203         mutex sync.Mutex
 204         calls []fakeCall
 205         reply func(http.ResponseWriter, fakeCall)
 206 }
 207
 208 // fakeCall is one call to a fake endpoint, numbered from 0.
 209 type fakeCall struct {
 210         number int
 211         header http.Header
 212         raw    []byte
 213         body   map[string]any
 214 }
 215
 216 func newFakeResponses(t *testing.T) *fakeResponses {
 217         t.Helper()
 218
 219         f := &fakeResponses{
 220                 reply: func(w http.ResponseWriter, call fakeCall) {
 221                         writeEvents(w, completedEvent([]any{messageItem("a line")}, 100, 0, 5))
 222                 },
 223         }
 224
 225         f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
 226         t.Cleanup(f.Close)
 227
 228         return f
 229 }
 230
 231 func (f *fakeResponses) serve(w http.ResponseWriter, r *http.Request) {
 232         raw, err := io.ReadAll(r.Body)
 233         if err != nil {
 234                 http.Error(w, err.Error(), http.StatusBadRequest)
 235                 return
 236         }
 237
 238         call := fakeCall{
 239                 header: r.Header.Clone(),
 240                 raw:    raw,
 241         }
 242
 243         err = json.Unmarshal(raw, &call.body)
 244         if err != nil {
 245                 http.Error(w, err.Error(), http.StatusBadRequest)
 246                 return
 247         }
 248
 249         f.mutex.Lock()
 250         call.number = len(f.calls)
 251         f.calls = append(f.calls, call)
 252         reply := f.reply
 253         f.mutex.Unlock()
 254
 255         reply(w, call)
 256 }
 257
 258 // Call returns call number n.
 259 func (f *fakeResponses) Call(n int) fakeCall {
 260         f.mutex.Lock()
 261         defer f.mutex.Unlock()
 262
 263         return f.calls[n]
 264 }
 265
 266 // Calls returns the number of calls.
 267 func (f *fakeResponses) Calls() int {
 268         f.mutex.Lock()
 269         defer f.mutex.Unlock()
 270
 271         return len(f.calls)
 272 }
 273
 274 func newTestClient(t *testing.T, server *fakeResponses) *Client {
 275         t.Helper()
 276
 277         return newClient(server.URL, apiKey(testKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
 278 }
 279
 280 func testRequest() Request {
 281         return Request{
 282                 Model:  "gpt-test",
 283                 Effort: "low",
 284                 Input:  []any{userMessage("hello")},
 285         }
 286 }
 287
 288 // writeEvents writes events as a server-sent event stream.
 289 func writeEvents(w http.ResponseWriter, events ...any) {
 290         w.Header().Set("Content-Type", "text/event-stream")
 291         for _, event := range events {
 292                 data, err := json.Marshal(event)
 293                 if err != nil {
 294                         panic(err)
 295                 }
 296
 297                 _, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data)
 298                 if err != nil {
 299                         panic(err)
 300                 }
 301         }
 302 }
 303
 304 func messageItem(text string) map[string]any {
 305         return map[string]any{
 306                 "type": "message",
 307                 "role": "assistant",
 308                 "content": []any{
 309                         map[string]any{"type": "output_text", "text": text},
 310                 },
 311         }
 312 }
 313
 314 func completedEvent(output []any, input, cached, out int) map[string]any {
 315         if output == nil {
 316                 output = []any{}
 317         }
 318
 319         return map[string]any{
 320                 "type": "response.completed",
 321                 "response": map[string]any{
 322                         "output": output,
 323                         "usage": map[string]any{
 324                                 "input_tokens":         input,
 325                                 "input_tokens_details": map[string]any{"cached_tokens": cached},
 326                                 "output_tokens":        out,
 327                         },
 328                 },
 329         }
 330 }
 331
 332 func readUsage(t *testing.T, path string) []usageRecord {
 333         t.Helper()
 334
 335         file, err := os.Open(path)
 336         require.NoError(t, err)
 337
 338         defer func() {
 339                 require.NoError(t, file.Close())
 340         }()
 341
 342         var records []usageRecord
 343         scanner := bufio.NewScanner(file)
 344         for scanner.Scan() {
 345                 var record usageRecord
 346                 require.NoError(t, json.Unmarshal(scanner.Bytes(), &record))
 347                 records = append(records, record)
 348         }
 349
 350         require.NoError(t, scanner.Err())
 351
 352         return records
 353 }