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

   1 package main
   2
   3 import (
   4         "io"
   5         "net/http"
   6         "os"
   7         "path/filepath"
   8         "testing"
   9         "time"
  10
  11         "github.com/stretchr/testify/require"
  12 )
  13
  14 func TestClientSendsTheFixedRequestFields(t *testing.T) {
  15         server := newModelFake(t)
  16         client := newTestModelClient(t, server)
  17
  18         _, err := client.Create(t.Context(), ModelRequest{
  19                 Model:        "gpt-test",
  20                 Effort:       "high",
  21                 Instructions: "You are a test.",
  22                 Tools: []ModelTool{{
  23                         Type:       "function",
  24                         Name:       "date",
  25                         Parameters: map[string]any{"type": "object"},
  26                 }},
  27                 Input: []any{modelUserMessage("<chat>\n</chat>\n", "hello")},
  28         })
  29         require.NoError(t, err)
  30
  31         call := server.Call(0)
  32         require.Equal(t, "Bearer "+modelTestKey, call.header.Get("Authorization"))
  33         require.Equal(t, "text/event-stream", call.header.Get("Accept"))
  34         require.Equal(t, "application/json", call.header.Get("Content-Type"))
  35
  36         body := call.body
  37         require.Equal(t, "gpt-test", body["model"])
  38         require.Equal(t, "You are a test.", body["instructions"])
  39         require.Equal(t, false, body["store"])
  40         require.Equal(t, true, body["stream"])
  41         require.Equal(t, []any{"reasoning.encrypted_content"}, body["include"])
  42         require.Equal(t, "agent-key", body["prompt_cache_key"])
  43         require.Equal(t, map[string]any{"effort": "high"}, body["reasoning"])
  44         require.NotContains(t, body, "max_output_tokens")
  45         require.NotContains(t, body, "metadata")
  46         require.Equal(t, []any{map[string]any{
  47                 "role": "user",
  48                 "content": []any{
  49                         map[string]any{"type": "input_text", "text": "<chat>\n</chat>\n"},
  50                         map[string]any{"type": "input_text", "text": "hello"},
  51                 },
  52         }}, body["input"])
  53
  54         require.Len(t, body["tools"], 1)
  55 }
  56
  57 func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {
  58         login := newChatGPTFake(t)
  59         auth := newTestChatGPTAuth(t, login)
  60         chatgptWriteTestTokens(t, auth, login.token(time.Hour), "refresh-0")
  61
  62         testCases := []struct {
  63                 name string
  64                 auth ModelAuthorizer
  65                 want string
  66         }{
  67                 {
  68                         name: "sends the session header to ChatGPT",
  69                         auth: auth,
  70                         want: "agent-key",
  71                 },
  72                 {
  73                         name: "omits the session header for the API",
  74                         auth: ModelAPIKey(modelTestKey),
  75                         want: "",
  76                 },
  77         }
  78
  79         for _, tc := range testCases {
  80                 t.Run(tc.name, func(t *testing.T) {
  81                         server := newModelFake(t)
  82                         client := newModelClient(server.URL, tc.auth, "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
  83                         _, err := client.Create(t.Context(), modelTestRequest())
  84                         require.NoError(t, err)
  85
  86                         call := server.Call(0)
  87                         require.Equal(t, tc.want, call.header.Get("session_id"))
  88                         require.Equal(t, "agent-key", call.body["prompt_cache_key"])
  89
  90                         // A subagent has its own key.
  91                         request := modelTestRequest()
  92                         request.Agent = "GreenOtter"
  93                         _, err = client.Create(t.Context(), request)
  94                         require.NoError(t, err)
  95
  96                         call = server.Call(1)
  97                         if tc.want != "" {
  98                                 require.Equal(t, "agent-key-GreenOtter", call.header.Get("session_id"))
  99                         }
 100
 101                         require.Equal(t, "agent-key-GreenOtter", call.body["prompt_cache_key"])
 102                 })
 103         }
 104 }
 105
 106 func TestClientReadsTheStream(t *testing.T) {
 107         server := newModelFake(t)
 108         server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
 109                 modelWriteEvents(w,
 110                         map[string]any{"type": "response.created", "response": map[string]any{}},
 111                         map[string]any{"type": "response.output_item.done", "item": map[string]any{"type": "reasoning", "encrypted_content": "opaque"}},
 112                         map[string]any{"type": "response.output_text.delta", "delta": "Hel"},
 113                         map[string]any{"type": "response.output_item.done", "item": modelMessageItem("Hello")},
 114                         map[string]any{"type": "response.output_item.done", "item": modelMessageItem(" world")},
 115                         modelCompletedEvent(nil, 1200, 1024, 30),
 116                 )
 117         }
 118
 119         response, err := newTestModelClient(t, server).Create(t.Context(), modelTestRequest())
 120         require.NoError(t, err)
 121         require.Equal(t, "Hello world", response.Text)
 122         require.Len(t, response.Output, 3)
 123         require.JSONEq(t, `{"type": "reasoning", "encrypted_content": "opaque"}`, string(response.Output[0]))
 124         require.Equal(t, ModelUsage{
 125                 Input:  1200,
 126                 Cached: 1024,
 127                 Output: 30,
 128         }, response.Usage)
 129 }
 130
 131 func TestClientTakesTheOutputOfTheCompletedResponse(t *testing.T) {
 132         server := newModelFake(t)
 133         server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
 134                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("only here")}, 10, 0, 2))
 135         }
 136
 137         response, err := newTestModelClient(t, server).Create(t.Context(), modelTestRequest())
 138         require.NoError(t, err)
 139         require.Equal(t, "only here", response.Text)
 140 }
 141
 142 func TestClientRecordsEachCall(t *testing.T) {
 143         server := newModelFake(t)
 144         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 145                 w.Header().Set("X-Codex-Primary-Used-Percent", "12")
 146                 w.Header().Set("X-Codex-Plan-Type", "plus")
 147                 w.Header().Set("X-Other", "not quota")
 148                 w.Header().Set("X-Codex-Turn-State", "opaque state")
 149
 150                 if call.number == 1 {
 151                         w.WriteHeader(http.StatusTooManyRequests)
 152
 153                         _, err := io.WriteString(w, `{"detail": "slow down"}`)
 154                         if err != nil {
 155                                 panic(err)
 156                         }
 157
 158                         return
 159                 }
 160
 161                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("ok")}, 500, 400, 7))
 162         }
 163
 164         client := newTestModelClient(t, server)
 165         request := modelTestRequest()
 166         request.Node = &NodeKey{2, 5}
 167         request.Agent = "GreenOtter"
 168         _, err := client.Create(t.Context(), request)
 169         require.NoError(t, err)
 170         require.NotContains(t, server.Call(0).body, "node")
 171         require.NotContains(t, string(server.Call(0).raw), "[2,5]")
 172         require.NotContains(t, server.Call(0).body, "agent")
 173
 174         _, err = client.Create(t.Context(), modelTestRequest())
 175         require.ErrorContains(t, err, "HTTP 429")
 176         require.ErrorContains(t, err, "slow down")
 177
 178         records := modelReadUsage(t, client.usage.path)
 179         require.Len(t, records, 2)
 180
 181         quota := map[string]string{
 182                 "x-codex-primary-used-percent": "12",
 183                 "x-codex-plan-type":            "plus",
 184         }
 185
 186         require.Equal(t, "gpt-test", records[0].Model)
 187         require.Equal(t, http.StatusOK, records[0].Status)
 188         require.Equal(t, 500, records[0].InputTokens)
 189         require.Equal(t, 400, records[0].CachedTokens)
 190         require.Equal(t, 7, records[0].OutputTokens)
 191         require.Equal(t, quota, records[0].Quota)
 192         require.Equal(t, &[2]int{2, 5}, records[0].Node)
 193         require.Equal(t, "GreenOtter", records[0].Agent)
 194         require.Equal(t, len("ok"), records[0].TextBytes)
 195         require.Greater(t, records[0].Seconds, 0.0)
 196         require.Empty(t, records[0].Error)
 197         require.NotEmpty(t, records[0].Date)
 198
 199         require.Nil(t, records[1].Node)
 200         require.Empty(t, records[1].Agent)
 201         require.Zero(t, records[1].TextBytes)
 202
 203         require.Equal(t, http.StatusTooManyRequests, records[1].Status)
 204         require.Equal(t, quota, records[1].Quota)
 205         require.Contains(t, records[1].Error, "HTTP 429")
 206
 207         data, err := os.ReadFile(client.usage.path)
 208         require.NoError(t, err)
 209         require.NotContains(t, string(data), modelTestKey)
 210
 211         info, err := os.Stat(client.usage.path)
 212         require.NoError(t, err)
 213         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 214 }
 215
 216 func TestClientRefusesCallsOverTheLimit(t *testing.T) {
 217         server := newModelFake(t)
 218         client := newTestModelClient(t, server)
 219         client.LimitCalls(2)
 220
 221         for range 2 {
 222                 require.False(t, client.LimitReached())
 223
 224                 _, err := client.Create(t.Context(), modelTestRequest())
 225                 require.NoError(t, err)
 226         }
 227
 228         require.True(t, client.LimitReached())
 229
 230         _, err := client.Create(t.Context(), modelTestRequest())
 231         require.ErrorIs(t, err, errModelCallLimit)
 232         require.Equal(t, 2, server.Calls())
 233         require.Len(t, modelReadUsage(t, client.usage.path), 2)
 234 }
 235
 236 func TestClientFailsWhenTheResponseFails(t *testing.T) {
 237         testCases := []struct {
 238                 name   string
 239                 events []any
 240                 want   string
 241         }{
 242                 {
 243                         name:   "reports the response failure cause",
 244                         events: []any{map[string]any{"type": "response.failed", "response": map[string]any{"error": map[string]any{"code": "server_error", "message": "busy"}}}},
 245                         want:   "server_error: busy",
 246                 },
 247                 {
 248                         name:   "reports why the response is incomplete",
 249                         events: []any{map[string]any{"type": "response.incomplete", "response": map[string]any{"incomplete_details": map[string]any{"reason": "content_filter"}}}},
 250                         want:   "content_filter",
 251                 },
 252                 {
 253                         name:   "reports the stream error cause",
 254                         events: []any{map[string]any{"type": "error", "code": "rate_limit", "message": "too many"}},
 255                         want:   "rate_limit: too many",
 256                 },
 257                 {
 258                         name:   "refuses an unfinished response stream",
 259                         events: []any{map[string]any{"type": "response.output_item.done", "item": modelMessageItem("cut")}},
 260                         want:   "ended before the response completed",
 261                 },
 262         }
 263
 264         for _, tc := range testCases {
 265                 t.Run(tc.name, func(t *testing.T) {
 266                         server := newModelFake(t)
 267                         server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
 268                                 modelWriteEvents(w, tc.events...)
 269                         }
 270
 271                         client := newTestModelClient(t, server)
 272                         _, err := client.Create(t.Context(), modelTestRequest())
 273                         require.ErrorContains(t, err, tc.want)
 274
 275                         records := modelReadUsage(t, client.usage.path)
 276                         require.Len(t, records, 1)
 277                         require.Contains(t, records[0].Error, tc.want)
 278                 })
 279         }
 280 }
 281
 282 func newTestModelClient(t *testing.T, server *ModelFake) *ModelClient {
 283         t.Helper()
 284
 285         return newModelClient(server.URL, ModelAPIKey(modelTestKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
 286 }