git.giacomolli.uk - onechat/blob - model_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         "time"
  15
  16         "github.com/stretchr/testify/require"
  17 )
  18
  19 // modelTestKey is a fake API key.
  20 const modelTestKey = "sk-test-not-a-real-key"
  21
  22 func TestClientSendsTheFixedRequestFields(t *testing.T) {
  23         server := newModelFake(t)
  24         client := newTestModelClient(t, server)
  25
  26         _, err := client.Create(t.Context(), ModelRequest{
  27                 Model:        "gpt-test",
  28                 Effort:       "high",
  29                 Instructions: "You are a test.",
  30                 Tools: []ModelTool{{
  31                         Type:       "function",
  32                         Name:       "date",
  33                         Parameters: map[string]any{"type": "object"},
  34                 }},
  35                 Input: []any{modelUserMessage("<chat>\n</chat>\n", "hello")},
  36         })
  37         require.NoError(t, err)
  38
  39         call := server.Call(0)
  40         require.Equal(t, "Bearer "+modelTestKey, call.header.Get("Authorization"))
  41         require.Equal(t, "text/event-stream", call.header.Get("Accept"))
  42         require.Equal(t, "application/json", call.header.Get("Content-Type"))
  43
  44         body := call.body
  45         require.Equal(t, "gpt-test", body["model"])
  46         require.Equal(t, "You are a test.", body["instructions"])
  47         require.Equal(t, false, body["store"])
  48         require.Equal(t, true, body["stream"])
  49         require.Equal(t, []any{"reasoning.encrypted_content"}, body["include"])
  50         require.Equal(t, "agent-key", body["prompt_cache_key"])
  51         require.Equal(t, map[string]any{"effort": "high"}, body["reasoning"])
  52         require.NotContains(t, body, "max_output_tokens")
  53         require.NotContains(t, body, "metadata")
  54         require.Equal(t, []any{map[string]any{
  55                 "role": "user",
  56                 "content": []any{
  57                         map[string]any{"type": "input_text", "text": "<chat>\n</chat>\n"},
  58                         map[string]any{"type": "input_text", "text": "hello"},
  59                 },
  60         }}, body["input"])
  61         require.Len(t, body["tools"], 1)
  62 }
  63
  64 func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {
  65         login := newChatGPTFake(t)
  66         auth := newTestChatGPTAuth(t, login)
  67         chatgptWriteTestTokens(t, auth, login.token(time.Hour), "refresh-0")
  68
  69         testCases := []struct {
  70                 name string
  71                 auth ModelAuthorizer
  72                 want string
  73         }{
  74                 {
  75                         name: "ChatGPT backend",
  76                         auth: auth,
  77                         want: "agent-key",
  78                 },
  79                 {
  80                         name: "OpenAI API",
  81                         auth: ModelAPIKey(modelTestKey),
  82                         want: "",
  83                 },
  84         }
  85
  86         for _, tc := range testCases {
  87                 t.Run(tc.name, func(t *testing.T) {
  88                         server := newModelFake(t)
  89                         client := newModelClient(server.URL, tc.auth, "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
  90                         _, err := client.Create(t.Context(), modelTestRequest())
  91                         require.NoError(t, err)
  92
  93                         call := server.Call(0)
  94                         require.Equal(t, tc.want, call.header.Get("session_id"))
  95                         require.Equal(t, "agent-key", call.body["prompt_cache_key"])
  96                 })
  97         }
  98 }
  99
 100 func TestClientReadsTheStream(t *testing.T) {
 101         server := newModelFake(t)
 102         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 103                 modelWriteEvents(w,
 104                         map[string]any{"type": "response.created", "response": map[string]any{}},
 105                         map[string]any{"type": "response.output_item.done", "item": map[string]any{"type": "reasoning", "encrypted_content": "opaque"}},
 106                         map[string]any{"type": "response.output_text.delta", "delta": "Hel"},
 107                         map[string]any{"type": "response.output_item.done", "item": modelMessageItem("Hello")},
 108                         map[string]any{"type": "response.output_item.done", "item": modelMessageItem(" world")},
 109                         modelCompletedEvent(nil, 1200, 1024, 30),
 110                 )
 111         }
 112
 113         response, err := newTestModelClient(t, server).Create(t.Context(), modelTestRequest())
 114         require.NoError(t, err)
 115         require.Equal(t, "Hello world", response.Text)
 116         require.Len(t, response.Output, 3)
 117         require.JSONEq(t, `{"type": "reasoning", "encrypted_content": "opaque"}`, string(response.Output[0]))
 118         require.Equal(t, ModelUsage{
 119                 Input:  1200,
 120                 Cached: 1024,
 121                 Output: 30,
 122         }, response.Usage)
 123 }
 124
 125 func TestClientTakesTheOutputOfTheCompletedResponse(t *testing.T) {
 126         server := newModelFake(t)
 127         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 128                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("only here")}, 10, 0, 2))
 129         }
 130
 131         response, err := newTestModelClient(t, server).Create(t.Context(), modelTestRequest())
 132         require.NoError(t, err)
 133         require.Equal(t, "only here", response.Text)
 134 }
 135
 136 func TestClientRecordsEachCall(t *testing.T) {
 137         server := newModelFake(t)
 138         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 139                 w.Header().Set("X-Codex-Primary-Used-Percent", "12")
 140                 w.Header().Set("X-Codex-Plan-Type", "plus")
 141                 w.Header().Set("X-Other", "not quota")
 142                 w.Header().Set("X-Codex-Turn-State", "opaque state")
 143                 if call.number == 1 {
 144                         w.WriteHeader(http.StatusTooManyRequests)
 145                         _, _ = io.WriteString(w, `{"detail": "slow down"}`)
 146
 147                         return
 148                 }
 149
 150                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("ok")}, 500, 400, 7))
 151         }
 152
 153         client := newTestModelClient(t, server)
 154         request := modelTestRequest()
 155         request.Node = &nodeKey{2, 5}
 156         _, err := client.Create(t.Context(), request)
 157         require.NoError(t, err)
 158         require.NotContains(t, server.Call(0).body, "node")
 159         require.NotContains(t, string(server.Call(0).raw), "[2,5]")
 160
 161         _, err = client.Create(t.Context(), modelTestRequest())
 162         require.ErrorContains(t, err, "HTTP 429")
 163         require.ErrorContains(t, err, "slow down")
 164
 165         records := modelReadUsage(t, client.usage.path)
 166         require.Len(t, records, 2)
 167
 168         quota := map[string]string{
 169                 "x-codex-primary-used-percent": "12",
 170                 "x-codex-plan-type":            "plus",
 171         }
 172         require.Equal(t, "gpt-test", records[0].Model)
 173         require.Equal(t, http.StatusOK, records[0].Status)
 174         require.Equal(t, 500, records[0].InputTokens)
 175         require.Equal(t, 400, records[0].CachedTokens)
 176         require.Equal(t, 7, records[0].OutputTokens)
 177         require.Equal(t, quota, records[0].Quota)
 178         require.Equal(t, &[2]int{2, 5}, records[0].Node)
 179         require.Equal(t, len("ok"), records[0].TextBytes)
 180         require.Greater(t, records[0].Seconds, 0.0)
 181         require.Empty(t, records[0].Error)
 182         require.NotEmpty(t, records[0].Date)
 183
 184         require.Nil(t, records[1].Node)
 185         require.Zero(t, records[1].TextBytes)
 186
 187         require.Equal(t, http.StatusTooManyRequests, records[1].Status)
 188         require.Equal(t, quota, records[1].Quota)
 189         require.Contains(t, records[1].Error, "HTTP 429")
 190
 191         data, err := os.ReadFile(client.usage.path)
 192         require.NoError(t, err)
 193         require.NotContains(t, string(data), modelTestKey)
 194
 195         info, err := os.Stat(client.usage.path)
 196         require.NoError(t, err)
 197         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 198 }
 199
 200 func TestClientRefusesCallsOverTheLimit(t *testing.T) {
 201         server := newModelFake(t)
 202         client := newTestModelClient(t, server)
 203         client.LimitCalls(2)
 204
 205         for range 2 {
 206                 require.False(t, client.LimitReached())
 207                 _, err := client.Create(t.Context(), modelTestRequest())
 208                 require.NoError(t, err)
 209         }
 210
 211         require.True(t, client.LimitReached())
 212         _, err := client.Create(t.Context(), modelTestRequest())
 213         require.ErrorIs(t, err, errModelCallLimit)
 214         require.Equal(t, 2, server.Calls())
 215         require.Len(t, modelReadUsage(t, client.usage.path), 2)
 216 }
 217
 218 func TestClientFailsWhenTheResponseFails(t *testing.T) {
 219         testCases := []struct {
 220                 name   string
 221                 events []any
 222                 want   string
 223         }{
 224                 {
 225                         name:   "failed response",
 226                         events: []any{map[string]any{"type": "response.failed", "response": map[string]any{"error": map[string]any{"code": "server_error", "message": "busy"}}}},
 227                         want:   "server_error: busy",
 228                 },
 229                 {
 230                         name:   "incomplete response",
 231                         events: []any{map[string]any{"type": "response.incomplete", "response": map[string]any{"incomplete_details": map[string]any{"reason": "content_filter"}}}},
 232                         want:   "content_filter",
 233                 },
 234                 {
 235                         name:   "error event",
 236                         events: []any{map[string]any{"type": "error", "code": "rate_limit", "message": "too many"}},
 237                         want:   "rate_limit: too many",
 238                 },
 239                 {
 240                         name:   "stream ends early",
 241                         events: []any{map[string]any{"type": "response.output_item.done", "item": modelMessageItem("cut")}},
 242                         want:   "ended before the response completed",
 243                 },
 244         }
 245
 246         for _, tc := range testCases {
 247                 t.Run(tc.name, func(t *testing.T) {
 248                         server := newModelFake(t)
 249                         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 250                                 modelWriteEvents(w, tc.events...)
 251                         }
 252
 253                         client := newTestModelClient(t, server)
 254                         _, err := client.Create(t.Context(), modelTestRequest())
 255                         require.ErrorContains(t, err, tc.want)
 256
 257                         records := modelReadUsage(t, client.usage.path)
 258                         require.Len(t, records, 1)
 259                         require.Contains(t, records[0].Error, tc.want)
 260                 })
 261         }
 262 }
 263
 264 // ModelFakeCall is one call to a fake endpoint, numbered from 0.
 265 type ModelFakeCall struct {
 266         number int
 267         header http.Header
 268         raw    []byte
 269         body   map[string]any
 270 }
 271
 272 // ModelFake is an httptest Responses endpoint. It keeps each call, and reply
 273 // writes the response. The default reply is one short message.
 274 type ModelFake struct {
 275         *httptest.Server
 276         mutex sync.Mutex
 277         calls []ModelFakeCall
 278         reply func(http.ResponseWriter, ModelFakeCall)
 279 }
 280
 281 // Call returns call number n.
 282 func (f *ModelFake) Call(n int) ModelFakeCall {
 283         f.mutex.Lock()
 284         defer f.mutex.Unlock()
 285
 286         return f.calls[n]
 287 }
 288
 289 // Calls returns the number of calls.
 290 func (f *ModelFake) Calls() int {
 291         f.mutex.Lock()
 292         defer f.mutex.Unlock()
 293
 294         return len(f.calls)
 295 }
 296
 297 func (f *ModelFake) serve(w http.ResponseWriter, r *http.Request) {
 298         raw, err := io.ReadAll(r.Body)
 299         if err != nil {
 300                 http.Error(w, err.Error(), http.StatusBadRequest)
 301                 return
 302         }
 303
 304         call := ModelFakeCall{
 305                 header: r.Header.Clone(),
 306                 raw:    raw,
 307         }
 308
 309         err = json.Unmarshal(raw, &call.body)
 310         if err != nil {
 311                 http.Error(w, err.Error(), http.StatusBadRequest)
 312                 return
 313         }
 314
 315         f.mutex.Lock()
 316         call.number = len(f.calls)
 317         f.calls = append(f.calls, call)
 318         reply := f.reply
 319         f.mutex.Unlock()
 320
 321         reply(w, call)
 322 }
 323
 324 func newModelFake(t *testing.T) *ModelFake {
 325         t.Helper()
 326
 327         f := &ModelFake{
 328                 reply: func(w http.ResponseWriter, call ModelFakeCall) {
 329                         modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
 330                 },
 331         }
 332
 333         f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
 334         t.Cleanup(f.Close)
 335
 336         return f
 337 }
 338
 339 func newTestModelClient(t *testing.T, server *ModelFake) *ModelClient {
 340         t.Helper()
 341
 342         return newModelClient(server.URL, ModelAPIKey(modelTestKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
 343 }
 344
 345 func modelTestRequest() ModelRequest {
 346         return ModelRequest{
 347                 Model:  "gpt-test",
 348                 Effort: "low",
 349                 Input:  []any{modelUserMessage("hello")},
 350         }
 351 }
 352
 353 // modelWriteEvents writes events as a server-sent event stream.
 354 func modelWriteEvents(w http.ResponseWriter, events ...any) {
 355         w.Header().Set("Content-Type", "text/event-stream")
 356         for _, event := range events {
 357                 data, err := json.Marshal(event)
 358                 if err != nil {
 359                         panic(err)
 360                 }
 361
 362                 _, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data)
 363                 if err != nil {
 364                         panic(err)
 365                 }
 366         }
 367 }
 368
 369 func modelMessageItem(text string) map[string]any {
 370         return map[string]any{
 371                 "type": "message",
 372                 "role": "assistant",
 373                 "content": []any{
 374                         map[string]any{"type": "output_text", "text": text},
 375                 },
 376         }
 377 }
 378
 379 func modelCompletedEvent(output []any, input, cached, out int) map[string]any {
 380         if output == nil {
 381                 output = []any{}
 382         }
 383
 384         return map[string]any{
 385                 "type": "response.completed",
 386                 "response": map[string]any{
 387                         "output": output,
 388                         "usage": map[string]any{
 389                                 "input_tokens":         input,
 390                                 "input_tokens_details": map[string]any{"cached_tokens": cached},
 391                                 "output_tokens":        out,
 392                         },
 393                 },
 394         }
 395 }
 396
 397 func modelReadUsage(t *testing.T, path string) []ModelUsageRecord {
 398         t.Helper()
 399
 400         file, err := os.Open(path)
 401         require.NoError(t, err)
 402
 403         defer func() {
 404                 require.NoError(t, file.Close())
 405         }()
 406
 407         var records []ModelUsageRecord
 408         scanner := bufio.NewScanner(file)
 409         for scanner.Scan() {
 410                 var record ModelUsageRecord
 411                 require.NoError(t, json.Unmarshal(scanner.Bytes(), &record))
 412                 records = append(records, record)
 413         }
 414
 415         require.NoError(t, scanner.Err())
 416
 417         return records
 418 }
 419
 420 func modelReasoningItem(id, encrypted string) map[string]any {
 421         return map[string]any{
 422                 "type":              "reasoning",
 423                 "id":                id,
 424                 "encrypted_content": encrypted,
 425                 "summary":           []any{},
 426         }
 427 }
 428
 429 func modelFunctionCallItem(id, callID, name, arguments string) map[string]any {
 430         return map[string]any{
 431                 "type":      "function_call",
 432                 "id":        id,
 433                 "call_id":   callID,
 434                 "name":      name,
 435                 "arguments": arguments,
 436         }
 437 }