From 4ee4f48ce3409cca5006d09eeee6e2043f761060 Mon Sep 17 00:00:00 2001 From: Biswajit Panda Date: Sun, 26 Jul 2026 06:15:01 +0530 Subject: [PATCH 1/2] fix: show budget-limit message on 429 instead of generic error --- .../api/static/widget/react/AgentChatApp.tsx | 21 ++++++++-- .../static/widget/react/shadcnAiElements.tsx | 6 ++- .../api/static/widget/styles/styles.ts | 2 + tests/agent_manager/test_api.py | 38 +++++++++++++++++++ 4 files changed, 61 insertions(+), 6 deletions(-) diff --git a/src/agent_manager/api/static/widget/react/AgentChatApp.tsx b/src/agent_manager/api/static/widget/react/AgentChatApp.tsx index 4bd932c1..53c658e4 100644 --- a/src/agent_manager/api/static/widget/react/AgentChatApp.tsx +++ b/src/agent_manager/api/static/widget/react/AgentChatApp.tsx @@ -32,6 +32,8 @@ type MessageEntry = { typing?: boolean; route?: string[]; tools?: ToolRecord[]; + /** When true the message represents a permanent error (e.g. budget exceeded). */ + error?: boolean; }; let nextMessageCounter = 0; @@ -150,6 +152,9 @@ export function AgentChatApp({ client, config, onAnswer, panelId, titleId }: Age setEntries((prev) => prev.map((entry) => (entry.id === id ? update(entry) : entry))); }, []); + const BUDGET_EXCEEDED_MESSAGE = + "This conversation has reached its context limit. Please start a new chat to continue."; + const submit = useCallback( async (text: string) => { const assistantId = nextMessageId("ai"); @@ -167,6 +172,16 @@ export function AgentChatApp({ client, config, onAnswer, panelId, titleId }: Age patchEntry(assistantId, (entry) => ({ ...entry, typing: false })); onAnswer(finalDetail); } catch (error) { + if (error instanceof AgentChatHttpError && error.status === 429) { + patchEntry(assistantId, () => ({ + id: assistantId, + role: "ai", + text: BUDGET_EXCEEDED_MESSAGE, + error: true, + })); + setSending(false); + return; + } try { const data = await sendToAgent(text); patchEntry(assistantId, () => ({ @@ -181,10 +196,7 @@ export function AgentChatApp({ client, config, onAnswer, panelId, titleId }: Age patchEntry(assistantId, () => ({ id: assistantId, role: "ai", - text: - error instanceof Error && error.message - ? `Something went wrong. Please try again.` - : "Something went wrong. Please try again.", + text: "Something went wrong. Please try again.", })); } } finally { @@ -256,6 +268,7 @@ export function AgentChatApp({ client, config, onAnswer, panelId, titleId }: Age key={entry.id} from={entry.role === "user" ? "user" : "assistant"} typing={entry.typing} + error={entry.error} > {entry.typing ? ( "..." diff --git a/src/agent_manager/api/static/widget/react/shadcnAiElements.tsx b/src/agent_manager/api/static/widget/react/shadcnAiElements.tsx index 44a0f963..a1f015ff 100644 --- a/src/agent_manager/api/static/widget/react/shadcnAiElements.tsx +++ b/src/agent_manager/api/static/widget/react/shadcnAiElements.tsx @@ -44,10 +44,11 @@ export function Message({ children, from, typing = false, -}: PropsWithChildren<{ from: "user" | "assistant"; typing?: boolean }>) { + error = false, +}: PropsWithChildren<{ from: "user" | "assistant"; typing?: boolean; error?: boolean }>) { return (
@@ -56,6 +57,7 @@ export function Message({ ); } + export function MessageContent({ children }: PropsWithChildren) { return
{children}
; } diff --git a/src/agent_manager/api/static/widget/styles/styles.ts b/src/agent_manager/api/static/widget/styles/styles.ts index b200e095..bde2d305 100644 --- a/src/agent_manager/api/static/widget/styles/styles.ts +++ b/src/agent_manager/api/static/widget/styles/styles.ts @@ -44,6 +44,8 @@ export function styles(config: AgentChatConfig): string { .msg { font-size: 14.5px; line-height: 1.55; word-wrap: break-word; white-space: pre-wrap; } .msg.ai { color: #18181b; max-width: 100%; } .msg.ai.typing { color: #a1a1aa; letter-spacing: 1px; } + .msg.ai.msg-error { color: #991b1b; background: #fff5f5; border: 1px solid #fecaca; + border-radius: 10px; padding: 8px 12px; } .msg.user { background: #f4f4f5; color: #18181b; border-radius: 18px; padding: 10px 14px; margin-left: auto; max-width: 88%; } .message-content { min-width: 0; } diff --git a/tests/agent_manager/test_api.py b/tests/agent_manager/test_api.py index 1e79b894..fe0301db 100644 --- a/tests/agent_manager/test_api.py +++ b/tests/agent_manager/test_api.py @@ -165,3 +165,41 @@ def test_create_accepts_stable_session_and_send_accepts_user(client: TestClient) sent = client.post("/conversations/sess-1/messages", json={"message": "hello", "user_id": "u1"}) assert sent.status_code == 200 + + +def test_send_returns_429_when_token_budget_exceeded() -> None: + """send_message returns 429 when the conversation token budget is exhausted.""" + app = FastAPI() + # max_tokens=1 means the budget is treated as exceeded as soon as any token is recorded. + # To trigger the guard we first send a real message (consumes tokens), then send a second one. + service = ConversationService(RecordingEngine(), MemoryRepository(), max_tokens=1) + app.state.service = service + app.include_router(router) + client = TestClient(app) + + cid = client.post("/conversations").json()["conversation_id"] + # First message succeeds and records token usage in the repository. + client.post(f"/conversations/{cid}/messages", json={"message": "first"}) + # Second message should be rejected because the budget is now exhausted. + response = client.post(f"/conversations/{cid}/messages", json={"message": "second"}) + + assert response.status_code == 429 + assert "budget" in response.json()["detail"] + + +def test_stream_returns_429_when_token_budget_exceeded() -> None: + """stream_message returns 429 when the conversation token budget is exhausted.""" + app = FastAPI() + service = ConversationService(RecordingEngine(), MemoryRepository(), max_tokens=1) + app.state.service = service + app.include_router(router) + client = TestClient(app) + + cid = client.post("/conversations").json()["conversation_id"] + # Consume the budget with a non-streaming send first. + client.post(f"/conversations/{cid}/messages", json={"message": "first"}) + # The streaming endpoint should now return 429. + response = client.post(f"/conversations/{cid}/messages/stream", json={"message": "second"}) + + assert response.status_code == 429 + assert "budget" in response.json()["detail"] From a6308ad60339c543b0b0ac8800869a6de957766e Mon Sep 17 00:00:00 2001 From: Biswajit Panda Date: Sun, 26 Jul 2026 17:23:09 +0530 Subject: [PATCH 2/2] Update conftest.py --- tests/agent_manager/conftest.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/agent_manager/conftest.py b/tests/agent_manager/conftest.py index efca65b3..a7ea19a4 100644 --- a/tests/agent_manager/conftest.py +++ b/tests/agent_manager/conftest.py @@ -30,7 +30,13 @@ async def run( self.prompts.append(message) self.histories.append(tuple(history)) self.contexts.append(context) - return RunResult(system_name="stub", visited=["agent"], answer=f"answer:{message}") + return RunResult( + system_name="stub", + visited=["agent"], + answer=f"answer:{message}", + input_tokens=5, + output_tokens=5, + ) async def stream( self,