Lay out turn requests for the prompt cache

[onechat] / responses.go

diff --git a/responses.go b/responses.go
index 9d6b75f6c05ce2e55335e77b44877b69ea56df5f..19ab6a1eda6bb54043f34f09c5eed89485bd8258 100644 (file)
--- a/responses.go
+++ b/responses.go
@@ -13,6 +13,7 @@ import (
        "path/filepath"
        "strings"
        "sync"
        "path/filepath"
        "strings"
        "sync"

+       "sync/atomic"

        "time"
 )
 
        "time"
 )
 


@@ -27,11 +28,18 @@ const (
        // backend, in canonical form.
        quotaPrefix = "X-Codex-"
 
        // 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
 )
 
        // 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 {
 // authorizer adds the credentials of a model call to the headers of its
 // request.
 type authorizer interface {


@@ -44,8 +52,10 @@ 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,
 // 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
 //
 // 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


@@ -57,17 +67,25 @@ type Client struct {
        cacheKey string
        usage    *usageLog
        http     *http.Client
        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
 }
 
 // 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
 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.
 }
 
 // Tool is the definition of a function tool.


@@ -179,13 +197,27 @@ func newClient(url string, auth authorizer, cacheKey, usagePath string) *Client
                usage: &usageLog{
                        path: usagePath,
                },
                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) {
 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,
        body, err := json.Marshal(requestBody{
                Model:        request.Model,
                Instructions: request.Instructions,


@@ -215,26 +247,39 @@ func (c *Client) Create(ctx context.Context, request Request) (Response, error)
                return Response{}, fmt.Errorf("authorize model call: %w", err)
        }
 
                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{
        record := usageRecord{

-               Date:  c.clock().Format(time.RFC3339Nano),
+               Date:  start.Format(time.RFC3339Nano),

                Model: request.Model,
        }
 
                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()
        }
 
        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.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.
 
        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)
        httpResponse, err := c.http.Do(httpRequest)
        if err != nil {
                return Response{}, fmt.Errorf("send model call: %w", err)


@@ -253,7 +298,7 @@ func (c *Client) send(httpRequest *http.Request, record *usageRecord) (response
                return Response{}, errors.Join(err, readErr)
        }
 
                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
 }
 
 // usageLog appends one JSON line for each model call to usage.jsonl. Calls can


@@ -263,15 +308,22 @@ type usageLog struct {
        path  string
 }
 
        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"`
 type usageRecord struct {
        Date         string            `json:"date"`

+       Seconds      float64           `json:"seconds"`

        Model        string            `json:"model"`
        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"`
        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"`
 }
        Quota        map[string]string `json:"quota,omitempty"`
        Error        string            `json:"error,omitempty"`
 }


@@ -333,11 +385,12 @@ func userMessage(texts ...string) any {
        }
 }
 
        }
 }
 

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


@@ -353,8 +406,9 @@ func quotaHeaders(header http.Header) map[string]string {
 
 // 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
 
 // 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
        var response Response
        reader := bufio.NewReader(stream)
        var data bytes.Buffer


@@ -372,7 +426,7 @@ func readStream(stream io.Reader) (Response, error) {
                                data.WriteByte('\n')
                        }
                } else if data.Len() > 0 {
                                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
                        data.Reset()
                        if eventErr != nil || done {
                                return response, eventErr


@@ -387,7 +441,7 @@ func readStream(stream io.Reader) (Response, error) {
 
 // readEvent adds one stream event to response, and tells if it ends the
 // stream.
 
 // 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 {
        var event streamEvent
        err := json.Unmarshal(data, &event)
        if err != nil {


@@ -395,6 +449,8 @@ func readEvent(response *Response, data []byte) (bool, error) {
        }
 
        switch event.Type {
        }
 
        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":
        case "response.output_item.done":
                response.Output = append(response.Output, event.Item)
        case "response.completed":