feat(cache): add streaming cache support for chat responses - #529
Conversation
Integrated caching support for streaming chat responses. Added utility functions `setStreamingCache`, `getStreamingCache`, and `generateStreamingCacheKey`. Enhanced chat logic to handle streaming response caching, including replaying cached chunks and saving streamed data. Improved logging with reconstructed content from cached data.
|
""" WalkthroughThe changes introduce a streaming cache mechanism for chat completions. The system now captures and stores each streamed chunk of a chat response, along with metadata, enabling replay of entire streaming responses from cache. This involves new cache key generation, chunked data structures, and updates to both the chat handler and cache utility functions. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant ChatHandler
participant Cache
participant SSE
Client->>ChatHandler: Request chat completion (streaming)
ChatHandler->>Cache: getStreamingCache(key)
alt Cache hit and completed
Cache-->>ChatHandler: Cached streaming data
ChatHandler->>SSE: Replay cached chunks to client
else Cache miss
ChatHandler->>SSE: Stream chunks to client (writeSSEAndCache)
ChatHandler->>Cache: setStreamingCache(key, chunks, metadata)
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~15–20 minutes Possibly related PRs
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
apps/gateway/src/lib/cache.ts (2)
255-260: Export the interface for better type safetyConsider exporting this interface to ensure type safety when used in
chat.tsand potential future consumers.-interface StreamingCacheChunk { +export interface StreamingCacheChunk { data: string; eventId: number; event?: string; timestamp: number; }
262-272: Export the interface and consider adding size validationExport this interface for type safety and consider adding validation to prevent excessive memory usage from large chunk arrays.
-interface StreamingCacheData { +export interface StreamingCacheData { chunks: StreamingCacheChunk[]; metadata: { model: string; provider: string; finishReason: string | null; totalChunks: number; duration: number; completed: boolean; }; }Consider implementing a maximum chunk limit or total size validation in
setStreamingCacheto prevent potential memory exhaustion from excessively large streaming responses.apps/gateway/src/chat/chat.ts (1)
1885-1911: Add null check for streamingCacheKeyThe function should validate that streamingCacheKey is not null before attempting to push chunks.
// Capture for streaming cache if enabled - if (cachingEnabled && streamingCacheKey) { + if (cachingEnabled && streamingCacheKey !== null) { streamingChunks.push({ data: sseData.data, eventId: sseData.id ? parseInt(sseData.id, 10) : eventId, event: sseData.event, timestamp: Date.now() - streamStartTime, }); }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
apps/gateway/src/chat/chat.ts(15 hunks)apps/gateway/src/lib/cache.ts(1 hunks)
🧰 Additional context used
📓 Path-based instructions (4)
**/*.{js,jsx,ts,tsx}
📄 CodeRabbit Inference Engine (.github/copilot-instructions.md)
Use localStorage instead of cookies for client-side data persistence
Files:
apps/gateway/src/chat/chat.tsapps/gateway/src/lib/cache.ts
**/*.{js,ts}
📄 CodeRabbit Inference Engine (.github/copilot-instructions.md)
**/*.{js,ts}: Use drizzle with the latest object syntax for database operations
For read queries, always usedb().query.<table>.findMany()ordb().query.<table>.findFirst()
Files:
apps/gateway/src/chat/chat.tsapps/gateway/src/lib/cache.ts
{apps/api,apps/gateway,packages/db}/**/*.ts
📄 CodeRabbit Inference Engine (CLAUDE.md)
{apps/api,apps/gateway,packages/db}/**/*.ts: Use Drizzle ORM with latest object syntax for database operations
For reads, usedb().query.<table>.findMany()ordb().query.<table>.findFirst()
Files:
apps/gateway/src/chat/chat.tsapps/gateway/src/lib/cache.ts
**/*.{ts,tsx}
📄 CodeRabbit Inference Engine (.cursor/rules/general.mdc)
Never use
as anyor: anyin TypeScript files.
Files:
apps/gateway/src/chat/chat.tsapps/gateway/src/lib/cache.ts
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: build / run
- GitHub Check: e2e / run
- GitHub Check: autofix
🔇 Additional comments (7)
apps/gateway/src/lib/cache.ts (2)
274-278: LGTM!Clean implementation that properly reuses the existing cache key generation logic.
297-310: LGTM!Well-implemented function that follows the established pattern with proper error handling and type safety.
apps/gateway/src/chat/chat.ts (5)
28-37: LGTM!Imports are correctly added for the new streaming cache functionality.
1642-1654: LGTM!Cache key generation correctly handles both streaming and non-streaming cases with a consistent payload structure.
1774-1825: LGTM!Non-streaming cache handling correctly maintains backward compatibility with proper logging.
2735-2757: LGTM!Streaming cache storage correctly saves completed streams with comprehensive metadata and proper error handling.
1982-1989: LGTM!All SSE writes have been consistently replaced with the caching-aware helper function, ensuring complete capture of streaming responses.
Also applies to: 2006-2023, 2072-2088, 2181-2184, 2233-2237, 2350-2360, 2396-2399, 2664-2674
| export async function setStreamingCache( | ||
| key: string, | ||
| data: StreamingCacheData, | ||
| expirationSeconds: number, | ||
| ): Promise<void> { | ||
| if (process.env.NODE_ENV === "test") { | ||
| // temp disable caching in test mode | ||
| return; | ||
| } | ||
|
|
||
| try { | ||
| await redisClient.set(key, JSON.stringify(data), "EX", expirationSeconds); | ||
| } catch (error) { | ||
| console.error("Error setting streaming cache:", error); | ||
| } | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add validation for data integrity and size limits
Consider validating the streaming cache data before storage to ensure data integrity and prevent memory issues.
export async function setStreamingCache(
key: string,
data: StreamingCacheData,
expirationSeconds: number,
): Promise<void> {
if (process.env.NODE_ENV === "test") {
// temp disable caching in test mode
return;
}
+ // Validate data structure
+ if (!data.chunks || !Array.isArray(data.chunks)) {
+ console.error("Invalid streaming cache data: chunks must be an array");
+ return;
+ }
+
+ // Limit maximum chunks to prevent memory issues (e.g., 10000 chunks)
+ const MAX_CHUNKS = 10000;
+ if (data.chunks.length > MAX_CHUNKS) {
+ console.error(`Streaming cache data exceeds maximum chunks limit: ${data.chunks.length} > ${MAX_CHUNKS}`);
+ return;
+ }
+
try {
await redisClient.set(key, JSON.stringify(data), "EX", expirationSeconds);
} catch (error) {
console.error("Error setting streaming cache:", error);
}
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| export async function setStreamingCache( | |
| key: string, | |
| data: StreamingCacheData, | |
| expirationSeconds: number, | |
| ): Promise<void> { | |
| if (process.env.NODE_ENV === "test") { | |
| // temp disable caching in test mode | |
| return; | |
| } | |
| try { | |
| await redisClient.set(key, JSON.stringify(data), "EX", expirationSeconds); | |
| } catch (error) { | |
| console.error("Error setting streaming cache:", error); | |
| } | |
| } | |
| export async function setStreamingCache( | |
| key: string, | |
| data: StreamingCacheData, | |
| expirationSeconds: number, | |
| ): Promise<void> { | |
| if (process.env.NODE_ENV === "test") { | |
| // temp disable caching in test mode | |
| return; | |
| } | |
| // Validate data structure | |
| if (!data.chunks || !Array.isArray(data.chunks)) { | |
| console.error("Invalid streaming cache data: chunks must be an array"); | |
| return; | |
| } | |
| // Limit maximum chunks to prevent memory issues (e.g., 10000 chunks) | |
| const MAX_CHUNKS = 10000; | |
| if (data.chunks.length > MAX_CHUNKS) { | |
| console.error( | |
| `Streaming cache data exceeds maximum chunks limit: ${data.chunks.length} > ${MAX_CHUNKS}` | |
| ); | |
| return; | |
| } | |
| try { | |
| await redisClient.set(key, JSON.stringify(data), "EX", expirationSeconds); | |
| } catch (error) { | |
| console.error("Error setting streaming cache:", error); | |
| } | |
| } |
🤖 Prompt for AI Agents
In apps/gateway/src/lib/cache.ts around lines 280 to 295, the setStreamingCache
function currently stores data without validation. Add checks before storing to
validate the data structure and ensure it does not exceed size limits to prevent
memory issues. Implement validation logic to verify required fields and data
types in the StreamingCacheData object and check the serialized data size before
calling redisClient.set, returning early or throwing an error if validation
fails.
Improved logic for accumulating and processing tool calls from streaming chat chunks. Updated tool call schema to record detailed data. Enhanced replay timing accuracy with original delays capped for long waits.
1 participant
Integrated caching support for streaming chat responses. Added utility functions
setStreamingCache,getStreamingCache, andgenerateStreamingCacheKey. Enhanced chat logic to handle streaming response caching, including replaying cached chunks and saving streamed data. Improved logging with reconstructed content from cached data.Summary by CodeRabbit
New Features
Bug Fixes
Chores