Skip to content

fix(responsecache): split stream cache entries and validate SSE#202

Merged
SantiagoDePolonia merged 3 commits intomainfrom
fix/stream-cache-separation
Apr 2, 2026
Merged

fix(responsecache): split stream cache entries and validate SSE#202
SantiagoDePolonia merged 3 commits intomainfrom
fix/stream-cache-separation

Conversation

@SantiagoDePolonia
Copy link
Copy Markdown
Contributor

@SantiagoDePolonia SantiagoDePolonia commented Apr 1, 2026

Summary

  • treat stream=true and stream=false as separate exact and semantic cache entries
  • persist valid raw SSE for streaming cache hits and replay it verbatim
  • validate SSE before caching so partial, malformed, or keepalive-only streams cannot poison cache
  • keep cache-hit usage accounting working by extracting usage from cached SSE bodies

Testing

  • go test ./internal/responsecache ./internal/usage ./internal/server

Summary by CodeRabbit

  • New Features

    • Streaming responses are cached separately from non-streaming responses; streaming cache stores and replays raw SSE bytes.
    • Usage extraction now supports cached streaming responses to produce usage entries.
  • Bug Fixes

    • Streaming vs non-streaming cache keying corrected so hits/misses behave independently and replay is byte-exact.
    • Invalid streaming payloads are rejected from cache population.
  • Tests

    • Added and updated tests validating streaming validation, caching separation, and usage extraction.

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai bot commented Apr 1, 2026

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: a6d35c34-c3b4-45a0-982c-a7379a5c5a7e

📥 Commits

Reviewing files that changed from the base of the PR and between 83b9b97 and 58b999c.

📒 Files selected for processing (3)
  • internal/responsecache/sse_validation.go
  • internal/responsecache/sse_validation_test.go
  • internal/responsecache/stream_cache.go

📝 Walkthrough

Walkthrough

This PR separates streaming and non-streaming cache entries, stores validated raw SSE bytes for streaming misses, replays those bytes verbatim on streaming hits, prevents caching of malformed SSE, and extends usage extraction to parse cached SSE bodies for token/usage metrics.

Changes

Cohort / File(s) Summary
Cache Keying & Semantic
internal/responsecache/semantic.go
Include Stream boolean in request parameter hash so streaming vs non-streaming produce distinct cache keys.
Stream Cache & Replay
internal/responsecache/stream_cache.go, internal/responsecache/simple.go
Streaming cache path now validates and stores raw SSE bytes and writes cached bytes verbatim on replay; removed reconstruction/rendering of streams and related helpers.
SSE Validation
internal/responsecache/sse_validation.go, internal/responsecache/sse_validation_test.go
New validation for cacheable SSE streams (validateCacheableSSE) plus table-driven tests covering valid/invalid cases and DONE termination rules.
Handler & Middleware Tests
internal/responsecache/handle_request_test.go, internal/responsecache/middleware_test.go, internal/responsecache/semantic_test.go
Updated and added tests asserting streaming/non-streaming cache separation, verbatim SSE hit replay, synthetic usage recording on streaming hits, and rejection of invalid SSE for cache population.
Usage Extraction
internal/usage/extractor.go, internal/usage/extractor_test.go
Extended extraction to parse cached SSE bodies via a stream observer (and static pricing resolver), plus tests validating token/usage extraction from SSE cached payloads.
Docs / Interface Note
internal/responsecache/responsecache.go
Updated docstring to reflect independent caching of streaming and non-streaming requests and verbatim SSE storage/replay.

Sequence Diagram(s)

mermaid
sequenceDiagram
participant Client
participant Middleware
participant Cache
participant Handler
participant UsageExtractor

Client->>Middleware: request (stream=true)
Middleware->>Cache: lookup(stream-key)
Cache-->>Middleware: MISS
Middleware->>Handler: invoke backend (streaming SSE)
Handler-->>Middleware: SSE bytes (rawStream)
Middleware->>Cache: validate & store rawStream as exact entry
Middleware->>Client: stream rawStream (SSE headers)

Note right of Cache: stored raw SSE bytes (exact)

Client->>Middleware: request (stream=true) repeat
Middleware->>Cache: lookup(stream-key)
Cache-->>Middleware: HIT (exact)
Middleware->>Client: write cached rawStream bytes verbatim

Note right of UsageExtractor: observes cached SSE when extracting usage
Middleware->>UsageExtractor: ExtractFromCachedResponseBody(cached rawStream)
UsageExtractor-->>Middleware: UsageEntry (tokens, provider)

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Possibly related PRs

Poem

🐰 I hopped through streams and crumbs of code,

Kept SSEs whole, not rebuilt or slowed.
Validated bytes, saved them true,
Tokens counted, providers too —
A tidy cache, in bunny-mode! 🥕

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 11.11% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately and concisely summarizes the main changes: splitting stream cache entries (stream=true/false as separate entries) and validating SSE payloads before caching.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/stream-cache-separation

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@internal/responsecache/sse_validation_test.go`:
- Around line 12-64: Add a new table-driven test case to the existing test cases
slice in sse_validation_test.go that verifies empty input is rejected: add an
entry with name "rejects empty input", raw: []byte("") and want: false so the
SSE validation logic (the same test harness used for the other cases) explicitly
asserts that an empty payload returns false.

In `@internal/responsecache/sse_validation.go`:
- Around line 28-40: The loop accepts events after a seen `[DONE]` because the
sawDone check runs after early-continue paths (comments/keepalives and
empty-data), so move the sawDone check so it runs immediately after calling
sseEventPayload(event) (i.e., right after "payload, hasData :=
sseEventPayload(event)") and before any continue; this ensures that once sawDone
is set (when payload equals cacheDonePayload) any subsequent event — including
comment/keepalive-only or another `[DONE]` — triggers an immediate false return.
- Around line 43-46: The current validation unmarshals into a map using
json.Unmarshal(payload, &decoded), allocating an unused map and incorrectly
rejecting valid non-object JSON; replace that logic with a syntax-only check
using json.Valid(payload) and remove the unused decoded variable so the function
simply returns false when json.Valid(payload) is false, otherwise proceeds as
valid. Ensure you reference the existing payload variable and remove the
json.Unmarshal call and decoded declaration, using json.Valid(payload) instead.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 41a69a42-06bf-431a-b7d8-14e1417b240c

📥 Commits

Reviewing files that changed from the base of the PR and between d5728e2 and 83b9b97.

📒 Files selected for processing (11)
  • internal/responsecache/handle_request_test.go
  • internal/responsecache/middleware_test.go
  • internal/responsecache/responsecache.go
  • internal/responsecache/semantic.go
  • internal/responsecache/semantic_test.go
  • internal/responsecache/simple.go
  • internal/responsecache/sse_validation.go
  • internal/responsecache/sse_validation_test.go
  • internal/responsecache/stream_cache.go
  • internal/usage/extractor.go
  • internal/usage/extractor_test.go

@SantiagoDePolonia SantiagoDePolonia merged commit cc42d4a into main Apr 2, 2026
14 checks passed
@SantiagoDePolonia SantiagoDePolonia deleted the fix/stream-cache-separation branch April 4, 2026 11:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant