[onechat] / responses.go

package main

import (
        "bufio"
        "bytes"
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "net/http"
        "os"
        "path/filepath"
        "strings"
        "sync"
        "time"
)

const (
        // chatgptURL is the Responses endpoint of the ChatGPT backend.
        chatgptURL = "https://chatgpt.com/backend-api/codex/responses"

        // apiURL is the Responses endpoint of the OpenAI API.
        apiURL = "https://api.openai.com/v1/responses"

        // 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"`
}

// modelEndpoint returns the URL and the authorizer of the endpoint in config.
// The ChatGPT backend uses the token file in dir. The OpenAI API uses the key in
// OPENAI_API_KEY, and there is no fallback from one to the other. All clients
// share the authorizer, so one refresh runs at a time.
func modelEndpoint(config Config, dir string) (string, authorizer, error) {
        if config.Endpoint == endpointAPI {
                key := os.Getenv("OPENAI_API_KEY")
                if key == "" {
                        return "", nil, errors.New("endpoint api needs OPENAI_API_KEY")
                }

                return apiURL, apiKey(key), nil
        }

        return chatgptURL, newChatGPTAuth(filepath.Join(dir, tokenFileName), authIssuer), nil
}

// 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
}