--- /dev/null
+package main
+
+import (
+ "bufio"
+ "bytes"
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "io"
+ "net/http"
+ "os"
+ "strings"
+ "sync"
+ "time"
+)
+
+const (
+ // quotaPrefix starts the names of the quota headers of the ChatGPT
+ // backend, in canonical form.
+ quotaPrefix = "X-Codex-"
+
+ // errorBytes is the most bytes of an error response that go into the
+ // error.
+ errorBytes = 4096
+)
+
+// authorizer adds the credentials of a model call to the headers of its
+// request.
+type authorizer interface {
+ Authorize(ctx context.Context, header http.Header) error
+}
+
+// apiKey authorizes calls to the OpenAI API with a key.
+type apiKey string
+
+// Client sends model calls to a Responses endpoint for one agent. The ChatGPT
+// backend and the OpenAI API take the same request body and differ only in the
+// URL and the credentials. Every call sends the fixed cache key of the agent,
+// so the turns and the compactions of the agent share one prompt cache. Every
+// call is recorded in the usage file.
+//
+// The request never holds max_output_tokens or metadata, because the ChatGPT
+// backend refuses them. It sets store to false, so the client sends the whole
+// input on each call, and it asks for the encrypted reasoning, so a later call
+// can send the reasoning items back.
+type Client struct {
+ url string
+ auth authorizer
+ cacheKey string
+ usage *usageLog
+ http *http.Client
+ clock func() time.Time
+}
+
+// Request is one model call. Input holds the input items in order, such as
+// userMessage values or output items of an earlier response.
+type Request struct {
+ Model string
+ Effort string
+ Instructions string
+ Tools []Tool
+ Input []any
+}
+
+// Tool is the definition of a function tool.
+type Tool struct {
+ Type string `json:"type"`
+ Name string `json:"name"`
+ Description string `json:"description"`
+ Parameters any `json:"parameters"`
+ Strict bool `json:"strict"`
+}
+
+// Response is the result of a model call. Output holds the output items as the
+// endpoint sent them. Text is the text of the message items.
+type Response struct {
+ Output []json.RawMessage
+ Text string
+ Usage Usage
+}
+
+// Usage gives the token counts of a call. Cached is the part of the input
+// tokens that came from the prompt cache.
+type Usage struct {
+ Input int
+ Cached int
+ Output int
+}
+
+// requestBody is the JSON body of a call.
+type requestBody struct {
+ Model string `json:"model"`
+ Instructions string `json:"instructions"`
+ Input []any `json:"input"`
+ Tools []Tool `json:"tools"`
+ Reasoning reasoning `json:"reasoning"`
+ Store bool `json:"store"`
+ Stream bool `json:"stream"`
+ Include []string `json:"include"`
+ PromptCacheKey string `json:"prompt_cache_key"`
+}
+
+type reasoning struct {
+ Effort string `json:"effort"`
+}
+
+// streamEvent holds the fields of the stream events that the client reads.
+type streamEvent struct {
+ Type string `json:"type"`
+ Item json.RawMessage `json:"item"`
+ Response *eventResponse `json:"response"`
+ Code string `json:"code"`
+ Message string `json:"message"`
+}
+
+type eventResponse struct {
+ Output []json.RawMessage `json:"output"`
+ Usage *wireUsage `json:"usage"`
+ Error *wireError `json:"error"`
+ IncompleteDetails *struct {
+ Reason string `json:"reason"`
+ } `json:"incomplete_details"`
+}
+
+type wireUsage struct {
+ InputTokens int `json:"input_tokens"`
+ OutputTokens int `json:"output_tokens"`
+ InputTokensDetails struct {
+ CachedTokens int `json:"cached_tokens"`
+ } `json:"input_tokens_details"`
+}
+
+type wireError struct {
+ Code string `json:"code"`
+ Message string `json:"message"`
+}
+
+// outputItem holds the fields of an output item that give its text.
+type outputItem struct {
+ Type string `json:"type"`
+ Content []struct {
+ Type string `json:"type"`
+ Text string `json:"text"`
+ } `json:"content"`
+}
+
+// newClient returns a client for the Responses endpoint at url. It records
+// each call in the usage file at usagePath.
+func newClient(url string, auth authorizer, cacheKey, usagePath string) *Client {
+ return &Client{
+ url: url,
+ auth: auth,
+ cacheKey: cacheKey,
+ usage: &usageLog{
+ path: usagePath,
+ },
+ http: &http.Client{},
+ clock: time.Now,
+ }
+}
+
+// Create sends the call and reads its event stream to the end.
+func (c *Client) Create(ctx context.Context, request Request) (Response, error) {
+ body, err := json.Marshal(requestBody{
+ Model: request.Model,
+ Instructions: request.Instructions,
+ Input: request.Input,
+ Tools: request.Tools,
+ Reasoning: reasoning{
+ Effort: request.Effort,
+ },
+ Store: false,
+ Stream: true,
+ Include: []string{"reasoning.encrypted_content"},
+ PromptCacheKey: c.cacheKey,
+ })
+ if err != nil {
+ return Response{}, fmt.Errorf("encode request: %w", err)
+ }
+
+ httpRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
+ if err != nil {
+ return Response{}, fmt.Errorf("create request: %w", err)
+ }
+
+ httpRequest.Header.Set("Content-Type", "application/json")
+ httpRequest.Header.Set("Accept", "text/event-stream")
+ err = c.auth.Authorize(ctx, httpRequest.Header)
+ if err != nil {
+ return Response{}, fmt.Errorf("authorize model call: %w", err)
+ }
+
+ record := usageRecord{
+ Date: c.clock().Format(time.RFC3339Nano),
+ Model: request.Model,
+ }
+
+ response, err := c.send(httpRequest, &record)
+ if err != nil {
+ record.Error = err.Error()
+ }
+
+ record.InputTokens = response.Usage.Input
+ record.CachedTokens = response.Usage.Cached
+ record.OutputTokens = response.Usage.Output
+
+ return response, errors.Join(err, c.usage.Write(record))
+}
+
+// send makes the HTTP call and reads its stream. It puts the status and the
+// quota headers into record.
+func (c *Client) send(httpRequest *http.Request, record *usageRecord) (response Response, err error) {
+ httpResponse, err := c.http.Do(httpRequest)
+ if err != nil {
+ return Response{}, fmt.Errorf("send model call: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, httpResponse.Body.Close())
+ }()
+
+ record.Status = httpResponse.StatusCode
+ record.Quota = quotaHeaders(httpResponse.Header)
+ if httpResponse.StatusCode != http.StatusOK {
+ text, readErr := io.ReadAll(io.LimitReader(httpResponse.Body, errorBytes))
+ err = fmt.Errorf("model call: HTTP %d: %s", httpResponse.StatusCode, strings.TrimSpace(string(text)))
+
+ return Response{}, errors.Join(err, readErr)
+ }
+
+ return readStream(httpResponse.Body)
+}
+
+// usageLog appends one JSON line for each model call to usage.jsonl. Calls can
+// run at once, so a mutex serializes the writes.
+type usageLog struct {
+ mutex sync.Mutex
+ path string
+}
+
+// usageRecord is one line of usage.jsonl. Status is the HTTP status, or 0 when
+// no response came. Quota holds the x-codex-* headers of the response.
+type usageRecord struct {
+ Date string `json:"date"`
+ Model string `json:"model"`
+ Status int `json:"status"`
+ InputTokens int `json:"input_tokens"`
+ CachedTokens int `json:"cached_tokens"`
+ OutputTokens int `json:"output_tokens"`
+ Quota map[string]string `json:"quota,omitempty"`
+ Error string `json:"error,omitempty"`
+}
+
+// Write appends record to the file with one write.
+func (u *usageLog) Write(record usageRecord) (err error) {
+ line, err := json.Marshal(record)
+ if err != nil {
+ return fmt.Errorf("encode usage: %w", err)
+ }
+
+ u.mutex.Lock()
+ defer u.mutex.Unlock()
+
+ file, err := os.OpenFile(u.path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
+ if err != nil {
+ return fmt.Errorf("open usage file: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, file.Close())
+ }()
+
+ _, err = file.Write(append(line, '\n'))
+ if err != nil {
+ return fmt.Errorf("write usage: %w", err)
+ }
+
+ return nil
+}
+
+// Authorize sets the key as the bearer token.
+func (k apiKey) Authorize(ctx context.Context, header http.Header) error {
+ header.Set("Authorization", "Bearer "+string(k))
+ return ctx.Err()
+}
+
+// userMessage returns a user input item with one text part for each text.
+func userMessage(texts ...string) any {
+ type part struct {
+ Type string `json:"type"`
+ Text string `json:"text"`
+ }
+
+ content := make([]part, len(texts))
+ for n, text := range texts {
+ content[n] = part{
+ Type: "input_text",
+ Text: text,
+ }
+ }
+
+ return struct {
+ Role string `json:"role"`
+ Content []part `json:"content"`
+ }{
+ Role: "user",
+ Content: content,
+ }
+}
+
+// quotaHeaders returns the x-codex-* headers, with lowercase names.
+func quotaHeaders(header http.Header) map[string]string {
+ var quota map[string]string
+ for name, values := range header {
+ if !strings.HasPrefix(name, quotaPrefix) || len(values) == 0 {
+ continue
+ }
+
+ if quota == nil {
+ quota = make(map[string]string)
+ }
+
+ quota[strings.ToLower(name)] = values[0]
+ }
+
+ return quota
+}
+
+// readStream reads the server-sent events of a call until the response
+// completes. The output items come from the output_item.done events, or from
+// the completed response when the stream sent none.
+func readStream(stream io.Reader) (Response, error) {
+ var response Response
+ reader := bufio.NewReader(stream)
+ var data bytes.Buffer
+ for {
+ line, err := reader.ReadBytes('\n')
+ if err != nil && !errors.Is(err, io.EOF) {
+ return response, fmt.Errorf("read event stream: %w", err)
+ }
+
+ line = bytes.TrimRight(line, "\r\n")
+ if len(line) > 0 {
+ value, found := bytes.CutPrefix(line, []byte("data:"))
+ if found {
+ data.Write(bytes.TrimPrefix(value, []byte(" ")))
+ data.WriteByte('\n')
+ }
+ } else if data.Len() > 0 {
+ done, eventErr := readEvent(&response, data.Bytes())
+ data.Reset()
+ if eventErr != nil || done {
+ return response, eventErr
+ }
+ }
+
+ if errors.Is(err, io.EOF) {
+ return response, errors.New("event stream ended before the response completed")
+ }
+ }
+}
+
+// readEvent adds one stream event to response, and tells if it ends the
+// stream.
+func readEvent(response *Response, data []byte) (bool, error) {
+ var event streamEvent
+ err := json.Unmarshal(data, &event)
+ if err != nil {
+ return true, fmt.Errorf("decode stream event: %w", err)
+ }
+
+ switch event.Type {
+ case "response.output_item.done":
+ response.Output = append(response.Output, event.Item)
+ case "response.completed":
+ if event.Response == nil {
+ return true, errors.New("completed event without a response")
+ }
+
+ if len(response.Output) == 0 {
+ response.Output = event.Response.Output
+ }
+
+ if event.Response.Usage != nil {
+ response.Usage = Usage{
+ Input: event.Response.Usage.InputTokens,
+ Cached: event.Response.Usage.InputTokensDetails.CachedTokens,
+ Output: event.Response.Usage.OutputTokens,
+ }
+ }
+
+ response.Text, err = outputText(response.Output)
+
+ return true, err
+ case "response.incomplete":
+ reason := "unknown"
+ if event.Response != nil && event.Response.IncompleteDetails != nil {
+ reason = event.Response.IncompleteDetails.Reason
+ }
+
+ return true, fmt.Errorf("response incomplete: %s", reason)
+ case "response.failed":
+ message := "no details"
+ if event.Response != nil && event.Response.Error != nil {
+ message = event.Response.Error.Code + ": " + event.Response.Error.Message
+ }
+
+ return true, fmt.Errorf("response failed: %s", message)
+ case "error":
+ return true, fmt.Errorf("stream error: %s: %s", event.Code, event.Message)
+ }
+
+ return false, nil
+}
+
+// outputText returns the text of the message items of output.
+func outputText(output []json.RawMessage) (string, error) {
+ var text strings.Builder
+ for _, raw := range output {
+ var item outputItem
+ err := json.Unmarshal(raw, &item)
+ if err != nil {
+ return "", fmt.Errorf("decode output item: %w", err)
+ }
+
+ if item.Type != "message" {
+ continue
+ }
+
+ for _, part := range item.Content {
+ if part.Type == "output_text" {
+ text.WriteString(part.Text)
+ }
+ }
+ }
+
+ return text.String(), nil
+}
--- /dev/null
+package main
+
+import (
+ "bufio"
+ "encoding/json"
+ "fmt"
+ "io"
+ "net/http"
+ "net/http/httptest"
+ "os"
+ "path/filepath"
+ "sync"
+ "testing"
+
+ "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("<chat>\n</chat>\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": "<chat>\n</chat>\n"},
+ map[string]any{"type": "input_text", "text": "hello"},
+ },
+ }}, body["input"])
+ require.Len(t, body["tools"], 1)
+}
+
+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")
+ 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)
+ _, err := client.Create(t.Context(), testRequest())
+ require.NoError(t, err)
+
+ _, 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.Empty(t, records[0].Error)
+ require.NotEmpty(t, records[0].Date)
+
+ 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
+}
projects / onechat / commitdiff