diff --git a/apps/desktop/src/main/main.ts b/apps/desktop/src/main/main.ts index 0fb686214f..ac1bc8449f 100644 --- a/apps/desktop/src/main/main.ts +++ b/apps/desktop/src/main/main.ts @@ -142,6 +142,7 @@ import { testConnection, } from '@maka/runtime'; import type { BotIncomingMessage, ToolArtifactRecorderInput } from '@maka/runtime'; +import type { ContextBudgetPolicy } from '@maka/runtime'; import { testProxyConnection } from '@maka/runtime/network/proxy-test'; import { fetchWeChatQrcode, pollWeChatQrcodeStatus } from './wechat-scan-login.js'; import { @@ -735,6 +736,7 @@ backends.register('ai-sdk', async (ctx) => { modelFactory: (input) => getAIModel({ ...input, fetch: modelFetch }), tools: builtinTools, providerOptions: buildProviderOptions(connection, model), + contextBudget: buildContextBudgetPolicy(connection), systemPrompt: ({ cwd }) => buildSystemPrompt(ctx.header, cwd), turnTailPrompt: ({ cwd }) => buildTurnTailPrompt(cwd), recordLlmCall: (event) => recordLlmCall({ repo: telemetryRepo, lookupPricing }, event), @@ -756,6 +758,38 @@ backends.register('ai-sdk', async (ctx) => { }); }); +function buildContextBudgetPolicy(connection: LlmConnection): ContextBudgetPolicy | undefined { + if (process.env.MAKA_CONTEXT_BUDGET === 'off') return undefined; + const maxHistoryEstimatedTokens = + parseOptionalPositiveInt(process.env.MAKA_CONTEXT_HISTORY_BUDGET_TOKENS) ?? + defaultHistoryBudgetTokens(connection); + if (maxHistoryEstimatedTokens === undefined) return undefined; + const maxHistoryTurns = parseOptionalPositiveInt(process.env.MAKA_CONTEXT_HISTORY_BUDGET_TURNS); + const minRecentTurns = parsePositiveInt(process.env.MAKA_CONTEXT_MIN_RECENT_TURNS, 2); + return { + name: 'desktop-default-history-budget', + maxHistoryEstimatedTokens, + ...(maxHistoryTurns !== undefined ? { maxHistoryTurns } : {}), + minRecentTurns, + }; +} + +function defaultHistoryBudgetTokens(connection: LlmConnection): number | undefined { + if (connection.providerType === 'deepseek') return undefined; + return 32_000; +} + +function parsePositiveInt(value: string | undefined, fallback: number): number { + const parsed = parseOptionalPositiveInt(value); + return parsed ?? fallback; +} + +function parseOptionalPositiveInt(value: string | undefined): number | undefined { + if (!value) return undefined; + const parsed = Number.parseInt(value, 10); + return Number.isFinite(parsed) && parsed > 0 ? parsed : undefined; +} + function buildSubscriptionModelFetch( connection: LlmConnection, sessionId: string, diff --git a/package.json b/package.json index 2339283dc6..d849f97a9f 100644 --- a/package.json +++ b/package.json @@ -19,7 +19,8 @@ "check:stale": "node scripts/check-stale-dist.mjs", "prepare:officecli": "node scripts/prepare-officecli.mjs", "check:officecli-bundle": "node scripts/check-officecli-bundle.mjs", - "check:release": "npm run check:stale && npm run check:officecli-bundle" + "check:release": "npm run check:stale && npm run check:officecli-bundle", + "cost:deepseek-baseline": "node scripts/deepseek-live-cost-baseline.mjs" }, "devDependencies": { "@types/node": "^25.0.0", diff --git a/packages/core/src/events.ts b/packages/core/src/events.ts index 0ffe76152a..f85b04d61f 100644 --- a/packages/core/src/events.ts +++ b/packages/core/src/events.ts @@ -10,7 +10,11 @@ */ import type { PermissionRequest, PermissionResponse, ToolCategory } from './permission.js'; -import type { PrefixChangeReason } from './usage-stats/types.js'; +import type { + ContextBudgetDiagnostic, + PrefixChangeReason, + PromptSegmentEstimate, +} from './usage-stats/types.js'; export const TOOL_OUTPUT_STREAMS = ['stdout', 'stderr'] as const; export const TOOL_OUTPUT_DELTA_MAX_CHARS = 8192; @@ -335,6 +339,8 @@ export interface TokenUsageEvent extends BaseEvent { contextRemaining?: number; prefixHash?: string; prefixChangeReason?: PrefixChangeReason; + promptSegments?: PromptSegmentEstimate[]; + contextBudget?: ContextBudgetDiagnostic; } export interface ErrorEvent extends BaseEvent { diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index b234f6baf4..18457ce34d 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -693,7 +693,10 @@ export { // usage-stats/types.ts export type { LlmCallRecord, + ContextBudgetDiagnostic, PricingConfig, + PromptSegmentEstimate, + PromptSegmentKind, TimeRange, ToolInvocationRecord, UsageBucket, diff --git a/packages/core/src/runtime-event.ts b/packages/core/src/runtime-event.ts index e110803944..0f5a9611c3 100644 --- a/packages/core/src/runtime-event.ts +++ b/packages/core/src/runtime-event.ts @@ -16,7 +16,11 @@ import type { AttachmentRef } from './events.js'; import type { PermissionRequest, PermissionResponse } from './permission.js'; -import type { PrefixChangeReason } from './usage-stats/types.js'; +import type { + ContextBudgetDiagnostic, + PrefixChangeReason, + PromptSegmentEstimate, +} from './usage-stats/types.js'; // ============================================================================ // Role / Author / Status @@ -180,6 +184,8 @@ export interface RuntimeEventTokenUsage { contextRemaining?: number; prefixHash?: string; prefixChangeReason?: PrefixChangeReason; + promptSegments?: PromptSegmentEstimate[]; + contextBudget?: ContextBudgetDiagnostic; } /** diff --git a/packages/core/src/session.ts b/packages/core/src/session.ts index 9b2ed36e70..5d81ddf03d 100644 --- a/packages/core/src/session.ts +++ b/packages/core/src/session.ts @@ -11,7 +11,11 @@ import type { AttachmentRef, ToolResultContent } from './events.js'; import type { PermissionMode } from './permission.js'; -import type { PrefixChangeReason } from './usage-stats/types.js'; +import type { + ContextBudgetDiagnostic, + PrefixChangeReason, + PromptSegmentEstimate, +} from './usage-stats/types.js'; export const SESSION_STATUSES = [ 'active', @@ -243,6 +247,8 @@ export interface TokenUsageMessage { costUsd?: number; prefixHash?: string; prefixChangeReason?: PrefixChangeReason; + promptSegments?: PromptSegmentEstimate[]; + contextBudget?: ContextBudgetDiagnostic; } export interface TurnStateMessage { diff --git a/packages/core/src/usage-stats/types.ts b/packages/core/src/usage-stats/types.ts index 84faa2a4f7..bac44e0731 100644 --- a/packages/core/src/usage-stats/types.ts +++ b/packages/core/src/usage-stats/types.ts @@ -72,6 +72,8 @@ export interface UsageLogRow { turnId?: string; prefixHash?: string; prefixChangeReason?: PrefixChangeReason; + promptSegments?: PromptSegmentEstimate[]; + contextBudget?: ContextBudgetDiagnostic; } export interface PricingConfig { @@ -117,6 +119,8 @@ export interface LlmCallRecord { startedAt: number; prefixHash?: string; prefixChangeReason?: PrefixChangeReason; + promptSegments?: PromptSegmentEstimate[]; + contextBudget?: ContextBudgetDiagnostic; } export type PrefixChangeReason = @@ -129,6 +133,35 @@ export type PrefixChangeReason = | 'stable' | 'unknown'; +export type PromptSegmentKind = + | 'system_prompt' + | 'tool_schema' + | 'prior_history' + | 'current_user' + | 'turn_tail'; + +export interface PromptSegmentEstimate { + kind: PromptSegmentKind; + chars: number; + estimatedTokens: number; + messageCount?: number; + eventCount?: number; + toolCount?: number; +} + +export interface ContextBudgetDiagnostic { + enabled: boolean; + policyName?: string; + maxHistoryEstimatedTokens?: number; + maxHistoryTurns?: number; + estimatedTokensBefore: number; + estimatedTokensAfter: number; + keptTurns: number; + droppedTurns: number; + keptEvents: number; + droppedEvents: number; +} + export interface ToolInvocationRecord { sessionId?: string; turnId?: string; diff --git a/packages/runtime/package.json b/packages/runtime/package.json index 5b53ea94a0..cbaee7f054 100644 --- a/packages/runtime/package.json +++ b/packages/runtime/package.json @@ -15,6 +15,7 @@ "./tool-output-delta": "./dist/tool-output-delta.js", "./stream-watchdog": "./dist/stream-watchdog.js", "./model-factory": "./dist/model-factory.js", + "./context-budget": "./dist/context-budget.js", "./test-connection": "./dist/test-connection.js", "./model-fetcher": "./dist/model-fetcher.js", "./materializer": "./dist/materializer.js", diff --git a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts index cb03788f18..e79e993066 100644 --- a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts +++ b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts @@ -24,6 +24,7 @@ import { canonicalizeToolSet, computeRequestShapeDiagnostic, } from '../request-shape.js'; +import { applyRuntimeEventContextBudget } from '../context-budget.js'; describe('AiSdkBackend model history', () => { test('prefers RuntimeEvent prior messages and appends current user once', async () => { @@ -800,6 +801,86 @@ describe('AiSdkBackend request-shape diagnostics', () => { }); }); +describe('AiSdkBackend context budget and prompt attribution', () => { + test('context budget keeps whole recent turns and drops older turns', () => { + const events = [ + runtimeTextEvent({ id: 'old-u', turnId: 'old', role: 'user', author: 'user', text: 'old user text' }), + runtimeTextEvent({ id: 'old-a', turnId: 'old', role: 'model', author: 'agent', text: 'old assistant text' }), + runtimeTextEvent({ id: 'new-u', turnId: 'new', role: 'user', author: 'user', text: 'new user text' }), + runtimeTextEvent({ id: 'new-a', turnId: 'new', role: 'model', author: 'agent', text: 'new assistant text' }), + ]; + + const budgeted = applyRuntimeEventContextBudget(events, { + name: 'test-budget', + maxHistoryEstimatedTokens: 1, + minRecentTurns: 1, + charsPerToken: 1, + }); + + assert.ok(budgeted); + assert.deepEqual([...new Set(budgeted.events.map((event) => event.turnId))], ['new']); + assert.equal(budgeted.diagnostic.droppedTurns, 1); + assert.equal(budgeted.diagnostic.keptTurns, 1); + assert.equal(budgeted.diagnostic.droppedEvents, 2); + }); + + test('usage events include prompt segments and context budget diagnostics', async () => { + const model = completionModel(); + const events: SessionEvent[] = []; + const backend = new AiSdkBackend({ + sessionId: 'session-1', + header: header(), + appendMessage: async () => {}, + connection: connection(), + apiKey: 'sk-test', + modelId: 'mock-model-id', + permissionEngine: new PermissionEngine({ newId: () => 'permission-id', now: () => 1 }), + modelFactory: () => model, + tools: [testTool('Read', z.object({ path: z.string() }))], + newId: idGenerator(), + now: monotonicClock(), + systemPrompt: 'durable system', + turnTailPrompt: 'volatile tail', + contextBudget: { + name: 'test-budget', + maxHistoryEstimatedTokens: 1, + minRecentTurns: 1, + charsPerToken: 1, + }, + }); + + for await (const event of backend.send({ + turnId: 'turn-current', + text: 'current user', + context: [], + runtimeContext: [ + runtimeTextEvent({ id: 'old-u', turnId: 'old', role: 'user', author: 'user', text: 'old user text' }), + runtimeTextEvent({ id: 'old-a', turnId: 'old', role: 'model', author: 'agent', text: 'old assistant text' }), + runtimeTextEvent({ id: 'new-u', turnId: 'new', role: 'user', author: 'user', text: 'new user text' }), + runtimeTextEvent({ id: 'new-a', turnId: 'new', role: 'model', author: 'agent', text: 'new assistant text' }), + ], + })) { + events.push(event); + } + + assert.deepEqual(compactPrompt(model), [ + { role: 'system', content: 'durable system' }, + { role: 'user', content: [{ type: 'text', text: 'new user text' }] }, + { role: 'assistant', content: [{ type: 'text', text: 'new assistant text' }] }, + { role: 'user', content: [{ type: 'text', text: 'current user\n\nvolatile tail' }] }, + ]); + const usage = events.find((event): event is Extract => + event.type === 'token_usage' + ); + assert.ok(usage); + assert.equal(usage.contextBudget?.policyName, 'test-budget'); + assert.equal(usage.contextBudget?.droppedTurns, 1); + assert.equal(usage.promptSegments?.some((segment) => segment.kind === 'prior_history'), true); + assert.equal(usage.promptSegments?.some((segment) => segment.kind === 'tool_schema'), true); + assert.equal(usage.promptSegments?.some((segment) => segment.kind === 'turn_tail'), true); + }); +}); + describe('AiSdkBackend RunTrace', () => { test('records turn, model, usage, and completion trace events without changing SessionEvents', async () => { const trace: RunTraceEvent[] = []; diff --git a/packages/runtime/src/ai-sdk-backend.ts b/packages/runtime/src/ai-sdk-backend.ts index c9029a56af..bc752cf660 100644 --- a/packages/runtime/src/ai-sdk-backend.ts +++ b/packages/runtime/src/ai-sdk-backend.ts @@ -39,7 +39,6 @@ import type { ErrorEvent, TextCompleteEvent, TokenUsageEvent, - AttachmentRef, } from '@maka/core/events'; import type { StoredMessage, @@ -57,6 +56,10 @@ import type { } from '@maka/core/backend-types'; import type { LlmConnection } from '@maka/core/llm-connections'; import type { LlmCallRecord, ToolInvocationRecord } from '@maka/core/usage-stats/types'; +import type { + ContextBudgetDiagnostic, + PromptSegmentEstimate, +} from '@maka/core/usage-stats/types'; import type { JSONValue, ModelMessage } from 'ai'; import { z } from 'zod'; @@ -92,8 +95,14 @@ import { import { canonicalizeToolSet, computeRequestShapeDiagnostic, + toolSchemaCharsForDiagnostics, type RequestShapeDiagnostic, } from './request-shape.js'; +import { + applyRuntimeEventContextBudget, + buildPromptSegmentEstimates, + type ContextBudgetPolicy, +} from './context-budget.js'; export { DEFAULT_PERMISSION_TIMEOUT_MS, @@ -177,6 +186,8 @@ export interface AiSdkBackendInput { turnTailPrompt?: string | ((context: SystemPromptContext) => string | undefined | Promise); /** Provider-native options passed through to ai-sdk. */ providerOptions?: Record; + /** Optional prior-history budget. Keeps whole turns to preserve tool-call/result pairs. */ + contextBudget?: ContextBudgetPolicy; /** Optional fire-and-forget telemetry hooks. Tool implementations remain unaware. */ recordLlmCall?: LlmTelemetryRecorder; recordToolInvocation?: ToolTelemetryRecorder; @@ -278,6 +289,8 @@ export class AiSdkBackend implements AgentBackend { let streamErrorClass: string | undefined; let rawFinishReason: string | undefined; let requestShapeForTelemetry: RequestShapeDiagnostic | undefined; + let promptSegmentsForTelemetry: PromptSegmentEstimate[] = []; + let contextBudgetForTelemetry: ContextBudgetDiagnostic | undefined; const trace = new RunTrace({ sessionId: this.sessionId, turnId, @@ -348,13 +361,25 @@ export class AiSdkBackend implements AgentBackend { const activeTools = canonicalTools.activeTools; const systemPrompt = await this.resolveSystemPrompt(); const turnTailPrompt = await this.resolveTurnTailPrompt(); + const currentUserContent = formatTextWithAttachmentRefs(input.text, input.attachments); const messages = [ ...priorReplay.messages, { role: 'user' as const, - content: this.buildUserContent(input.text, input.attachments, turnTailPrompt), + content: this.appendTurnTailPrompt(currentUserContent, turnTailPrompt), }, ]; + const promptSegments = buildPromptSegmentEstimates({ + systemPrompt, + toolSchemaChars: toolSchemaCharsForDiagnostics(canonicalTools.providerTools, activeTools), + toolCount: canonicalTools.providerTools.length, + priorMessages: priorReplay.messages, + priorRuntimeEventCount: priorReplay.runtimeEventCount, + currentUserContent, + turnTailPrompt, + }); + promptSegmentsForTelemetry = promptSegments; + contextBudgetForTelemetry = priorReplay.contextBudget; const requestShape = computeRequestShapeDiagnostic({ connection: this.input.connection, modelId: this.input.modelId, @@ -369,6 +394,8 @@ export class AiSdkBackend implements AgentBackend { trace.modelStreamStarted(activeTools, { prefixHash: requestShape.prefixHash, prefixChangeReason: requestShape.prefixChangeReason, + promptSegments, + ...(priorReplay.contextBudget ? { contextBudget: priorReplay.contextBudget } : {}), }); const result = await this.modelAdapter.startStream({ @@ -478,6 +505,8 @@ export class AiSdkBackend implements AgentBackend { ...(tokenUsage.cacheWriteInputTokens > 0 ? { cacheCreation: tokenUsage.cacheWriteInputTokens } : {}), prefixHash: requestShape.prefixHash, prefixChangeReason: requestShape.prefixChangeReason, + promptSegments, + ...(priorReplay.contextBudget ? { contextBudget: priorReplay.contextBudget } : {}), }; await this.input.appendMessage(tu).catch(() => {}); queue.push({ @@ -497,6 +526,8 @@ export class AiSdkBackend implements AgentBackend { ...(tokenUsage.cacheWriteInputTokens > 0 ? { cacheCreation: tokenUsage.cacheWriteInputTokens } : {}), prefixHash: requestShape.prefixHash, prefixChangeReason: requestShape.prefixChangeReason, + promptSegments, + ...(priorReplay.contextBudget ? { contextBudget: priorReplay.contextBudget } : {}), } satisfies TokenUsageEvent); } } catch { @@ -571,6 +602,8 @@ export class AiSdkBackend implements AgentBackend { prefixHash: requestShapeForTelemetry.prefixHash, prefixChangeReason: requestShapeForTelemetry.prefixChangeReason, } : {}), + ...(promptSegmentsForTelemetry.length > 0 ? { promptSegments: promptSegmentsForTelemetry } : {}), + ...(contextBudgetForTelemetry !== undefined ? { contextBudget: contextBudgetForTelemetry } : {}), }); queue.close(); } @@ -646,6 +679,8 @@ export class AiSdkBackend implements AgentBackend { messages: ModelMessage[]; gate: RuntimeEventReplayFallbackGate | 'stored_message_projection'; diagnostics: RuntimeEventModelReplayPlan['diagnostics']; + runtimeEventCount?: number; + contextBudget?: ContextBudgetDiagnostic; } { const projectedMessages = this.materializePriorMessages( input.context.filter((message) => message.turnId !== input.turnId), @@ -653,12 +688,24 @@ export class AiSdkBackend implements AgentBackend { if (!input.runtimeContext) { return { messages: projectedMessages, gate: 'stored_message_projection', diagnostics: [] }; } + const budgeted = applyRuntimeEventContextBudget( + input.runtimeContext.filter((event) => event.turnId !== input.turnId), + this.input.contextBudget, + ); + const runtimeContext = budgeted?.events + ?? input.runtimeContext.filter((event) => event.turnId !== input.turnId); const plan = buildRuntimeEventModelReplayPlan( - input.runtimeContext.filter((event) => event.turnId !== input.turnId), + runtimeContext, ); if (plan.items.length === 0) { - return { messages: projectedMessages, gate: 'stored_message_projection', diagnostics: plan.diagnostics }; + return { + messages: projectedMessages, + gate: 'stored_message_projection', + diagnostics: plan.diagnostics, + runtimeEventCount: runtimeContext.length, + ...(budgeted ? { contextBudget: budgeted.diagnostic } : {}), + }; } if (hasBlockingReplayDiagnostics(plan)) { @@ -666,6 +713,8 @@ export class AiSdkBackend implements AgentBackend { messages: projectedMessages, gate: 'runtime_replay_unsupported_semantics', diagnostics: plan.diagnostics, + runtimeEventCount: runtimeContext.length, + ...(budgeted ? { contextBudget: budgeted.diagnostic } : {}), }; } @@ -674,6 +723,8 @@ export class AiSdkBackend implements AgentBackend { messages: plan.textMessages, gate: 'runtime_replay_text_only', diagnostics: plan.diagnostics, + runtimeEventCount: runtimeContext.length, + ...(budgeted ? { contextBudget: budgeted.diagnostic } : {}), }; } @@ -682,6 +733,8 @@ export class AiSdkBackend implements AgentBackend { messages: projectedMessages, gate: 'runtime_replay_unsupported_semantics', diagnostics: plan.diagnostics, + runtimeEventCount: runtimeContext.length, + ...(budgeted ? { contextBudget: budgeted.diagnostic } : {}), }; } @@ -689,6 +742,8 @@ export class AiSdkBackend implements AgentBackend { messages: this.materializeRuntimeReplayPlan(plan), gate: 'runtime_replay_provider_native', diagnostics: plan.diagnostics, + runtimeEventCount: runtimeContext.length, + ...(budgeted ? { contextBudget: budgeted.diagnostic } : {}), }; } @@ -758,9 +813,8 @@ export class AiSdkBackend implements AgentBackend { return out; } - /** Build the user content payload for the current turn (text + attachment refs). */ - private buildUserContent(text: string, attachments?: AttachmentRef[], turnTailPrompt?: string): string { - const content = formatTextWithAttachmentRefs(text, attachments); + /** Append provider-visible volatile turn facts after the durable user content. */ + private appendTurnTailPrompt(content: string, turnTailPrompt?: string): string { if (!turnTailPrompt) return content; return `${content}\n\n${turnTailPrompt}`; } diff --git a/packages/runtime/src/ai-sdk-flow.ts b/packages/runtime/src/ai-sdk-flow.ts index 0aafcf1fa9..75d034ee8d 100644 --- a/packages/runtime/src/ai-sdk-flow.ts +++ b/packages/runtime/src/ai-sdk-flow.ts @@ -319,6 +319,8 @@ export function mapSessionEventToRuntimeEvent( ...(event.prefixChangeReason !== undefined ? { prefixChangeReason: event.prefixChangeReason } : {}), + ...(event.promptSegments !== undefined ? { promptSegments: event.promptSegments } : {}), + ...(event.contextBudget !== undefined ? { contextBudget: event.contextBudget } : {}), }, }, }; diff --git a/packages/runtime/src/context-budget.ts b/packages/runtime/src/context-budget.ts new file mode 100644 index 0000000000..6540f710ef --- /dev/null +++ b/packages/runtime/src/context-budget.ts @@ -0,0 +1,201 @@ +import type { ModelMessage } from 'ai'; +import type { RuntimeEvent } from '@maka/core/runtime-event'; +import type { + ContextBudgetDiagnostic, + PromptSegmentEstimate, +} from '@maka/core/usage-stats/types'; + +export interface ContextBudgetPolicy { + name?: string; + /** + * Approximate max model-visible prior-history tokens. This is an estimate + * used for shaping, not provider billing. + */ + maxHistoryEstimatedTokens?: number; + /** Hard cap on prior turns retained for model replay. */ + maxHistoryTurns?: number; + /** Keep at least this many recent turns even if the token estimate exceeds the cap. */ + minRecentTurns?: number; + /** Estimate conversion. Defaults to 4 chars/token, intentionally conservative for mixed text. */ + charsPerToken?: number; +} + +export interface BudgetedRuntimeContext { + events: RuntimeEvent[]; + diagnostic: ContextBudgetDiagnostic; +} + +export interface PromptSegmentInput { + systemPrompt?: string; + toolSchemaChars: number; + toolCount: number; + priorMessages: readonly ModelMessage[]; + priorRuntimeEventCount?: number; + currentUserContent: string; + turnTailPrompt?: string; + charsPerToken?: number; +} + +export function applyRuntimeEventContextBudget( + events: readonly RuntimeEvent[], + policy: ContextBudgetPolicy | undefined, +): BudgetedRuntimeContext | undefined { + const enabled = Boolean(policy?.maxHistoryEstimatedTokens || policy?.maxHistoryTurns); + if (!enabled) return undefined; + const charsPerToken = policy?.charsPerToken ?? 4; + const maxTokens = finitePositive(policy?.maxHistoryEstimatedTokens); + const maxTurns = finitePositive(policy?.maxHistoryTurns); + const minRecentTurns = Math.max(0, Math.floor(policy?.minRecentTurns ?? 1)); + const turnGroups = groupEventsByTurn(events, charsPerToken); + const estimatedTokensBefore = estimateRuntimeEventsTokens(events, charsPerToken); + + const keptTurnIds = new Set(); + let keptTokens = 0; + for (let index = turnGroups.length - 1; index >= 0; index -= 1) { + const group = turnGroups[index]!; + const nextTurnCount = keptTurnIds.size + 1; + const mustKeep = nextTurnCount <= minRecentTurns; + const wouldExceedTurns = maxTurns !== undefined && nextTurnCount > maxTurns; + const wouldExceedTokens = + maxTokens !== undefined && + keptTokens > 0 && + keptTokens + group.estimatedTokens > maxTokens; + if (!mustKeep && (wouldExceedTurns || wouldExceedTokens)) break; + keptTurnIds.add(group.turnId); + keptTokens += group.estimatedTokens; + } + + const keptEvents = events.filter((event) => keptTurnIds.has(event.turnId)); + const diagnostic: ContextBudgetDiagnostic = { + enabled: true, + ...(policy?.name ? { policyName: policy.name } : {}), + ...(maxTokens !== undefined ? { maxHistoryEstimatedTokens: maxTokens } : {}), + ...(maxTurns !== undefined ? { maxHistoryTurns: maxTurns } : {}), + estimatedTokensBefore, + estimatedTokensAfter: estimateRuntimeEventsTokens(keptEvents, charsPerToken), + keptTurns: keptTurnIds.size, + droppedTurns: Math.max(0, turnGroups.length - keptTurnIds.size), + keptEvents: keptEvents.length, + droppedEvents: Math.max(0, events.length - keptEvents.length), + }; + return { events: keptEvents, diagnostic }; +} + +export function buildPromptSegmentEstimates(input: PromptSegmentInput): PromptSegmentEstimate[] { + const charsPerToken = input.charsPerToken ?? 4; + return [ + segment('system_prompt', input.systemPrompt?.length ?? 0, charsPerToken), + { + ...segment('tool_schema', input.toolSchemaChars, charsPerToken), + toolCount: input.toolCount, + }, + { + ...segment('prior_history', estimateModelMessagesChars(input.priorMessages), charsPerToken), + messageCount: input.priorMessages.length, + ...(input.priorRuntimeEventCount !== undefined ? { eventCount: input.priorRuntimeEventCount } : {}), + }, + segment('current_user', input.currentUserContent.length, charsPerToken), + segment('turn_tail', input.turnTailPrompt?.length ?? 0, charsPerToken), + ]; +} + +export function estimateModelMessagesChars(messages: readonly ModelMessage[]): number { + return messages.reduce((total, message) => total + estimateModelMessageChars(message), 0); +} + +export function estimateRuntimeEventsTokens( + events: readonly RuntimeEvent[], + charsPerToken = 4, +): number { + const chars = events.reduce((total, event) => total + estimateRuntimeEventChars(event), 0); + return estimateTokens(chars, charsPerToken); +} + +export function estimateTokens(chars: number, charsPerToken = 4): number { + if (chars <= 0) return 0; + return Math.ceil(chars / Math.max(1, charsPerToken)); +} + +function groupEventsByTurn(events: readonly RuntimeEvent[], charsPerToken: number): Array<{ + turnId: string; + estimatedTokens: number; +}> { + const order: string[] = []; + const byTurn = new Map(); + for (const event of events) { + const key = event.turnId || ''; + const group = byTurn.get(key); + if (group) group.push(event); + else { + order.push(key); + byTurn.set(key, [event]); + } + } + return order.map((turnId) => ({ + turnId, + estimatedTokens: estimateRuntimeEventsTokens(byTurn.get(turnId) ?? [], charsPerToken), + })); +} + +function estimateRuntimeEventChars(event: RuntimeEvent): number { + let total = 0; + const content = event.content; + if (content?.kind === 'text' || content?.kind === 'thinking') total += content.text.length; + else if (content?.kind === 'function_call') total += content.name.length + stableJsonLength(content.args); + else if (content?.kind === 'function_response') total += content.name.length + stableJsonLength(content.result); + else if (content?.kind === 'error') total += content.message.length; + return total; +} + +function estimateModelMessageChars(message: ModelMessage): number { + const raw = message as unknown as { content?: unknown }; + return estimateContentChars(raw.content); +} + +function estimateContentChars(content: unknown): number { + if (typeof content === 'string') return content.length; + if (Array.isArray(content)) { + return content.reduce((total, part) => total + estimatePartChars(part), 0); + } + return stableJsonLength(content); +} + +function estimatePartChars(part: unknown): number { + if (!part || typeof part !== 'object') return stableJsonLength(part); + const value = part as Record; + let total = 0; + for (const key of ['text', 'toolName', 'toolCallId'] as const) { + if (typeof value[key] === 'string') total += value[key].length; + } + for (const key of ['input', 'output'] as const) { + if (value[key] !== undefined) total += stableJsonLength(value[key]); + } + return total; +} + +function segment( + kind: PromptSegmentEstimate['kind'], + chars: number, + charsPerToken: number, +): PromptSegmentEstimate { + return { + kind, + chars, + estimatedTokens: estimateTokens(chars, charsPerToken), + }; +} + +function stableJsonLength(value: unknown): number { + if (value === undefined) return 0; + try { + return JSON.stringify(value)?.length ?? 0; + } catch { + return String(value).length; + } +} + +function finitePositive(value: number | undefined): number | undefined { + return typeof value === 'number' && Number.isFinite(value) && value > 0 + ? Math.floor(value) + : undefined; +} diff --git a/packages/runtime/src/index.ts b/packages/runtime/src/index.ts index 57367ce628..53923f1189 100644 --- a/packages/runtime/src/index.ts +++ b/packages/runtime/src/index.ts @@ -69,6 +69,18 @@ export type { StreamWatchdogInput, StreamWatchdogPhase, StreamWatchdogTimeout } export { getAIModel, buildProviderOptions } from './model-factory.js'; export type { ModelFactoryInput as GetAIModelInput } from './model-factory.js'; +export { + applyRuntimeEventContextBudget, + buildPromptSegmentEstimates, + estimateModelMessagesChars, + estimateRuntimeEventsTokens, + estimateTokens, +} from './context-budget.js'; +export type { + BudgetedRuntimeContext, + ContextBudgetPolicy, + PromptSegmentInput, +} from './context-budget.js'; export { testConnection } from './test-connection.js'; export { fetchProviderModels } from './model-fetcher.js'; diff --git a/packages/runtime/src/request-shape.ts b/packages/runtime/src/request-shape.ts index 84319b98c9..b3f50db234 100644 --- a/packages/runtime/src/request-shape.ts +++ b/packages/runtime/src/request-shape.ts @@ -63,7 +63,7 @@ export function computeRequestShapeDiagnostic( providerOptionsHash: stableHash(input.providerOptions ?? {}), toolSchemaHash: stableHash({ activeTools: [...input.activeTools], - providerTools: input.providerTools.map(toolShapeForHash), + providerTools: input.providerTools.map(toolShapeForDiagnostics), }), historyProjectionHash: stableHash(input.priorMessages.map(messageShapeForHash)), }; @@ -75,6 +75,16 @@ export function computeRequestShapeDiagnostic( }; } +export function toolSchemaCharsForDiagnostics( + providerTools: readonly MakaTool[], + activeTools: readonly string[], +): number { + return stableStringify({ + activeTools: [...activeTools], + providerTools: providerTools.map(toolShapeForDiagnostics), + }).length; +} + export function stableHash(value: unknown): string { return `sha256:${createHash('sha256').update(stableStringify(value)).digest('hex')}`; } @@ -96,7 +106,7 @@ function classifyPrefixChange( return 'stable'; } -function toolShapeForHash(tool: MakaTool): unknown { +function toolShapeForDiagnostics(tool: MakaTool): unknown { return { name: tool.name, description: tool.description, diff --git a/packages/runtime/src/run-trace.ts b/packages/runtime/src/run-trace.ts index f0332ea633..c3f906292d 100644 --- a/packages/runtime/src/run-trace.ts +++ b/packages/runtime/src/run-trace.ts @@ -1,5 +1,9 @@ import { generalizedErrorMessage } from '@maka/core/redaction'; -import type { PrefixChangeReason } from '@maka/core/usage-stats/types'; +import type { + ContextBudgetDiagnostic, + PrefixChangeReason, + PromptSegmentEstimate, +} from '@maka/core/usage-stats/types'; export type RunTracePhase = 'turn' | 'model' | 'tool' | 'permission' | 'abort' | 'usage'; @@ -93,7 +97,12 @@ export class RunTrace { modelStreamStarted( activeTools: readonly string[], - prefix?: { prefixHash: string; prefixChangeReason: PrefixChangeReason }, + prefix?: { + prefixHash: string; + prefixChangeReason: PrefixChangeReason; + promptSegments?: PromptSegmentEstimate[]; + contextBudget?: ContextBudgetDiagnostic; + }, ): void { this.emit('model', 'model_stream_started', 'Model stream started', { activeTools: [...activeTools], diff --git a/packages/runtime/src/runtime-event-read-model.ts b/packages/runtime/src/runtime-event-read-model.ts index d233c50391..60ff5763fd 100644 --- a/packages/runtime/src/runtime-event-read-model.ts +++ b/packages/runtime/src/runtime-event-read-model.ts @@ -491,6 +491,8 @@ function projectTokenUsage( ...(usage.costUsd !== undefined ? { costUsd: usage.costUsd } : {}), ...(usage.prefixHash !== undefined ? { prefixHash: usage.prefixHash } : {}), ...(usage.prefixChangeReason !== undefined ? { prefixChangeReason: usage.prefixChangeReason } : {}), + ...(usage.promptSegments !== undefined ? { promptSegments: usage.promptSegments } : {}), + ...(usage.contextBudget !== undefined ? { contextBudget: usage.contextBudget } : {}), }); return true; } @@ -815,6 +817,8 @@ function semanticMessage(message: StoredMessage): unknown { costUsd: message.costUsd, prefixHash: message.prefixHash, prefixChangeReason: message.prefixChangeReason, + promptSegments: message.promptSegments, + contextBudget: message.contextBudget, }; case 'turn_state': return { diff --git a/packages/storage/src/telemetry-repo.ts b/packages/storage/src/telemetry-repo.ts index 957c7684fc..7742dbaa69 100644 --- a/packages/storage/src/telemetry-repo.ts +++ b/packages/storage/src/telemetry-repo.ts @@ -157,6 +157,8 @@ class FileTelemetryRepo implements TelemetryRepo { ...(row.turnId ? { turnId: row.turnId } : {}), ...(row.prefixHash ? { prefixHash: row.prefixHash } : {}), ...(row.prefixChangeReason ? { prefixChangeReason: row.prefixChangeReason } : {}), + ...(row.promptSegments ? { promptSegments: row.promptSegments } : {}), + ...(row.contextBudget ? { contextBudget: row.contextBudget } : {}), } satisfies UsageLogRow)); return { rows: rows.slice(offset, offset + limit), total: rows.length }; } diff --git a/scripts/deepseek-live-cost-baseline.mjs b/scripts/deepseek-live-cost-baseline.mjs new file mode 100644 index 0000000000..92006579ff --- /dev/null +++ b/scripts/deepseek-live-cost-baseline.mjs @@ -0,0 +1,292 @@ +#!/usr/bin/env node +import { mkdir, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join, resolve } from 'node:path'; +import { randomUUID } from 'node:crypto'; +import { + AiSdkBackend, + BackendRegistry, + PermissionEngine, + SessionManager, + buildBuiltinTools, + buildProviderOptions, + computeCost, + createDefaultPermissionEngineDeps, + getAIModel, + getBuiltinPricing, +} from '../packages/runtime/dist/index.js'; +import { + createAgentRunStore, + createRuntimeEventStore, + createSessionStore, +} from '../packages/storage/dist/index.js'; + +const apiKey = process.env.DEEPSEEK_API_KEY; +if (!apiKey) { + throw new Error('DEEPSEEK_API_KEY is not set. Source your local secret env before running this script.'); +} + +const repoRoot = resolve(new URL('..', import.meta.url).pathname); +const outputRoot = resolve(process.env.MAKA_COST_BASELINE_OUTPUT ?? join(tmpdir(), 'maka-deepseek-cost-baseline')); +const runId = new Date().toISOString().replace(/[:.]/g, '-'); +const workspaceRoot = join(outputRoot, runId, 'workspace'); +await mkdir(workspaceRoot, { recursive: true }); + +const model = process.env.MAKA_COST_BASELINE_MODEL ?? 'deepseek-chat'; +const turnCount = parsePositiveInt(process.env.MAKA_COST_BASELINE_TURNS, 10); +const toolMode = process.env.MAKA_COST_BASELINE_TOOLS ?? 'none'; +const seed = process.env.MAKA_COST_BASELINE_SEED ?? runId; +const cwd = resolve(process.env.MAKA_COST_BASELINE_CWD ?? repoRoot); +const contextBudget = buildContextBudgetPolicy(); +const stablePolicyLines = parsePositiveInt(process.env.MAKA_COST_BASELINE_STABLE_POLICY_LINES, 140); +const payloadLines = parsePositiveInt(process.env.MAKA_COST_BASELINE_PAYLOAD_LINES, 70); + +const sessionStore = createSessionStore(workspaceRoot); +const runStore = createAgentRunStore(workspaceRoot); +const runtimeEventStore = createRuntimeEventStore(workspaceRoot); +const permissionEngine = new PermissionEngine(createDefaultPermissionEngineDeps()); +const backends = new BackendRegistry(); +const llmRecords = []; +const runTraceEvents = []; +const tools = toolMode === 'builtin' ? buildBuiltinTools() : []; + +const connection = { + slug: 'deepseek-live-cost-baseline', + name: 'DeepSeek live cost baseline', + providerType: 'deepseek', + baseUrl: 'https://api.deepseek.com', + defaultModel: model, + enabled: true, + createdAt: Date.now(), + updatedAt: Date.now(), +}; + +const durablePrefix = [ + 'You are a concise Maka runtime cost baseline assistant.', + `Baseline seed: ${seed}`, + 'Always answer exactly: OK', + 'The following stable policy block is intentionally repeated to make provider prefix caching observable.', + '', + Array.from({ length: stablePolicyLines }, (_, index) => + `Stable policy line ${String(index + 1).padStart(3, '0')}: preserve the durable system prefix, avoid unnecessary wording churn, and keep responses short.`, + ).join('\n'), + '', +].join('\n'); + +function turnTailPrompt() { + return [ + '', + `cwd: ${cwd}`, + `git_branch: ${process.env.MAKA_COST_BASELINE_BRANCH ?? 'unknown'}`, + `calendar_date: ${process.env.MAKA_COST_BASELINE_DATE ?? new Date().toISOString().slice(0, 10)}`, + '', + ].join('\n'); +} + +backends.register('ai-sdk', async (ctx) => + new AiSdkBackend({ + sessionId: ctx.sessionId, + header: { ...ctx.header, model }, + appendMessage: (message) => ctx.store.appendMessage(ctx.sessionId, message), + connection, + apiKey, + modelId: model, + permissionEngine, + modelFactory: getAIModel, + tools, + providerOptions: buildProviderOptions(connection, model), + contextBudget, + systemPrompt: durablePrefix, + turnTailPrompt, + recordLlmCall: (record) => llmRecords.push(record), + recordRunTrace: (event) => runTraceEvents.push(event), + newId: randomUUID, + now: Date.now, + maxSteps: 1, + streamConnectTimeoutMs: 30_000, + streamIdleTimeoutMs: 120_000, + }), +); + +const manager = new SessionManager({ + store: sessionStore, + runStore, + runtimeEventStore, + backends, + newId: randomUUID, + now: Date.now, +}); +const session = await manager.createSession({ + cwd, + backend: 'ai-sdk', + llmConnectionSlug: connection.slug, + model, + permissionMode: 'explore', + name: 'DeepSeek live cost baseline', +}); + +const turns = []; +const repeatedPayload = Array.from({ length: payloadLines }, (_, index) => + `baseline fact ${String(index + 1).padStart(2, '0')}: this stable user payload is repeated to expose how much new-tail text becomes cache miss.`, +).join('\n'); + +for (let i = 1; i <= turnCount; i += 1) { + const turnId = `cost-turn-${String(i).padStart(2, '0')}`; + const text = [ + `Turn ${i}. Answer exactly OK.`, + repeatedPayload, + `Unique turn marker: ${String(i).padStart(2, '0')}.`, + ].join('\n'); + const events = []; + const startedAt = Date.now(); + for await (const event of manager.sendMessage(session.id, { turnId, text })) { + events.push(event); + } + const finishedAt = Date.now(); + const usageEvent = events.find((event) => event.type === 'token_usage'); + const completeEvent = events.find((event) => event.type === 'complete'); + const errorEvent = events.find((event) => event.type === 'error'); + const llmRecord = llmRecords.at(-1); + const cost = llmRecord + ? computeCost( + { + inputTokens: llmRecord.inputTokens, + outputTokens: llmRecord.outputTokens, + cacheHitInputTokens: llmRecord.cacheHitInputTokens, + cacheMissInputTokens: llmRecord.cacheMissInputTokens, + cacheWriteInputTokens: llmRecord.cacheWriteInputTokens, + }, + getBuiltinPricing(`${connection.providerType}:${model}`), + ) + : undefined; + turns.push({ + turn: i, + turnId, + durationMs: finishedAt - startedAt, + eventCount: events.length, + status: errorEvent ? 'error' : 'ok', + stopReason: completeEvent?.stopReason, + prefixChangeReason: usageEvent?.prefixChangeReason, + prefixHash: usageEvent?.prefixHash, + input: usageEvent?.input ?? llmRecord?.inputTokens, + cacheHitInput: usageEvent?.cacheHitInput ?? llmRecord?.cacheHitInputTokens, + cacheMissInput: usageEvent?.cacheMissInput ?? llmRecord?.cacheMissInputTokens, + output: usageEvent?.output ?? llmRecord?.outputTokens, + total: usageEvent?.total, + estimatedCostUsd: cost?.totalCost, + promptSegments: usageEvent?.promptSegments ?? llmRecord?.promptSegments, + contextBudget: usageEvent?.contextBudget ?? llmRecord?.contextBudget, + errorReason: errorEvent?.reason, + }); +} + +const totals = turns.reduce((acc, turn) => { + acc.input += turn.input ?? 0; + acc.cacheHitInput += turn.cacheHitInput ?? 0; + acc.cacheMissInput += turn.cacheMissInput ?? 0; + acc.output += turn.output ?? 0; + acc.estimatedCostUsd += turn.estimatedCostUsd ?? 0; + return acc; +}, { input: 0, cacheHitInput: 0, cacheMissInput: 0, output: 0, estimatedCostUsd: 0 }); + +const report = { + sourceRef: process.env.MAKA_COST_BASELINE_SOURCE_REF ?? 'local-build', + repoRoot, + workspaceRoot, + model, + seed, + toolMode, + toolCount: tools.length, + turnCount, + stablePolicyLines, + payloadLines, + contextBudget, + sessionId: session.id, + totals, + turns, + runTracePrefixEvents: runTraceEvents + .filter((event) => event.data?.prefixHash || event.data?.prefixChangeReason) + .map((event) => ({ + phase: event.phase, + type: event.type, + prefixHash: event.data?.prefixHash, + prefixChangeReason: event.data?.prefixChangeReason, + promptSegments: event.data?.promptSegments, + contextBudget: event.data?.contextBudget, + })), +}; + +const outputDir = join(outputRoot, runId); +await mkdir(outputDir, { recursive: true }); +const jsonPath = join(outputDir, 'deepseek-live-cost-baseline.json'); +const markdownPath = join(outputDir, 'deepseek-live-cost-baseline.md'); +await writeFile(jsonPath, `${JSON.stringify(report, null, 2)}\n`, 'utf8'); +await writeFile(markdownPath, renderMarkdown(report, jsonPath), 'utf8'); +console.log(JSON.stringify({ jsonPath, markdownPath, totals, turnCount, toolMode, contextBudget }, null, 2)); + +function buildContextBudgetPolicy() { + if (process.env.MAKA_CONTEXT_BUDGET === 'off') return undefined; + const maxHistoryEstimatedTokens = parseOptionalPositiveInt( + process.env.MAKA_CONTEXT_HISTORY_BUDGET_TOKENS, + ); + if (maxHistoryEstimatedTokens === undefined) return undefined; + const maxHistoryTurns = parseOptionalPositiveInt(process.env.MAKA_CONTEXT_HISTORY_BUDGET_TURNS); + return { + name: process.env.MAKA_CONTEXT_BUDGET_NAME ?? 'cost-baseline-history-budget', + maxHistoryEstimatedTokens, + minRecentTurns: parsePositiveInt(process.env.MAKA_CONTEXT_MIN_RECENT_TURNS, 2), + ...(maxHistoryTurns !== undefined ? { maxHistoryTurns } : {}), + }; +} + +function parsePositiveInt(value, fallback) { + return parseOptionalPositiveInt(value) ?? fallback; +} + +function parseOptionalPositiveInt(value) { + if (!value) return undefined; + const parsed = Number.parseInt(value, 10); + return Number.isFinite(parsed) && parsed > 0 ? parsed : undefined; +} + +function renderMarkdown(report, jsonPath) { + const lines = [ + '# DeepSeek Live Cost Baseline', + '', + `JSON: \`${jsonPath}\``, + `Model: \`${report.model}\``, + `Turns: ${report.turnCount}`, + `Tools: ${report.toolMode} (${report.toolCount})`, + `Stable policy lines: ${report.stablePolicyLines}`, + `Payload lines: ${report.payloadLines}`, + `Context budget: ${report.contextBudget ? JSON.stringify(report.contextBudget) : 'off'}`, + '', + '## Totals', + '', + `- input: ${report.totals.input}`, + `- cacheHitInput: ${report.totals.cacheHitInput}`, + `- cacheMissInput: ${report.totals.cacheMissInput}`, + `- output: ${report.totals.output}`, + `- estimatedCostUsd: ${report.totals.estimatedCostUsd}`, + '', + '## Turns', + '', + '| turn | input | hit | miss | output | reason | prior history est | budget after |', + '| ---: | ---: | ---: | ---: | ---: | --- | ---: | ---: |', + ]; + for (const turn of report.turns) { + const prior = turn.promptSegments?.find((segment) => segment.kind === 'prior_history'); + lines.push([ + `| ${turn.turn}`, + turn.input ?? 0, + turn.cacheHitInput ?? 0, + turn.cacheMissInput ?? 0, + turn.output ?? 0, + turn.prefixChangeReason ?? '', + prior?.estimatedTokens ?? 0, + turn.contextBudget?.estimatedTokensAfter ?? 0, + ].join(' | ') + ' |'); + } + lines.push(''); + return `${lines.join('\n')}\n`; +}