Give each subagent its own prompt cache key

A call of a subagent sends the cache key of the agent with "-" and the
name of the subagent, in the body and in the session_id header. With the
key of the agent, the second request of a subagent in a live check read
0 of 4,435 input tokens from the prompt cache, although it starts with
the whole first request, while the agent made calls with the same key at
the same time. With its own key, the second request of a subagent in the
same check read 4,608 of 4,761 tokens from the cache. The calls of the
agent keep the key of the agent.

Issue: b1ac698

diff --git a/model.go b/model.go
index f15be221e2d494e977595d3437d274029eb85552..61dbc5829a4bbfd4fc360ec9d67a68704697e8ed 100644 (file)
--- a/model.go
+++ b/model.go
@@ -138,9 +138,11 @@ func (k ModelAPIKey) Authorize(ctx context.Context, header http.Header) error {
 // 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.
 // 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.

-// 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.
+// A call of a subagent sends the key of the agent with "-" and the name of the
+// subagent, so that the calls of a subagent do not share a cache entry with the
+// calls of the agent that run at the same time. 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


@@ -173,6 +175,11 @@ func (c *ModelClient) Create(ctx context.Context, request ModelRequest) (ModelRe
                return ModelResponse{}, errModelCallLimit
        }
 
                return ModelResponse{}, errModelCallLimit
        }
 

+       cacheKey := c.cacheKey
+       if request.Agent != "" {
+               cacheKey += "-" + request.Agent
+       }
+

        body, err := json.Marshal(ModelRequestBody{
                Model:        request.Model,
                Instructions: request.Instructions,
        body, err := json.Marshal(ModelRequestBody{
                Model:        request.Model,
                Instructions: request.Instructions,


@@ -184,7 +191,7 @@ func (c *ModelClient) Create(ctx context.Context, request ModelRequest) (ModelRe
                Store:          false,
                Stream:         true,
                Include:        []string{"reasoning.encrypted_content"},
                Store:          false,
                Stream:         true,
                Include:        []string{"reasoning.encrypted_content"},

-               PromptCacheKey: c.cacheKey,
+               PromptCacheKey: cacheKey,

        })
        if err != nil {
                return ModelResponse{}, fmt.Errorf("encode request: %w", err)
        })
        if err != nil {
                return ModelResponse{}, fmt.Errorf("encode request: %w", err)


@@ -204,7 +211,7 @@ func (c *ModelClient) Create(ctx context.Context, request ModelRequest) (ModelRe
 
        _, chatgpt := c.auth.(*ChatGPTAuth)
        if chatgpt {
 
        _, chatgpt := c.auth.(*ChatGPTAuth)
        if chatgpt {

-               httpRequest.Header.Set("session_id", c.cacheKey)
+               httpRequest.Header.Set("session_id", cacheKey)

        }
 
        start := time.Now()
        }
 
        start := time.Now()


@@ -304,8 +311,9 @@ type ModelOutputItem struct {
 // modelUserMessage 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 or next. Agent names the subagent that makes the call, or is
 // modelUserMessage 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 or next. Agent names the subagent that makes the call, or is

-// empty for a call of the agent itself. Node, Turn, and Agent go only to the
-// usage file, never to the endpoint.
+// empty for a call of the agent itself. Node, Turn, and Agent go to the usage
+// file, never to the endpoint as fields. Agent also gives the call the cache
+// key of its subagent.

 type ModelRequest struct {
        Model        string
        Effort       string
 type ModelRequest struct {
        Model        string
        Effort       string

diff --git a/model_test.go b/model_test.go
index 16946a061dfe25447e34bded0128e079732d0523..11a7d00bcddf7ad8f68373675a8428111fa81fb4 100644 (file)
--- a/model_test.go
+++ b/model_test.go
@@ -85,6 +85,19 @@ func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {
                        call := server.Call(0)
                        require.Equal(t, tc.want, call.header.Get("session_id"))
                        require.Equal(t, "agent-key", call.body["prompt_cache_key"])
                        call := server.Call(0)
                        require.Equal(t, tc.want, call.header.Get("session_id"))
                        require.Equal(t, "agent-key", call.body["prompt_cache_key"])

+
+                       // A subagent has its own key.
+                       request := modelTestRequest()
+                       request.Agent = "GreenOtter"
+                       _, err = client.Create(t.Context(), request)
+                       require.NoError(t, err)
+
+                       call = server.Call(1)
+                       if tc.want != "" {
+                               require.Equal(t, "agent-key-GreenOtter", call.header.Get("session_id"))
+                       }
+
+                       require.Equal(t, "agent-key-GreenOtter", call.body["prompt_cache_key"])

                })
        }
 }
                })
        }
 }


@@ -150,7 +163,7 @@ func TestClientRecordsEachCall(t *testing.T) {
        require.NoError(t, err)
        require.NotContains(t, server.Call(0).body, "node")
        require.NotContains(t, string(server.Call(0).raw), "[2,5]")
        require.NoError(t, err)
        require.NotContains(t, server.Call(0).body, "node")
        require.NotContains(t, string(server.Call(0).raw), "[2,5]")

-       require.NotContains(t, string(server.Call(0).raw), "GreenOtter")
+       require.NotContains(t, server.Call(0).body, "agent")

 
        _, err = client.Create(t.Context(), modelTestRequest())
        require.ErrorContains(t, err, "HTTP 429")
 
        _, err = client.Create(t.Context(), modelTestRequest())
        require.ErrorContains(t, err, "HTTP 429")