diff --git a/backend/app/agent/router.py b/backend/app/agent/router.py index 5321a95b..55b986b1 100644 --- a/backend/app/agent/router.py +++ b/backend/app/agent/router.py @@ -472,6 +472,13 @@ def fmt(payload: dict) -> str: total = getattr(usage_data, "total_tokens", None) if total is not None: usage_payload["totalTokens"] = total + # Forward prompt-cache hit count so the client can observe cache + # effectiveness (OpenRouter-native accounting; may be absent on + # providers without caching). + finish_details = getattr(usage_data, "prompt_tokens_details", None) + finish_cached = getattr(finish_details, "cached_tokens", None) if finish_details else None + if finish_cached is not None: + usage_payload["cachedTokens"] = finish_cached finish_event["messageMetadata"] = {"usage": usage_payload} yield fmt(finish_event) diff --git a/frontend/src/db/index.ts b/frontend/src/db/index.ts index b21f39f4..fe7453a0 100644 --- a/frontend/src/db/index.ts +++ b/frontend/src/db/index.ts @@ -109,6 +109,12 @@ export interface ThreadSummaryRecord { threadId: string summary: string coversThroughMessageId: string + /** + * Position of the cut in the history at summary time. Lets compaction trim by + * index when the id can't be found (paginated/restored threads), instead of + * falling back to sending the full untrimmed history. Optional for back-compat. + */ + coversThroughIndex?: number tokenBudget: number createdAt: number } diff --git a/frontend/src/features/agent/application/useZoberChat.ts b/frontend/src/features/agent/application/useZoberChat.ts index d61ea190..fee84b6d 100644 --- a/frontend/src/features/agent/application/useZoberChat.ts +++ b/frontend/src/features/agent/application/useZoberChat.ts @@ -32,10 +32,12 @@ import { estimateTokens, normalizeMessagesForBackend, PAGE_SIZE, - TOKEN_BUDGET, + pruneToFit, + readUsageTokens, + USABLE, } from '@/features/agent/lib/agent-utils' import { buildPrompt, resolveThreadId } from '@/features/agent/lib/context-assembler' -import { buildHistoryToStore, maybeRunBackgroundSummary } from '@/features/agent/lib/context-assembler/background-summary' +import { buildHistoryToStore, maybeCompact } from '@/features/agent/lib/context-assembler/background-summary' import { computeLessonExhaustion } from '@/features/agent/lib/context-assembler/exhaustion' import { ToolExecutor } from '@/features/agent/lib/tools/executor' import { @@ -51,8 +53,6 @@ import { getEffectiveDueItems } from '@/shared/lib/skillSessionProgress' const MAX_TOOL_ROUNDS_LESSON = 5 const MAX_TOOL_ROUNDS_GLOBAL = 5 const MAX_INPUT_CHARS = 8000 -const TOKEN_BUDGET_SOFT = 0.8 -const TOKEN_BUDGET_HARD = 1.0 type AgentActionsDispatch = (action: AgentAction) => void @@ -125,6 +125,9 @@ export function useZoberChat(args: ZoberChatArgs) { const toolCallCountRef = useRef(0) const errorCountRef = useRef(0) const exercisesThisSessionRef = useRef(0) + // Real token usage from the last completed turn (when the backend reports it), + // used as the primary overflow signal for compaction; undefined → fall back to estimate. + const lastUsageTokensRef = useRef(undefined) // Live context built into a ref to avoid stale closures in transport const ctxRef = useRef(null) @@ -310,19 +313,23 @@ export function useZoberChat(args: ZoberChatArgs) { const builtPrompt = ctx ? buildPrompt(ctx) : '' const includeTools = !ctx?.lesson?.exhausted - - const projectedTokens = estimateTokens(finalMessages) + estimateTextTokens(builtPrompt) - - if (projectedTokens > TOKEN_BUDGET_HARD * TOKEN_BUDGET) { - throw new Error('Conversation too long. Please start a new chat.') - } - if (projectedTokens > TOKEN_BUDGET_SOFT * TOKEN_BUDGET) { - console.warn(`[useZoberChat] Approaching context limit: ${projectedTokens} / ${TOKEN_BUDGET}`) + const systemTokens = estimateTextTokens(builtPrompt) + + // No hard block. Idle compaction (maybeCompact, post-response) is the + // primary sizing mechanism; here we apply the LLM-free pruneToFit + // backstop so a send is never refused. With a 1M window this rarely fires. + let outgoing = finalMessages + let projectedTokens = estimateTokens(outgoing) + systemTokens + if (projectedTokens > USABLE) { + outgoing = pruneToFit(outgoing, USABLE - systemTokens) + projectedTokens = estimateTokens(outgoing) + systemTokens + if (projectedTokens > USABLE) + console.warn(`[useZoberChat] still over budget after prune: ${projectedTokens} / ${USABLE}`) } return { body: { - messages: finalMessages, + messages: outgoing, system_prompt: builtPrompt, openrouter_api_key: apiKey || null, tools: includeTools ? getToolDefinitions(toolPool) : [], @@ -404,6 +411,12 @@ export function useZoberChat(args: ZoberChatArgs) { console.error('Agent chat error:', err) toast.error(err.message || 'Unknown error') }, + onFinish({ message }) { + // Capture real usage if the backend streams it on message metadata, so + // compaction keys off actual token counts (opencode-style) rather than the + // CJK estimate. undefined when absent → maybeCompact falls back to estimate. + lastUsageTokensRef.current = readUsageTokens(message) + }, }) // Load persisted thread from IDB and hydrate useChat state via setMessages. @@ -486,7 +499,9 @@ export function useZoberChat(args: ZoberChatArgs) { exercisesCompleted: exercisesThisSessionRef.current, }) } - void maybeRunBackgroundSummary(db, threadId, fullHistory, apiKey, API_BASE, locale) + // Post-response, idle: compact when the turn reached the usable budget. + // Prefers real usage from this turn; falls back to the CJK estimate. + void maybeCompact(db, threadId, fullHistory, apiKey, API_BASE, locale, lastUsageTokensRef.current) })() }, [status, messages, db, narrowed, threadId, apiKey, locale]) diff --git a/frontend/src/features/agent/lib/agent-utils.simulate.test.ts b/frontend/src/features/agent/lib/agent-utils.simulate.test.ts index 1f790f7e..5133b4a4 100644 --- a/frontend/src/features/agent/lib/agent-utils.simulate.test.ts +++ b/frontend/src/features/agent/lib/agent-utils.simulate.test.ts @@ -45,7 +45,7 @@ describe('agent Utils Compaction Pipeline', () => { expect(hasText, 'The text response was lost').toBe(true) }) - it('lobotomy Bug: compactForTokenBudget should never stub GUIDANCE_TOOLS', () => { + it('lobotomy Bug: the normalize pipeline should never stub GUIDANCE_TOOLS', () => { // Simulate a message way over budget to force compaction of old messages const heavyString = 'x'.repeat(64_000 * 5) diff --git a/frontend/src/features/agent/lib/agent-utils.test.ts b/frontend/src/features/agent/lib/agent-utils.test.ts index 51391453..cd3f5721 100644 --- a/frontend/src/features/agent/lib/agent-utils.test.ts +++ b/frontend/src/features/agent/lib/agent-utils.test.ts @@ -1,6 +1,6 @@ import type { UIMessage } from 'ai' import { describe, expect, it } from 'vitest' -import { compactForTokenBudget, estimateTokens, normalizeMessagesForBackend } from '@/features/agent/lib/agent-utils' +import { estimateTokens, isOverflow, normalizeMessagesForBackend, pruneToFit, readUsageTokens, USABLE } from '@/features/agent/lib/agent-utils' /** * Build a UIMessage for tests. Parts often carry extra fields (toolName, args) @@ -436,17 +436,56 @@ describe('estimateTokens', () => { }) }) -describe('compactForTokenBudget', () => { +describe('readUsageTokens', () => { + it('reads canonical camelCase metadata.usage.totalTokens', () => { + expect(readUsageTokens({ metadata: { usage: { totalTokens: 22889 } } })).toBe(22889) + }) + + it('accepts snake_case total_tokens (raw OpenRouter shape)', () => { + expect(readUsageTokens({ metadata: { usage: { total_tokens: 22889 } } })).toBe(22889) + }) + + it('falls back to promptTokens when total is absent', () => { + expect(readUsageTokens({ metadata: { usage: { promptTokens: 22213 } } })).toBe(22213) + expect(readUsageTokens({ metadata: { usage: { prompt_tokens: 22213 } } })).toBe(22213) + }) + + it('reads usage placed directly on metadata', () => { + expect(readUsageTokens({ metadata: { totalTokens: 100 } })).toBe(100) + }) + + it('returns undefined when absent or non-numeric (→ estimate fallback)', () => { + expect(readUsageTokens({})).toBeUndefined() + expect(readUsageTokens(null)).toBeUndefined() + expect(readUsageTokens({ metadata: {} })).toBeUndefined() + expect(readUsageTokens({ metadata: { usage: { totalTokens: 0 } } })).toBeUndefined() + expect(readUsageTokens({ metadata: { usage: { totalTokens: 'x' } } })).toBeUndefined() + }) +}) + +describe('isOverflow', () => { + it('true iff tokens reach the usable budget', () => { + expect(isOverflow(USABLE - 1)).toBe(false) + expect(isOverflow(USABLE)).toBe(true) + expect(isOverflow(USABLE + 1)).toBe(true) + }) + + it('honours an explicit budget', () => { + expect(isOverflow(100, 200)).toBe(false) + expect(isOverflow(200, 200)).toBe(true) + }) +}) + +describe('pruneToFit', () => { it('returns messages unchanged when under budget', () => { const messages = [ msg({ id: '1', role: 'user', parts: [{ type: 'text', text: 'Hello' }] }), msg({ id: '2', role: 'assistant', parts: [{ type: 'text', text: 'Hi' }] }), ] - const result = compactForTokenBudget(messages, 10000, 6) - expect(result).toEqual(messages) + expect(pruneToFit(messages, 10000, 6)).toEqual(messages) }) - it('stubs tool results in older messages when over budget', () => { + it('stubs tool results in older messages when over budget, keeps protected tail full', () => { const messages = [ msg({ id: '1', @@ -462,24 +501,34 @@ describe('compactForTokenBudget', () => { msg({ id: '3', role: 'assistant', parts: [{ type: 'text', text: 'response' }] }), msg({ id: '4', role: 'user', parts: [{ type: 'text', text: 'next' }] }), ] - const result = compactForTokenBudget(messages, 500, 2) - // Tail kept verbatim + const result = pruneToFit(messages, 500, 2) + // Never deletes — length unchanged + expect(result.length).toBe(messages.length) + // Protected tail kept verbatim expect(result.at(-1)).toEqual(messages.at(-1)) expect(result.at(-2)).toEqual(messages.at(-2)) - // Old tool result should be stubbed - const oldPart = part(result[0] as UIMessage) - expect(oldPart.output).not.toContain('x'.repeat(100)) + // Old tool result stubbed + expect(part(result[0] as UIMessage).output).not.toContain('x'.repeat(100)) }) - it('drops oldest messages if still over budget after stubbing', () => { - const messages = Array.from({ length: 20 }, (_, i) => msg({ - id: `${i}`, - role: i % 2 === 0 ? 'user' as const : 'assistant' as const, - parts: [{ type: 'text', text: 'a'.repeat(500) }], - })) - const result = compactForTokenBudget(messages, 1000, 4) - expect(result.length).toBeLessThan(messages.length) - expect(result.at(-1)!.id).toBe('19') + it('never stubs guidance tools', () => { + const messages = [ + msg({ + id: '1', + role: 'assistant', + parts: [{ + type: 'tool-get_core_guidelines', + toolName: 'get_core_guidelines', + state: 'output-available', + output: 'CORE RULES '.repeat(500), + }], + }), + msg({ id: '2', role: 'user', parts: [{ type: 'text', text: 'hi' }] }), + msg({ id: '3', role: 'assistant', parts: [{ type: 'text', text: 'hello' }] }), + ] + const result = pruneToFit(messages, 100, 1) + expect(result.length).toBe(messages.length) + expect(part(result[0] as UIMessage).output).toContain('CORE RULES') }) it('preserves user and assistant text in older messages', () => { @@ -490,34 +539,15 @@ describe('compactForTokenBudget', () => { role: 'assistant', parts: [ { type: 'text', text: 'important answer' }, - { - type: 'tool-get_vocabulary', - toolName: 'get_vocabulary', - state: 'output-available', - output: 'x'.repeat(3000), - }, + { type: 'tool-get_vocabulary', toolName: 'get_vocabulary', state: 'output-available', output: 'x'.repeat(3000) }, ], }), msg({ id: '3', role: 'user', parts: [{ type: 'text', text: 'follow-up' }] }), ] - const result = compactForTokenBudget(messages, 500, 1) - // User text preserved + const result = pruneToFit(messages, 500, 1) expect(part(result[0] as UIMessage).text).toBe('important question') - // Assistant text preserved, tool result stubbed const assistantParts = (result[1] as any).parts expect(assistantParts.find((p: any) => p.type === 'text')?.text).toBe('important answer') - }) - - it('does not start on a tool-role message', () => { - const messages = [ - msg({ id: '0', role: 'assistant' as any, parts: [{ type: 'tool-x', toolName: 'x', state: 'output-available', output: 'y' }] }), - msg({ id: '1', role: 'user', parts: [{ type: 'text', text: 'hi' }] }), - msg({ id: '2', role: 'assistant', parts: [{ type: 'text', text: 'hello' }] }), - ] - // Budget so low that message 0 would be dropped - const result = compactForTokenBudget(messages, 100, 2) - if (result.length < messages.length) { - expect(result[0]?.role).not.toBe('tool') - } + expect(assistantParts.find((p: any) => p.type?.startsWith('tool-'))?.output).not.toContain('x'.repeat(100)) }) }) diff --git a/frontend/src/features/agent/lib/agent-utils.ts b/frontend/src/features/agent/lib/agent-utils.ts index b6f56093..2bc8e115 100644 --- a/frontend/src/features/agent/lib/agent-utils.ts +++ b/frontend/src/features/agent/lib/agent-utils.ts @@ -11,6 +11,12 @@ const DATA_TOOLS = new Set([ 'get_vocabulary', 'get_progress_summary', 'recall_memory', + // NOTE: `search_document` is deliberately NOT here. Dedup keys by tool name and + // keeps only the latest occurrence — but each search_document call returns + // DIFFERENT passages for a different query. Deduping would rewrite earlier + // retrievals to {status:'superseded'}, so an agent that makes several queries in + // one turn reads its own results as "no results / not in knowledge base". + // Old RAG payloads are freed only by compaction (aging out) / pruneToFit (overflow). ]) // ── Types ── @@ -40,34 +46,75 @@ export function toolName(p: ToolPart): string { /** * [STAGE 0: Helper] - * Rough approximation of token count: ~4 characters per token. - * While not perfectly accurate for all models, it provides a stable heuristic for - * budget-clearing decisions without needing a heavy tokenizer in the browser. - * - * Example: "Hello" (5 chars) -> ~1.25 tokens - * @param messages - The array of UIMessages to calculate tokens for. - * @returns The estimated token count. + * CJK-aware token estimate. Latin text runs ~4 chars/token, but CJK codepoints + * (Han, kana, Hangul) are ~1-2 tokens *each* — char/4 under-counts Mandarin by + * 3-4×, which is the dominant content here. We count CJK at ~1.7 tokens and + * everything else at char/4. This is only the FALLBACK signal; real usage from + * the model response is preferred when available. */ +function isCjkCodepoint(c: number): boolean { + return ( + (c >= 0x4E00 && c <= 0x9FFF) // CJK Unified Ideographs + || (c >= 0x3400 && c <= 0x4DBF) // Extension A + || (c >= 0xF900 && c <= 0xFAFF) // Compatibility Ideographs + || (c >= 0x3040 && c <= 0x30FF) // Hiragana + Katakana + || (c >= 0xAC00 && c <= 0xD7A3) // Hangul syllables + ) +} + +function textTokens(s: string): number { + let cjk = 0 + let other = 0 + for (const ch of s) { + if (isCjkCodepoint(ch.codePointAt(0)!)) + cjk++ + else + other++ + } + return cjk * 1.7 + other / 4 +} + export function estimateTokens(messages: UIMessage[]): number { if (messages.length === 0) return 0 - let chars = 0 + let tokens = 0 for (const msg of messages) { for (const part of msg.parts) { if (part.type === 'text') - chars += (part.text ?? '').length + tokens += textTokens(part.text ?? '') else if (isToolPart(part) && part.output != null) - chars += typeof part.output === 'string' ? part.output.length : JSON.stringify(part.output).length - chars += 20 // per-part overhead (role, type, toolName) + tokens += textTokens(typeof part.output === 'string' ? part.output : JSON.stringify(part.output)) + tokens += 5 // per-part overhead (role, type, toolName) } - chars += 30 // per-message overhead + tokens += 8 // per-message overhead } - return Math.ceil(chars / 4) + return Math.ceil(tokens) } -/** Cheap char-based token estimate for a single string (no UIMessage wrapping). */ +/** CJK-aware token estimate for a single string (no UIMessage wrapping). */ export function estimateTextTokens(text: string): number { - return Math.ceil(text.length / 4) + return Math.ceil(textTokens(text)) +} + +/** + * Extract the real token count for the last turn from a UI message's metadata, + * used as the primary overflow signal (the CJK estimate is the fallback). + * + * CONTRACT — the backend must stream this on the assistant message metadata: + * metadata: { usage: { totalTokens, promptTokens?, cachedTokens? } } + * camelCase is canonical; snake_case (`total_tokens` / `prompt_tokens`) is also + * accepted so a backend that forwards OpenRouter's raw shape still works. + * + * Prefers `totalTokens` (this turn's input+output ≈ next turn's context), then + * `promptTokens`. Returns undefined when absent → caller falls back to estimate. + */ +export function readUsageTokens(message: unknown): number | undefined { + const meta = (message as { metadata?: Record } | null)?.metadata + if (!meta) + return undefined + const u = (meta.usage ?? meta) as Record + const v = u.totalTokens ?? u.total_tokens ?? u.promptTokens ?? u.prompt_tokens + return typeof v === 'number' && v > 0 ? v : undefined } /** @@ -318,80 +365,54 @@ function deduplicateDataToolResults(messages: UIMessage[]): UIMessage[] { }) } -export const TOKEN_BUDGET = 64_000 -export const VERBATIM_TAIL = 15 +// ── Token budget (mirrors opencode's overflow.ts `usable`) ── +// The agent runs on deepseek-v4-flash (1M context). USABLE reserves room for the +// model's output so a full-context request never gets truncated server-side. +export const MODEL_CONTEXT_WINDOW = 1_000_000 +export const RESERVE = 20_000 +export const USABLE = MODEL_CONTEXT_WINDOW - RESERVE + +// No-LLM prune fallback protects this many trailing messages (active context). +export const PROTECT_RECENT_MESSAGES = 15 + +/** opencode `isOverflow`: a turn's token count has reached the usable budget. */ +export function isOverflow(tokens: number, budget: number = USABLE): boolean { + return tokens >= budget +} /** - * [STAGE 7] - * The Final Defense. If history still exceeds the token budget (e.g. 64k), - * we aggressively truncate the past. - * - * 1. Preserves a sliding window (VERBATIM_TAIL=15) at the bottom. - * 2. In messages ABOVE that window, we stub ALL tool outputs. - * 3. Never deletes Knowledge Tools (Whitelist) to prevent "lobotomy". - * - * Before: 70,000 tokens (Crashes API) - * After: 60,000 tokens (Safe) + * LLM-free prune fallback. Stubs tool outputs in OLDER messages (beyond the + * protected recent window) to reclaim tokens WITHOUT summarizing. Never deletes + * messages, never touches GUIDANCE_TOOLS (rules must survive), and never touches + * the protected tail — so active context incl. live `search_document` passages + * stays full. `compact()` is the primary sizing mechanism; this is the no-network + * backstop when `/api/summarize` is slow or unavailable. */ -export function compactForTokenBudget( +export function pruneToFit( messages: UIMessage[], - budget: number = TOKEN_BUDGET, - verbatimTail: number = VERBATIM_TAIL, + budget: number = USABLE, + protectRecent: number = PROTECT_RECENT_MESSAGES, ): UIMessage[] { if (messages.length === 0 || estimateTokens(messages) <= budget) return messages - const splitAt = Math.max(0, messages.length - verbatimTail) - const tail = messages.slice(splitAt) - const older = messages.slice(0, splitAt) - - // Stub tool result content in older messages - const compacted = older.map((msg) => { - if (msg.role !== 'assistant') + const splitAt = Math.max(0, messages.length - protectRecent) + return messages.map((msg, i) => { + if (i >= splitAt || msg.role !== 'assistant') return msg - if (!msg.parts.some(p => - isToolPart(p) && (p.state === 'output-available' || p.state === 'output-error'), - )) { + if (!msg.parts.some(p => isToolPart(p) && p.state === 'output-available')) return msg - } return { ...msg, parts: msg.parts.map((p) => { - if (!isToolPart(p)) + if (!isToolPart(p) || p.state !== 'output-available') return p - if (p.state === 'output-available') { - // EXEMPT: Never stub guidance/knowledge tools, or the agent lobotomizes itself - if (GUIDANCE_TOOLS.has(toolName(p))) { - return p - } - return { ...p, output: `[${toolName(p)} result omitted]` } - } - return p + if (GUIDANCE_TOOLS.has(toolName(p))) + return p + return { ...p, output: `[${toolName(p)} result omitted]` } }), } }) - - const result = [...compacted, ...tail] - - // Drop oldest if still over budget - while (result.length > verbatimTail && estimateTokens(result) > budget) { - const dropIndex = result.findIndex((msg, i) => { - // Stop searching if we hit the guarded verbatim tail - if (i >= result.length - verbatimTail) - return false - // Protect guidance tools from being deleted - if (msg.role === 'assistant' && msg.parts.some(p => isToolPart(p) && GUIDANCE_TOOLS.has(toolName(p)))) { - return false - } - return true - }) - - if (dropIndex === -1) - break - result.splice(dropIndex, 1) - } - - return result } // ── Public API ── @@ -441,23 +462,20 @@ export function compactVocab(e: { id: string, word: string, romanization?: strin * [ Stage 6: deduplicateDataToolResults ] * Similar to Stage 5, but for temporary Data. If the agent calls `get_vocabulary` at minute 1, * and `get_vocabulary` again at minute 20, the Minute 1 data is stale. - * This finds older data fetches and stubs them: + * This finds older data fetches and stubs them (keeping the latest full): * - Output: `{ status: "superseded" }` + * `search_document` is included here so stale RAG duplicates are freed while the + * latest result stays intact (the answer's source-of-truth is never trimmed). * - * [ Stage 7: compactForTokenBudget ] - * The Last Resort. If the history is STILL heavily over the budget (e.g. > 64,000 tokens), - * it aggressively prunes older chat history. - * - Step A: It ring-fences the `VERBATIM_TAIL` (the last 15 messages) so immediate context is safe. - * - Step B: In messages older than the Tail, it blanks out ALL normal tool outputs. - * - Step C: It loops to literally delete the oldest messages one by one (`slice`) until - * budget is met. - * - CRITICAL: It actively searches for and EXEMPTS knowledge tools (e.g., `get_core_guidelines`) - * from deletion so the agent never forgets its core instructions. + * Sizing is NOT done here anymore. Overflow is handled by `compact()` + * (background-summary.ts) — summarize old turns, keep the recent tail full — with + * `pruneToFit` as the LLM-free backstop. Mirrors opencode: this pipeline is the + * per-send cleanup; compaction owns the budget. * - * Diagram of a heavily compressed history sent to the LLM: + * Diagram of a normalized history sent to the LLM: * * [User] "How does this app work?" - * [Asst] <-- Protected by Stage 7 + * [Asst] <-- Guidance always kept * [User] "What's my vocab?" * [Asst] <-- Deduplicated by Stage 6 * [User] "Wait show me again." @@ -473,6 +491,5 @@ export function normalizeMessagesForBackend(messages: UIMessage[]): UIMessage[] result = summarizeRenderOutputs(result) result = compressStaleGuidance(result) result = deduplicateDataToolResults(result) - result = compactForTokenBudget(result) return result } diff --git a/frontend/src/features/agent/lib/context-assembler/background-summary.ts b/frontend/src/features/agent/lib/context-assembler/background-summary.ts index 3a212c84..4872cc29 100644 --- a/frontend/src/features/agent/lib/context-assembler/background-summary.ts +++ b/frontend/src/features/agent/lib/context-assembler/background-summary.ts @@ -1,20 +1,55 @@ import type { UIMessage } from '@ai-sdk/react' import type { ShadowLearnDB, ThreadSummaryRecord } from '@/db' import { getLatestSummary, getThread, putThreadSummary, saveThreadMessages } from '@/db' -import { estimateTokens, TOKEN_BUDGET } from '@/features/agent/lib/agent-utils' +import { estimateTokens, isOverflow } from '@/features/agent/lib/agent-utils' -const COMPACTION_TRIGGER_RATIO = 0.7 -const MIN_MESSAGES_BEFORE_SUMMARY = 40 +// Mirrors opencode compaction.ts: keep the recent tail verbatim, summarize the rest. +const TAIL_TURNS = 2 +const MIN_PRESERVE_RECENT_TOKENS = 2_000 +const MAX_PRESERVE_RECENT_TOKENS = 8_000 + +// Structured but TUTOR-shaped (not opencode's coding-agent template). Passed to +// /api/summarize so the backend produces a summary that preserves teaching +// continuity — what was taught, what the learner struggles with, what drill is mid-flight. +export const TUTOR_SUMMARY_TEMPLATE = `Summarise the tutoring conversation so far using EXACTLY this Markdown structure (keep the headings, fill each with bullets or "(none)"): +## Topics Covered +## Grammar Points Explained +## Vocabulary Touched +## Mistake Patterns +## Pending Drill or Exercise State +## Open Questions` const inFlight = new Set() /** - * Given the full message history and an existing summary, return what to persist to IDB. - * Collapses everything up to and including `coversThroughMessageId` into a single synthetic - * assistant message containing the summary, keeping the messages after the cut visible. - * The cut point is set at the second-to-last assistant message, so the most recent - * user→assistant exchange is always preserved as real messages. - * When no summary exists, returns `fullHistory` unchanged. + * Index where the preserved tail begins (inclusive). A *turn* is a user-bounded + * exchange: a user message + every downstream message it produced (assistant text + * + all tool roundtrips). Keep the last TAIL_TURNS such turns, capped at + * MAX_PRESERVE_RECENT_TOKENS (floor MIN). Everything before is summarised. + */ +export function selectTailStart(messages: UIMessage[]): number { + let tokens = 0 + let userSeen = 0 + let tailStart = messages.length + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i].role === 'user') { + userSeen++ + if (userSeen > TAIL_TURNS && tokens >= MIN_PRESERVE_RECENT_TOKENS) + break + if (tokens >= MAX_PRESERVE_RECENT_TOKENS) + break + } + tokens += estimateTokens([messages[i]]) + tailStart = i + } + return tailStart +} + +/** + * Collapse everything up to and including `coversThroughMessageId` into a single + * synthetic assistant message holding the summary, keeping later messages visible. + * Falls back to the stored index when the id can't be found (paginated/restored + * threads) — never silently returns the full untrimmed history. */ export function buildHistoryToStore( fullHistory: UIMessage[], @@ -23,13 +58,16 @@ export function buildHistoryToStore( if (!summary) return fullHistory - const cutIdx = fullHistory.findIndex(m => m.id === summary.coversThroughMessageId) + let cutIdx = fullHistory.findIndex(m => m.id === summary.coversThroughMessageId) + if (cutIdx < 0 && typeof summary.coversThroughIndex === 'number' && summary.coversThroughIndex < fullHistory.length) { + console.warn('[buildHistoryToStore] coversThroughMessageId not found — trimming by stored index', summary.coversThroughIndex) + cutIdx = summary.coversThroughIndex + } if (cutIdx < 0) { - console.warn('[buildHistoryToStore] coversThroughMessageId not found in history — returning full history. Summary may be stale.', summary.coversThroughMessageId) + console.warn('[buildHistoryToStore] cannot locate cut point — returning full history. Summary may be stale.', summary.coversThroughMessageId) return fullHistory } - const postSummaryMessages = fullHistory.slice(cutIdx + 1) return [ { id: 'compaction-assistant', @@ -37,100 +75,116 @@ export function buildHistoryToStore( content: summary.summary, parts: [{ type: 'text', text: summary.summary }], } as UIMessage, - ...postSummaryMessages, + ...fullHistory.slice(cutIdx + 1), ] } -export async function maybeRunBackgroundSummary( +/** + * opencode-style compaction: summarise older turns into a structured Compaction + * message, keep the recent tail verbatim (full tool outputs incl. live RAG), + * persist. Returns true if it compacted. THROWS on summarize failure so the + * synchronous send-path caller can fall back to a prune; the idle caller + * (`maybeCompact`) swallows + logs. + */ +export async function compact( db: ShadowLearnDB, threadId: string, messages: UIMessage[], apiKey: string, apiBase: string, locale: string, -): Promise { - if (messages.length < MIN_MESSAGES_BEFORE_SUMMARY) - return +): Promise { + const tailStart = selectTailStart(messages) + if (tailStart <= 0) + return false // everything fits in the preserved tail — nothing older to summarise - const last = messages.at(-1)! - const flightKey = `${threadId}:${last.id}` - if (inFlight.has(flightKey)) - return - inFlight.add(flightKey) + const cutMsg = messages[tailStart - 1] + const older = messages.slice(0, tailStart) - let previous: Awaited> + const resp = await fetch(`${apiBase}/api/summarize`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + messages: older.map((m: any) => ({ role: m.role, parts: m.parts })), + template: TUTOR_SUMMARY_TEMPLATE, + openrouter_api_key: apiKey || null, + locale, + }), + }) + if (!resp.ok) { + console.warn('[compact] /api/summarize failed', resp.status) + throw new Error(`summarize failed: ${resp.status}`) + } + + let parsed: { summary?: string } try { - previous = await getLatestSummary(db, threadId) + parsed = await resp.json() } - catch { - inFlight.delete(flightKey) - return + catch (e) { + console.warn('[compact] could not parse summarize response', e) + throw new Error('summarize: invalid JSON') + } + if (!parsed.summary) { + console.warn('[compact] summarize returned no summary') + throw new Error('summarize: empty summary') } - const coveredIdx = previous - ? messages.findIndex(m => m.id === previous!.coversThroughMessageId) - : -1 - const uncovered = messages.slice(coveredIdx + 1) - - if (uncovered.length < MIN_MESSAGES_BEFORE_SUMMARY) { - inFlight.delete(flightKey) - return + const newSummary: ThreadSummaryRecord = { + threadId, + summary: parsed.summary, + coversThroughMessageId: cutMsg.id, + coversThroughIndex: tailStart - 1, + tokenBudget: estimateTokens(messages), + createdAt: Date.now(), } - if (estimateTokens(uncovered) < COMPACTION_TRIGGER_RATIO * TOKEN_BUDGET) { - inFlight.delete(flightKey) + await putThreadSummary(db, newSummary) + + // Rewrite stored history compacted (re-fetch in case messages arrived since). + const thread = await getThread(db, threadId) + if (thread) + await saveThreadMessages(db, threadId, buildHistoryToStore(thread.messages, newSummary), thread.surface, thread.ownerId) + + return true +} + +/** + * Idle/post-response trigger. Compacts only when the turn's token count has + * reached the usable budget (opencode `isOverflow`). `tokens` should be the real + * usage reported by the model when available; falls back to the CJK-aware estimate. + * Errors are swallowed (logged) — this runs in an effect and must not throw. + */ +export async function maybeCompact( + db: ShadowLearnDB, + threadId: string, + messages: UIMessage[], + apiKey: string, + apiBase: string, + locale: string, + tokens?: number, +): Promise { + const count = tokens ?? estimateTokens(messages) + if (!isOverflow(count)) return - } - // Cut after the second-to-last assistant message so the most recent - // user→assistant exchange stays visible after compaction. - const assistantIdxs: number[] = [] - for (let i = messages.length - 1; i >= 0 && assistantIdxs.length < 2; i--) { - if (messages[i].role === 'assistant') - assistantIdxs.push(i) - } - const cutIdx = assistantIdxs.length >= 2 ? assistantIdxs[1] : (assistantIdxs[0] ?? messages.length - 1) - const cutMsg = messages[cutIdx] + const last = messages.at(-1) + if (!last) + return + const key = `${threadId}:${last.id}` + if (inFlight.has(key)) + return + inFlight.add(key) try { - const resp = await fetch(`${apiBase}/api/summarize`, { - method: 'POST', - headers: { 'content-type': 'application/json' }, - body: JSON.stringify({ - messages: messages.slice(0, cutIdx + 1).map((m: any) => ({ role: m.role, parts: m.parts })), - openrouter_api_key: apiKey || null, - locale, - }), - }) - if (!resp.ok) + // Avoid re-summarising the same already-covered range. + const previous = await getLatestSummary(db, threadId) + if (previous && previous.coversThroughMessageId === messages[selectTailStart(messages) - 1]?.id) return - - let parsedJson: { summary?: string } - try { parsedJson = await resp.json() } - catch { return } - if (!parsedJson.summary) - return - - const newSummary: ThreadSummaryRecord = { - threadId, - summary: parsedJson.summary, - coversThroughMessageId: cutMsg.id, - tokenBudget: estimateTokens(messages), - createdAt: Date.now(), - } - await putThreadSummary(db, newSummary) - - // Write compacted form so future loads start clean. - // Fetch current thread in case new messages arrived after our snapshot. - const thread = await getThread(db, threadId) - if (thread) { - const toStore = buildHistoryToStore(thread.messages, newSummary) - await saveThreadMessages(db, threadId, toStore, thread.surface, thread.ownerId) - } + await compact(db, threadId, messages, apiKey, apiBase, locale) } - catch { - // Background failure silent by design. + catch (e) { + console.warn('[maybeCompact] compaction failed', e) } finally { - inFlight.delete(flightKey) + inFlight.delete(key) } } diff --git a/frontend/src/features/agent/lib/prompt/sections.ts b/frontend/src/features/agent/lib/prompt/sections.ts index f6f41c29..881fba36 100644 --- a/frontend/src/features/agent/lib/prompt/sections.ts +++ b/frontend/src/features/agent/lib/prompt/sections.ts @@ -38,6 +38,16 @@ export function toDateOnly(iso?: string): string { return (iso ?? new Date().toISOString()).slice(0, 10) } +export function learningResourceProtocolBlock(): string { + return tag('learning_resource_protocol', [ + 'Learning tips and vocabulary methods = call `search_document` FIRST, then ground the answer in the returned passages.', + '- **TRIGGER:** any question about HOW to learn Chinese — memorizing words, building vocabulary, planning a study schedule, retention techniques, effective daily habits, or any request for a recommended video/lesson on a specific learning topic (e.g. "how do I memorize tones", "recommend a video about spaced repetition", "how should I plan my study day").', + '- **DO NOT answer from training data alone.** The knowledge base contains curated content from language-learning experts; prefer retrieved passages over generic advice.', + '- **NO REDUNDANT CALLS:** if a prior `search_document` result in this conversation already covers the topic, reuse those passages.', + '- **SHOW THE SOURCE URL** when one appears in the returned passages — link the learner directly to the original video.', + ]) +} + // ── Static blocks (no per-session data → cacheable) ────────────────────────── export function roleBlock(surface: PromptSurface): string { @@ -86,12 +96,14 @@ export function instructionsBlock(surface: PromptSurface): string { '- Skip filler and preamble. Use one sentence when possible.', '- Use save_memory() to remember important user preferences or observations.', '- **Call `recall_memory()` proactively when the user asks about their goals, preferences, history, or learning context** — do not rely solely on the Memory Summary above.', - '- Answer general learning questions, including grammar (see ``). Only DECLINE questions about a specific YouTube lesson\'s content or lesson-specific actions (exercises, shadowing) — point the user into the lesson for those. Boundary: "what does 把 mean / how does 了 work" = general grammar, ANSWER it; "what does the 3rd sentence of my Daily Conversation lesson mean / start a drill" = lesson-specific, REDIRECT into the lesson.', - '- For grammar, follow `` below — never `get_skill_guide`.', + '- Answer general learning questions, including grammar (see ``) and learning strategies (see ``). Only DECLINE questions about a specific YouTube lesson\'s content or lesson-specific actions (exercises, shadowing) — point the user into the lesson for those. Boundary: "what does 把 mean / how does 了 work" = general grammar, ANSWER it; "what does the 3rd sentence of my Daily Conversation lesson mean / start a drill" = lesson-specific, REDIRECT into the lesson.', + '- For grammar, follow `` — never `get_skill_guide`.', + '- For learning tips, study planning, vocabulary methods, or video recommendations, follow `` — call `search_document` first.', '- If asked for tips, advice, or a topic covered in core guidelines or skill guides, ALWAYS use get_core_guidelines() or get_skill_guide() to provide accurate info.', '- Do not re-call `get_core_guidelines` or `get_skill_guide` if already loaded this session.', ]), grammarProtocolBlock(), + learningResourceProtocolBlock(), ]) } return compose([ @@ -100,7 +112,8 @@ export function instructionsBlock(surface: PromptSurface): string { '- Skip filler and preamble. Use one sentence when possible.', '- **Call `get_core_guidelines()` at session start — loads SLA principles, feedback templates, and session protocols.**', '- **ALWAYS call `get_skill_guide({ skill })` BEFORE giving advice, tips, or answering "how-to" questions about specific skills (tones, pronunciation, vocabulary, listening, speaking, characters).**', - '- For grammar, follow `` below — never `get_skill_guide`.', + '- For grammar, follow `` — never `get_skill_guide`.', + '- For learning tips, study planning, vocabulary methods, or video recommendations, follow `` — call `search_document` first.', '- Chain tools when needed, but always end with a user-visible response.', '- Use get_study_context (composite) before suggesting exercises — it covers all data in one call.', '- Save important user observations with save_memory().', @@ -108,6 +121,7 @@ export function instructionsBlock(surface: PromptSurface): string { '- Do not call `get_vocabulary` without a specific purpose — avoid speculative data fetching.', ]), grammarProtocolBlock(), + learningResourceProtocolBlock(), tag('exercise_rendering_rules', [ 'STRICT RULES — exercises MUST be rendered via tools, never as chat text.', '- **NEVER write exercise questions as plain text in the chat.** Exercises MUST always be rendered via `render_study_session`.', diff --git a/frontend/src/features/agent/lib/tools/data/searchDocument.ts b/frontend/src/features/agent/lib/tools/data/searchDocument.ts index 6d6eb49d..b2441c09 100644 --- a/frontend/src/features/agent/lib/tools/data/searchDocument.ts +++ b/frontend/src/features/agent/lib/tools/data/searchDocument.ts @@ -28,11 +28,14 @@ export async function executeSearchDocument( export const searchDocumentTool = buildTool({ name: 'search_document', - description: 'Search the knowledge base and return relevant verbatim passages to ground your answer. The knowledge base holds two kinds of documents: (1) the ShadowLearn app user manual — how to use features of the app; (2) a grammar-point reference compiled from well-known language-learning YouTube channels. Call this for questions like "how do I use shadowing mode?" (manual) or "explain the 把 construction" / "when do I use this grammar point?" (grammar reference). Pass a natural-language question, not keywords. Ground your answer in the returned passages and cite the source document; do not add facts beyond them.', + description: 'Search the knowledge base and return relevant verbatim passages to ground your answer. The knowledge base holds three kinds of documents: (1) the ShadowLearn app user manual — how to use features of the app; (2) a grammar-point reference compiled from well-known language-learning YouTube channels; (3) learning strategy content covering vocabulary acquisition methods, memorization techniques, study scheduling, and effective learning habits. Call this for: "how do I use shadowing mode?" (manual), "explain the 把 construction" (grammar), or "how do I memorize Chinese words / plan my study day / learn vocabulary effectively?" (learning strategies). Also call this when the user asks for a recommended video or lesson on any of these topics. Pass a natural-language question, not keywords. Ground your answer in the returned passages and cite the source; do not add facts beyond them.', inputSchema: SearchDocumentSchema, isConcurrencySafe: () => true, isReadOnly: () => true, - maxResultSizeChars: 100_000, - searchHint: 'search app manual, how to use feature, grammar point explanation, construction pattern language reference', + // RAG passages are the source of truth the agent fetched to answer; never chop + // them at produce-time. Context is managed downstream by compaction (stale + // passages are pruned only after being summarized away, never the active one). + maxResultSizeChars: Number.MAX_SAFE_INTEGER, + searchHint: 'search app manual, how to use feature, grammar point explanation, construction pattern language reference, vocabulary memorization methods, study planning, learning tips, effective Chinese learning strategies', execute: async (input, _context) => executeSearchDocument(input as SearchDocumentArgs), }) diff --git a/frontend/src/shared/ui/ai-elements/prompt-input.tsx b/frontend/src/shared/ui/ai-elements/prompt-input.tsx index 821f198c..2e010558 100644 --- a/frontend/src/shared/ui/ai-elements/prompt-input.tsx +++ b/frontend/src/shared/ui/ai-elements/prompt-input.tsx @@ -1253,6 +1253,7 @@ export function PromptInputSubmit({ const handleClick = useCallback( (e: any) => { + console.warn('[PromptInputSubmit] click', { status, isGenerating, hasOnStop: !!onStop }) if (isGenerating && onStop) { e.preventDefault() onStop() @@ -1260,7 +1261,7 @@ export function PromptInputSubmit({ } onClick?.(e) }, - [isGenerating, onStop, onClick], + [isGenerating, onStop, onClick, status], ) return ( diff --git a/frontend/tests/background-summary.test.ts b/frontend/tests/background-summary.test.ts index 7a6a0c4b..4c308e3d 100644 --- a/frontend/tests/background-summary.test.ts +++ b/frontend/tests/background-summary.test.ts @@ -1,28 +1,101 @@ +import type { UIMessage } from '@ai-sdk/react' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { getLatestSummary, initDB } from '@/db' -import { maybeRunBackgroundSummary } from '@/features/agent/lib/context-assembler/background-summary' +import { USABLE } from '@/features/agent/lib/agent-utils' +import { buildHistoryToStore, compact, maybeCompact, selectTailStart } from '@/features/agent/lib/context-assembler/background-summary' import 'fake-indexeddb/auto' -describe('maybeRunBackgroundSummary', () => { +function bigMsgs(n: number): UIMessage[] { + return Array.from({ length: n }, (_, i) => ({ + id: `m${i}`, + role: i % 2 === 0 ? 'user' : 'assistant', + parts: [{ type: 'text', text: 'x'.repeat(4000) }], + })) as any +} + +function textMsg(id: string): UIMessage { + return ({ id, role: 'user', parts: [{ type: 'text', text: id }] }) as any +} + +describe('selectTailStart', () => { + it('keeps recent turns and leaves older content to summarize', () => { + const start = selectTailStart(bigMsgs(40)) + expect(start).toBeGreaterThan(0) + expect(start).toBeLessThan(40) + }) + + it('returns 0 when everything fits in the tail', () => { + expect(selectTailStart([textMsg('1'), textMsg('2')])).toBe(0) + }) +}) + +describe('maybeCompact', () => { beforeEach(() => { (globalThis as any).indexedDB = new (globalThis as any).IDBFactory() }) afterEach(() => vi.restoreAllMocks()) - it('skips when message count below threshold', async () => { + it('skips when not over budget (real usage below USABLE)', async () => { const db = await initDB() const spy = vi.spyOn(globalThis, 'fetch').mockResolvedValue(new Response()) - await maybeRunBackgroundSummary(db, 'tid', [], 'k', 'http://x') + await maybeCompact(db, 'tid', bigMsgs(40), 'k', 'http://x', 'en', 100) expect(spy).not.toHaveBeenCalled() db.close() }) - it('persists summary on successful JSON response', async () => { + it('compacts + persists when over budget, recording the cut index', async () => { const db = await initDB() - vi.spyOn(globalThis, 'fetch').mockResolvedValue(new Response(JSON.stringify({ summary: 'hi' }), { headers: { 'content-type': 'application/json' } })) - // 50 large messages so estimateTokens > 70% of TOKEN_BUDGET (64_000) - const msgs = Array.from({ length: 50 }, (_, i) => ({ id: `m${i}`, role: 'user' as const, parts: [{ type: 'text', text: 'x'.repeat(4000) }] })) - await maybeRunBackgroundSummary(db, 'tid', msgs as any, 'k', 'http://x') + vi.spyOn(globalThis, 'fetch').mockResolvedValue( + new Response(JSON.stringify({ summary: 'hi' }), { headers: { 'content-type': 'application/json' } }), + ) + await maybeCompact(db, 'tid', bigMsgs(40), 'k', 'http://x', 'en', USABLE) const s = await getLatestSummary(db, 'tid') expect(s?.summary).toBe('hi') + expect(typeof s?.coversThroughIndex).toBe('number') + db.close() + }) +}) + +describe('compact', () => { + beforeEach(() => { (globalThis as any).indexedDB = new (globalThis as any).IDBFactory() }) + afterEach(() => vi.restoreAllMocks()) + + it('throws on summarize failure so the send-path can fall back to prune', async () => { + const db = await initDB() + vi.spyOn(globalThis, 'fetch').mockResolvedValue(new Response('', { status: 500 })) + await expect(compact(db, 'tid', bigMsgs(40), 'k', 'http://x', 'en')).rejects.toThrow() db.close() }) + + it('returns false when nothing older than the tail', async () => { + const db = await initDB() + const spy = vi.spyOn(globalThis, 'fetch').mockResolvedValue(new Response()) + const did = await compact(db, 'tid', [textMsg('1')], 'k', 'http://x', 'en') + expect(did).toBe(false) + expect(spy).not.toHaveBeenCalled() + db.close() + }) +}) + +describe('buildHistoryToStore', () => { + const summary = (over: Partial) => + ({ threadId: 't', summary: 'S', coversThroughMessageId: 'b', tokenBudget: 0, createdAt: 0, ...over }) as any + + it('collapses everything up to the cut id into the summary message', () => { + const full = [textMsg('a'), textMsg('b'), textMsg('c')] + const out = buildHistoryToStore(full, summary({ coversThroughMessageId: 'b' })) + expect(out[0].id).toBe('compaction-assistant') + expect(out.slice(1).map(m => m.id)).toEqual(['c']) + }) + + it('falls back to the stored index when the id is missing', () => { + const full = [textMsg('a'), textMsg('b'), textMsg('c')] + const out = buildHistoryToStore(full, summary({ coversThroughMessageId: 'GONE', coversThroughIndex: 0 })) + expect(out[0].id).toBe('compaction-assistant') + expect(out.slice(1).map(m => m.id)).toEqual(['b', 'c']) + }) + + it('returns full history when neither id nor index resolves', () => { + const full = [textMsg('a'), textMsg('b')] + const out = buildHistoryToStore(full, summary({ coversThroughMessageId: 'GONE', coversThroughIndex: undefined })) + expect(out).toEqual(full) + }) })