diff --git a/examples/starter/.env.example b/examples/starter/.env.example index 5e8733aa..232aed1d 100644 --- a/examples/starter/.env.example +++ b/examples/starter/.env.example @@ -21,3 +21,10 @@ OPENAI_BASE_URL=http://host.docker.internal:11434/v1 # Optional — read by the bank_name resolver. # BANK_NAME=Northwind Bank + +# Optional — a per-conversation *cumulative* token budget: every turn's input + +# output tokens, summed over the life of the conversation. This is not the +# model's context window. Once a conversation has spent this many tokens in +# total, further turns are refused with a 429; the widget shows how full the +# budget is so the user sees it coming. +# CONTEXT_MAX_TOKENS=2000 diff --git a/src/agent_manager/api/routes.py b/src/agent_manager/api/routes.py index 445b2f9b..cbcfa0e3 100644 --- a/src/agent_manager/api/routes.py +++ b/src/agent_manager/api/routes.py @@ -19,6 +19,7 @@ SendMessageRequest, SendMessageResponse, StreamEventOut, + TokenBudgetResponse, ToolRecord, ) from agent_manager.application import ( @@ -50,6 +51,20 @@ async def list_messages(conversation_id: str, service: Service) -> list[MessageO return [MessageOut(role=m.role, content=m.content, created_at=m.created_at) for m in msgs] +@router.get("/conversations/{conversation_id}/usage", response_model=TokenBudgetResponse) +async def get_usage(conversation_id: str, service: Service) -> TokenBudgetResponse: + try: + usage = await service.usage(conversation_id) + except ConversationNotFound as exc: + raise HTTPException(status_code=404, detail="conversation not found") from exc + return TokenBudgetResponse( + used_tokens=usage.used_tokens, + max_tokens=usage.max_tokens, + percent=usage.percent, + severity=usage.severity, + ) + + @router.post("/conversations/{conversation_id}/messages", response_model=SendMessageResponse) async def send_message( conversation_id: str, body: SendMessageRequest, service: Service diff --git a/src/agent_manager/api/schemas.py b/src/agent_manager/api/schemas.py index 3777ee5c..1c2452a2 100644 --- a/src/agent_manager/api/schemas.py +++ b/src/agent_manager/api/schemas.py @@ -10,7 +10,7 @@ from pydantic import BaseModel -from agent_manager.domain import Role +from agent_manager.domain import BudgetSeverity, Role class CreateConversationRequest(BaseModel): @@ -49,6 +49,13 @@ class SendMessageResponse(BaseModel): used_tools: list[ToolRecord] +class TokenBudgetResponse(BaseModel): + used_tokens: int + max_tokens: int | None = None + percent: float = 0.0 + severity: BudgetSeverity = BudgetSeverity.NORMAL + + class StreamEventOut(BaseModel): type: str content: str | None = None diff --git a/src/agent_manager/api/static/demo.html b/src/agent_manager/api/static/demo.html index db9313bc..852e0161 100644 --- a/src/agent_manager/api/static/demo.html +++ b/src/agent_manager/api/static/demo.html @@ -21,16 +21,15 @@

Drop-in chat widget

your product, paste these two lines — pointing at your backend:

<script type="module" src="https://your-backend/widget.js"></script>
-<agent-chat title="Home Assistant" color="#2563eb"></agent-chat>
+<agent-chat title="Support" color="#18181b"></agent-chat> -

↘ The launcher is in the bottom-right corner. Click it and try - “turn on the kitchen lights”, then “now turn it off”.

+

↘ The launcher is in the bottom-right corner. Click it and ask a question.

