package main
import (
"bufio"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"sync"
"testing"
"time"
"github.com/stretchr/testify/require"
)
// testKey is a fake API key.
const testKey = "sk-test-not-a-real-key"
func TestClientSendsTheFixedRequestFields(t *testing.T) {
server := newFakeResponses(t)
client := newTestClient(t, server)
_, err := client.Create(t.Context(), Request{
Model: "gpt-test",
Effort: "high",
Instructions: "You are a test.",
Tools: []Tool{{
Type: "function",
Name: "date",
Parameters: map[string]any{"type": "object"},
}},
Input: []any{userMessage("\n\n", "hello")},
})
require.NoError(t, err)
call := server.Call(0)
require.Equal(t, "Bearer "+testKey, call.header.Get("Authorization"))
require.Equal(t, "text/event-stream", call.header.Get("Accept"))
require.Equal(t, "application/json", call.header.Get("Content-Type"))
body := call.body
require.Equal(t, "gpt-test", body["model"])
require.Equal(t, "You are a test.", body["instructions"])
require.Equal(t, false, body["store"])
require.Equal(t, true, body["stream"])
require.Equal(t, []any{"reasoning.encrypted_content"}, body["include"])
require.Equal(t, "agent-key", body["prompt_cache_key"])
require.Equal(t, map[string]any{"effort": "high"}, body["reasoning"])
require.NotContains(t, body, "max_output_tokens")
require.NotContains(t, body, "metadata")
require.Equal(t, []any{map[string]any{
"role": "user",
"content": []any{
map[string]any{"type": "input_text", "text": "\n\n"},
map[string]any{"type": "input_text", "text": "hello"},
},
}}, body["input"])
require.Len(t, body["tools"], 1)
}
func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {
login := newFakeAuth(t)
auth := newTestAuth(t, login)
writeTestTokens(t, auth, login.token(time.Hour), "refresh-0")
testCases := []struct {
name string
auth authorizer
want string
}{
{
name: "ChatGPT backend",
auth: auth,
want: "agent-key",
},
{
name: "OpenAI API",
auth: apiKey(testKey),
want: "",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
server := newFakeResponses(t)
client := newClient(server.URL, tc.auth, "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
_, err := client.Create(t.Context(), testRequest())
require.NoError(t, err)
call := server.Call(0)
require.Equal(t, tc.want, call.header.Get("session_id"))
require.Equal(t, "agent-key", call.body["prompt_cache_key"])
})
}
}
func TestClientReadsTheStream(t *testing.T) {
server := newFakeResponses(t)
server.reply = func(w http.ResponseWriter, call fakeCall) {
writeEvents(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"}},
map[string]any{"type": "response.output_text.delta", "delta": "Hel"},
map[string]any{"type": "response.output_item.done", "item": messageItem("Hello")},
map[string]any{"type": "response.output_item.done", "item": messageItem(" world")},
completedEvent(nil, 1200, 1024, 30),
)
}
response, err := newTestClient(t, server).Create(t.Context(), testRequest())
require.NoError(t, err)
require.Equal(t, "Hello world", response.Text)
require.Len(t, response.Output, 3)
require.JSONEq(t, `{"type": "reasoning", "encrypted_content": "opaque"}`, string(response.Output[0]))
require.Equal(t, Usage{
Input: 1200,
Cached: 1024,
Output: 30,
}, response.Usage)
}
func TestClientTakesTheOutputOfTheCompletedResponse(t *testing.T) {
server := newFakeResponses(t)
server.reply = func(w http.ResponseWriter, call fakeCall) {
writeEvents(w, completedEvent([]any{messageItem("only here")}, 10, 0, 2))
}
response, err := newTestClient(t, server).Create(t.Context(), testRequest())
require.NoError(t, err)
require.Equal(t, "only here", response.Text)
}
func TestClientRecordsEachCall(t *testing.T) {
server := newFakeResponses(t)
server.reply = func(w http.ResponseWriter, call fakeCall) {
w.Header().Set("X-Codex-Primary-Used-Percent", "12")
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
}
writeEvents(w, completedEvent([]any{messageItem("ok")}, 500, 400, 7))
}
client := newTestClient(t, server)
request := testRequest()
request.Node = &nodeKey{2, 5}
_, 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]")
_, err = client.Create(t.Context(), testRequest())
require.ErrorContains(t, err, "HTTP 429")
require.ErrorContains(t, err, "slow down")
records := readUsage(t, client.usage.path)
require.Len(t, records, 2)
quota := map[string]string{
"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, 400, records[0].CachedTokens)
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, 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.Zero(t, records[1].TextBytes)
require.Equal(t, http.StatusTooManyRequests, records[1].Status)
require.Equal(t, quota, records[1].Quota)
require.Contains(t, records[1].Error, "HTTP 429")
data, err := os.ReadFile(client.usage.path)
require.NoError(t, err)
require.NotContains(t, string(data), testKey)
info, err := os.Stat(client.usage.path)
require.NoError(t, err)
require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
}
func TestClientFailsWhenTheResponseFails(t *testing.T) {
testCases := []struct {
name string
events []any
want string
}{
{
name: "failed response",
events: []any{map[string]any{"type": "response.failed", "response": map[string]any{"error": map[string]any{"code": "server_error", "message": "busy"}}}},
want: "server_error: busy",
},
{
name: "incomplete response",
events: []any{map[string]any{"type": "response.incomplete", "response": map[string]any{"incomplete_details": map[string]any{"reason": "content_filter"}}}},
want: "content_filter",
},
{
name: "error event",
events: []any{map[string]any{"type": "error", "code": "rate_limit", "message": "too many"}},
want: "rate_limit: too many",
},
{
name: "stream ends early",
events: []any{map[string]any{"type": "response.output_item.done", "item": messageItem("cut")}},
want: "ended before the response completed",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
server := newFakeResponses(t)
server.reply = func(w http.ResponseWriter, call fakeCall) {
writeEvents(w, tc.events...)
}
client := newTestClient(t, server)
_, err := client.Create(t.Context(), testRequest())
require.ErrorContains(t, err, tc.want)
records := readUsage(t, client.usage.path)
require.Len(t, records, 1)
require.Contains(t, records[0].Error, tc.want)
})
}
}
// fakeResponses is an httptest Responses endpoint. It keeps each call, and
// reply writes the response. The default reply is one short message.
type fakeResponses struct {
*httptest.Server
mutex sync.Mutex
calls []fakeCall
reply func(http.ResponseWriter, fakeCall)
}
// fakeCall is one call to a fake endpoint, numbered from 0.
type fakeCall struct {
number int
header http.Header
raw []byte
body map[string]any
}
func newFakeResponses(t *testing.T) *fakeResponses {
t.Helper()
f := &fakeResponses{
reply: func(w http.ResponseWriter, call fakeCall) {
writeEvents(w, completedEvent([]any{messageItem("a line")}, 100, 0, 5))
},
}
f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
t.Cleanup(f.Close)
return f
}
func (f *fakeResponses) 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 := fakeCall{
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)
}
// Call returns call number n.
func (f *fakeResponses) Call(n int) fakeCall {
f.mutex.Lock()
defer f.mutex.Unlock()
return f.calls[n]
}
// Calls returns the number of calls.
func (f *fakeResponses) Calls() int {
f.mutex.Lock()
defer f.mutex.Unlock()
return len(f.calls)
}
func newTestClient(t *testing.T, server *fakeResponses) *Client {
t.Helper()
return newClient(server.URL, apiKey(testKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
}
func testRequest() Request {
return Request{
Model: "gpt-test",
Effort: "low",
Input: []any{userMessage("hello")},
}
}
// writeEvents writes events as a server-sent event stream.
func writeEvents(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 messageItem(text string) map[string]any {
return map[string]any{
"type": "message",
"role": "assistant",
"content": []any{
map[string]any{"type": "output_text", "text": text},
},
}
}
func completedEvent(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 readUsage(t *testing.T, path string) []usageRecord {
t.Helper()
file, err := os.Open(path)
require.NoError(t, err)
defer func() {
require.NoError(t, file.Close())
}()
var records []usageRecord
scanner := bufio.NewScanner(file)
for scanner.Scan() {
var record usageRecord
require.NoError(t, json.Unmarshal(scanner.Bytes(), &record))
records = append(records, record)
}
require.NoError(t, scanner.Err())
return records
}