diff --git a/.changeset/context-token-budget.md b/.changeset/context-token-budget.md new file mode 100644 index 00000000..3b2e540f --- /dev/null +++ b/.changeset/context-token-budget.md @@ -0,0 +1,7 @@ +--- +"@open-codesign/shared": patch +"@open-codesign/core": patch +"@open-codesign/i18n": patch +--- + +Fit each generation run into the model's context window. Before the request, the system prompt, tool definitions and the prompt are measured in tokens, and a quarter of the window (at most the 32,000 tokens a run requests) stays free for the reply; the request asks for no more output than that, so models with a window under 128k tokens request less than 32,000. If those parts do not fit, the run stops with an error that names the window and points to the context window setting in Settings → Models, instead of sending a request the model would truncate. Chat history keeps its newest messages that fit what remains. diff --git a/packages/core/src/agent.test.ts b/packages/core/src/agent.test.ts index 8c45d431..2e5a24c2 100644 --- a/packages/core/src/agent.test.ts +++ b/packages/core/src/agent.test.ts @@ -645,7 +645,7 @@ describe('generateViaAgent()', () => { }); }); - it('gives the agent model the configured context window', async () => { + it('gives the agent model the configured context window and a matching output limit', async () => { scriptedAgent = { assistantText: RESPONSE_WITH_ARTIFACT }; await generateViaAgent({ prompt: 'design a dashboard', @@ -658,6 +658,49 @@ describe('generateViaAgent()', () => { }); expect(agentCalls[0]?.options.initialState?.model?.contextWindow).toBe(32_768); + // A quarter of the window, so the 32,000-token default does not overrun it. + expect(agentCalls[0]?.options.initialState?.model?.maxTokens).toBe(8_192); + }); + + it('fails before prompting when the run does not fit the context window', async () => { + scriptedAgent = { assistantText: RESPONSE_WITH_ARTIFACT }; + await expect( + generateViaAgent({ + prompt: 'design a dashboard', + history: [], + model: { provider: 'custom-lmstudio', modelId: 'qwen3.6-35b-a3b' }, + apiKey: 'sk-test', + baseUrl: 'http://127.0.0.1:1234/v1', + wire: 'openai-chat', + contextWindow: 2_048, + }), + ).rejects.toMatchObject({ code: 'GENERATION_CONTEXT_EXCEEDED' }); + expect(agentCalls).toHaveLength(0); + }); + + it('starts the agent with the newest history that fits the context window', async () => { + scriptedAgent = { assistantText: RESPONSE_WITH_ARTIFACT }; + const history = Array.from({ length: 6 }, (_, index) => ({ + role: index % 2 === 0 ? ('user' as const) : ('assistant' as const), + content: String(index).repeat(16_000), + })); + await generateViaAgent({ + prompt: 'design a dashboard', + history, + model: { provider: 'custom-lmstudio', modelId: 'qwen3.6-35b-a3b' }, + apiKey: 'sk-test', + baseUrl: 'http://127.0.0.1:1234/v1', + wire: 'openai-chat', + contextWindow: 32_768, + }); + + const messages = (agentCalls[0]?.options.initialState?.messages ?? []) as Array<{ + role: string; + }>; + expect(messages.length).toBeGreaterThan(0); + expect(messages.length).toBeLessThan(history.length); + expect(messages[0]?.role).toBe('user'); + expect(JSON.stringify(messages.at(-1))).toContain('5'.repeat(100)); }); it("keeps pi-agent-core's default stream when no request timeout is configured", async () => { diff --git a/packages/core/src/agent.ts b/packages/core/src/agent.ts index bbf2a47d..a6dab1e5 100644 --- a/packages/core/src/agent.ts +++ b/packages/core/src/agent.ts @@ -62,6 +62,7 @@ import { } from '@open-codesign/shared'; import type { TSchema } from '@sinclair/typebox'; import type { ActiveRunMessages } from './active-messages.js'; +import { planContextBudget } from './context-budget.js'; import { buildTransformContext } from './context-prune.js'; import { resolveContextWindow } from './context-window.js'; import { remapProviderError } from './errors.js'; @@ -1375,7 +1376,19 @@ async function generateViaAgentInternal( return retryAgent; }; - let agent = createRetryAgent(historyAsAgentMessages); + const { history: budgetedHistory, ...budget } = planContextBudget({ + contextWindow: piModel.contextWindow, + maxOutputTokens: piModel.maxTokens, + systemPrompt: augmentedSystemPrompt, + tools, + prompt: userContent, + promptImages: promptImages.length, + history: historyAsAgentMessages, + }); + log.info('[generate] step=context_budget', { ...ctx, ...budget }); + // Request no more output than the budget reserves, so input and output fit the window. + piModel.maxTokens = budget.outputReserve; + let agent = createRetryAgent(budgetedHistory); const admitRun = async (run: () => Promise): Promise => { assertNotCancelled(); const target = agent; diff --git a/packages/core/src/context-budget.test.ts b/packages/core/src/context-budget.test.ts new file mode 100644 index 00000000..6ea25349 --- /dev/null +++ b/packages/core/src/context-budget.test.ts @@ -0,0 +1,92 @@ +import type { AgentMessage } from '@mariozechner/pi-agent-core'; +import { describe, expect, it } from 'vitest'; +import { planContextBudget } from './context-budget.js'; + +const user = (chars: number): AgentMessage => ({ + role: 'user', + content: 'u'.repeat(chars), + timestamp: 0, +}); +const assistant = (chars: number): AgentMessage => + ({ + role: 'assistant', + content: [{ type: 'text', text: 'a'.repeat(chars) }], + timestamp: 0, + }) as unknown as AgentMessage; + +const base = { + maxOutputTokens: 32_000, + systemPrompt: 's'.repeat(4_000), + tools: [], + prompt: 'p'.repeat(400), + promptImages: 0, +}; + +describe('planContextBudget', () => { + it('keeps all history and reserves the requested output in a large window', () => { + const history = [user(400), assistant(400)]; + const budget = planContextBudget({ ...base, contextWindow: 200_000, history }); + + expect(budget.outputReserve).toBe(32_000); + expect(budget.systemTokens).toBe(1_000); + expect(budget.promptTokens).toBe(100); + expect(budget.history).toEqual(history); + expect(budget.droppedHistory).toBe(0); + }); + + it('reserves a quarter of a small window for the reply', () => { + expect(planContextBudget({ ...base, contextWindow: 8_192, history: [] }).outputReserve).toBe( + 2_048, + ); + }); + + it('drops the oldest history that does not fit and starts at a user message', () => { + // 8,192 - 2,048 reserve - 1,101 fixed (system, empty tool list, prompt) leaves 5,043 tokens. + const history = [user(8_000), assistant(8_000), user(8_000), assistant(8_000)]; + const budget = planContextBudget({ ...base, contextWindow: 8_192, history }); + + expect(budget.history).toEqual(history.slice(2)); + expect(budget.historyTokens).toBe(4_000); + expect(budget.droppedHistory).toBe(2); + }); + + it('drops a leading assistant message left by the cut', () => { + // 3,000 tokens of history fit: the last assistant and user messages plus + // the assistant before them, which cannot open the history. + const history = [user(4_000), assistant(4_000), user(4_000), assistant(4_000)]; + const budget = planContextBudget({ ...base, contextWindow: 5_600, history }); + + expect(budget.history).toEqual(history.slice(2)); + expect(budget.droppedHistory).toBe(2); + }); + + it('counts tool definitions by name, description and parameters', () => { + const tool = { + name: 'read', + description: 'd'.repeat(3_935), + parameters: { type: 'object' }, + label: 'l'.repeat(4_000), + }; + // The serialized definition is 4,000 characters; `label` is not sent to the model. + expect( + planContextBudget({ ...base, contextWindow: 200_000, tools: [tool], history: [] }).toolTokens, + ).toBe(1_000); + // 2,400 - 600 reserve leaves 1,800: room for the system prompt and prompt, not the tool. + expect(() => planContextBudget({ ...base, contextWindow: 2_400, history: [] })).not.toThrow(); + expect(() => + planContextBudget({ ...base, contextWindow: 2_400, tools: [tool], history: [] }), + ).toThrow(expect.objectContaining({ code: 'GENERATION_CONTEXT_EXCEEDED' })); + }); + + it('counts prompt images and fails before the request when the fixed parts do not fit', () => { + expect(() => + planContextBudget({ ...base, contextWindow: 4_096, promptImages: 2, history: [] }), + ).toThrow( + expect.objectContaining({ + code: 'GENERATION_CONTEXT_EXCEEDED', + message: + "This run needs about 3501 tokens before any chat history, plus 1024 for the reply, but the model's context window is 4096 tokens.", + }), + ); + }); +}); diff --git a/packages/core/src/context-budget.ts b/packages/core/src/context-budget.ts new file mode 100644 index 00000000..f018f7d4 --- /dev/null +++ b/packages/core/src/context-budget.ts @@ -0,0 +1,81 @@ +import type { AgentMessage } from '@mariozechner/pi-agent-core'; +import { estimateTokens } from '@mariozechner/pi-coding-agent'; +import { CodesignError, ERROR_CODES } from '@open-codesign/shared'; + +/** pi's estimate for one image (`estimateTokens` counts images in tool results as 1,200 tokens). */ +const IMAGE_TOKENS = 1_200; + +export interface ContextBudgetInput { + contextWindow: number; + /** The most output the run requests from the model. */ + maxOutputTokens: number; + systemPrompt: string; + tools: ReadonlyArray<{ name: string; description: string; parameters: unknown }>; + /** The run's user message, which also carries the brief, memory and attachments as text. */ + prompt: string; + promptImages: number; + history: AgentMessage[]; +} + +export interface ContextBudget { + contextWindow: number; + outputReserve: number; + systemTokens: number; + toolTokens: number; + promptTokens: number; + historyTokens: number; + history: AgentMessage[]; + droppedHistory: number; +} + +/** + * Fits a run into the model's context window, in tokens estimated the way pi + * does (characters / 4). A quarter of the window, at most the requested + * output, stays free for the reply. The system prompt, tool definitions and + * prompt must fit, or the run fails before the request; history keeps its + * newest messages that fit what remains, starting at a user message. + */ +export function planContextBudget(input: ContextBudgetInput): ContextBudget { + const outputReserve = Math.min(input.maxOutputTokens, Math.floor(input.contextWindow / 4)); + const systemTokens = Math.ceil(input.systemPrompt.length / 4); + const toolTokens = Math.ceil( + JSON.stringify( + input.tools.map(({ name, description, parameters }) => ({ name, description, parameters })), + ).length / 4, + ); + const promptTokens = + estimateTokens({ role: 'user', content: input.prompt, timestamp: 0 }) + + input.promptImages * IMAGE_TOKENS; + const fixedTokens = systemTokens + toolTokens + promptTokens; + const available = input.contextWindow - outputReserve - fixedTokens; + if (available < 0) { + throw new CodesignError( + `This run needs about ${fixedTokens} tokens before any chat history, plus ${outputReserve} for the reply, but the model's context window is ${input.contextWindow} tokens.`, + ERROR_CODES.GENERATION_CONTEXT_EXCEEDED, + ); + } + + let start = input.history.length; + let historyTokens = 0; + while (start > 0) { + const tokens = estimateTokens(input.history[start - 1] as AgentMessage); + if (historyTokens + tokens > available) break; + historyTokens += tokens; + start -= 1; + } + while (start < input.history.length && input.history[start]?.role !== 'user') { + historyTokens -= estimateTokens(input.history[start] as AgentMessage); + start += 1; + } + const history = input.history.slice(start); + return { + contextWindow: input.contextWindow, + outputReserve, + systemTokens, + toolTokens, + promptTokens, + historyTokens, + history, + droppedHistory: start, + }; +} diff --git a/packages/i18n/src/locales/en.json b/packages/i18n/src/locales/en.json index e8a7aa5e..75e5de4f 100644 --- a/packages/i18n/src/locales/en.json +++ b/packages/i18n/src/locales/en.json @@ -1240,6 +1240,7 @@ "reasoningPolicy": "The endpoint rejected reasoning metadata for this model. Lower the reasoning depth and try again.", "modelIdShape": "The selected model id includes a provider prefix this endpoint does not accept.", "generationTimeout": "The generation reached the configured timeout before the model finished.", + "contextWindowExceeded": "The system prompt, tools and this request need more tokens than the model's context window allows.", "relayStreamingBug": "The gateway may mishandle OpenAI Responses API SSE events (older sub2api / claude2api / anyrouter builds cut the stream short).", "transportInterrupted": "The provider connection ended before the turn completed. This is usually a gateway timeout or network interruption.", "referenceUrlInvalid": "The reference URL is invalid or not supported.", @@ -1256,6 +1257,7 @@ "addV1": "Add /v1", "waitAndRetry": "Wait and retry", "adjustGenerationTimeout": "Adjust generation timeout", + "adjustContextWindow": "Check the context window", "checkNetwork": "Check network / VPN", "checkVpn": "Check VPN / firewall", "reportBug": "Report this bug", @@ -1523,6 +1525,7 @@ "INPUT_EMPTY_HTML": "Existing HTML is required for this operation.", "INPUT_UNSUPPORTED_MODE": "This generation mode is not supported.", "GENERATION_TIMEOUT": "Generation timed out. Try a shorter prompt or increase the timeout in Settings.", + "GENERATION_CONTEXT_EXCEEDED": "This run does not fit the model's context window. Shorten the request or its attachments, set the provider's context window in Settings → Models, or use a model with a larger window.", "GENERATION_INCOMPLETE": "The agent stopped before completing verification. Continue the run to finish the workspace artifact.", "WORKSPACE_MISSING": "This design has no workspace bound. Reopen or recreate the design.", "ARTIFACT_PROTOCOL_INVALID": "The generated artifact did not follow the required protocol.", diff --git a/packages/i18n/src/locales/es.json b/packages/i18n/src/locales/es.json index 8d7453cd..f592e08a 100644 --- a/packages/i18n/src/locales/es.json +++ b/packages/i18n/src/locales/es.json @@ -1061,6 +1061,7 @@ "reasoningPolicy": "El endpoint rechazó los metadatos de razonamiento para este modelo. Reduce la profundidad de razonamiento e inténtalo de nuevo.", "modelIdShape": "El id del modelo seleccionado incluye un prefijo de proveedor que este endpoint no acepta.", "generationTimeout": "La generación alcanzó el tiempo límite configurado antes de que el modelo terminara.", + "contextWindowExceeded": "El prompt del sistema, las herramientas y esta solicitud necesitan más tokens de los que permite la ventana de contexto del modelo.", "relayStreamingBug": "La pasarela puede manejar mal los eventos SSE de la API OpenAI Responses (versiones antiguas de sub2api / claude2api / anyrouter cortan el flujo prematuramente).", "transportInterrupted": "La conexión con el proveedor terminó antes de completar el turno. Normalmente es un timeout de la pasarela o una interrupción de red.", "referenceUrlInvalid": "La URL de referencia no es válida o no es compatible.", @@ -1077,6 +1078,7 @@ "addV1": "Agregar /v1", "waitAndRetry": "Esperar y reintentar", "adjustGenerationTimeout": "Ajustar tiempo límite de generación", + "adjustContextWindow": "Revisar la ventana de contexto", "checkNetwork": "Revisar red / VPN", "checkVpn": "Revisar VPN / cortafuegos", "reportBug": "Reportar este error", @@ -1346,6 +1348,7 @@ "INPUT_EMPTY_HTML": "Se requiere el HTML existente para esta operación.", "INPUT_UNSUPPORTED_MODE": "Este modo de generación no es compatible.", "GENERATION_TIMEOUT": "La generación agotó el tiempo de espera. Prueba con una instrucción más corta o aumenta el tiempo de espera en Configuración.", + "GENERATION_CONTEXT_EXCEEDED": "Esta ejecución no cabe en la ventana de contexto del modelo. Acorta la solicitud o sus adjuntos, configura la ventana de contexto del proveedor en Ajustes → Modelos o usa un modelo con una ventana mayor.", "GENERATION_INCOMPLETE": "El agente se detuvo antes de completar la verificación. Continúa la ejecución para terminar el artefacto del espacio de trabajo.", "WORKSPACE_MISSING": "Este diseño no tiene un espacio de trabajo vinculado. Vuelve a abrirlo o recréalo.", "ARTIFACT_PROTOCOL_INVALID": "El artefacto generado no siguió el protocolo requerido.", diff --git a/packages/i18n/src/locales/pt-BR.json b/packages/i18n/src/locales/pt-BR.json index 86487a59..a4578d65 100644 --- a/packages/i18n/src/locales/pt-BR.json +++ b/packages/i18n/src/locales/pt-BR.json @@ -1053,6 +1053,7 @@ "reasoningPolicy": "O endpoint rejeitou metadados de reasoning para este modelo. Reduza a profundidade de reasoning e tente novamente.", "modelIdShape": "O id do modelo selecionado inclui um prefixo de provedor que este endpoint não aceita.", "generationTimeout": "A geração atingiu o tempo limite configurado antes de o modelo terminar.", + "contextWindowExceeded": "O prompt do sistema, as ferramentas e esta solicitação precisam de mais tokens do que a janela de contexto do modelo permite.", "relayStreamingBug": "O gateway pode tratar incorretamente eventos SSE da OpenAI Responses API e cortar o stream cedo.", "transportInterrupted": "A conexão com o provedor terminou antes da conclusão. Normalmente é timeout do gateway ou interrupção de rede.", "referenceUrlInvalid": "A URL de referência é inválida ou não é suportada.", @@ -1069,6 +1070,7 @@ "addV1": "Adicionar /v1", "waitAndRetry": "Aguardar e tentar novamente", "adjustGenerationTimeout": "Ajustar tempo limite da geração", + "adjustContextWindow": "Verificar a janela de contexto", "checkNetwork": "Verificar rede / VPN", "checkVpn": "Verificar VPN / firewall", "reportBug": "Reportar este bug", @@ -1336,6 +1338,7 @@ "INPUT_EMPTY_HTML": "Um HTML existente é obrigatório para esta operação.", "INPUT_UNSUPPORTED_MODE": "Este modo de geração não é suportado.", "GENERATION_TIMEOUT": "A geração excedeu o tempo limite. Tente um prompt mais curto ou aumente o tempo em Configurações.", + "GENERATION_CONTEXT_EXCEEDED": "Esta execução não cabe na janela de contexto do modelo. Encurte a solicitação ou os anexos, defina a janela de contexto do provedor em Configurações → Modelos ou use um modelo com uma janela maior.", "GENERATION_INCOMPLETE": "O agente parou antes de concluir a verificação. Continue a execução para finalizar o artefato do espaço de trabalho.", "WORKSPACE_MISSING": "Este design não tem um espaço de trabalho vinculado. Reabra ou recrie o design.", "ARTIFACT_PROTOCOL_INVALID": "O artefato gerado não seguiu o protocolo obrigatório.", diff --git a/packages/i18n/src/locales/zh-CN.json b/packages/i18n/src/locales/zh-CN.json index eb1dc727..1b9bbfd6 100644 --- a/packages/i18n/src/locales/zh-CN.json +++ b/packages/i18n/src/locales/zh-CN.json @@ -1236,6 +1236,7 @@ "reasoningPolicy": "端点拒绝了该模型的 reasoning 元数据。请降低 reasoning 深度后重试。", "modelIdShape": "当前模型 ID 带有这个端点不接受的 Provider 前缀。", "generationTimeout": "生成已达到当前配置的超时时间，但模型还未完成。", + "contextWindowExceeded": "系统提示、工具定义和本次请求需要的 token 超过了模型的上下文窗口。", "relayStreamingBug": "网关可能错误处理了 OpenAI Responses API 的 SSE 事件（老版本 sub2api / claude2api / anyrouter 会把流提前截断）。", "transportInterrupted": "Provider 连接在本轮完成前中断，通常是网关超时或网络中断。", "referenceUrlInvalid": "Reference URL 无效或不受支持。", @@ -1252,6 +1253,7 @@ "addV1": "添加 /v1", "waitAndRetry": "等待后重试", "adjustGenerationTimeout": "调整生成超时时间", + "adjustContextWindow": "检查上下文窗口", "checkNetwork": "检查网络 / VPN", "checkVpn": "检查 VPN / 防火墙", "reportBug": "报告此 Bug", @@ -1519,6 +1521,7 @@ "INPUT_EMPTY_HTML": "此操作需要已有的 HTML 内容。", "INPUT_UNSUPPORTED_MODE": "不支持该生成模式。", "GENERATION_TIMEOUT": "生成超时，请缩短 prompt 或在设置中增加超时时间。", + "GENERATION_CONTEXT_EXCEEDED": "这次生成超出了模型的上下文窗口。请缩短请求或附件，在设置 → 模型里填写服务商的上下文窗口，或换用窗口更大的模型。", "GENERATION_INCOMPLETE": "Agent 在完成校验前停止了。请继续运行，以完成工作区 artifact。", "WORKSPACE_MISSING": "此设计未绑定工作区，请重新打开或重新创建设计。", "ARTIFACT_PROTOCOL_INVALID": "生成的 artifact 未遵循所需协议。", diff --git a/packages/shared/src/diagnostics.test.ts b/packages/shared/src/diagnostics.test.ts index f6afc5fd..8ce035be 100644 --- a/packages/shared/src/diagnostics.test.ts +++ b/packages/shared/src/diagnostics.test.ts @@ -258,6 +258,20 @@ describe('diagnoseGenerateFailure', () => { expect(result[0]?.suggestedFix?.settingsTab).toBe('advanced'); }); + it('points context window overflows at the provider context window setting', () => { + const result = diagnoseGenerateFailure({ + ...ctx, + code: 'GENERATION_CONTEXT_EXCEEDED', + message: + "This run needs about 9000 tokens before any chat history, plus 2048 for the reply, but the model's context window is 8192 tokens.", + }); + + expect(result[0]?.category).toBe('context-window-exceeded'); + expect(result[0]?.cause).toBe('diagnostics.cause.contextWindowExceeded'); + expect(result[0]?.suggestedFix?.kind).toBe('openSettings'); + expect(result[0]?.suggestedFix?.settingsTab).toBe('models'); + }); + it('maps reference URL errors by CodesignError code before provider heuristics', () => { const result = diagnoseGenerateFailure({ ...ctx, diff --git a/packages/shared/src/diagnostics.ts b/packages/shared/src/diagnostics.ts index bf3008f0..5c098391 100644 --- a/packages/shared/src/diagnostics.ts +++ b/packages/shared/src/diagnostics.ts @@ -32,6 +32,7 @@ export type DiagnosticCategory = | 'relay-stream-cutoff' | 'gateway-waf-blocked' | 'generation-timeout' + | 'context-window-exceeded' | 'model-discovery-degraded' | 'transport-interrupted' | 'upstream-server-error' @@ -352,6 +353,21 @@ export function diagnoseGenerateFailure(ctx: GenerateFailureContext): Diagnostic ]; } + if (code === 'GENERATION_CONTEXT_EXCEEDED') { + return [ + h({ + cause: 'diagnostics.cause.contextWindowExceeded', + category: 'context-window-exceeded', + severity: 'error', + suggestedFix: { + kind: 'openSettings', + label: 'diagnostics.fix.adjustContextWindow', + settingsTab: 'models', + }, + }), + ]; + } + if (code === 'REFERENCE_URL_UNSUPPORTED') { return [ h({ diff --git a/packages/shared/src/error-codes.ts b/packages/shared/src/error-codes.ts index de8b172c..a0bc6598 100644 --- a/packages/shared/src/error-codes.ts +++ b/packages/shared/src/error-codes.ts @@ -39,6 +39,7 @@ export const ERROR_CODES = { INPUT_EMPTY_HTML: 'INPUT_EMPTY_HTML', INPUT_UNSUPPORTED_MODE: 'INPUT_UNSUPPORTED_MODE', GENERATION_TIMEOUT: 'GENERATION_TIMEOUT', + GENERATION_CONTEXT_EXCEEDED: 'GENERATION_CONTEXT_EXCEEDED', GENERATION_INCOMPLETE: 'GENERATION_INCOMPLETE', WORKSPACE_MISSING: 'WORKSPACE_MISSING', ARTIFACT_PROTOCOL_INVALID: 'ARTIFACT_PROTOCOL_INVALID', @@ -229,6 +230,12 @@ export const ERROR_CODE_DESCRIPTIONS: Record