diff --git a/src/agent_manager/api/static/widget.js b/src/agent_manager/api/static/widget.js index a69e7455..3b15819c 100644 --- a/src/agent_manager/api/static/widget.js +++ b/src/agent_manager/api/static/widget.js @@ -52141,6 +52141,19 @@ var AgentChatClient = class { used_tools: Array.isArray(data.used_tools) ? data.used_tools : void 0 }; } + async getUsage(conversationId) { + const response = await fetch(`${this.endpoint}/conversations/${conversationId}/usage`); + if (!response.ok) { + throw new AgentChatHttpError(response.status); + } + const data = await response.json(); + return { + used_tokens: Number(data.used_tokens) || 0, + max_tokens: data.max_tokens == null ? null : Number(data.max_tokens), + percent: Number(data.percent) || 0, + severity: data.severity ?? "normal" + }; + } async *streamMessage(conversationId, message) { const response = await fetch(`${this.endpoint}/conversations/${conversationId}/messages/stream`, { method: "POST", @@ -53051,7 +53064,19 @@ function useConversation(client, endpoint) { return []; } }, [client, endpoint]); - return (0, import_react9.useMemo)(() => ({ send, stream, loadHistory }), [send, stream, loadHistory]); + const loadUsage = (0, import_react9.useCallback)(async () => { + const stored = getStoredConversationId(endpoint); + if (!stored) return null; + try { + return await client.getUsage(stored); + } catch { + return null; + } + }, [client, endpoint]); + return (0, import_react9.useMemo)( + () => ({ send, stream, loadHistory, loadUsage }), + [send, stream, loadHistory, loadUsage] + ); } // src/agent_manager/api/static/widget/react/AgentChatApp.tsx @@ -53072,14 +53097,19 @@ function AgentChatApp({ client, config, onAnswer, panelId, titleId }) { const [loaded, setLoaded] = (0, import_react10.useState)(false); const [sending, setSending] = (0, import_react10.useState)(false); const [entries, setEntries] = (0, import_react10.useState)([]); + const [usage, setUsage] = (0, import_react10.useState)(null); const launcherRef = (0, import_react10.useRef)(null); const inputRef = (0, import_react10.useRef)(null); + const refreshUsage = (0, import_react10.useCallback)(async () => { + setUsage(await conversation.loadUsage()); + }, [conversation]); const loadHistory = (0, import_react10.useCallback)(async () => { if (loaded) return; setLoaded(true); const history = await conversation.loadHistory(); if (history.length) setEntries(history.map(toEntry)); - }, [conversation, loaded]); + await refreshUsage(); + }, [conversation, loaded, refreshUsage]); (0, import_react10.useEffect)(() => { if (inline) void loadHistory(); }, [inline, loadHistory]); @@ -53134,9 +53164,10 @@ function AgentChatApp({ client, config, onAnswer, panelId, titleId }) { await sendWithoutStreaming(text10, pending.id); } finally { setSending(false); + void refreshUsage(); } }, - [conversation, onAnswer, replaceEntry, sendWithoutStreaming] + [conversation, onAnswer, refreshUsage, replaceEntry, sendWithoutStreaming] ); const toggle = () => void (open ? closeChat() : openChat()); return /* @__PURE__ */ (0, import_jsx_runtime4.jsxs)( @@ -53184,7 +53215,10 @@ function AgentChatApp({ client, config, onAnswer, panelId, titleId }) { } ), /* @__PURE__ */ (0, import_jsx_runtime4.jsxs)(PromptInputFooter, { children: [ - /* @__PURE__ */ (0, import_jsx_runtime4.jsx)("span", { className: "prompt-hint", children: "Enter to send \xB7 Shift+Enter for a new line" }), + /* @__PURE__ */ (0, import_jsx_runtime4.jsxs)("div", { className: "footer-start", children: [ + usage ? /* @__PURE__ */ (0, import_jsx_runtime4.jsx)(BudgetMeter, { usage }) : null, + /* @__PURE__ */ (0, import_jsx_runtime4.jsx)("span", { className: "prompt-hint", children: "Enter to send \xB7 Shift+Enter for a new line" }) + ] }), /* @__PURE__ */ (0, import_jsx_runtime4.jsx)(PromptInputSubmit, { disabled: sending }) ] }) ] }), @@ -53280,6 +53314,71 @@ function Welcome({ title }) { /* @__PURE__ */ (0, import_jsx_runtime4.jsx)("p", { className: "welcome-title", children: title }) ] }); } +var RING_SIZE = 18; +var RING_STROKE = 2.5; +var RING_RADIUS = (RING_SIZE - RING_STROKE) / 2; +var RING_CIRCUMFERENCE = 2 * Math.PI * RING_RADIUS; +function BudgetMeter({ usage }) { + const { used_tokens: used, max_tokens: max, percent, severity } = usage; + if (!max) return null; + const center = RING_SIZE / 2; + const ring = { cx: center, cy: center, r: RING_RADIUS, fill: "none", strokeWidth: RING_STROKE }; + const rounded = Math.round(percent); + return /* @__PURE__ */ (0, import_jsx_runtime4.jsxs)( + "span", + { + className: `budget-meter ${severity}`, + role: "img", + "aria-label": `Token budget ${rounded}% used`, + tabIndex: 0, + children: [ + /* @__PURE__ */ (0, import_jsx_runtime4.jsxs)( + "svg", + { + className: "budget-ring", + width: RING_SIZE, + height: RING_SIZE, + viewBox: `0 0 ${RING_SIZE} ${RING_SIZE}`, + "aria-hidden": true, + children: [ + /* @__PURE__ */ (0, import_jsx_runtime4.jsx)("circle", { className: "budget-ring-track", ...ring }), + /* @__PURE__ */ (0, import_jsx_runtime4.jsx)( + "circle", + { + className: "budget-ring-value", + ...ring, + strokeLinecap: "round", + strokeDasharray: RING_CIRCUMFERENCE, + strokeDashoffset: RING_CIRCUMFERENCE * (1 - percent / 100) + } + ) + ] + } + ), + /* @__PURE__ */ (0, import_jsx_runtime4.jsxs)("span", { className: "budget-percent", children: [ + rounded, + "%" + ] }), + /* @__PURE__ */ (0, import_jsx_runtime4.jsxs)("span", { className: "budget-popover", role: "tooltip", children: [ + /* @__PURE__ */ (0, import_jsx_runtime4.jsxs)("span", { className: "budget-popover-head", children: [ + /* @__PURE__ */ (0, import_jsx_runtime4.jsx)("span", { children: "Token budget" }), + /* @__PURE__ */ (0, import_jsx_runtime4.jsxs)("span", { className: "budget-popover-count", children: [ + formatTokens(Math.min(used, max)), + " of ", + formatTokens(max) + ] }) + ] }), + /* @__PURE__ */ (0, import_jsx_runtime4.jsx)("span", { className: "budget-bar", children: /* @__PURE__ */ (0, import_jsx_runtime4.jsx)("span", { className: "budget-bar-fill", style: { width: `${percent}%` } }) }) + ] }) + ] + } + ); +} +function formatTokens(tokens) { + if (tokens >= 1e6) return `${(tokens / 1e6).toFixed(1).replace(/\.0$/, "")}M`; + if (tokens >= 1e3) return `${(tokens / 1e3).toFixed(1).replace(/\.0$/, "")}k`; + return `${tokens}`; +} function avatarStyle(avatar) { if (!avatar) return void 0; return { backgroundImage: `url("${avatar.replace(/"/g, "%22")}")` }; @@ -53428,6 +53527,34 @@ function styles(config) { .prompt-footer { grid-column: 1 / -1; display: flex; align-items: center; justify-content: space-between; gap: 10px; color: #a1a1aa; font-size: 11.5px; } .prompt-hint { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } + .footer-start { display: flex; align-items: center; gap: 8px; min-width: 0; } + .budget-meter { position: relative; display: inline-flex; align-items: center; gap: 5px; + flex: 0 0 auto; color: #71717a; font-size: 11.5px; cursor: default; outline: none; } + .budget-ring { transform: rotate(-90deg); } + .budget-ring-track { stroke: #e4e4e7; } + .budget-ring-value { stroke: #3f3f46; + transition: stroke-dashoffset .3s ease, stroke .3s ease; } + .budget-percent { font-variant-numeric: tabular-nums; } + .budget-meter.warning .budget-ring-value { stroke: #f59e0b; } + .budget-meter.warning .budget-percent { color: #b45309; } + .budget-meter.critical .budget-ring-value { stroke: #ef4444; } + .budget-meter.critical .budget-percent { color: #b91c1c; } + .budget-popover { position: absolute; bottom: calc(100% + 8px); left: 0; width: 200px; + background: #fff; color: #18181b; border: 1px solid #e4e4e7; border-radius: 10px; + padding: 10px 12px; box-shadow: 0 10px 25px rgba(0,0,0,.12); + opacity: 0; transform: translateY(4px); pointer-events: none; + transition: opacity .15s ease, transform .15s ease; z-index: 5; } + .budget-meter:hover .budget-popover, .budget-meter:focus-visible .budget-popover { + opacity: 1; transform: none; } + .budget-popover-head { display: flex; align-items: baseline; justify-content: space-between; + gap: 12px; font-size: 12px; font-weight: 600; white-space: nowrap; } + .budget-popover-count { color: #71717a; font-weight: 400; font-variant-numeric: tabular-nums; } + .budget-bar { display: block; margin-top: 8px; height: 4px; background: #f4f4f5; + border-radius: 999px; overflow: hidden; } + .budget-bar-fill { display: block; height: 100%; background: #3f3f46; + border-radius: 999px; transition: width .3s ease; } + .budget-meter.warning .budget-bar-fill { background: #f59e0b; } + .budget-meter.critical .budget-bar-fill { background: #ef4444; } .powered { text-align: center; padding: 0 14px 10px; color: #a1a1aa; font-size: 11px; letter-spacing: .01em; } @media (prefers-reduced-motion: reduce) { @@ -53442,6 +53569,7 @@ function styles(config) { .panel.open .composer { animation: none; } .welcome { animation: none; } .msg-action svg { animation: none; } + .budget-ring-value, .budget-bar-fill, .budget-popover { transition: none; } } @media (max-width: 480px) { .panel:not(.inline) { width: 100vw; height: 100dvh; max-height: 100dvh; diff --git a/src/agent_manager/api/static/widget/api/AgentChatClient.ts b/src/agent_manager/api/static/widget/api/AgentChatClient.ts index 2566307b..20cc8568 100644 --- a/src/agent_manager/api/static/widget/api/AgentChatClient.ts +++ b/src/agent_manager/api/static/widget/api/AgentChatClient.ts @@ -1,4 +1,4 @@ -import type { ChatMessage, SendMessageResponse, StreamEvent } from "../types"; +import type { ChatMessage, TokenBudget, SendMessageResponse, StreamEvent } from "../types"; export class AgentChatHttpError extends Error { constructor(readonly status: number) { @@ -44,6 +44,20 @@ export class AgentChatClient { }; } + async getUsage(conversationId: string): Promise { + const response = await fetch(`${this.endpoint}/conversations/${conversationId}/usage`); + if (!response.ok) { + throw new AgentChatHttpError(response.status); + } + const data = await response.json(); + return { + used_tokens: Number(data.used_tokens) || 0, + max_tokens: data.max_tokens == null ? null : Number(data.max_tokens), + percent: Number(data.percent) || 0, + severity: data.severity ?? "normal", + }; + } + async *streamMessage(conversationId: string, message: string): AsyncGenerator { const response = await fetch(`${this.endpoint}/conversations/${conversationId}/messages/stream`, { method: "POST", diff --git a/src/agent_manager/api/static/widget/react/AgentChatApp.tsx b/src/agent_manager/api/static/widget/react/AgentChatApp.tsx index 0428b3eb..d2e0f7c8 100644 --- a/src/agent_manager/api/static/widget/react/AgentChatApp.tsx +++ b/src/agent_manager/api/static/widget/react/AgentChatApp.tsx @@ -6,6 +6,7 @@ import type { AgentChatAnswerDetail, AgentChatConfig, ChatMessage, + TokenBudget, MessageEntry, ToolRecord, } from "../types"; @@ -55,15 +56,21 @@ export function AgentChatApp({ client, config, onAnswer, panelId, titleId }: Age const [loaded, setLoaded] = useState(false); const [sending, setSending] = useState(false); const [entries, setEntries] = useState([]); + const [usage, setUsage] = useState(null); const launcherRef = useRef(null); const inputRef = useRef(null); + const refreshUsage = useCallback(async () => { + setUsage(await conversation.loadUsage()); + }, [conversation]); + const loadHistory = useCallback(async () => { if (loaded) return; setLoaded(true); const history = await conversation.loadHistory(); if (history.length) setEntries(history.map(toEntry)); - }, [conversation, loaded]); + await refreshUsage(); + }, [conversation, loaded, refreshUsage]); useEffect(() => { if (inline) void loadHistory(); @@ -125,9 +132,10 @@ export function AgentChatApp({ client, config, onAnswer, panelId, titleId }: Age await sendWithoutStreaming(text, pending.id); } finally { setSending(false); + void refreshUsage(); } }, - [conversation, onAnswer, replaceEntry, sendWithoutStreaming], + [conversation, onAnswer, refreshUsage, replaceEntry, sendWithoutStreaming], ); const toggle = () => void (open ? closeChat() : openChat()); @@ -187,7 +195,10 @@ export function AgentChatApp({ client, config, onAnswer, panelId, titleId }: Age placeholder="Message..." /> - Enter to send · Shift+Enter for a new line +
+ {usage ? : null} + Enter to send · Shift+Enter for a new line +
@@ -337,6 +348,65 @@ function Welcome({ title }: { title: string }) { ); } +const RING_SIZE = 18; +const RING_STROKE = 2.5; +const RING_RADIUS = (RING_SIZE - RING_STROKE) / 2; +const RING_CIRCUMFERENCE = 2 * Math.PI * RING_RADIUS; + +/** Cumulative tokens this conversation has spent, not the current context size. */ +function BudgetMeter({ usage }: { usage: TokenBudget }) { + const { used_tokens: used, max_tokens: max, percent, severity } = usage; + if (!max) return null; + + const center = RING_SIZE / 2; + const ring = { cx: center, cy: center, r: RING_RADIUS, fill: "none", strokeWidth: RING_STROKE }; + const rounded = Math.round(percent); + + return ( + + + + + + {rounded}% + + + Token budget + + {formatTokens(Math.min(used, max))} of {formatTokens(max)} + + + + + + + + ); +} + +function formatTokens(tokens: number): string { + if (tokens >= 1_000_000) return `${(tokens / 1_000_000).toFixed(1).replace(/\.0$/, "")}M`; + if (tokens >= 1_000) return `${(tokens / 1_000).toFixed(1).replace(/\.0$/, "")}k`; + return `${tokens}`; +} + function avatarStyle(avatar: string) { if (!avatar) return undefined; return { backgroundImage: `url("${avatar.replace(/"/g, "%22")}")` }; diff --git a/src/agent_manager/api/static/widget/react/useConversation.ts b/src/agent_manager/api/static/widget/react/useConversation.ts index 9d4c0338..3fab9540 100644 --- a/src/agent_manager/api/static/widget/react/useConversation.ts +++ b/src/agent_manager/api/static/widget/react/useConversation.ts @@ -6,12 +6,13 @@ import { removeStoredConversationId, setStoredConversationId, } from "../storage/conversationStorage"; -import type { ChatMessage, SendMessageResponse, StreamEvent } from "../types"; +import type { ChatMessage, TokenBudget, SendMessageResponse, StreamEvent } from "../types"; export interface Conversation { send(text: string): Promise; stream(text: string): AsyncGenerator; loadHistory(): Promise; + loadUsage(): Promise; } const isMissingConversation = (error: unknown): boolean => @@ -69,5 +70,18 @@ export function useConversation(client: AgentChatClient, endpoint: string): Conv } }, [client, endpoint]); - return useMemo(() => ({ send, stream, loadHistory }), [send, stream, loadHistory]); + const loadUsage = useCallback(async () => { + const stored = getStoredConversationId(endpoint); + if (!stored) return null; + try { + return await client.getUsage(stored); + } catch { + return null; + } + }, [client, endpoint]); + + return useMemo( + () => ({ send, stream, loadHistory, loadUsage }), + [send, stream, loadHistory, loadUsage], + ); } diff --git a/src/agent_manager/api/static/widget/styles/styles.ts b/src/agent_manager/api/static/widget/styles/styles.ts index e461fd16..ba034801 100644 --- a/src/agent_manager/api/static/widget/styles/styles.ts +++ b/src/agent_manager/api/static/widget/styles/styles.ts @@ -137,6 +137,34 @@ export function styles(config: AgentChatConfig): string { .prompt-footer { grid-column: 1 / -1; display: flex; align-items: center; justify-content: space-between; gap: 10px; color: #a1a1aa; font-size: 11.5px; } .prompt-hint { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } + .footer-start { display: flex; align-items: center; gap: 8px; min-width: 0; } + .budget-meter { position: relative; display: inline-flex; align-items: center; gap: 5px; + flex: 0 0 auto; color: #71717a; font-size: 11.5px; cursor: default; outline: none; } + .budget-ring { transform: rotate(-90deg); } + .budget-ring-track { stroke: #e4e4e7; } + .budget-ring-value { stroke: #3f3f46; + transition: stroke-dashoffset .3s ease, stroke .3s ease; } + .budget-percent { font-variant-numeric: tabular-nums; } + .budget-meter.warning .budget-ring-value { stroke: #f59e0b; } + .budget-meter.warning .budget-percent { color: #b45309; } + .budget-meter.critical .budget-ring-value { stroke: #ef4444; } + .budget-meter.critical .budget-percent { color: #b91c1c; } + .budget-popover { position: absolute; bottom: calc(100% + 8px); left: 0; width: 200px; + background: #fff; color: #18181b; border: 1px solid #e4e4e7; border-radius: 10px; + padding: 10px 12px; box-shadow: 0 10px 25px rgba(0,0,0,.12); + opacity: 0; transform: translateY(4px); pointer-events: none; + transition: opacity .15s ease, transform .15s ease; z-index: 5; } + .budget-meter:hover .budget-popover, .budget-meter:focus-visible .budget-popover { + opacity: 1; transform: none; } + .budget-popover-head { display: flex; align-items: baseline; justify-content: space-between; + gap: 12px; font-size: 12px; font-weight: 600; white-space: nowrap; } + .budget-popover-count { color: #71717a; font-weight: 400; font-variant-numeric: tabular-nums; } + .budget-bar { display: block; margin-top: 8px; height: 4px; background: #f4f4f5; + border-radius: 999px; overflow: hidden; } + .budget-bar-fill { display: block; height: 100%; background: #3f3f46; + border-radius: 999px; transition: width .3s ease; } + .budget-meter.warning .budget-bar-fill { background: #f59e0b; } + .budget-meter.critical .budget-bar-fill { background: #ef4444; } .powered { text-align: center; padding: 0 14px 10px; color: #a1a1aa; font-size: 11px; letter-spacing: .01em; } @media (prefers-reduced-motion: reduce) { @@ -151,6 +179,7 @@ export function styles(config: AgentChatConfig): string { .panel.open .composer { animation: none; } .welcome { animation: none; } .msg-action svg { animation: none; } + .budget-ring-value, .budget-bar-fill, .budget-popover { transition: none; } } @media (max-width: 480px) { .panel:not(.inline) { width: 100vw; height: 100dvh; max-height: 100dvh; diff --git a/src/agent_manager/api/static/widget/types.ts b/src/agent_manager/api/static/widget/types.ts index 7eaa9b0b..92c06267 100644 --- a/src/agent_manager/api/static/widget/types.ts +++ b/src/agent_manager/api/static/widget/types.ts @@ -28,6 +28,16 @@ export interface ChatMessage { created_at?: string; } +export type BudgetSeverity = "normal" | "warning" | "critical"; + +/** Cumulative tokens spent by a conversation against its configured budget. */ +export interface TokenBudget { + used_tokens: number; + max_tokens: number | null; + percent: number; + severity: BudgetSeverity; +} + export interface MessageEntry { id: string; role: "user" | "ai"; diff --git a/src/agent_manager/application/service.py b/src/agent_manager/application/service.py index e044c2c4..b5fdc8f0 100644 --- a/src/agent_manager/application/service.py +++ b/src/agent_manager/application/service.py @@ -17,7 +17,13 @@ from agent_engine.runtime.hooks import RunContext from agent_engine.runtime.streaming import RunStreamEvent from agent_manager.application.context import build_history -from agent_manager.domain import ConversationMessage, Message, Repository, Role +from agent_manager.domain import ( + ConversationMessage, + Message, + Repository, + Role, + TokenBudgetUsage, +) class ConversationNotFound(Exception): @@ -76,6 +82,11 @@ async def history(self, conversation_id: str) -> list[Message]: await self._require(conversation_id) return await self._repository.list_messages(conversation_id) + async def usage(self, conversation_id: str) -> TokenBudgetUsage: + await self._require(conversation_id) + used = await self._repository.get_token_usage(conversation_id) + return TokenBudgetUsage.from_totals(used, self._max_tokens) + async def send( self, conversation_id: str, text: str, *, user_id: str | None = None ) -> RunResult: diff --git a/src/agent_manager/domain/__init__.py b/src/agent_manager/domain/__init__.py index e8857caa..49ec6188 100644 --- a/src/agent_manager/domain/__init__.py +++ b/src/agent_manager/domain/__init__.py @@ -1,17 +1,20 @@ """Domain layer: value objects and ports. Pure Python, no frameworks.""" from agent_manager.domain.models import ( + BudgetSeverity, ConversationContext, ConversationMessage, ConversationSession, ConversationSnapshot, Message, Role, + TokenBudgetUsage, User, ) from agent_manager.domain.repository import Repository __all__ = [ + "BudgetSeverity", "ConversationContext", "ConversationMessage", "ConversationSession", @@ -19,5 +22,6 @@ "Message", "Repository", "Role", + "TokenBudgetUsage", "User", ] diff --git a/src/agent_manager/domain/models.py b/src/agent_manager/domain/models.py index 361de7d9..ffe353f3 100644 --- a/src/agent_manager/domain/models.py +++ b/src/agent_manager/domain/models.py @@ -7,6 +7,9 @@ from enum import StrEnum from typing import Any +BUDGET_WARNING_PERCENT = 65.0 +BUDGET_CRITICAL_PERCENT = 85.0 + class Role(StrEnum): USER = "user" @@ -94,3 +97,40 @@ class ConversationContext: message_count: int source: str snapshot: ConversationSnapshot | None = None + + +class BudgetSeverity(StrEnum): + NORMAL = "normal" + WARNING = "warning" + CRITICAL = "critical" + + +@dataclass(frozen=True) +class TokenBudgetUsage: + """How much of a conversation's lifetime token budget has been spent. + + This is cumulative consumption — every turn's input + output tokens summed + over the whole conversation — measured against `context_max_tokens`, the + budget the service enforces (see `ConversationTokenBudgetExceeded`). It is + deliberately *not* the size of the context window currently sent to the + model: history is re-sent each turn, so the same message is counted again + every time it is included. + """ + + used_tokens: int + max_tokens: int | None + percent: float + severity: BudgetSeverity + + @classmethod + def from_totals(cls, used_tokens: int, max_tokens: int | None) -> TokenBudgetUsage: + if not max_tokens: + return cls(used_tokens, max_tokens, 0.0, BudgetSeverity.NORMAL) + percent = min(used_tokens / max_tokens * 100, 100.0) + if percent >= BUDGET_CRITICAL_PERCENT: + severity = BudgetSeverity.CRITICAL + elif percent >= BUDGET_WARNING_PERCENT: + severity = BudgetSeverity.WARNING + else: + severity = BudgetSeverity.NORMAL + return cls(used_tokens, max_tokens, percent, severity) diff --git a/tests/agent_manager/test_api.py b/tests/agent_manager/test_api.py index 1e79b894..694566e3 100644 --- a/tests/agent_manager/test_api.py +++ b/tests/agent_manager/test_api.py @@ -14,6 +14,7 @@ from agent_engine.runtime.streaming import RunStreamEvent from agent_manager.api.routes import router from agent_manager.application import ConversationService +from agent_manager.domain import TokenBudgetUsage from agent_manager.infrastructure.persistence.memory_repository import MemoryRepository from tests.agent_manager.conftest import RecordingEngine @@ -43,6 +44,59 @@ def test_create_send_history_round_trip(client: TestClient) -> None: def test_unknown_conversation_returns_404(client: TestClient) -> None: assert client.get("/conversations/nope/messages").status_code == 404 assert client.post("/conversations/nope/messages", json={"message": "x"}).status_code == 404 + assert client.get("/conversations/nope/usage").status_code == 404 + + +def test_usage_reports_null_budget_when_unset(client: TestClient) -> None: + cid = client.post("/conversations").json()["conversation_id"] + assert client.get(f"/conversations/{cid}/usage").json() == { + "used_tokens": 0, + "max_tokens": None, + "percent": 0.0, + "severity": "normal", + } + + +def test_usage_reports_cumulative_tokens_and_severity_against_budget() -> None: + class TokenEngine(RecordingEngine): + async def run( + self, + message: str, + *, + history: Sequence[ChatMessage] = (), + context: RunContext | None = None, + ) -> RunResult: + return RunResult( + system_name="stub", + visited=["agent"], + answer="ok", + input_tokens=600, + output_tokens=100, + ) + + app = FastAPI() + app.state.service = ConversationService(TokenEngine(), MemoryRepository(), max_tokens=1000) + app.include_router(router) + client = TestClient(app) + + cid = client.post("/conversations").json()["conversation_id"] + client.post(f"/conversations/{cid}/messages", json={"message": "hi"}) + + body = client.get(f"/conversations/{cid}/usage").json() + assert body["used_tokens"] == 700 + assert body["max_tokens"] == 1000 + assert body["percent"] == pytest.approx(70.0) + assert body["severity"] == "warning" + + +def test_token_budget_severity_thresholds() -> None: + assert TokenBudgetUsage.from_totals(0, None).severity == "normal" + assert TokenBudgetUsage.from_totals(640, 1000).severity == "normal" + # Both thresholds are inclusive: exactly 65% warns, exactly 85% is critical. + assert TokenBudgetUsage.from_totals(650, 1000).severity == "warning" + assert TokenBudgetUsage.from_totals(849, 1000).severity == "warning" + assert TokenBudgetUsage.from_totals(850, 1000).severity == "critical" + assert TokenBudgetUsage.from_totals(5000, 1000).percent == 100.0 class _SubAgentEngine(Engine): diff --git a/tests/e2e/widget.spec.ts b/tests/e2e/widget.spec.ts index 4b3be11c..c499eb12 100644 --- a/tests/e2e/widget.spec.ts +++ b/tests/e2e/widget.spec.ts @@ -432,6 +432,50 @@ test("backend error renders a user-friendly message", async ({ page }) => { await expect.poll(() => shadowText(page, ".messages")).toContain("Something went wrong. Please try again."); }); +test("budget meter shows cumulative token usage against the budget after a turn", async ({ page }) => { + await mockConversationApi(page); + await page.route("**/conversations/*/usage", async (route: Route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + used_tokens: 900, + max_tokens: 1000, + percent: 90, + severity: "critical", + }), + }); + }); + await page.goto("/widget-demo.html"); + await shadowClick(page, ".launcher"); + + await shadowFill(page, ".input", "hello"); + await page.keyboard.press("Enter"); + + await expect.poll(() => shadowText(page, ".messages")).toContain("Echo: hello"); + await expect.poll(() => shadowExists(page, ".budget-meter")).toBe(true); + await expect.poll(() => shadowText(page, ".budget-percent")).toBe("90%"); + await expect.poll(() => shadowClassContains(page, ".budget-meter", "critical")).toBe(true); +}); + +test("budget meter stays hidden when no budget is configured", async ({ page }) => { + await mockConversationApi(page); + await page.route("**/conversations/*/usage", async (route: Route) => { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ used_tokens: 0, max_tokens: null, percent: 0, severity: "normal" }), + }); + }); + await page.goto("/widget-demo.html"); + await shadowClick(page, ".launcher"); + await shadowFill(page, ".input", "hello"); + await page.keyboard.press("Enter"); + + await expect.poll(() => shadowText(page, ".messages")).toContain("Echo: hello"); + await expect.poll(() => shadowExists(page, ".budget-meter")).toBe(false); +}); + test("stale stored conversation is replaced before sending to the agent", async ({ page }) => { const calls = await mockConversationApiWithStaleConversation(page); await page.goto("/widget-demo.html");