diff --git a/internal/responsecache/handle_request_test.go b/internal/responsecache/handle_request_test.go index d509c5471..b6303318c 100644 --- a/internal/responsecache/handle_request_test.go +++ b/internal/responsecache/handle_request_test.go @@ -2,6 +2,7 @@ package responsecache import ( "bytes" + "encoding/json" "net/http" "net/http/httptest" "testing" @@ -261,3 +262,630 @@ func TestHandleRequest_ExactHitWritesSyntheticUsageEntry(t *testing.T) { t.Fatalf("unexpected tokens: %+v", entry) } } + +func TestHandleRequest_CacheControlNoCacheBypassesAllLayers(t *testing.T) { + store := cache.NewMapStore() + defer store.Close() + + emb := &mockEmbedder{vector: []float32{1, 0, 0}} + vecStore := NewMapVecStore() + semCfg := config.SemanticCacheConfig{ + Enabled: true, + SimilarityThreshold: 0.90, + TTL: 3600, + MaxConversationMessages: 10, + } + + m := &ResponseCacheMiddleware{ + simple: newSimpleCacheMiddleware(store, time.Hour, nil), + semantic: newSemanticCacheMiddleware(emb, vecStore, semCfg, nil), + } + + body := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"handle-request-no-cache"}]}`) + e := echo.New() + handlerCalls := 0 + + run := func(cacheControl string) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + if cacheControl != "" { + req.Header.Set("Cache-Control", cacheControl) + } + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + if err := m.HandleRequest(c, body, func() error { + handlerCalls++ + return c.JSON(http.StatusOK, map[string]int{"n": handlerCalls}) + }); err != nil { + t.Fatalf("HandleRequest: %v", err) + } + return rec + } + + rec1 := run("") + if got := rec1.Header().Get("X-Cache"); got != "" { + t.Fatalf("first request should miss cache, got X-Cache=%q", got) + } + + m.simple.wg.Wait() + m.semantic.wg.Wait() + + rec2 := run("no-cache") + if got := rec2.Header().Get("X-Cache"); got != "" { + t.Fatalf("no-cache request should bypass cache, got X-Cache=%q", got) + } + if !bytes.Contains(rec2.Body.Bytes(), []byte(`"n":2`)) { + t.Fatalf("no-cache response body = %q, want fresh handler response", rec2.Body.String()) + } + + rec3 := run("") + if got := rec3.Header().Get("X-Cache"); got != "HIT (exact)" { + t.Fatalf("follow-up request should still hit original cache entry, got X-Cache=%q", got) + } + if !bytes.Contains(rec3.Body.Bytes(), []byte(`"n":1`)) { + t.Fatalf("cached response body = %q, want original cached payload", rec3.Body.String()) + } + if handlerCalls != 2 { + t.Fatalf("expected handler to run exactly twice, got %d calls", handlerCalls) + } +} + +func TestHandleRequest_StreamingMissPopulatesExactCacheAcrossModes(t *testing.T) { + store := cache.NewMapStore() + defer store.Close() + + m := &ResponseCacheMiddleware{ + simple: newSimpleCacheMiddleware(store, time.Hour, nil), + } + + streamBody := []byte(`{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"cache-streaming-cross-mode"}]}`) + jsonBody := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"cache-streaming-cross-mode"}]}`) + e := echo.New() + handlerCalls := 0 + + run := func(body []byte) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + plan := &core.ExecutionPlan{ + Mode: core.ExecutionModeTranslated, + ProviderType: "openai", + Resolution: &core.RequestModelResolution{ + ResolvedSelector: core.ModelSelector{Provider: "openai", Model: "gpt-4"}, + }, + } + c.SetRequest(req.WithContext(core.WithExecutionPlan(req.Context(), plan))) + if err := m.HandleRequest(c, body, func() error { + handlerCalls++ + c.Response().Header().Set("Content-Type", "text/event-stream") + c.Response().WriteHeader(http.StatusOK) + _, _ = c.Response().Write([]byte("data: {\"id\":\"chatcmpl-stream-cache\",\"object\":\"chat.completion.chunk\",\"created\":1234567890,\"model\":\"gpt-4\",\"provider\":\"openai\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n")) + _, _ = c.Response().Write([]byte("data: {\"id\":\"chatcmpl-stream-cache\",\"object\":\"chat.completion.chunk\",\"created\":1234567890,\"model\":\"gpt-4\",\"provider\":\"openai\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" world\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":2,\"total_tokens\":13}}\n\n")) + _, _ = c.Response().Write([]byte("data: [DONE]\n\n")) + return nil + }); err != nil { + t.Fatalf("HandleRequest: %v", err) + } + return rec + } + + rec1 := run(streamBody) + if got := rec1.Header().Get("X-Cache"); got != "" { + t.Fatalf("streaming miss should not be cache hit, got X-Cache=%q", got) + } + if got := rec1.Header().Get("Content-Type"); got != "text/event-stream" { + t.Fatalf("streaming miss Content-Type = %q, want text/event-stream", got) + } + if handlerCalls != 1 { + t.Fatalf("expected 1 handler invocation after streaming miss, got %d", handlerCalls) + } + + m.simple.wg.Wait() + + rec2 := run(jsonBody) + if got := rec2.Header().Get("X-Cache"); got != "HIT (exact)" { + t.Fatalf("non-streaming follow-up should hit exact cache, got X-Cache=%q", got) + } + if got := rec2.Header().Get("Content-Type"); got != "application/json" { + t.Fatalf("non-streaming hit Content-Type = %q, want application/json", got) + } + if !bytes.Contains(rec2.Body.Bytes(), []byte(`"content":"Hello world"`)) { + t.Fatalf("non-streaming cache hit body = %q, want reconstructed JSON response", rec2.Body.String()) + } + if handlerCalls != 1 { + t.Fatalf("non-streaming exact hit should not call handler again, got %d calls", handlerCalls) + } + + rec3 := run(streamBody) + if got := rec3.Header().Get("X-Cache"); got != "HIT (exact)" { + t.Fatalf("streaming follow-up should hit exact cache, got X-Cache=%q", got) + } + if got := rec3.Header().Get("Content-Type"); got != "text/event-stream" { + t.Fatalf("streaming hit Content-Type = %q, want text/event-stream", got) + } + if !bytes.Contains(rec3.Body.Bytes(), []byte("Hello world")) || !bytes.Contains(rec3.Body.Bytes(), []byte("[DONE]")) { + t.Fatalf("streaming cache hit body = %q, want synthesized SSE with content and [DONE]", rec3.Body.String()) + } + if handlerCalls != 1 { + t.Fatalf("streaming exact hit should not call handler again, got %d calls", handlerCalls) + } +} + +func TestReconstructStreamingResponse_PreservesChatReasoningContent(t *testing.T) { + raw := []byte( + "data: {\"id\":\"chatcmpl-reasoning\",\"object\":\"chat.completion.chunk\",\"created\":1234567890,\"model\":\"claude-sonnet\",\"provider\":\"anthropic\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"reasoning_content\":\"think first\"},\"finish_reason\":null}]}\n\n" + + "data: {\"id\":\"chatcmpl-reasoning\",\"object\":\"chat.completion.chunk\",\"created\":1234567890,\"model\":\"claude-sonnet\",\"provider\":\"anthropic\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"final answer\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":4,\"total_tokens\":14}}\n\n" + + "data: [DONE]\n\n", + ) + + cached, ok := reconstructStreamingResponse("/v1/chat/completions", raw, streamResponseDefaults{ + Model: "claude-sonnet", + Provider: "anthropic", + }) + if !ok { + t.Fatal("expected streamed chat response to reconstruct successfully") + } + if !bytes.Contains(cached, []byte(`"reasoning_content":"think first"`)) { + t.Fatalf("reconstructed chat response = %q, want reasoning_content preserved", string(cached)) + } + + replay, err := renderCachedChatStream([]byte(`{"model":"claude-sonnet","stream":true}`), cached) + if err != nil { + t.Fatalf("renderCachedChatStream() error = %v", err) + } + if !bytes.Contains(replay, []byte(`"reasoning_content":"think first"`)) { + t.Fatalf("cached chat replay = %q, want reasoning_content delta", string(replay)) + } + if bytes.Contains(replay, []byte(`"usage"`)) { + t.Fatalf("cached chat replay without include_usage = %q, did not expect usage chunk", string(replay)) + } + + replayWithUsage, err := renderCachedChatStream([]byte(`{"model":"claude-sonnet","stream":true,"stream_options":{"include_usage":true}}`), cached) + if err != nil { + t.Fatalf("renderCachedChatStream(include_usage) error = %v", err) + } + if !bytes.Contains(replayWithUsage, []byte(`"usage"`)) { + t.Fatalf("cached chat replay with include_usage = %q, want usage chunk", string(replayWithUsage)) + } +} + +func TestRenderCachedChatStream_EmitsStandaloneUsageChunk(t *testing.T) { + raw := []byte( + "data: {\"id\":\"chatcmpl-usage\",\"object\":\"chat.completion.chunk\",\"created\":1234567890,\"model\":\"gpt-4o-mini\",\"provider\":\"openai\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n" + + "data: {\"id\":\"chatcmpl-usage\",\"object\":\"chat.completion.chunk\",\"created\":1234567890,\"model\":\"gpt-4o-mini\",\"provider\":\"openai\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" world\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":2,\"total_tokens\":13}}\n\n" + + "data: [DONE]\n\n", + ) + + cached, ok := reconstructStreamingResponse("/v1/chat/completions", raw, streamResponseDefaults{ + Model: "gpt-4o-mini", + Provider: "openai", + }) + if !ok { + t.Fatal("expected streamed chat response to reconstruct successfully") + } + + replay, err := renderCachedChatStream([]byte(`{"model":"gpt-4o-mini","stream":true,"stream_options":{"include_usage":true}}`), cached) + if err != nil { + t.Fatalf("renderCachedChatStream() error = %v", err) + } + + var events []map[string]any + parseSSEJSONEvents(replay, func(event map[string]any) { + events = append(events, event) + }) + if len(events) != 2 { + t.Fatalf("len(events) = %d, want 2 chat events before [DONE]", len(events)) + } + + firstChoices, ok := events[0]["choices"].([]any) + if !ok || len(firstChoices) != 1 { + t.Fatalf("first event choices = %#v, want len=1", events[0]["choices"]) + } + if _, ok := events[0]["usage"]; ok { + t.Fatalf("first event should not carry usage, got %#v", events[0]["usage"]) + } + + secondChoices, ok := events[1]["choices"].([]any) + if !ok || len(secondChoices) != 0 { + t.Fatalf("usage event choices = %#v, want empty slice", events[1]["choices"]) + } + usage, ok := events[1]["usage"].(map[string]any) + if !ok { + t.Fatalf("usage event usage = %#v, want object", events[1]["usage"]) + } + if got, ok := jsonNumberToInt(usage["total_tokens"]); !ok || got != 13 { + t.Fatalf("usage.total_tokens = %#v, want 13", usage["total_tokens"]) + } +} + +func TestReconstructStreamingResponse_PreservesChatLogprobs(t *testing.T) { + raw := []byte( + "data: {\"id\":\"chatcmpl-logprobs\",\"object\":\"chat.completion.chunk\",\"created\":1234567890,\"model\":\"gpt-4o-mini\",\"provider\":\"openai\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"Hello\"},\"logprobs\":null,\"finish_reason\":null}]}\n\n" + + "data: {\"id\":\"chatcmpl-logprobs\",\"object\":\"chat.completion.chunk\",\"created\":1234567890,\"model\":\"gpt-4o-mini\",\"provider\":\"openai\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" world\"},\"logprobs\":null,\"finish_reason\":\"stop\"}]}\n\n" + + "data: [DONE]\n\n", + ) + + cached, ok := reconstructStreamingResponse("/v1/chat/completions", raw, streamResponseDefaults{ + Model: "gpt-4o-mini", + Provider: "openai", + }) + if !ok { + t.Fatal("expected streamed chat response to reconstruct successfully") + } + if !bytes.Contains(cached, []byte(`"logprobs":null`)) { + t.Fatalf("reconstructed chat response = %q, want choice.logprobs preserved", string(cached)) + } + + replay, err := renderCachedChatStream([]byte(`{"model":"gpt-4o-mini","stream":true}`), cached) + if err != nil { + t.Fatalf("renderCachedChatStream() error = %v", err) + } + if !bytes.Contains(replay, []byte(`"logprobs":null`)) { + t.Fatalf("cached chat replay = %q, want choice.logprobs preserved", string(replay)) + } +} + +func TestReconstructStreamingResponse_PreservesResponsesReasoningText(t *testing.T) { + raw := []byte( + "event: response.created\n" + + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_reasoning_build\",\"object\":\"response\",\"created_at\":1234567890,\"model\":\"grok-4\",\"provider\":\"xai\",\"status\":\"in_progress\",\"output\":[]}}\n\n" + + "event: response.output_item.added\n" + + "data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"id\":\"rs_1\",\"type\":\"reasoning\",\"status\":\"in_progress\",\"summary\":[]}}\n\n" + + "event: response.reasoning_text.delta\n" + + "data: {\"type\":\"response.reasoning_text.delta\",\"item_id\":\"rs_1\",\"output_index\":0,\"content_index\":0,\"delta\":\"step by\"}\n\n" + + "event: response.reasoning_text.delta\n" + + "data: {\"type\":\"response.reasoning_text.delta\",\"item_id\":\"rs_1\",\"output_index\":0,\"content_index\":1,\"delta\":\"step\"}\n\n" + + "event: response.output_item.done\n" + + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"rs_1\",\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[]}}\n\n" + + "event: response.completed\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_reasoning_build\",\"object\":\"response\",\"created_at\":1234567890,\"model\":\"grok-4\",\"provider\":\"xai\",\"status\":\"completed\",\"output\":[{\"id\":\"rs_1\",\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[]}]}}\n\n" + + "data: [DONE]\n\n", + ) + + cached, ok := reconstructStreamingResponse("/v1/responses", raw, streamResponseDefaults{ + Model: "grok-4", + Provider: "xai", + }) + if !ok { + t.Fatal("expected streamed responses payload to reconstruct successfully") + } + + var response map[string]any + if err := json.Unmarshal(cached, &response); err != nil { + t.Fatalf("json.Unmarshal(cached) error = %v", err) + } + output, ok := response["output"].([]any) + if !ok || len(output) != 1 { + t.Fatalf("reconstructed output = %#v, want len=1", response["output"]) + } + item, ok := output[0].(map[string]any) + if !ok { + t.Fatalf("reconstructed output[0] = %#v, want object", output[0]) + } + if _, ok := item["content"]; ok { + t.Fatalf("reasoning item should not use content field, got %#v", item["content"]) + } + summary, ok := item["summary"].([]any) + if !ok || len(summary) != 2 { + t.Fatalf("reconstructed reasoning summary = %#v, want len=2", item["summary"]) + } + wantTexts := []string{"step by", "step"} + for i, wantText := range wantTexts { + part, ok := summary[i].(map[string]any) + if !ok { + t.Fatalf("reconstructed summary[%d] = %#v, want object", i, summary[i]) + } + if got, _ := part["type"].(string); got != "reasoning_text" { + t.Fatalf("reconstructed summary part type = %q, want reasoning_text", got) + } + if got, _ := part["text"].(string); got != wantText { + t.Fatalf("reconstructed summary part text = %q, want %q", got, wantText) + } + } + + replay, err := renderCachedResponsesStream([]byte(`{"model":"grok-4","stream":true}`), cached) + if err != nil { + t.Fatalf("renderCachedResponsesStream() error = %v", err) + } + var reasoningDeltas []map[string]any + parseSSEJSONEvents(replay, func(event map[string]any) { + if eventType, _ := event["type"].(string); eventType == "response.reasoning_text.delta" { + reasoningDeltas = append(reasoningDeltas, event) + } + }) + if len(reasoningDeltas) != 2 { + t.Fatalf("cached responses replay = %q, want 2 reasoning_text delta events", string(replay)) + } + for i, deltaEvent := range reasoningDeltas { + if got, _ := deltaEvent["delta"].(string); got != wantTexts[i] { + t.Fatalf("reasoning delta text = %q, want %q", got, wantTexts[i]) + } + if got, _ := deltaEvent["item_id"].(string); got != "rs_1" { + t.Fatalf("reasoning delta item_id = %q, want rs_1", got) + } + if got, ok := jsonNumberToInt(deltaEvent["output_index"]); !ok || got != 0 { + t.Fatalf("reasoning delta output_index = %#v, want 0", deltaEvent["output_index"]) + } + if got, ok := jsonNumberToInt(deltaEvent["content_index"]); !ok || got != i { + t.Fatalf("reasoning delta content_index = %#v, want %d", deltaEvent["content_index"], i) + } + } +} + +func TestReconstructStreamingResponse_HonorsResponsesTextDeltaLocators(t *testing.T) { + raw := []byte( + "event: response.created\n" + + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_text_locator\",\"object\":\"response\",\"created_at\":1234567890,\"model\":\"gpt-4o-mini\",\"provider\":\"openai\",\"status\":\"in_progress\",\"output\":[{\"id\":\"rs_1\",\"type\":\"reasoning\",\"status\":\"in_progress\",\"summary\":[]}]}}\n\n" + + "event: response.output_text.delta\n" + + "data: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_1\",\"output_index\":1,\"content_index\":0,\"delta\":\"final\"}\n\n" + + "event: response.output_text.delta\n" + + "data: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_1\",\"output_index\":1,\"content_index\":1,\"delta\":\"answer\"}\n\n" + + "event: response.completed\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_text_locator\",\"object\":\"response\",\"created_at\":1234567890,\"model\":\"gpt-4o-mini\",\"provider\":\"openai\",\"status\":\"completed\",\"output\":[{\"id\":\"rs_1\",\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[{\"type\":\"reasoning_text\",\"text\":\"step by step\"}]},{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"status\":\"completed\",\"content\":[]}]}}\n\n" + + "data: [DONE]\n\n", + ) + + cached, ok := reconstructStreamingResponse("/v1/responses", raw, streamResponseDefaults{ + Model: "gpt-4o-mini", + Provider: "openai", + }) + if !ok { + t.Fatal("expected streamed responses payload to reconstruct successfully") + } + + var response map[string]any + if err := json.Unmarshal(cached, &response); err != nil { + t.Fatalf("json.Unmarshal(cached) error = %v", err) + } + output, ok := response["output"].([]any) + if !ok || len(output) != 2 { + t.Fatalf("reconstructed output = %#v, want len=2", response["output"]) + } + reasoningItem, ok := output[0].(map[string]any) + if !ok { + t.Fatalf("reconstructed output[0] = %#v, want object", output[0]) + } + if _, ok := reasoningItem["content"]; ok { + t.Fatalf("reasoning item should not use content field, got %#v", reasoningItem["content"]) + } + messageItem, ok := output[1].(map[string]any) + if !ok { + t.Fatalf("reconstructed output[1] = %#v, want object", output[1]) + } + content, ok := messageItem["content"].([]any) + if !ok || len(content) != 2 { + t.Fatalf("reconstructed message content = %#v, want len=2", messageItem["content"]) + } + wantTexts := []string{"final", "answer"} + for i, wantText := range wantTexts { + messagePart, ok := content[i].(map[string]any) + if !ok { + t.Fatalf("reconstructed message content[%d] = %#v, want object", i, content[i]) + } + if got, _ := messagePart["text"].(string); got != wantText { + t.Fatalf("reconstructed message text = %q, want %q", got, wantText) + } + } + + replay, err := renderCachedResponsesStream([]byte(`{"model":"gpt-4o-mini","stream":true}`), cached) + if err != nil { + t.Fatalf("renderCachedResponsesStream() error = %v", err) + } + var textDeltas []map[string]any + parseSSEJSONEvents(replay, func(event map[string]any) { + if eventType, _ := event["type"].(string); eventType == "response.output_text.delta" { + textDeltas = append(textDeltas, event) + } + }) + if len(textDeltas) != 2 { + t.Fatalf("cached responses replay = %q, want 2 output_text delta events", string(replay)) + } + for i, deltaEvent := range textDeltas { + if got, _ := deltaEvent["delta"].(string); got != wantTexts[i] { + t.Fatalf("text delta text = %q, want %q", got, wantTexts[i]) + } + if got, _ := deltaEvent["item_id"].(string); got != "msg_1" { + t.Fatalf("text delta item_id = %q, want msg_1", got) + } + if got, ok := jsonNumberToInt(deltaEvent["output_index"]); !ok || got != 1 { + t.Fatalf("text delta output_index = %#v, want 1", deltaEvent["output_index"]) + } + if got, ok := jsonNumberToInt(deltaEvent["content_index"]); !ok || got != i { + t.Fatalf("text delta content_index = %#v, want %d", deltaEvent["content_index"], i) + } + } +} + +func TestRenderCachedResponsesStream_PreservesReasoningTextDeltas(t *testing.T) { + cached := []byte(`{ + "id":"resp_reasoning", + "object":"response", + "created_at":1234567890, + "model":"grok-4", + "provider":"xai", + "status":"completed", + "output":[ + { + "id":"rs_1", + "type":"reasoning", + "status":"completed", + "summary":[{"type":"reasoning_text","text":"step by step"}] + }, + { + "id":"msg_1", + "type":"message", + "role":"assistant", + "status":"completed", + "content":[{"type":"output_text","text":"final answer"}] + } + ] + }`) + + replay, err := renderCachedResponsesStream([]byte(`{"model":"grok-4","stream":true}`), cached) + if err != nil { + t.Fatalf("renderCachedResponsesStream() error = %v", err) + } + if !bytes.Contains(replay, []byte("event: response.reasoning_text.delta")) { + t.Fatalf("cached responses replay = %q, want reasoning_text delta event", string(replay)) + } + if !bytes.Contains(replay, []byte("step by step")) { + t.Fatalf("cached responses replay = %q, want reasoning delta text", string(replay)) + } + if !bytes.Contains(replay, []byte("event: response.output_text.delta")) { + t.Fatalf("cached responses replay = %q, want output_text delta event", string(replay)) + } +} + +func TestRenderCachedResponsesStream_FunctionCallAddedItemOmitsArguments(t *testing.T) { + cached := []byte(`{ + "id":"resp_function_call", + "object":"response", + "created_at":1234567890, + "model":"gpt-4o-mini", + "provider":"openai", + "status":"completed", + "output":[ + { + "id":"fc_1", + "type":"function_call", + "status":"completed", + "call_id":"call_1", + "name":"lookup_weather", + "arguments":"{\"city\":\"Warsaw\"}" + } + ] + }`) + + replay, err := renderCachedResponsesStream([]byte(`{"model":"gpt-4o-mini","stream":true}`), cached) + if err != nil { + t.Fatalf("renderCachedResponsesStream() error = %v", err) + } + + var addedItem map[string]any + var argDelta map[string]any + var argDone map[string]any + parseSSEJSONEvents(replay, func(event map[string]any) { + switch eventType, _ := event["type"].(string); eventType { + case "response.output_item.added": + addedItem, _ = event["item"].(map[string]any) + case "response.function_call_arguments.delta": + argDelta = event + case "response.function_call_arguments.done": + argDone = event + } + }) + + if addedItem == nil { + t.Fatalf("cached responses replay = %q, want output_item.added event", string(replay)) + } + if _, ok := addedItem["arguments"]; ok { + t.Fatalf("added item arguments = %#v, want omitted", addedItem["arguments"]) + } + if argDelta == nil || argDone == nil { + t.Fatalf("cached responses replay = %q, want function_call_arguments delta and done events", string(replay)) + } + if got, _ := argDelta["delta"].(string); got != `{"city":"Warsaw"}` { + t.Fatalf("arguments delta = %q, want full arguments", got) + } + if got, _ := argDone["arguments"].(string); got != `{"city":"Warsaw"}` { + t.Fatalf("arguments done = %q, want full arguments", got) + } +} + +func TestReconstructStreamingResponse_PreservesResponsesTerminalEvents(t *testing.T) { + tests := []struct { + name string + eventName string + status string + terminalResponse string + assertTerminal func(*testing.T, map[string]any) + }{ + { + name: "failed", + eventName: "response.failed", + status: "failed", + terminalResponse: `{"id":"resp_failed","object":"response","created_at":1234567890,"model":"gpt-4o-mini","provider":"openai","status":"failed","error":{"code":"boom","message":"upstream failed"},"metadata":{"trace":"abc"},"output":[]}`, + assertTerminal: func(t *testing.T, response map[string]any) { + t.Helper() + errMap, ok := response["error"].(map[string]any) + if !ok { + t.Fatalf("terminal response error = %#v, want object", response["error"]) + } + if got, _ := errMap["code"].(string); got != "boom" { + t.Fatalf("terminal response error.code = %q, want boom", got) + } + }, + }, + { + name: "incomplete", + eventName: "response.incomplete", + status: "incomplete", + terminalResponse: `{"id":"resp_incomplete","object":"response","created_at":1234567890,"model":"gpt-4o-mini","provider":"openai","status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"metadata":{"trace":"def"},"output":[]}`, + assertTerminal: func(t *testing.T, response map[string]any) { + t.Helper() + details, ok := response["incomplete_details"].(map[string]any) + if !ok { + t.Fatalf("terminal response incomplete_details = %#v, want object", response["incomplete_details"]) + } + if got, _ := details["reason"].(string); got != "max_output_tokens" { + t.Fatalf("terminal response incomplete_details.reason = %q, want max_output_tokens", got) + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + raw := []byte( + "event: response.created\n" + + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_terminal\",\"object\":\"response\",\"created_at\":1234567890,\"model\":\"gpt-4o-mini\",\"provider\":\"openai\",\"status\":\"in_progress\",\"output\":[]}}\n\n" + + "event: " + tt.eventName + "\n" + + "data: {\"type\":\"" + tt.eventName + "\",\"response\":" + tt.terminalResponse + "}\n\n" + + "data: [DONE]\n\n", + ) + + cached, ok := reconstructStreamingResponse("/v1/responses", raw, streamResponseDefaults{ + Model: "gpt-4o-mini", + Provider: "openai", + }) + if !ok { + t.Fatal("expected streamed responses payload to reconstruct successfully") + } + + var cachedResponse map[string]any + if err := json.Unmarshal(cached, &cachedResponse); err != nil { + t.Fatalf("json.Unmarshal(cached) error = %v", err) + } + if got, _ := cachedResponse["status"].(string); got != tt.status { + t.Fatalf("cached response status = %q, want %q", got, tt.status) + } + tt.assertTerminal(t, cachedResponse) + + replay, err := renderCachedResponsesStream([]byte(`{"model":"gpt-4o-mini","stream":true}`), cached) + if err != nil { + t.Fatalf("renderCachedResponsesStream() error = %v", err) + } + + var terminalEvent map[string]any + parseSSEJSONEvents(replay, func(event map[string]any) { + if eventType, _ := event["type"].(string); eventType == tt.eventName { + terminalEvent = event + } + }) + if terminalEvent == nil { + t.Fatalf("cached responses replay = %q, want terminal event %s", string(replay), tt.eventName) + } + response, ok := terminalEvent["response"].(map[string]any) + if !ok { + t.Fatalf("terminal event response = %#v, want object", terminalEvent["response"]) + } + if got, _ := response["status"].(string); got != tt.status { + t.Fatalf("terminal event status = %q, want %q", got, tt.status) + } + tt.assertTerminal(t, response) + }) + } +} diff --git a/internal/responsecache/middleware_test.go b/internal/responsecache/middleware_test.go index 86046b63b..9993e833b 100644 --- a/internal/responsecache/middleware_test.go +++ b/internal/responsecache/middleware_test.go @@ -200,27 +200,87 @@ func TestHashRequest_ModeChangesKey(t *testing.T) { } } -func TestSimpleCacheMiddleware_SkipsStreaming(t *testing.T) { +func TestHashRequest_StreamIncludeUsageChangesKey(t *testing.T) { + base := []byte(`{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + withUsage := []byte(`{"model":"gpt-4","stream":true,"stream_options":{"include_usage":true},"messages":[{"role":"user","content":"hi"}]}`) + plan := &core.ExecutionPlan{ + Mode: core.ExecutionModeTranslated, + ProviderType: "openai", + Resolution: &core.RequestModelResolution{ + ResolvedSelector: core.ModelSelector{Provider: "openai", Model: "gpt-4"}, + }, + } + + first := hashRequest("/v1/chat/completions", base, plan) + second := hashRequest("/v1/chat/completions", withUsage, plan) + + if first == second { + t.Fatal("stream_options.include_usage should affect the exact cache key") + } +} + +func TestSimpleCacheMiddleware_SharesCacheAcrossStreamingAndNonStreaming(t *testing.T) { store := cache.NewMapStore() defer store.Close() mw := NewResponseCacheMiddlewareWithStore(store, time.Hour) e := echo.New() + installResolvedExecutionPlan(e, "openai", "gpt-4") e.Use(mw.Middleware()) callCount := 0 e.POST("/v1/chat/completions", func(c *echo.Context) error { callCount++ - return c.JSON(http.StatusOK, map[string]string{"n": "1"}) + return c.JSON(http.StatusOK, &core.ChatResponse{ + ID: "chatcmpl-shared-cache", + Object: "chat.completion", + Model: "gpt-4", + Provider: "openai", + Created: 1234567890, + Choices: []core.Choice{ + { + Index: 0, + Message: core.ResponseMessage{ + Role: "assistant", + Content: "shared cached response", + }, + FinishReason: "stop", + }, + }, + Usage: core.Usage{ + PromptTokens: 9, + CompletionTokens: 3, + TotalTokens: 12, + }, + }) }) - body := []byte(`{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"hi"}]}`) - for range 2 { - req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) - req.Header.Set("Content-Type", "application/json") - rec := httptest.NewRecorder() - e.ServeHTTP(rec, req) + nonStreamingBody := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}`) + streamingBody := []byte(`{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + + req1 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(nonStreamingBody)) + req1.Header.Set("Content-Type", "application/json") + rec1 := httptest.NewRecorder() + e.ServeHTTP(rec1, req1) + if rec1.Header().Get("X-Cache") != "" { + t.Fatalf("first request should miss cache, got X-Cache=%q", rec1.Header().Get("X-Cache")) } - if callCount != 2 { - t.Fatalf("streaming requests should not be cached, handler called %d times", callCount) + + mw.simple.wg.Wait() + + req2 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(streamingBody)) + req2.Header.Set("Content-Type", "application/json") + rec2 := httptest.NewRecorder() + e.ServeHTTP(rec2, req2) + if got := rec2.Header().Get("X-Cache"); got != "HIT (exact)" { + t.Fatalf("streaming request should reuse cached full response, got X-Cache=%q", got) + } + if got := rec2.Header().Get("Content-Type"); got != "text/event-stream" { + t.Fatalf("streaming cache hit Content-Type = %q, want text/event-stream", got) + } + if !bytes.Contains(rec2.Body.Bytes(), []byte("shared cached response")) || !bytes.Contains(rec2.Body.Bytes(), []byte("[DONE]")) { + t.Fatalf("streaming cache hit body = %q, want synthesized SSE", rec2.Body.String()) + } + if callCount != 1 { + t.Fatalf("expected streaming replay to avoid second handler call, got %d calls", callCount) } } diff --git a/internal/responsecache/responsecache.go b/internal/responsecache/responsecache.go index 9b66120ac..6f24efa6d 100644 --- a/internal/responsecache/responsecache.go +++ b/internal/responsecache/responsecache.go @@ -94,7 +94,8 @@ func (m *ResponseCacheMiddleware) Middleware() echo.MiddlewareFunc { // HandleRequest runs the full dual-layer cache check (exact then semantic) for a // translated inference request that has already been guardrail-patched. // body is the final patched request bytes; next is the real LLM call. -// Returns true if the request was served from cache. +// Streaming misses are reconstructed into full JSON before storage; streaming +// hits replay that stored JSON as synthetic SSE. func (m *ResponseCacheMiddleware) HandleRequest(c *echo.Context, body []byte, next func() error) error { if m == nil { return next() diff --git a/internal/responsecache/semantic.go b/internal/responsecache/semantic.go index 8dae23e3e..01b292a4d 100644 --- a/internal/responsecache/semantic.go +++ b/internal/responsecache/semantic.go @@ -64,10 +64,6 @@ func (m *semanticCacheMiddleware) Handle(c *echo.Context, body []byte, next func return next() } - if isStreamingRequest(path, body) { - return next() - } - ctx := c.Request().Context() plan := core.GetExecutionPlan(ctx) @@ -105,20 +101,20 @@ func (m *semanticCacheMiddleware) Handle(c *echo.Context, body []byte, next func } if len(results) > 0 && float64(results[0].Score) >= threshold { - auditlog.EnrichEntryWithCacheType(c, CacheTypeSemantic) - c.Response().Header().Set("Content-Type", "application/json") - c.Response().Header().Set("X-Cache", "HIT (semantic)") - c.Response().WriteHeader(http.StatusOK) - _, _ = c.Response().Write(results[0].Response) - if m.hitRecorder != nil { - m.hitRecorder(c, results[0].Response, CacheTypeSemantic) + replayErr := writeCachedResponse(c, path, body, results[0].Response, CacheTypeSemantic) + if replayErr == nil { + auditlog.EnrichEntryWithCacheType(c, CacheTypeSemantic) + if m.hitRecorder != nil { + m.hitRecorder(c, results[0].Response, CacheTypeSemantic) + } + slog.Info("semantic cache hit", + "path", path, + "score", results[0].Score, + "request_id", c.Request().Header.Get("X-Request-ID"), + ) + return nil } - slog.Info("semantic cache hit", - "path", path, - "score", results[0].Score, - "request_id", c.Request().Header.Get("X-Request-ID"), - ) - return nil + slog.Warn("semantic cache replay failed", "path", path, "err", replayErr) } capture := &responseCapture{ @@ -137,8 +133,11 @@ func (m *semanticCacheMiddleware) Handle(c *echo.Context, body []byte, next func if core.GetFallbackUsed(c.Request().Context()) { return nil } - - data := bytes.Clone(capture.body.Bytes()) + data, ok := capture.cachedBody(path, streamResponseDefaultsFromContext(c.Request().Context()), c.Response().Header().Get("Content-Type")) + if !ok { + slog.Warn("semantic cache: failed to reconstruct cacheable response body", "path", path) + return nil + } ttl := time.Duration(m.cfg.TTL) * time.Second if ttl == 0 { ttl = time.Hour @@ -346,13 +345,13 @@ func extractTextFromContent(content any) string { // (e.g. "/v1/chat/completions") and isolates entries across distinct endpoints. func computeParamsHash(body []byte, endpointPath string, plan *core.ExecutionPlan, guardrailsHash, embedderIdentity string) string { var req struct { - Model string `json:"model"` - Temperature *float64 `json:"temperature"` - TopP *float64 `json:"top_p"` - MaxTokens *int `json:"max_tokens"` - Tools []map[string]any `json:"tools"` - ResponseFormat any `json:"response_format"` - Stream bool `json:"stream"` + Model string `json:"model"` + Temperature *float64 `json:"temperature"` + TopP *float64 `json:"top_p"` + MaxTokens *int `json:"max_tokens"` + Tools []map[string]any `json:"tools"` + ResponseFormat any `json:"response_format"` + StreamOptions *core.StreamOptions `json:"stream_options"` } _ = json.Unmarshal(body, &req) @@ -397,7 +396,10 @@ func computeParamsHash(body []byte, endpointPath string, plan *core.ExecutionPla } h.Write([]byte{0}) - h.Write([]byte(strconv.FormatBool(req.Stream))) + if streamOptions := normalizeStreamOptionsForCache(req.StreamOptions); streamOptions != nil { + soJSON, _ := json.Marshal(streamOptions) + h.Write(soJSON) + } h.Write([]byte{0}) h.Write([]byte(guardrailsHash)) @@ -512,9 +514,8 @@ func ShouldSkipExactCache(req *http.Request) bool { return strings.EqualFold(req.Header.Get("X-Cache-Type"), CacheTypeSemantic) } -// ShouldSkipAllCache reports whether caching must be bypassed for this request -// (X-Cache-Control: no-store or Cache-Control containing no-store), matching -// shouldSkipSemanticCache / shouldSkipCache semantics for the no-store directive. +// ShouldSkipAllCache reports whether caching must be bypassed for this request, +// matching the exact-cache middleware semantics for no-cache and no-store. func ShouldSkipAllCache(req *http.Request) bool { if strings.EqualFold(req.Header.Get("X-Cache-Control"), "no-store") { return true @@ -526,7 +527,7 @@ func ShouldSkipAllCache(req *http.Request) bool { directives := strings.Split(strings.ToLower(cc), ",") for _, d := range directives { d = strings.TrimSpace(d) - if d == "no-store" { + if d == "no-cache" || d == "no-store" { return true } } diff --git a/internal/responsecache/semantic_test.go b/internal/responsecache/semantic_test.go index 081460b72..cc286eeb4 100644 --- a/internal/responsecache/semantic_test.go +++ b/internal/responsecache/semantic_test.go @@ -207,6 +207,25 @@ func TestSemanticCacheMiddleware_ParamsHashIsolation_Temperature(t *testing.T) { } } +func TestComputeParamsHash_StreamIncludeUsageChangesHash(t *testing.T) { + base := []byte(`{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"same prompt"}]}`) + withUsage := []byte(`{"model":"gpt-4","stream":true,"stream_options":{"include_usage":true},"messages":[{"role":"user","content":"same prompt"}]}`) + plan := &core.ExecutionPlan{ + Mode: core.ExecutionModeTranslated, + ProviderType: "openai", + Resolution: &core.RequestModelResolution{ + ResolvedSelector: core.ModelSelector{Provider: "openai", Model: "gpt-4"}, + }, + } + + first := computeParamsHash(base, "/v1/chat/completions", plan, "", "") + second := computeParamsHash(withUsage, "/v1/chat/completions", plan, "", "") + + if first == second { + t.Fatal("stream_options.include_usage should affect semantic params_hash") + } +} + func TestSemanticCacheMiddleware_GuardrailsHashIsolation(t *testing.T) { m, store, emb := newTestSemanticMiddleware(0.90, 10, false) emb.vector = []float32{1, 0, 0} @@ -284,15 +303,70 @@ func TestSemanticCacheMiddleware_ExcludeSystemPrompt(t *testing.T) { } } -func TestSemanticCacheMiddleware_StreamingSkipped(t *testing.T) { +func TestSemanticCacheMiddleware_StreamingMissPopulatesSemanticCacheAcrossModes(t *testing.T) { m, store, _ := newTestSemanticMiddleware(0.90, 10, false) + streamBody := []byte(`{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"semantic-stream-cache"}]}`) + jsonBody := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"semantic-stream-cache"}]}`) + e := echo.New() + handlerCalls := 0 + + run := func(body []byte) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + if err := m.Handle(c, body, func() error { + handlerCalls++ + c.Response().Header().Set("Content-Type", "text/event-stream") + c.Response().WriteHeader(http.StatusOK) + _, _ = c.Response().Write([]byte("data: {\"id\":\"chatcmpl-semantic-stream\",\"object\":\"chat.completion.chunk\",\"created\":1234567890,\"model\":\"gpt-4\",\"provider\":\"openai\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"Semantic\"},\"finish_reason\":null}]}\n\n")) + _, _ = c.Response().Write([]byte("data: {\"id\":\"chatcmpl-semantic-stream\",\"object\":\"chat.completion.chunk\",\"created\":1234567890,\"model\":\"gpt-4\",\"provider\":\"openai\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" cache\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":2,\"total_tokens\":9}}\n\n")) + _, _ = c.Response().Write([]byte("data: [DONE]\n\n")) + return nil + }); err != nil { + t.Fatalf("Handle error: %v", err) + } + return rec + } + + rec1 := run(streamBody) + if got := rec1.Header().Get("X-Cache"); got != "" { + t.Fatalf("streaming miss should not be cache hit, got X-Cache=%q", got) + } - body := []byte(`{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"hi"}]}`) - serveSemanticRequest(t, m, body, "") m.wg.Wait() - if store.Len() != 0 { - t.Fatal("streaming requests should be skipped by semantic cache") + if store.Len() != 1 { + t.Fatalf("expected streaming miss to populate semantic cache, got %d entries", store.Len()) + } + + rec2 := run(jsonBody) + if got := rec2.Header().Get("X-Cache"); got != "HIT (semantic)" { + t.Fatalf("non-streaming follow-up should hit semantic cache, got X-Cache=%q", got) + } + if got := rec2.Header().Get("Content-Type"); got != "application/json" { + t.Fatalf("non-streaming semantic hit Content-Type = %q, want application/json", got) + } + if !bytes.Contains(rec2.Body.Bytes(), []byte(`"content":"Semantic cache"`)) { + t.Fatalf("semantic cache hit body = %q, want reconstructed JSON response", rec2.Body.String()) + } + if handlerCalls != 1 { + t.Fatalf("semantic hit should not call handler again, got %d calls", handlerCalls) + } + + rec3 := run(streamBody) + if got := rec3.Header().Get("X-Cache"); got != "HIT (semantic)" { + t.Fatalf("streaming follow-up should hit semantic cache, got X-Cache=%q", got) + } + if got := rec3.Header().Get("Content-Type"); got != "text/event-stream" { + t.Fatalf("streaming semantic hit Content-Type = %q, want text/event-stream", got) + } + if !bytes.Contains(rec3.Body.Bytes(), []byte("Semantic cache")) || !bytes.Contains(rec3.Body.Bytes(), []byte("[DONE]")) { + t.Fatalf("streaming semantic hit body = %q, want synthesized SSE", rec3.Body.String()) + } + if handlerCalls != 1 { + t.Fatalf("streaming semantic hit should not call handler again, got %d calls", handlerCalls) } } @@ -462,6 +536,14 @@ func TestShouldSkipAllCache_CacheControlNoStore(t *testing.T) { } } +func TestShouldSkipAllCache_CacheControlNoCache(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + req.Header.Set("Cache-Control", "private, no-cache, max-age=0") + if !ShouldSkipAllCache(req) { + t.Fatal("expected ShouldSkipAllCache for Cache-Control: no-cache") + } +} + func TestSemanticCacheMiddleware_HitMarksAuditEntryCacheType(t *testing.T) { m, _, _ := newTestSemanticMiddleware(0.90, 10, false) body := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"semantic-cache-type"}]}`) diff --git a/internal/responsecache/simple.go b/internal/responsecache/simple.go index 569cfad46..da3d7373c 100644 --- a/internal/responsecache/simple.go +++ b/internal/responsecache/simple.go @@ -5,6 +5,7 @@ import ( "context" "crypto/sha256" "encoding/hex" + "encoding/json" "io" "log/slog" "net/http" @@ -80,9 +81,6 @@ func (m *simpleCacheMiddleware) Middleware() echo.MiddlewareFunc { if !cacheable { return next(c) } - if isStreamingRequest(path, body) { - return next(c) - } plan := core.GetExecutionPlan(c.Request().Context()) if shouldSkipCacheForExecutionPlan(plan) { return next(c) @@ -110,11 +108,11 @@ func (m *simpleCacheMiddleware) TryHit(c *echo.Context, body []byte) (bool, erro return false, nil } if len(cached) > 0 { + if err := writeCachedResponse(c, path, body, cached, CacheTypeExact); err != nil { + slog.Warn("response cache replay failed", "path", path, "cache_type", CacheTypeExact, "err", err) + return false, nil + } auditlog.EnrichEntryWithCacheType(c, CacheTypeExact) - c.Response().Header().Set("Content-Type", "application/json") - c.Response().Header().Set("X-Cache", "HIT (exact)") - c.Response().WriteHeader(http.StatusOK) - _, _ = c.Response().Write(cached) if m.hitRecorder != nil { m.hitRecorder(c, cached, CacheTypeExact) } @@ -148,7 +146,11 @@ func (m *simpleCacheMiddleware) StoreAfter(c *echo.Context, body []byte, next fu if core.GetFallbackUsed(c.Request().Context()) { return nil } - data := bytes.Clone(capture.body.Bytes()) + data, ok := capture.cachedBody(path, streamResponseDefaultsFromContext(c.Request().Context()), c.Response().Header().Get("Content-Type")) + if !ok { + slog.Warn("response cache: failed to reconstruct cacheable response body", "path", path) + return nil + } m.enqueueWrite(cacheWriteJob{key: key, data: data}) } return nil @@ -282,7 +284,7 @@ func hashRequest(path string, body []byte, plan *core.ExecutionPlan) string { h.Write([]byte(plan.ResolvedQualifiedModel())) h.Write([]byte{0}) } - h.Write(body) + h.Write(cacheKeyRequestBody(path, body)) return hex.EncodeToString(h.Sum(nil)) } @@ -292,6 +294,21 @@ type responseCapture struct { status int } +func (r *responseCapture) cachedBody(path string, defaults streamResponseDefaults, contentType string) ([]byte, bool) { + if r == nil || r.body == nil || r.body.Len() == 0 { + return nil, false + } + + raw := bytes.Clone(r.body.Bytes()) + if isEventStreamContentType(contentType) { + return reconstructStreamingResponse(path, raw, defaults) + } + if !json.Valid(raw) { + return nil, false + } + return raw, true +} + func (r *responseCapture) WriteHeader(code int) { r.status = code r.ResponseWriter.WriteHeader(code) diff --git a/internal/responsecache/stream_cache.go b/internal/responsecache/stream_cache.go new file mode 100644 index 000000000..605587829 --- /dev/null +++ b/internal/responsecache/stream_cache.go @@ -0,0 +1,1326 @@ +package responsecache + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "sort" + "strings" + + "github.com/labstack/echo/v5" + + "gomodel/internal/auditlog" + "gomodel/internal/core" +) + +var ( + cacheLFEventBoundary = []byte("\n\n") + cacheCRLFEventBoundary = []byte("\r\n\r\n") + cacheDataPrefix = []byte("data:") + cacheDonePayload = []byte("[DONE]") +) + +type streamResponseDefaults struct { + Model string + Provider string +} + +type chatToolCallState struct { + Index int + ID string + Type string + Name string + Arguments strings.Builder +} + +type chatChoiceState struct { + Index int + Role string + Content strings.Builder + Reasoning strings.Builder + FinishReason string + Logprobs json.RawMessage + HasLogprobs bool + ToolCalls map[int]*chatToolCallState +} + +type chatStreamCacheBuilder struct { + defaults streamResponseDefaults + seen bool + ID string + Model string + Provider string + Object string + SystemFingerprint string + Created int64 + Usage map[string]any + Choices map[int]*chatChoiceState +} + +type responsesOutputState struct { + Index int + Item map[string]any + + TextParts map[int]*strings.Builder + ReasoningParts map[int]*strings.Builder + Arguments strings.Builder + HasArgs bool +} + +type responsesStreamCacheBuilder struct { + defaults streamResponseDefaults + seen bool + Response map[string]any + ID string + Object string + Model string + Provider string + Status string + CreatedAt int64 + Usage map[string]any + Error map[string]any + Output map[int]*responsesOutputState + ItemIDs map[string]int + AssistantIndex int + HasAssistant bool + ReasoningIndex int + HasReasoning bool +} + +func cacheKeyRequestBody(path string, body []byte) []byte { + switch path { + case "/v1/chat/completions": + req, err := core.DecodeChatRequest(body, nil) + if err != nil || req == nil { + return body + } + req.Stream = false + req.StreamOptions = normalizeStreamOptionsForCache(req.StreamOptions) + normalized, err := json.Marshal(req) + if err != nil { + return body + } + return normalized + case "/v1/responses": + req, err := core.DecodeResponsesRequest(body, nil) + if err != nil || req == nil { + return body + } + req.Stream = false + req.StreamOptions = normalizeStreamOptionsForCache(req.StreamOptions) + normalized, err := json.Marshal(req) + if err != nil { + return body + } + return normalized + default: + return body + } +} + +func streamResponseDefaultsFromContext(ctx context.Context) streamResponseDefaults { + var defaults streamResponseDefaults + + plan := core.GetExecutionPlan(ctx) + if plan == nil { + return defaults + } + + defaults.Provider = strings.TrimSpace(plan.ProviderType) + if plan.Resolution != nil { + if defaults.Provider == "" { + defaults.Provider = strings.TrimSpace(plan.Resolution.ResolvedSelector.Provider) + } + defaults.Model = strings.TrimSpace(plan.Resolution.ResolvedSelector.Model) + } + + return defaults +} + +func isEventStreamContentType(contentType string) bool { + if contentType == "" { + return false + } + mediaType := strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0])) + return mediaType == "text/event-stream" +} + +func writeCachedResponse(c *echo.Context, path string, requestBody, cached []byte, cacheType string) error { + if isStreamingRequest(path, requestBody) { + streamBody, err := renderCachedStream(path, requestBody, cached) + if err != nil { + return err + } + auditlog.EnrichEntryWithStream(c, true) + c.Response().Header().Set("Content-Type", "text/event-stream") + c.Response().Header().Set("Cache-Control", "no-cache") + c.Response().Header().Set("Connection", "keep-alive") + c.Response().Header().Set("X-Cache", "HIT ("+cacheType+")") + c.Response().WriteHeader(http.StatusOK) + _, _ = c.Response().Write(streamBody) + return nil + } + + c.Response().Header().Set("Content-Type", "application/json") + c.Response().Header().Set("X-Cache", "HIT ("+cacheType+")") + c.Response().WriteHeader(http.StatusOK) + _, _ = c.Response().Write(cached) + return nil +} + +func renderCachedStream(path string, requestBody, cached []byte) ([]byte, error) { + switch path { + case "/v1/chat/completions": + return renderCachedChatStream(requestBody, cached) + case "/v1/responses": + return renderCachedResponsesStream(requestBody, cached) + default: + return nil, errors.New("cached streaming replay is not supported for this path") + } +} + +func renderCachedChatStream(requestBody, cached []byte) ([]byte, error) { + var resp core.ChatResponse + if err := json.Unmarshal(cached, &resp); err != nil { + return nil, err + } + + var out bytes.Buffer + includeUsage := streamIncludeUsageRequested("/v1/chat/completions", requestBody) + usage := chatUsageMap(resp.Usage) + if !includeUsage { + usage = nil + } + for _, choice := range resp.Choices { + delta := map[string]any{} + role := strings.TrimSpace(choice.Message.Role) + if role == "" { + role = "assistant" + } + delta["role"] = role + + if content := core.ExtractTextContent(choice.Message.Content); content != "" { + delta["content"] = content + } + if reasoning := chatReasoningContent(choice.Message); reasoning != "" { + delta["reasoning_content"] = reasoning + } + if len(choice.Message.ToolCalls) > 0 { + delta["tool_calls"] = renderChatToolCalls(choice.Message.ToolCalls) + } + + renderedChoice := map[string]any{ + "index": choice.Index, + "delta": delta, + "finish_reason": choice.FinishReason, + } + if len(choice.Logprobs) > 0 { + renderedChoice["logprobs"] = choice.Logprobs + } + + chunk := map[string]any{ + "id": resp.ID, + "object": "chat.completion.chunk", + "model": resp.Model, + "choices": []map[string]any{renderedChoice}, + } + if resp.Created != 0 { + chunk["created"] = resp.Created + } + if resp.Provider != "" { + chunk["provider"] = resp.Provider + } + if resp.SystemFingerprint != "" { + chunk["system_fingerprint"] = resp.SystemFingerprint + } + if err := appendSSEJSONEvent(&out, "", chunk); err != nil { + return nil, err + } + } + + if usage != nil { + chunk := map[string]any{ + "id": resp.ID, + "object": "chat.completion.chunk", + "model": resp.Model, + "choices": []map[string]any{}, + "usage": usage, + } + if resp.Created != 0 { + chunk["created"] = resp.Created + } + if resp.Provider != "" { + chunk["provider"] = resp.Provider + } + if err := appendSSEJSONEvent(&out, "", chunk); err != nil { + return nil, err + } + } + + out.WriteString("data: [DONE]\n\n") + return out.Bytes(), nil +} + +func renderCachedResponsesStream(requestBody, cached []byte) ([]byte, error) { + var resp map[string]any + if err := json.Unmarshal(cached, &resp); err != nil { + return nil, err + } + + var out bytes.Buffer + includeUsage := streamIncludeUsageRequested("/v1/responses", requestBody) + responseWithUsage := cloneJSONMap(resp) + if !includeUsage { + delete(responseWithUsage, "usage") + } + + respID, _ := responseWithUsage["id"].(string) + respObject, _ := responseWithUsage["object"].(string) + respModel, _ := responseWithUsage["model"].(string) + respProvider, _ := responseWithUsage["provider"].(string) + respCreatedAt := responseWithUsage["created_at"] + created := map[string]any{ + "id": respID, + "object": nonEmpty(respObject, "response"), + "status": "in_progress", + "model": respModel, + "provider": respProvider, + "created_at": respCreatedAt, + } + if err := appendSSEJSONEvent(&out, "response.created", map[string]any{ + "type": "response.created", + "response": created, + }); err != nil { + return nil, err + } + + output, _ := responseWithUsage["output"].([]any) + for i, itemAny := range output { + itemMap, ok := itemAny.(map[string]any) + if !ok { + continue + } + itemID, _ := itemMap["id"].(string) + added := responsesAddedItem(itemMap) + if err := appendSSEJSONEvent(&out, "response.output_item.added", map[string]any{ + "type": "response.output_item.added", + "item": added, + "output_index": i, + }); err != nil { + return nil, err + } + if err := appendResponsesItemDeltaEvents(&out, itemMap, itemID, i); err != nil { + return nil, err + } + done := cloneJSONMap(itemMap) + if _, ok := done["status"]; !ok || done["status"] == "" { + done["status"] = "completed" + } + if err := appendSSEJSONEvent(&out, "response.output_item.done", map[string]any{ + "type": "response.output_item.done", + "item": done, + "output_index": i, + }); err != nil { + return nil, err + } + } + + terminalEventName := responsesTerminalEventName(responseWithUsage) + if err := appendSSEJSONEvent(&out, terminalEventName, map[string]any{ + "type": terminalEventName, + "response": responseWithUsage, + }); err != nil { + return nil, err + } + out.WriteString("data: [DONE]\n\n") + return out.Bytes(), nil +} + +func reconstructStreamingResponse(path string, raw []byte, defaults streamResponseDefaults) ([]byte, bool) { + switch path { + case "/v1/chat/completions": + builder := &chatStreamCacheBuilder{ + defaults: defaults, + Choices: make(map[int]*chatChoiceState), + } + parseSSEJSONEvents(raw, builder.OnJSONEvent) + return builder.Build() + case "/v1/responses": + builder := &responsesStreamCacheBuilder{ + defaults: defaults, + Output: make(map[int]*responsesOutputState), + ItemIDs: make(map[string]int), + } + parseSSEJSONEvents(raw, builder.OnJSONEvent) + return builder.Build() + default: + return nil, false + } +} + +func parseSSEJSONEvents(raw []byte, onJSON func(map[string]any)) { + for len(raw) > 0 { + idx, sepLen := nextCacheEventBoundary(raw) + event := raw + if idx != -1 { + event = raw[:idx] + raw = raw[idx+sepLen:] + } else { + raw = nil + } + + payload, ok := parseCacheEventJSON(event) + if !ok { + if idx == -1 { + break + } + continue + } + onJSON(payload) + if idx == -1 { + break + } + } +} + +func parseCacheEventJSON(event []byte) (map[string]any, bool) { + lines := bytes.Split(event, []byte("\n")) + payloadLines := make([][]byte, 0, len(lines)) + for _, line := range lines { + data, ok := parseCacheDataLine(line) + if !ok { + continue + } + payloadLines = append(payloadLines, data) + } + if len(payloadLines) == 0 { + return nil, false + } + + jsonData := bytes.Join(payloadLines, []byte("\n")) + if bytes.Equal(jsonData, cacheDonePayload) { + return nil, false + } + + var payload map[string]any + if err := json.Unmarshal(jsonData, &payload); err != nil { + return nil, false + } + return payload, true +} + +func nextCacheEventBoundary(data []byte) (idx int, sepLen int) { + lfIdx := bytes.Index(data, cacheLFEventBoundary) + crlfIdx := bytes.Index(data, cacheCRLFEventBoundary) + + switch { + case lfIdx == -1: + if crlfIdx == -1 { + return -1, 0 + } + return crlfIdx, len(cacheCRLFEventBoundary) + case crlfIdx == -1 || lfIdx < crlfIdx: + return lfIdx, len(cacheLFEventBoundary) + default: + return crlfIdx, len(cacheCRLFEventBoundary) + } +} + +func parseCacheDataLine(line []byte) ([]byte, bool) { + line = bytes.TrimSuffix(line, []byte("\r")) + if !bytes.HasPrefix(line, cacheDataPrefix) { + return nil, false + } + payload := bytes.TrimPrefix(line, cacheDataPrefix) + if len(payload) > 0 && payload[0] == ' ' { + payload = payload[1:] + } + return payload, true +} + +func (b *chatStreamCacheBuilder) OnJSONEvent(event map[string]any) { + if b == nil { + return + } + b.seen = true + + if id, ok := event["id"].(string); ok && id != "" { + b.ID = id + } + if model, ok := event["model"].(string); ok && model != "" { + b.Model = model + } + if provider, ok := event["provider"].(string); ok && provider != "" { + b.Provider = provider + } + if object, ok := event["object"].(string); ok && object != "" { + b.Object = object + } + if fingerprint, ok := event["system_fingerprint"].(string); ok && fingerprint != "" { + b.SystemFingerprint = fingerprint + } + if created, ok := jsonNumberToInt64(event["created"]); ok { + b.Created = created + } + if usage, ok := event["usage"].(map[string]any); ok { + b.Usage = cloneJSONMap(usage) + } + + choices, ok := event["choices"].([]any) + if !ok { + return + } + for _, choiceAny := range choices { + choiceMap, ok := choiceAny.(map[string]any) + if !ok { + continue + } + index, ok := jsonNumberToInt(choiceMap["index"]) + if !ok { + index = len(b.Choices) + } + state := b.choice(index) + if finish, ok := choiceMap["finish_reason"].(string); ok && finish != "" { + state.FinishReason = finish + } + if logprobs, ok := choiceMap["logprobs"]; ok { + raw, err := json.Marshal(logprobs) + if err == nil { + state.Logprobs = raw + state.HasLogprobs = true + } + } + + delta, ok := choiceMap["delta"].(map[string]any) + if !ok { + continue + } + if role, ok := delta["role"].(string); ok && role != "" { + state.Role = role + } + if content, ok := delta["content"].(string); ok && content != "" { + _, _ = state.Content.WriteString(content) + } + if reasoning, ok := delta["reasoning_content"].(string); ok && reasoning != "" { + _, _ = state.Reasoning.WriteString(reasoning) + } + + toolCalls, ok := delta["tool_calls"].([]any) + if !ok { + continue + } + for _, toolAny := range toolCalls { + toolMap, ok := toolAny.(map[string]any) + if !ok { + continue + } + toolIndex, ok := jsonNumberToInt(toolMap["index"]) + if !ok { + toolIndex = len(state.ToolCalls) + } + toolState := state.toolCall(toolIndex) + if id, ok := toolMap["id"].(string); ok && id != "" { + toolState.ID = id + } + if typ, ok := toolMap["type"].(string); ok && typ != "" { + toolState.Type = typ + } + function, ok := toolMap["function"].(map[string]any) + if !ok { + continue + } + if name, ok := function["name"].(string); ok && name != "" { + toolState.Name = name + } + if arguments, ok := function["arguments"].(string); ok && arguments != "" { + _, _ = toolState.Arguments.WriteString(arguments) + } + } + } +} + +func (b *chatStreamCacheBuilder) Build() ([]byte, bool) { + if b == nil || !b.seen { + return nil, false + } + + choiceIndexes := make([]int, 0, len(b.Choices)) + for index := range b.Choices { + choiceIndexes = append(choiceIndexes, index) + } + sort.Ints(choiceIndexes) + + choices := make([]map[string]any, 0, len(choiceIndexes)) + for _, index := range choiceIndexes { + state := b.Choices[index] + message := map[string]any{ + "role": nonEmpty(state.Role, "assistant"), + } + + content := state.Content.String() + toolCalls := buildChatToolCalls(state.ToolCalls) + switch { + case content != "": + message["content"] = content + case len(toolCalls) > 0: + message["content"] = nil + default: + message["content"] = "" + } + if len(toolCalls) > 0 { + message["tool_calls"] = toolCalls + } + if reasoning := state.Reasoning.String(); reasoning != "" { + message["reasoning_content"] = reasoning + } + + choice := map[string]any{ + "index": index, + "message": message, + "finish_reason": state.FinishReason, + } + if state.HasLogprobs { + choice["logprobs"] = state.Logprobs + } + choices = append(choices, choice) + } + + response := map[string]any{ + "id": b.ID, + "object": "chat.completion", + "model": nonEmpty(b.Model, b.defaults.Model), + "choices": choices, + } + if provider := nonEmpty(b.Provider, b.defaults.Provider); provider != "" { + response["provider"] = provider + } + if b.Created != 0 { + response["created"] = b.Created + } + if b.SystemFingerprint != "" { + response["system_fingerprint"] = b.SystemFingerprint + } + if b.Usage != nil { + response["usage"] = b.Usage + } + + data, err := json.Marshal(response) + if err != nil { + return nil, false + } + return data, true +} + +func (b *chatStreamCacheBuilder) choice(index int) *chatChoiceState { + state, ok := b.Choices[index] + if ok { + return state + } + state = &chatChoiceState{ + Index: index, + ToolCalls: make(map[int]*chatToolCallState), + } + b.Choices[index] = state + return state +} + +func (c *chatChoiceState) toolCall(index int) *chatToolCallState { + state, ok := c.ToolCalls[index] + if ok { + return state + } + state = &chatToolCallState{Index: index} + c.ToolCalls[index] = state + return state +} + +func (b *responsesStreamCacheBuilder) OnJSONEvent(event map[string]any) { + if b == nil { + return + } + b.seen = true + + eventType, _ := event["type"].(string) + switch eventType { + case "response.created", "response.completed", "response.failed", "response.incomplete", "response.done": + response, ok := event["response"].(map[string]any) + if !ok { + return + } + b.captureResponseMetadata(response) + if output, ok := response["output"].([]any); ok { + for index, itemAny := range output { + itemMap, ok := itemAny.(map[string]any) + if !ok { + continue + } + b.output(index).SetItem(itemMap) + if itemID, _ := itemMap["id"].(string); itemID != "" { + b.ItemIDs[itemID] = index + } + if itemType, _ := itemMap["type"].(string); itemType == "message" { + if role, _ := itemMap["role"].(string); role == "assistant" { + b.AssistantIndex = index + b.HasAssistant = true + } + } else if itemType == "reasoning" { + b.ReasoningIndex = index + b.HasReasoning = true + } + } + } + case "response.output_item.added", "response.output_item.done": + index, ok := jsonNumberToInt(event["output_index"]) + if !ok { + return + } + item, ok := event["item"].(map[string]any) + if !ok { + return + } + state := b.output(index) + state.SetItem(item) + if itemID, _ := item["id"].(string); itemID != "" { + b.ItemIDs[itemID] = index + } + if itemType, _ := item["type"].(string); itemType == "message" { + if role, _ := item["role"].(string); role == "assistant" { + b.AssistantIndex = index + b.HasAssistant = true + } + } else if itemType == "reasoning" { + b.ReasoningIndex = index + b.HasReasoning = true + } + case "response.output_text.delta": + delta, _ := event["delta"].(string) + if delta == "" { + return + } + contentIndex, _ := jsonNumberToInt(event["content_index"]) + index, ok := b.lookupOutputIndex(event) + if !ok { + index = 0 + if b.HasAssistant { + index = b.AssistantIndex + } + } + b.rememberOutputLocator(event, index) + b.AssistantIndex = index + b.HasAssistant = true + b.output(index).AppendText(contentIndex, delta) + case "response.reasoning_text.delta": + delta, _ := event["delta"].(string) + if delta == "" { + return + } + contentIndex, _ := jsonNumberToInt(event["content_index"]) + outputIndex, hasOutputIndex := jsonNumberToInt(event["output_index"]) + index, ok := b.lookupOutputIndex(event) + if !ok { + index = b.ensureReasoningOutputIndex(outputIndex, hasOutputIndex) + } + b.rememberOutputLocator(event, index) + b.ReasoningIndex = index + b.HasReasoning = true + b.output(index).AppendReasoning(contentIndex, delta) + case "response.function_call_arguments.delta": + index, ok := b.lookupOutputIndex(event) + if !ok { + return + } + delta, _ := event["delta"].(string) + if delta == "" { + return + } + b.output(index).AppendArguments(delta) + case "response.function_call_arguments.done": + index, ok := b.lookupOutputIndex(event) + if !ok { + return + } + arguments, _ := event["arguments"].(string) + b.output(index).SetArguments(arguments) + } +} + +func (b *responsesStreamCacheBuilder) Build() ([]byte, bool) { + if b == nil || !b.seen { + return nil, false + } + + indexes := make([]int, 0, len(b.Output)) + for index := range b.Output { + indexes = append(indexes, index) + } + sort.Ints(indexes) + + output := make([]map[string]any, 0, len(indexes)) + for _, index := range indexes { + item := b.Output[index].BuildItem() + if len(item) == 0 { + continue + } + output = append(output, item) + } + + response := cloneJSONMap(b.Response) + if response == nil { + response = map[string]any{ + "id": b.ID, + "object": nonEmpty(b.Object, "response"), + "created_at": b.CreatedAt, + "model": nonEmpty(b.Model, b.defaults.Model), + "status": nonEmpty(b.Status, "completed"), + } + if provider := nonEmpty(b.Provider, b.defaults.Provider); provider != "" { + response["provider"] = provider + } + if b.Usage != nil { + response["usage"] = b.Usage + } + if b.Error != nil { + response["error"] = b.Error + } + } + response["output"] = output + if _, ok := response["id"]; !ok { + response["id"] = b.ID + } + if _, ok := response["object"]; !ok { + response["object"] = nonEmpty(b.Object, "response") + } + if _, ok := response["created_at"]; !ok && b.CreatedAt != 0 { + response["created_at"] = b.CreatedAt + } + if _, ok := response["model"]; !ok { + response["model"] = nonEmpty(b.Model, b.defaults.Model) + } + if _, ok := response["status"]; !ok { + response["status"] = nonEmpty(b.Status, "completed") + } + if provider := nonEmpty(b.Provider, b.defaults.Provider); provider != "" { + if _, ok := response["provider"]; !ok { + response["provider"] = provider + } + } + if b.Usage != nil { + if _, ok := response["usage"]; !ok { + response["usage"] = b.Usage + } + } + if b.Error != nil { + if _, ok := response["error"]; !ok { + response["error"] = b.Error + } + } + + data, err := json.Marshal(response) + if err != nil { + return nil, false + } + return data, true +} + +func (b *responsesStreamCacheBuilder) captureResponseMetadata(response map[string]any) { + b.Response = cloneJSONMap(response) + if id, ok := response["id"].(string); ok && id != "" { + b.ID = id + } + if object, ok := response["object"].(string); ok && object != "" { + b.Object = object + } + if model, ok := response["model"].(string); ok && model != "" { + b.Model = model + } + if provider, ok := response["provider"].(string); ok && provider != "" { + b.Provider = provider + } + if status, ok := response["status"].(string); ok && status != "" { + b.Status = status + } + if createdAt, ok := jsonNumberToInt64(response["created_at"]); ok { + b.CreatedAt = createdAt + } + if usage, ok := response["usage"].(map[string]any); ok { + b.Usage = cloneJSONMap(usage) + } + if errMap, ok := response["error"].(map[string]any); ok { + b.Error = cloneJSONMap(errMap) + } +} + +func (b *responsesStreamCacheBuilder) output(index int) *responsesOutputState { + state, ok := b.Output[index] + if ok { + return state + } + state = &responsesOutputState{Index: index} + b.Output[index] = state + return state +} + +func (b *responsesStreamCacheBuilder) lookupOutputIndex(event map[string]any) (int, bool) { + if index, ok := jsonNumberToInt(event["output_index"]); ok { + return index, true + } + itemID, _ := event["item_id"].(string) + index, ok := b.ItemIDs[itemID] + return index, ok +} + +func (b *responsesStreamCacheBuilder) rememberOutputLocator(event map[string]any, index int) { + itemID, _ := event["item_id"].(string) + if itemID == "" { + return + } + b.ItemIDs[itemID] = index +} + +func (b *responsesStreamCacheBuilder) ensureReasoningOutputIndex(outputIndex int, hasOutputIndex bool) int { + if hasOutputIndex { + b.ReasoningIndex = outputIndex + b.HasReasoning = true + return outputIndex + } + if b.HasReasoning { + return b.ReasoningIndex + } + index := len(b.Output) + b.ReasoningIndex = index + b.HasReasoning = true + return index +} + +func (s *responsesOutputState) SetItem(item map[string]any) { + s.Item = cloneJSONMap(item) +} + +func (s *responsesOutputState) AppendText(contentIndex int, delta string) { + if delta == "" { + return + } + part := s.textPart(contentIndex) + _, _ = part.WriteString(delta) +} + +func (s *responsesOutputState) AppendReasoning(contentIndex int, delta string) { + if delta == "" { + return + } + part := s.reasoningPart(contentIndex) + _, _ = part.WriteString(delta) +} + +func (s *responsesOutputState) AppendArguments(delta string) { + if delta == "" { + return + } + _, _ = s.Arguments.WriteString(delta) + s.HasArgs = true +} + +func (s *responsesOutputState) SetArguments(arguments string) { + s.Arguments = strings.Builder{} + _, _ = s.Arguments.WriteString(arguments) + s.HasArgs = true +} + +func (s *responsesOutputState) textPart(contentIndex int) *strings.Builder { + if s.TextParts == nil { + s.TextParts = make(map[int]*strings.Builder) + } + part, ok := s.TextParts[contentIndex] + if ok { + return part + } + part = &strings.Builder{} + s.TextParts[contentIndex] = part + return part +} + +func (s *responsesOutputState) reasoningPart(contentIndex int) *strings.Builder { + if s.ReasoningParts == nil { + s.ReasoningParts = make(map[int]*strings.Builder) + } + part, ok := s.ReasoningParts[contentIndex] + if ok { + return part + } + part = &strings.Builder{} + s.ReasoningParts[contentIndex] = part + return part +} + +func (s *responsesOutputState) BuildItem() map[string]any { + item := cloneJSONMap(s.Item) + if item == nil { + item = make(map[string]any) + } + + itemType, _ := item["type"].(string) + if len(s.TextParts) > 0 { + if itemType == "" { + itemType = "message" + item["type"] = itemType + } + if item["role"] == nil { + item["role"] = "assistant" + } + item["content"] = buildResponsesContentParts(item["content"], "output_text", s.TextParts) + } + if len(s.ReasoningParts) > 0 { + if itemType == "" { + itemType = "reasoning" + item["type"] = itemType + } + targetField := "content" + if itemType == "reasoning" { + targetField = "summary" + } else { + if _, ok := item["summary"].([]any); ok { + targetField = "summary" + } + } + item[targetField] = buildResponsesContentParts(item[targetField], "reasoning_text", s.ReasoningParts) + } + if s.HasArgs { + if itemType == "" { + itemType = "function_call" + item["type"] = itemType + } + item["arguments"] = s.Arguments.String() + } + if _, ok := item["status"]; !ok || item["status"] == "" { + item["status"] = "completed" + } + + return item +} + +func buildResponsesContentParts(existing any, partType string, parts map[int]*strings.Builder) []map[string]any { + if len(parts) == 0 { + return nil + } + + existingParts, _ := existing.([]any) + maxIndex := len(existingParts) - 1 + for index := range parts { + if index > maxIndex { + maxIndex = index + } + } + + built := make([]map[string]any, 0, maxIndex+1) + for index := 0; index <= maxIndex; index++ { + existingPart, existingOK := cloneJSONPart(existingParts, index) + partBuilder, hasPart := parts[index] + + switch { + case hasPart: + if existingPart == nil { + existingPart = make(map[string]any) + } + existingPart["type"] = partType + existingPart["text"] = partBuilder.String() + built = append(built, existingPart) + case existingOK: + built = append(built, existingPart) + } + } + + return built +} + +func cloneJSONPart(parts []any, index int) (map[string]any, bool) { + if index < 0 || index >= len(parts) { + return nil, false + } + part, ok := parts[index].(map[string]any) + if !ok { + return nil, false + } + return cloneJSONMap(part), true +} + +func buildChatToolCalls(states map[int]*chatToolCallState) []map[string]any { + if len(states) == 0 { + return nil + } + + indexes := make([]int, 0, len(states)) + for index := range states { + indexes = append(indexes, index) + } + sort.Ints(indexes) + + toolCalls := make([]map[string]any, 0, len(indexes)) + for _, index := range indexes { + state := states[index] + toolCall := map[string]any{ + "id": state.ID, + "type": nonEmpty(state.Type, "function"), + "index": index, + "function": map[string]any{ + "name": state.Name, + "arguments": state.Arguments.String(), + }, + } + toolCalls = append(toolCalls, toolCall) + } + return toolCalls +} + +func renderChatToolCalls(toolCalls []core.ToolCall) []map[string]any { + if len(toolCalls) == 0 { + return nil + } + rendered := make([]map[string]any, 0, len(toolCalls)) + for index, toolCall := range toolCalls { + rendered = append(rendered, map[string]any{ + "index": index, + "id": toolCall.ID, + "type": nonEmpty(toolCall.Type, "function"), + "function": map[string]any{ + "name": toolCall.Function.Name, + "arguments": toolCall.Function.Arguments, + }, + }) + } + return rendered +} + +func normalizeStreamOptionsForCache(src *core.StreamOptions) *core.StreamOptions { + if src == nil || !src.IncludeUsage { + return nil + } + cloned := *src + return &cloned +} + +func streamIncludeUsageRequested(path string, requestBody []byte) bool { + switch path { + case "/v1/chat/completions": + req, err := core.DecodeChatRequest(requestBody, nil) + if err != nil || req == nil { + return false + } + return normalizeStreamOptionsForCache(req.StreamOptions) != nil + case "/v1/responses": + req, err := core.DecodeResponsesRequest(requestBody, nil) + if err != nil || req == nil { + return false + } + return normalizeStreamOptionsForCache(req.StreamOptions) != nil + default: + return false + } +} + +func chatReasoningContent(message core.ResponseMessage) string { + raw := message.ExtraFields.Lookup("reasoning_content") + if len(raw) == 0 { + return "" + } + var reasoning string + if err := json.Unmarshal(raw, &reasoning); err != nil { + return "" + } + return reasoning +} + +func responsesAddedItem(item map[string]any) map[string]any { + added := cloneJSONMap(item) + if added == nil { + return nil + } + added["status"] = "in_progress" + delete(added, "arguments") + if _, ok := added["content"].([]any); ok { + added["content"] = []any{} + } + if _, ok := added["summary"].([]any); ok { + added["summary"] = []any{} + } + return added +} + +func responsesTerminalEventName(response map[string]any) string { + status, _ := response["status"].(string) + switch status { + case "failed": + return "response.failed" + case "incomplete": + return "response.incomplete" + default: + return "response.completed" + } +} + +func appendResponsesItemDeltaEvents(out *bytes.Buffer, item map[string]any, itemID string, outputIndex int) error { + if out == nil || item == nil { + return nil + } + + if arguments, ok := item["arguments"].(string); ok && arguments != "" { + if err := appendSSEJSONEvent(out, "response.function_call_arguments.delta", map[string]any{ + "type": "response.function_call_arguments.delta", + "item_id": itemID, + "output_index": outputIndex, + "delta": arguments, + }); err != nil { + return err + } + if err := appendSSEJSONEvent(out, "response.function_call_arguments.done", map[string]any{ + "type": "response.function_call_arguments.done", + "item_id": itemID, + "output_index": outputIndex, + "arguments": arguments, + }); err != nil { + return err + } + } + + for _, key := range []string{"content", "summary"} { + parts, ok := item[key].([]any) + if !ok { + continue + } + for contentIndex, partAny := range parts { + part, ok := partAny.(map[string]any) + if !ok { + continue + } + eventName, payload, ok := responsesContentDeltaEvent(part, itemID, outputIndex, contentIndex) + if !ok { + continue + } + if err := appendSSEJSONEvent(out, eventName, payload); err != nil { + return err + } + } + } + + return nil +} + +func responsesContentDeltaEvent(part map[string]any, itemID string, outputIndex, contentIndex int) (string, map[string]any, bool) { + partType, _ := part["type"].(string) + text, _ := part["text"].(string) + if partType == "" || text == "" { + return "", nil, false + } + + var eventName string + switch partType { + case "output_text": + eventName = "response.output_text.delta" + case "reasoning_text": + eventName = "response.reasoning_text.delta" + default: + return "", nil, false + } + + payload := map[string]any{ + "type": eventName, + "delta": text, + "output_index": outputIndex, + "content_index": contentIndex, + } + if itemID != "" { + payload["item_id"] = itemID + } + + return eventName, payload, true +} + +func chatUsageMap(usage core.Usage) map[string]any { + if usage.PromptTokens == 0 && + usage.CompletionTokens == 0 && + usage.TotalTokens == 0 && + usage.PromptTokensDetails == nil && + usage.CompletionTokensDetails == nil && + len(usage.RawUsage) == 0 { + return nil + } + result, err := toJSONMap(usage) + if err != nil { + return nil + } + return result +} + +func appendSSEJSONEvent(out *bytes.Buffer, eventName string, payload any) error { + data, err := json.Marshal(payload) + if err != nil { + return err + } + if eventName != "" { + out.WriteString("event: ") + out.WriteString(eventName) + out.WriteByte('\n') + } + out.WriteString("data: ") + out.Write(data) + out.WriteString("\n\n") + return nil +} + +func toJSONMap(value any) (map[string]any, error) { + data, err := json.Marshal(value) + if err != nil { + return nil, err + } + var result map[string]any + if err := json.Unmarshal(data, &result); err != nil { + return nil, err + } + return result, nil +} + +func cloneJSONMap(src map[string]any) map[string]any { + if src == nil { + return nil + } + dst := make(map[string]any, len(src)) + for key, value := range src { + dst[key] = value + } + return dst +} + +func jsonNumberToInt(value any) (int, bool) { + switch v := value.(type) { + case float64: + return int(v), true + case int: + return v, true + case int64: + return int(v), true + default: + return 0, false + } +} + +func jsonNumberToInt64(value any) (int64, bool) { + switch v := value.(type) { + case float64: + return int64(v), true + case int: + return int64(v), true + case int64: + return v, true + default: + return 0, false + } +} + +func nonEmpty(value, fallback string) string { + value = strings.TrimSpace(value) + if value != "" { + return value + } + return strings.TrimSpace(fallback) +} diff --git a/internal/server/translated_inference_service.go b/internal/server/translated_inference_service.go index ee89fd461..78b300732 100644 --- a/internal/server/translated_inference_service.go +++ b/internal/server/translated_inference_service.go @@ -149,9 +149,10 @@ func handleTranslatedInference[R any]( return handleWithCache(s, c, req, isStream(req), plan, dispatch) } -// handleWithCache injects the guardrails hash into context, then either routes the -// request through the dual-layer response cache (non-streaming) or calls dispatch -// directly (streaming). R is the post-patch request type. +// handleWithCache injects the guardrails hash into context, then routes the +// request through the dual-layer response cache when caching is enabled. +// Streaming requests are stored as full responses and replayed as SSE on hits. +// R is the post-patch request type. func handleWithCache[R any]( s *translatedInferenceService, c *echo.Context, @@ -169,7 +170,7 @@ func handleWithCache[R any]( c.SetRequest(c.Request().WithContext(ctx)) } - if s.responseCache != nil && !stream && (plan == nil || plan.CacheEnabled()) { + if s.responseCache != nil && (plan == nil || plan.CacheEnabled()) { body, marshalErr := marshalRequestBody(req) if marshalErr != nil { slog.Debug("marshalRequestBody failed", "err", marshalErr)