Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions .changeset/context-token-budget.md
Original file line number Diff line number Diff line change
@@ -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.
45 changes: 44 additions & 1 deletion packages/core/src/agent.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand All @@ -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 () => {
Expand Down
15 changes: 14 additions & 1 deletion packages/core/src/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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<void>): Promise<void> => {
assertNotCancelled();
const target = agent;
Expand Down
92 changes: 92 additions & 0 deletions packages/core/src/context-budget.test.ts
Original file line number Diff line number Diff line change
@@ -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.",
}),
);
});
});
81 changes: 81 additions & 0 deletions packages/core/src/context-budget.ts
Original file line number Diff line number Diff line change
@@ -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,
};
}
3 changes: 3 additions & 0 deletions packages/i18n/src/locales/en.json
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand All @@ -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",
Expand Down Expand Up @@ -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.",
Expand Down
3 changes: 3 additions & 0 deletions packages/i18n/src/locales/es.json
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand All @@ -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",
Expand Down Expand Up @@ -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.",
Expand Down
3 changes: 3 additions & 0 deletions packages/i18n/src/locales/pt-BR.json
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand All @@ -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",
Expand Down Expand Up @@ -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.",
Expand Down
3 changes: 3 additions & 0 deletions packages/i18n/src/locales/zh-CN.json
Original file line number Diff line number Diff line change
Expand Up @@ -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 无效或不受支持。",
Expand All @@ -1252,6 +1253,7 @@
"addV1": "添加 /v1",
"waitAndRetry": "等待后重试",
"adjustGenerationTimeout": "调整生成超时时间",
"adjustContextWindow": "检查上下文窗口",
"checkNetwork": "检查网络 / VPN",
"checkVpn": "检查 VPN / 防火墙",
"reportBug": "报告此 Bug",
Expand Down Expand Up @@ -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 未遵循所需协议。",
Expand Down
Loading
Loading