"io"
"net/http"
"os"
+ "path/filepath"
"strings"
"sync"
+ "sync/atomic"
"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-"
+ // turnStateHeader is an x-codex-* header that holds an opaque state of
+ // the backend, not quota.
+ turnStateHeader = "X-Codex-Turn-State"
+
// errorBytes is the most bytes of an error response that go into the
// error.
errorBytes = 4096
)
+// errCallLimit tells that a client sent its limit of calls.
+var errCallLimit = errors.New("call limit reached")
+
// authorizer adds the credentials of a model call to the headers of its
// request.
type authorizer interface {
// 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.
+// so the turns and the compactions of the agent share one prompt cache. The
+// ChatGPT backend caches only with the key in the session_id header, so a call
+// to it sends the key there and in the body. 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
cacheKey string
usage *usageLog
http *http.Client
- clock func() time.Time
+ limit int64
+ calls atomic.Int64
}
// Request is one model call. Input holds the input items in order, such as
-// userMessage values or output items of an earlier response.
+// userMessage values or output items of an earlier response. Node names the
+// node that a compaction builds, or is nil. Turn names the request of a turn:
+// first, next, or warmup. Node and Turn go only to the usage file, never to the
+// endpoint. A call with StopAtOutput ends at the first output item, after the
+// endpoint has read the whole input, and gives no output and no usage.
type Request struct {
Model string
Effort string
Instructions string
Tools []Tool
Input []any
+ Node *nodeKey
+ Turn string
+ StopAtOutput bool
}
// Tool is the definition of a function tool.
} `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 {
usage: &usageLog{
path: usagePath,
},
- http: &http.Client{},
- clock: time.Now,
+ http: &http.Client{},
}
}
-// Create sends the call and reads its event stream to the end.
+// LimitCalls makes the client refuse each call after the first limit calls.
+func (c *Client) LimitCalls(limit int64) {
+ c.limit = limit
+}
+
+// LimitReached tells if the client refuses further calls.
+func (c *Client) LimitReached() bool {
+ return c.limit > 0 && c.calls.Load() >= c.limit
+}
+
+// Create sends the call and reads its event stream to the end. Over the call
+// limit it sends nothing and returns errCallLimit.
func (c *Client) Create(ctx context.Context, request Request) (Response, error) {
+ if c.calls.Add(1) > c.limit && c.limit > 0 {
+ return Response{}, errCallLimit
+ }
+
body, err := json.Marshal(requestBody{
Model: request.Model,
Instructions: request.Instructions,
return Response{}, fmt.Errorf("authorize model call: %w", err)
}
+ _, chatgpt := c.auth.(*chatgptAuth)
+ if chatgpt {
+ httpRequest.Header.Set("session_id", c.cacheKey)
+ }
+
+ start := time.Now()
record := usageRecord{
- Date: c.clock().Format(time.RFC3339Nano),
+ Date: start.Format(time.RFC3339Nano),
Model: request.Model,
}
- response, err := c.send(httpRequest, &record)
+ if request.Node != nil {
+ record.Node = &[2]int{request.Node.level, request.Node.index}
+ }
+
+ record.Turn = request.Turn
+ response, err := c.send(httpRequest, &record, request.StopAtOutput)
if err != nil {
record.Error = err.Error()
}
+ record.Seconds = time.Since(start).Seconds()
record.InputTokens = response.Usage.Input
record.CachedTokens = response.Usage.Cached
record.OutputTokens = response.Usage.Output
+ record.TextBytes = len(response.Text)
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) {
+func (c *Client) send(httpRequest *http.Request, record *usageRecord, stopAtOutput bool) (response Response, err error) {
httpResponse, err := c.http.Do(httpRequest)
if err != nil {
return Response{}, fmt.Errorf("send model call: %w", err)
return Response{}, errors.Join(err, readErr)
}
- return readStream(httpResponse.Body)
+ return readStream(httpResponse.Body, stopAtOutput)
}
// usageLog appends one JSON line for each model call to usage.jsonl. Calls can
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.
+// usageRecord is one line of usage.jsonl. Date is the start of the call, and
+// Seconds its length. Node is the node [l, i] of a compaction. Status is the
+// HTTP status, or 0 when no response came. TextBytes is the size of the text of
+// the reply. Turn names the request of a turn. Quota holds the x-codex-* headers of the response, without the
+// turn state.
type usageRecord struct {
Date string `json:"date"`
+ Seconds float64 `json:"seconds"`
Model string `json:"model"`
+ Node *[2]int `json:"node,omitempty"`
+ Turn string `json:"turn,omitempty"`
Status int `json:"status"`
InputTokens int `json:"input_tokens"`
CachedTokens int `json:"cached_tokens"`
OutputTokens int `json:"output_tokens"`
+ TextBytes int `json:"text_bytes"`
Quota map[string]string `json:"quota,omitempty"`
Error string `json:"error,omitempty"`
}
}
}
-// quotaHeaders returns the x-codex-* headers, with lowercase names.
+// quotaHeaders returns the x-codex-* headers, with lowercase names, without the
+// turn state.
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 {
+ if !strings.HasPrefix(name, quotaPrefix) || name == turnStateHeader || len(values) == 0 {
continue
}
// 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) {
+// the completed response when the stream sent none. With stopAtOutput it ends
+// at the first output item.
+func readStream(stream io.Reader, stopAtOutput bool) (Response, error) {
var response Response
reader := bufio.NewReader(stream)
var data bytes.Buffer
data.WriteByte('\n')
}
} else if data.Len() > 0 {
- done, eventErr := readEvent(&response, data.Bytes())
+ done, eventErr := readEvent(&response, data.Bytes(), stopAtOutput)
data.Reset()
if eventErr != nil || done {
return response, eventErr
// readEvent adds one stream event to response, and tells if it ends the
// stream.
-func readEvent(response *Response, data []byte) (bool, error) {
+func readEvent(response *Response, data []byte, stopAtOutput bool) (bool, error) {
var event streamEvent
err := json.Unmarshal(data, &event)
if err != nil {
}
switch event.Type {
+ case "response.output_item.added":
+ return stopAtOutput, nil
case "response.output_item.done":
response.Output = append(response.Output, event.Item)
case "response.completed":