package main
import (
- "bufio"
- "encoding/json"
- "fmt"
"io"
"net/http"
- "net/http/httptest"
"os"
"path/filepath"
- "sync"
"testing"
"time"
"github.com/stretchr/testify/require"
)
-// modelTestKey is a fake API key.
-const modelTestKey = "sk-test-not-a-real-key"
-
func TestClientSendsTheFixedRequestFields(t *testing.T) {
server := newModelFake(t)
client := newTestModelClient(t, server)
call := server.Call(0)
require.Equal(t, tc.want, call.header.Get("session_id"))
require.Equal(t, "agent-key", call.body["prompt_cache_key"])
+
+ // A subagent has its own key.
+ request := modelTestRequest()
+ request.Agent = "GreenOtter"
+ _, err = client.Create(t.Context(), request)
+ require.NoError(t, err)
+
+ call = server.Call(1)
+ if tc.want != "" {
+ require.Equal(t, "agent-key-GreenOtter", call.header.Get("session_id"))
+ }
+
+ require.Equal(t, "agent-key-GreenOtter", call.body["prompt_cache_key"])
})
}
}
func TestClientReadsTheStream(t *testing.T) {
server := newModelFake(t)
- server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
+ server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
modelWriteEvents(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"}},
func TestClientTakesTheOutputOfTheCompletedResponse(t *testing.T) {
server := newModelFake(t)
- server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
+ server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("only here")}, 10, 0, 2))
}
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
client := newTestModelClient(t, server)
request := modelTestRequest()
request.Node = &NodeKey{2, 5}
+ request.Agent = "GreenOtter"
_, err := client.Create(t.Context(), request)
require.NoError(t, err)
require.NotContains(t, server.Call(0).body, "node")
require.NotContains(t, string(server.Call(0).raw), "[2,5]")
+ require.NotContains(t, server.Call(0).body, "agent")
_, err = client.Create(t.Context(), modelTestRequest())
require.ErrorContains(t, err, "HTTP 429")
"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, 7, records[0].OutputTokens)
require.Equal(t, quota, records[0].Quota)
require.Equal(t, &[2]int{2, 5}, records[0].Node)
+ require.Equal(t, "GreenOtter", records[0].Agent)
require.Equal(t, len("ok"), records[0].TextBytes)
require.Greater(t, records[0].Seconds, 0.0)
require.Empty(t, records[0].Error)
require.NotEmpty(t, records[0].Date)
require.Nil(t, records[1].Node)
+ require.Empty(t, records[1].Agent)
require.Zero(t, records[1].TextBytes)
require.Equal(t, http.StatusTooManyRequests, records[1].Status)
for range 2 {
require.False(t, client.LimitReached())
+
_, err := client.Create(t.Context(), modelTestRequest())
require.NoError(t, err)
}
require.True(t, client.LimitReached())
+
_, err := client.Create(t.Context(), modelTestRequest())
require.ErrorIs(t, err, errModelCallLimit)
require.Equal(t, 2, server.Calls())
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
server := newModelFake(t)
- server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
+ server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
modelWriteEvents(w, tc.events...)
}
}
}
-// ModelFakeCall is one call to a fake endpoint, numbered from 0.
-type ModelFakeCall struct {
- number int
- header http.Header
- raw []byte
- body map[string]any
-}
-
-// ModelFake is an httptest Responses endpoint. It keeps each call, and reply
-// writes the response. The default reply is one short message.
-type ModelFake struct {
- *httptest.Server
- mutex sync.Mutex
- calls []ModelFakeCall
- reply func(http.ResponseWriter, ModelFakeCall)
-}
-
-// Call returns call number n.
-func (f *ModelFake) Call(n int) ModelFakeCall {
- f.mutex.Lock()
- defer f.mutex.Unlock()
-
- return f.calls[n]
-}
-
-// Calls returns the number of calls.
-func (f *ModelFake) Calls() int {
- f.mutex.Lock()
- defer f.mutex.Unlock()
-
- return len(f.calls)
-}
-
-func (f *ModelFake) 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 := ModelFakeCall{
- 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)
-}
-
-func newModelFake(t *testing.T) *ModelFake {
- t.Helper()
-
- f := &ModelFake{
- reply: func(w http.ResponseWriter, call ModelFakeCall) {
- modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
- },
- }
-
- f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
- t.Cleanup(f.Close)
-
- return f
-}
-
func newTestModelClient(t *testing.T, server *ModelFake) *ModelClient {
t.Helper()
return newModelClient(server.URL, ModelAPIKey(modelTestKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
}
-
-func modelTestRequest() ModelRequest {
- return ModelRequest{
- Model: "gpt-test",
- Effort: "low",
- Input: []any{modelUserMessage("hello")},
- }
-}
-
-// modelWriteEvents writes events as a server-sent event stream.
-func modelWriteEvents(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 modelMessageItem(text string) map[string]any {
- return map[string]any{
- "type": "message",
- "role": "assistant",
- "content": []any{
- map[string]any{"type": "output_text", "text": text},
- },
- }
-}
-
-func modelCompletedEvent(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 modelReadUsage(t *testing.T, path string) []ModelUsageRecord {
- t.Helper()
-
- file, err := os.Open(path)
- require.NoError(t, err)
-
- defer func() {
- require.NoError(t, file.Close())
- }()
-
- var records []ModelUsageRecord
- scanner := bufio.NewScanner(file)
- for scanner.Scan() {
- var record ModelUsageRecord
- require.NoError(t, json.Unmarshal(scanner.Bytes(), &record))
- records = append(records, record)
- }
-
- require.NoError(t, scanner.Err())
-
- return records
-}
-
-func modelReasoningItem(id, encrypted string) map[string]any {
- return map[string]any{
- "type": "reasoning",
- "id": id,
- "encrypted_content": encrypted,
- "summary": []any{},
- }
-}
-
-func modelFunctionCallItem(id, callID, name, arguments string) map[string]any {
- return map[string]any{
- "type": "function_call",
- "id": id,
- "call_id": callID,
- "name": name,
- "arguments": arguments,
- }
-}
projects / onechat / blobdiff