Send the cache key in the session_id header

Each call to the ChatGPT backend sends the cache key of the agent in the
session_id header too, and keeps prompt_cache_key in the body. With the
key in the body only, the probe got no cached token, also on five tries
of one node with the same prefix. The same request sent twice with the
header got 1,664 of 1,756 input tokens from the cache on the second
call. Calls to the OpenAI API do not send the header.

Issue: b14d537

diff --git a/responses.go b/responses.go
index cdc1558195e159cd111f0e4968f8e9f5d9aea941..096439f5cc9564d92f8fcfed7636ae60960c91a6 100644 (file)
--- a/responses.go
+++ b/responses.go
@@ -48,8 +48,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,
-// 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
@@ -217,6 +219,11 @@ func (c *Client) Create(ctx context.Context, request Request) (Response, error)
                return Response{}, fmt.Errorf("authorize model call: %w", err)
        }
 
+       _, chatgpt := c.auth.(*chatgptAuth)
+       if chatgpt {
+               httpRequest.Header.Set("session_id", c.cacheKey)
+       }
+
        record := usageRecord{
                Date:  time.Now().Format(time.RFC3339Nano),
                Model: request.Model,

diff --git a/responses_test.go b/responses_test.go
index 577e6111e5b2c3e80135768237cc1c9ef59b0219..de928bd426ccd0f32f9ddfa0274d448fe1cb183a 100644 (file)
--- a/responses_test.go
+++ b/responses_test.go
@@ -11,6 +11,7 @@ import (
        "path/filepath"
        "sync"
        "testing"
+       "time"
 
        "github.com/stretchr/testify/require"
 )
@@ -60,6 +61,42 @@ func TestClientSendsTheFixedRequestFields(t *testing.T) {
        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) {