diff --git a/llm/transformer/openai/responses/cache_anchor.go b/llm/transformer/openai/responses/cache_anchor.go new file mode 100644 index 000000000..215577369 --- /dev/null +++ b/llm/transformer/openai/responses/cache_anchor.go @@ -0,0 +1,95 @@ +package responses + +import ( + "crypto/sha256" + "encoding/hex" + "strconv" + + "github.com/looplj/axonhub/llm" +) + +// anchorMaxBytesPerUnit bounds how much each content unit (a message's plain +// string content, or a single field of a content part) contributes to the +// hash. The budget is per unit rather than global so that a large instruction +// prefix can never starve the first user message out of the fingerprint. +const anchorMaxBytesPerUnit = 4096 + +// conversationAnchor derives a stable fingerprint for the conversation a +// request belongs to: the contiguous leading system/developer messages plus +// the first user message. Later turns of the same conversation keep this head +// unchanged, while sibling conversations multiplexed over one client session +// (e.g. Claude Code subagents sharing a session_id) diverge in their first +// user message. Combining this anchor with the session ID therefore yields a +// per-conversation prompt_cache_key instead of a per-session one. +func conversationAnchor(messages []llm.Message) string { + h := sha256.New() + + // Each unit is hashed as :. The length prefix keeps + // concatenated units unambiguous and still distinguishes oversized units + // (e.g. two base64 images sharing their first 4 KiB) by their size. + writeUnit := func(s string) { + h.Write([]byte(strconv.Itoa(len(s)))) + h.Write([]byte{':'}) + + if len(s) > anchorMaxBytesPerUnit { + s = s[:anchorMaxBytesPerUnit] + } + + h.Write([]byte(s)) + } + + hashed := false + + for _, msg := range messages { + if msg.Role != "system" && msg.Role != "developer" && msg.Role != "user" { + break + } + + h.Write([]byte(msg.Role)) + h.Write([]byte{0x00}) + + if msg.Content.Content != nil { + writeUnit(*msg.Content.Content) + } + + for _, part := range msg.Content.MultipleContent { + h.Write([]byte(part.Type)) + h.Write([]byte{0x1f}) + + switch { + case part.Text != nil: + writeUnit(*part.Text) + case part.ImageURL != nil: + writeUnit(part.ImageURL.URL) + case part.VideoURL != nil: + writeUnit(part.VideoURL.URL) + case part.Document != nil: + writeUnit(part.Document.MIMEType) + writeUnit(part.Document.URL) + case part.InputAudio != nil: + writeUnit(part.InputAudio.Format) + writeUnit(part.InputAudio.Data) + case part.Compact != nil: + writeUnit(part.Compact.ID) + writeUnit(part.Compact.EncryptedContent) + } + + h.Write([]byte{0x1f}) + } + + h.Write([]byte{0x1e}) + + hashed = true + + // The first user message ends the stable conversation head. + if msg.Role == "user" { + break + } + } + + if !hashed { + return "" + } + + return hex.EncodeToString(h.Sum(nil))[:16] +} diff --git a/llm/transformer/openai/responses/outbound.go b/llm/transformer/openai/responses/outbound.go index 390a61b55..7ed2736cd 100644 --- a/llm/transformer/openai/responses/outbound.go +++ b/llm/transformer/openai/responses/outbound.go @@ -274,6 +274,13 @@ func (t *OutboundTransformer) TransformRequest(ctx context.Context, llmReq *llm. if lo.FromPtr(payload.PromptCacheKey) == "" { if sessionID, ok := shared.GetSessionID(ctx); ok { + // A session may multiplex several concurrent conversations + // (e.g. Claude Code subagents); scope the cache key to the + // conversation so they do not evict each other upstream. + if anchor := conversationAnchor(llmReq.Messages); anchor != "" { + sessionID = sessionID + "-" + anchor + } + payload.PromptCacheKey = lo.ToPtr(sessionID) } } diff --git a/llm/transformer/openai/responses/outbound_test.go b/llm/transformer/openai/responses/outbound_test.go index 9a4ac497d..7a32d3710 100644 --- a/llm/transformer/openai/responses/outbound_test.go +++ b/llm/transformer/openai/responses/outbound_test.go @@ -1050,7 +1050,87 @@ func TestOutboundTransformer_TransformRequest_UsesSharedSessionIDAsPromptCacheKe err = json.Unmarshal(httpReq.Body, &payload) require.NoError(t, err) require.NotNil(t, payload.PromptCacheKey) - require.Equal(t, "shared-session-123", *payload.PromptCacheKey) + require.Equal(t, "shared-session-123-"+conversationAnchor(req.Messages), *payload.PromptCacheKey) +} + +func TestOutboundTransformer_TransformRequest_PromptCacheKeyScopedPerConversation(t *testing.T) { + transformer, err := NewOutboundTransformer("https://api.openai.com", "test-api-key") + require.NoError(t, err) + + ctx := shared.WithSessionID(context.Background(), "shared-session-123") + + newReq := func(firstUser string, extraTurns ...llm.Message) *llm.Request { + messages := []llm.Message{ + {Role: "system", Content: llm.MessageContent{Content: lo.ToPtr("You are an agent.")}}, + {Role: "user", Content: llm.MessageContent{Content: lo.ToPtr(firstUser)}}, + } + messages = append(messages, extraTurns...) + + return &llm.Request{Model: "gpt-5.4", Messages: messages} + } + + cacheKey := func(req *llm.Request) string { + httpReq, err := transformer.TransformRequest(ctx, req) + require.NoError(t, err) + + var payload Request + + require.NoError(t, json.Unmarshal(httpReq.Body, &payload)) + require.NotNil(t, payload.PromptCacheKey) + + return *payload.PromptCacheKey + } + + // Later turns of the same conversation keep the same cache key. + turn1 := cacheKey(newReq("task A")) + turn2 := cacheKey(newReq("task A", + llm.Message{Role: "assistant", Content: llm.MessageContent{Content: lo.ToPtr("working")}}, + llm.Message{Role: "user", Content: llm.MessageContent{Content: lo.ToPtr("continue")}}, + )) + require.Equal(t, turn1, turn2) + + // Sibling conversations in the same session get distinct cache keys. + require.NotEqual(t, turn1, cacheKey(newReq("task B"))) + + // Client-provided keys are preserved untouched. + explicit := newReq("task A") + explicit.PromptCacheKey = lo.ToPtr("client-key") + require.Equal(t, "client-key", cacheKey(explicit)) + + // A large shared instruction prefix must not starve the first user + // message out of the fingerprint: sibling conversations still get + // distinct keys. + largeSystem := strings.Repeat("shared instructions. ", 2048) + largeReq := func(firstUser string) *llm.Request { + return &llm.Request{ + Model: "gpt-5.4", + Messages: []llm.Message{ + {Role: "system", Content: llm.MessageContent{Content: lo.ToPtr(largeSystem)}}, + {Role: "user", Content: llm.MessageContent{Content: lo.ToPtr(firstUser)}}, + }, + } + } + require.NotEqual(t, cacheKey(largeReq("task A")), cacheKey(largeReq("task B"))) + + // Non-text content contributes to the fingerprint: first user messages + // that differ only by an image part get distinct keys. + imageReq := func(imageURL string) *llm.Request { + return &llm.Request{ + Model: "gpt-5.4", + Messages: []llm.Message{ + {Role: "user", Content: llm.MessageContent{ + MultipleContent: []llm.MessageContentPart{ + {Type: "text", Text: lo.ToPtr("describe this image")}, + {Type: "image_url", ImageURL: &llm.ImageURL{URL: imageURL}}, + }, + }}, + }, + } + } + require.NotEqual(t, + cacheKey(imageReq("https://example.com/a.png")), + cacheKey(imageReq("https://example.com/b.png")), + ) } func TestOutboundTransformer_TransformResponse(t *testing.T) {