diff --git a/api/chat_history.py b/api/chat_history.py new file mode 100644 index 00000000..0fa27a65 --- /dev/null +++ b/api/chat_history.py @@ -0,0 +1,81 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Chat history response shaping — keep GET /chat under Lambda's 6MB limit. + +Pure functions (no boto3 / Powertools) so they are unit-testable outside +the Lambda bundle. Three defenses, applied in order by get_chat: + +1. strip_tool_result_content — toolResult bodies are only needed by the + agent (which reads AgentCore Memory directly); the UI needs status only. +2. truncate_tool_inputs — toolUse.input can embed entire slide JSON bodies. + The UI renders only short scalar fields (slide_id, purpose, instruction, + slide_groups), so long strings are cut. +3. cap_messages_size — backstop: drop oldest messages until the serialized + payload fits, keeping the most recent conversation visible. +""" + +import json +from typing import Any, Dict, List, Tuple + +# Per-string cap inside toolUse.input. Generous enough that slide_groups +# (JSON-encoded string parsed by ComposeCard) survives intact in practice, +# while slide JSON bodies (tens of KB each) are slashed. +MAX_INPUT_STR = 8_000 + +# Serialized messages budget — headroom below the 6MB Lambda response limit +# for the JSON envelope and transport overhead. +MAX_MESSAGES_BYTES = 4_500_000 + +TRUNCATION_MARKER = "…[truncated]" + + +def strip_tool_result_content(messages: List[Dict]) -> None: + """Empty toolResult content blocks in place — frontend needs status only.""" + for msg in messages: + if msg.get("role") == "user" and isinstance(msg.get("content"), list): + for block in msg["content"]: + if isinstance(block, dict) and "toolResult" in block: + block["toolResult"]["content"] = [] + + +def _truncate_value(value: Any) -> Any: + """Recursively cap string lengths inside a toolUse input value.""" + if isinstance(value, str) and len(value) > MAX_INPUT_STR: + return value[:MAX_INPUT_STR] + TRUNCATION_MARKER + if isinstance(value, dict): + return {k: _truncate_value(v) for k, v in value.items()} + if isinstance(value, list): + return [_truncate_value(v) for v in value] + return value + + +def truncate_tool_inputs(messages: List[Dict]) -> None: + """Cap long strings inside toolUse.input blocks in place.""" + for msg in messages: + if not isinstance(msg.get("content"), list): + continue + for block in msg["content"]: + if isinstance(block, dict) and "toolUse" in block: + tool_use = block["toolUse"] + if isinstance(tool_use.get("input"), dict): + tool_use["input"] = _truncate_value(tool_use["input"]) + + +def cap_messages_size( + messages: List[Dict], max_bytes: int = MAX_MESSAGES_BYTES, +) -> Tuple[List[Dict], bool]: + """Drop oldest messages until the serialized payload fits max_bytes. + + Returns: + (messages, truncated) — truncated is True if any message was dropped. + """ + sizes = [len(json.dumps(m, ensure_ascii=False).encode("utf-8")) for m in messages] + total = sum(sizes) + if total <= max_bytes: + return messages, False + + start = 0 + while start < len(messages) - 1 and total > max_bytes: + total -= sizes[start] + start += 1 + return messages[start:], True diff --git a/api/index.py b/api/index.py index 24c4bcfd..6913327f 100644 --- a/api/index.py +++ b/api/index.py @@ -26,6 +26,11 @@ from aws_lambda_powertools.utilities.typing import LambdaContext from boto3.dynamodb.conditions import Key from authz import authorize +from chat_history import ( + cap_messages_size, + strip_tool_result_content, + truncate_tool_inputs, +) from common import get_user_id, now_iso, presigned_url from shared.schema import ( deck_pk, deck_sk, shared_pk, fav_sk, upload_sk, @@ -1253,14 +1258,18 @@ def get_chat(session_id: str) -> Dict[str, Any]: # Strands stores events in reverse chronological order messages.reverse() - # Strip toolResult content — frontend only needs status for ToolCard display. - # Agent reads from Amazon Bedrock AgentCore Memory directly, not via this API. - # This prevents Lambda 6MB response limit errors on long conversations. - for msg in messages: - if msg.get("role") == "user" and isinstance(msg.get("content"), list): - for block in msg["content"]: - if isinstance(block, dict) and "toolResult" in block: - block["toolResult"]["content"] = [] + # Keep the response under Lambda's 6MB limit (see chat_history.py): + # toolResult bodies and giant toolUse inputs are UI-irrelevant; the + # agent reads AgentCore Memory directly, not via this API. + strip_tool_result_content(messages) + truncate_tool_inputs(messages) + messages, truncated = cap_messages_size(messages) + if truncated: + logger.warning( + "Chat history for session %s exceeded size budget — oldest messages dropped", + session_id, + ) + return {"messages": messages, "truncated": True} except Exception as e: logger.warning("Failed to read chat history from AgentCore Memory: %s", e) diff --git a/tests/test_chat_history.py b/tests/test_chat_history.py new file mode 100644 index 00000000..0a11e2b0 --- /dev/null +++ b/tests/test_chat_history.py @@ -0,0 +1,104 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Tests for api.chat_history — GET /chat 6MB response shaping.""" + +import json + +from api.chat_history import ( + MAX_INPUT_STR, + TRUNCATION_MARKER, + cap_messages_size, + strip_tool_result_content, + truncate_tool_inputs, +) + + +def _tool_use_msg(input_value: dict) -> dict: + return { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "t1", "name": "write_slide", "input": input_value}}, + ], + } + + +class TestStripToolResultContent: + def test_strips_content_keeps_status(self): + messages = [{ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "t1", + "status": "success", + "content": [{"text": "x" * 100_000}], + }, + }], + }] + strip_tool_result_content(messages) + tr = messages[0]["content"][0]["toolResult"] + assert tr["content"] == [] + assert tr["status"] == "success" + + def test_ignores_plain_text_messages(self): + messages = [{"role": "user", "content": "hello"}] + strip_tool_result_content(messages) + assert messages[0]["content"] == "hello" + + +class TestTruncateToolInputs: + def test_long_string_truncated_with_marker(self): + messages = [_tool_use_msg({"slide_json": "x" * (MAX_INPUT_STR + 1000)})] + truncate_tool_inputs(messages) + value = messages[0]["content"][0]["toolUse"]["input"]["slide_json"] + assert len(value) == MAX_INPUT_STR + len(TRUNCATION_MARKER) + assert value.endswith(TRUNCATION_MARKER) + + def test_short_ui_fields_untouched(self): + input_value = { + "slide_id": "intro", + "purpose": "Compose the intro slide", + "slide_groups": json.dumps([{"slugs": ["a", "b"], "instruction": "short"}]), + } + messages = [_tool_use_msg(dict(input_value))] + truncate_tool_inputs(messages) + assert messages[0]["content"][0]["toolUse"]["input"] == input_value + + def test_nested_structures_truncated(self): + messages = [_tool_use_msg({ + "groups": [{"instruction": "y" * (MAX_INPUT_STR + 1)}], + })] + truncate_tool_inputs(messages) + value = messages[0]["content"][0]["toolUse"]["input"]["groups"][0]["instruction"] + assert value.endswith(TRUNCATION_MARKER) + + def test_non_string_values_untouched(self): + messages = [_tool_use_msg({"count": 5, "flag": True, "empty": None})] + truncate_tool_inputs(messages) + assert messages[0]["content"][0]["toolUse"]["input"] == { + "count": 5, "flag": True, "empty": None, + } + + +class TestCapMessagesSize: + def test_under_budget_unchanged(self): + messages = [{"role": "user", "content": "hi"}] * 3 + result, truncated = cap_messages_size(messages, max_bytes=10_000) + assert result == messages + assert truncated is False + + def test_drops_oldest_first(self): + messages = [ + {"role": "user", "content": f"msg-{i}: " + "x" * 100} for i in range(10) + ] + result, truncated = cap_messages_size(messages, max_bytes=500) + assert truncated is True + assert result # newest messages survive + assert result[-1] == messages[-1] + assert result[0] != messages[0] + + def test_always_keeps_last_message(self): + messages = [{"role": "user", "content": "x" * 1000}] * 2 + result, truncated = cap_messages_size(messages, max_bytes=10) + assert truncated is True + assert len(result) == 1 + assert result[0] is messages[-1] diff --git a/web-ui/src/components/chat/ChatPanel.tsx b/web-ui/src/components/chat/ChatPanel.tsx index f6e07610..a90e30dc 100644 --- a/web-ui/src/components/chat/ChatPanel.tsx +++ b/web-ui/src/components/chat/ChatPanel.tsx @@ -16,21 +16,21 @@ import { buildAttachedMarkers } from "@/lib/attachmentMarker" import { generateSessionId, setAgentConfig } from "@/services/agentCoreService" import { getChatHistory, patchDeck } from "@/services/deckService" import type { UploadedFile } from "@/services/uploadService" -import { useChatStream, type Message, type ToolUseCallbackData } from "@/hooks/useChatStream" +import { useChatStream, type ToolUseCallbackData } from "@/hooks/useChatStream" import { ChatInput, type ChatInputHandle } from "./ChatInput" import { ChatMessage, ToolUse } from "./ChatMessage" -import { McpStatusBar, McpServerStatus } from "./McpStatusBar" +import { McpStatusBar } from "./McpStatusBar" import { FileDropZone } from "./FileDropZone" import { useIsMobile } from "@/hooks/UseMobile" import { Send, ChevronRight } from "lucide-react" import { ModeSelector } from "./ModeSelector" import { usePreferences } from "@/hooks/usePreferences" import { notifyError } from "@/lib/errors" +import { toast } from "sonner" import { isLocalHistoryFormat, parseLocalHistory, parseCloudHistory } from "./chatHistory" interface ChatPanelProps { deckId: string - deckName?: string chatSessionId?: string slideSlugs?: string[] onDeckCreated?: (deckId: string) => void @@ -43,14 +43,14 @@ export interface ChatPanelHandle { insertAtCursor: (text: string) => void } -export const ChatPanel = forwardRef(function ChatPanel({ deckId, deckName, chatSessionId, slideSlugs, onDeckCreated, onPreviewInvalidated, onWorkflowPhase }, ref) { +export const ChatPanel = forwardRef(function ChatPanel({ deckId, chatSessionId, slideSlugs, onDeckCreated, onPreviewInvalidated, onWorkflowPhase }, ref) { // --- Session --- const [sessionId, setSessionId] = useState(() => { if (chatSessionId) return chatSessionId if (deckId === "new") return generateSessionId() return deckId.padEnd(36, "0") }) - useEffect(() => { if (chatSessionId && chatSessionId !== sessionId) setSessionId(chatSessionId) }, [chatSessionId]) + useEffect(() => { if (chatSessionId) setSessionId((prev) => (chatSessionId !== prev ? chatSessionId : prev)) }, [chatSessionId]) // --- Config --- const [configLoaded, setConfigLoaded] = useState(false) @@ -174,7 +174,10 @@ export const ChatPanel = forwardRef(function Ch if (!sessionId) return setHistoryLoading(true) try { - const history = await getChatHistory(sessionId, idToken ?? "", deckId || undefined) + const { messages: history, truncated } = await getChatHistory(sessionId, idToken ?? "", deckId || undefined) + if (truncated) { + toast.info("Older messages in this conversation were omitted to keep loading fast.") + } if (history.length > 0) { // Local mode: .chat.json is already in ChatPanel's internal format if (IS_LOCAL && isLocalHistoryFormat(history)) { diff --git a/web-ui/src/components/chat/ChatPanelShell.tsx b/web-ui/src/components/chat/ChatPanelShell.tsx index e87fa2bd..b22c8ad0 100644 --- a/web-ui/src/components/chat/ChatPanelShell.tsx +++ b/web-ui/src/components/chat/ChatPanelShell.tsx @@ -211,7 +211,6 @@ export function ChatPanelShell({ key={`a-${panelAKey}`} ref={panelAVisible ? chatRef : undefined} deckId="new" - deckName="New Deck" slideSlugs={panelAOwnsCurrentDeck ? (slideSlugs || []) : []} onDeckCreated={handlePanelADeckCreated} onPreviewInvalidated={onPreviewInvalidated} @@ -226,7 +225,6 @@ export function ChatPanelShell({ key={`b-${panelBKey}`} ref={panelBVisible ? chatRef : undefined} deckId={deckId!} - deckName={deckName || undefined} chatSessionId={chatSessionId} slideSlugs={slideSlugs || []} onDeckCreated={handlePanelBDeckCreated} diff --git a/web-ui/src/services/deckService.ts b/web-ui/src/services/deckService.ts index 363bc9be..45ec26c7 100644 --- a/web-ui/src/services/deckService.ts +++ b/web-ui/src/services/deckService.ts @@ -195,35 +195,41 @@ export interface ChatMessage { timestamp: number } +export interface ChatHistoryResult { + messages: ChatMessage[] + /** True when the backend dropped oldest messages to fit the response size limit. */ + truncated: boolean +} + /** * Fetch chat history for a session. * * @param sessionId - Conversation session ID * @param idToken - Cognito ID token - * @returns Array of chat messages sorted by timestamp + * @returns Messages sorted by timestamp, plus a truncation flag */ -export async function getChatHistory(sessionId: string, idToken: string, deckId?: string): Promise { +export async function getChatHistory(sessionId: string, idToken: string, deckId?: string): Promise { if (IS_LOCAL) { - if (!deckId || deckId === "new") return [] + if (!deckId || deckId === "new") return { messages: [], truncated: false } // Load session context + saved messages const res = await fetch("/api/agent/load", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ sessionId, deckId }), }) - if (!res.ok) return [] + if (!res.ok) return { messages: [], truncated: false } const data = await res.json() - return data.messages || [] + return { messages: data.messages || [], truncated: false } } const base = await getApiBaseUrl() const response = await fetch(`${base}chat/${sessionId}`, { headers: { Authorization: `Bearer ${idToken}` }, }) - if (!response.ok) return [] + if (!response.ok) return { messages: [], truncated: false } const data = await response.json() - return data.messages || [] + return { messages: data.messages || [], truncated: data.truncated === true } } /**