diff --git a/docs-web/architecture/high-concurrency-orchestration.md b/docs-web/architecture/high-concurrency-orchestration.md index 88db8df8bb..17899d5a7e 100644 --- a/docs-web/architecture/high-concurrency-orchestration.md +++ b/docs-web/architecture/high-concurrency-orchestration.md @@ -101,6 +101,12 @@ the dashboard, and interactive replies. - Incremental Codex telemetry follows the native thread id emitted by the current exec stream (or the exact requested resume id) and reads only that rollout from the paired runtime volume. A previous invocation's newer file cannot replace the persisted continuation identity. +- Codex rollout parsing skips any single JSONL record above 2 MiB, truncates retained transcript + fields, and keeps the newest 256 conversation item groups. Oversized generated assets or command + output cannot exhaust the server heap, and parsing resumes at the next record. +- Jules full-history telemetry is serialized process-wide, joins duplicate in-flight session reads, + tokenizes text in slices of at most 64 KiB, and releases raw activities before SQLite + reconciliation. Wide hosted-session syncs therefore cannot multiply large patch payloads in heap. - Antigravity sends `--conversation` only for provider-native ids. Logical/workspace continuation markers use `--continue` inside the isolated paired runtime volume instead. diff --git a/docs-web/content/docs/architecture-high-concurrency-orchestration.mdx b/docs-web/content/docs/architecture-high-concurrency-orchestration.mdx index 88db8df8bb..17899d5a7e 100644 --- a/docs-web/content/docs/architecture-high-concurrency-orchestration.mdx +++ b/docs-web/content/docs/architecture-high-concurrency-orchestration.mdx @@ -101,6 +101,12 @@ the dashboard, and interactive replies. - Incremental Codex telemetry follows the native thread id emitted by the current exec stream (or the exact requested resume id) and reads only that rollout from the paired runtime volume. A previous invocation's newer file cannot replace the persisted continuation identity. +- Codex rollout parsing skips any single JSONL record above 2 MiB, truncates retained transcript + fields, and keeps the newest 256 conversation item groups. Oversized generated assets or command + output cannot exhaust the server heap, and parsing resumes at the next record. +- Jules full-history telemetry is serialized process-wide, joins duplicate in-flight session reads, + tokenizes text in slices of at most 64 KiB, and releases raw activities before SQLite + reconciliation. Wide hosted-session syncs therefore cannot multiply large patch payloads in heap. - Antigravity sends `--conversation` only for provider-native ids. Logical/workspace continuation markers use `--continue` inside the isolated paired runtime volume instead. diff --git a/docs/architecture/high-concurrency-orchestration.md b/docs/architecture/high-concurrency-orchestration.md index 1d02baf764..7f3e5a648a 100644 --- a/docs/architecture/high-concurrency-orchestration.md +++ b/docs/architecture/high-concurrency-orchestration.md @@ -212,6 +212,10 @@ Codex rollout transport binds every read to the native thread id emitted by the `codex exec --json` stream (or the exact requested resume id); it never treats an unrelated newest rollout in a reused runtime home as the invocation's identity. Parsing retains a byte cursor, handles split UTF-8/JSONL records, caps work per poll, and resets on source rotation or truncation. +The parser discards any single JSONL record above 2 MiB, bounds retained message/tool fields, and +keeps only the newest 256 conversation item groups. Large generated assets or command output can +therefore remain in the provider-owned rollout without exhausting the server heap; normalized usage, +session identity, and later records continue to be processed. Claude transport also reads appended bytes. Qwen mutable JSON files and the Antigravity SQLite source use a coherent full read only after their cheap metadata changes; unchanged polls do not copy or parse them. @@ -234,7 +238,12 @@ Structured invocation messages reconcile by stable ordinal in one SQLite transac Text-only completion fallback remains append-only so retry and audit messages are not removed. Hosted Jules activity-to-message sync uses the same atomic suffix reconciliation instead of clearing -and reinserting the invocation transcript on each poll. +and reinserting the invocation transcript on each poll. Full Jules histories are fetched and +tokenized through a one-at-a-time process queue, duplicate in-flight requests for the same session +join the existing work, and tokenizer input is sliced to at most 64 KiB. The raw activity array is +released after bounded messages and numeric usage are derived, before SQLite reconciliation begins. +This keeps wide hosted-session synchronization from multiplying large patch and media payloads in +the Node.js heap. Streaming provider activity is buffered for 250 ms or 50 source records, then adjacent records from the same originator are compacted into bounded 16 KiB rows before one batch transaction. This avoids diff --git a/docs/architecture/usage-telemetry-and-stats.md b/docs/architecture/usage-telemetry-and-stats.md index 7eccbf3821..2eed4937ee 100644 --- a/docs/architecture/usage-telemetry-and-stats.md +++ b/docs/architecture/usage-telemetry-and-stats.md @@ -126,6 +126,12 @@ initial target. Code UX derives the rollout date from that id and reads only lookup remains only as a compatibility fallback when a legacy caller has no native id, so prior sessions in the same runtime home cannot overwrite the persisted continuation identity. +Live rollout parsing is bounded independently of rollout file size. A JSONL record larger than +2 MiB is skipped until its terminating newline, per-turn text/tool payloads are truncated with an +explicit marker, and only the newest 256 conversation item groups remain in memory. This protects +the runtime from generated binary assets and extreme command output while preserving cumulative +usage snapshots, native session identity, and every later well-formed record. + ### Qwen Code Qwen Code runs via its OpenAI-compatible request/response logging (`enableOpenAILoggingDir`), written to a directory that is reset at the start of every run so usage aggregation only ever sums the current invocation's own log files — unlike Codex/OpenCode, there is no cross-run cumulative counter to isolate. @@ -221,7 +227,18 @@ Stats pricing still prefers configured model-pricing overrides and catalogue tok ### Jules -Jules does not expose a compatible native token contract. Instead of excluding it, Code UX computes **estimated** tokens for Jules by accumulating input and output characters divided by 4 (the characters-per-token heuristic). +Jules does not expose a compatible native token contract. Instead of excluding it, Code UX computes +**estimated** tokens from the cumulative activity stream. The estimator models replayed input +context, generated messages, progress/tool turns, and added lines from patch artifacts. It uses the +`cl100k_base` tokenizer in slices of at most 64 KiB so a multi-megabyte patch cannot create one +unbounded tokenizer allocation. + +Full-conversation usage synchronization is process-wide serialized. Calls for the same session join +the existing in-flight sync, while different sessions wait on a one-at-a-time queue. Once bounded +invocation messages and numeric usage have been derived, Code UX releases the raw remote activity +array before reconciling SQLite rows. This prevents a wide set of hosted sessions from retaining and +tokenizing multiple large histories concurrently without changing the persisted message or usage +contract. During live synchronization (`syncLiveInvocation`), expected 404 responses indicating that a session or activity stream is unavailable are handled gracefully: they are logged at the debug level and skipped to avoid spamming the logs with warnings. For terminal sync (`calculateAndSaveUsageForTask`), the system is conservative: if the session returns a 404, it skips creating a new usage record to prevent saving "fake" empty records unless an existing prompt or usage record is already present to allow safe estimation. diff --git a/src/domain/jules/jules-usage-estimator.ts b/src/domain/jules/jules-usage-estimator.ts index cf5166b7c4..6fc9523880 100644 --- a/src/domain/jules/jules-usage-estimator.ts +++ b/src/domain/jules/jules-usage-estimator.ts @@ -38,6 +38,11 @@ export const JULES_SYSTEM_PROMPT_TOKENS = 800; * compaction Jules performs and prevents quadratic blow-up on long sessions. */ export const JULES_CONTEXT_TOKEN_CAP = 200_000; +/** Maximum string slice passed to the tokenizer at once. `js-tiktoken` + * materializes token arrays and regex matches, so feeding it a multi-megabyte + * patch in one call can temporarily consume gigabytes of V8 heap. */ +export const JULES_TOKENIZER_CHUNK_CHARS = 64 * 1024; + /** Fallback tokens-per-added-line when a diff isn't available but PR git stats are. */ export const JULES_TOKENS_PER_ADDED_LINE = 12; @@ -80,6 +85,72 @@ export function extractAddedDiffLines(unidiffPatch: string): string { return added.join("\n"); } +/** Tokenizes large values in bounded slices. Jules usage is already an + * estimate, and the tiny BPE-boundary variance is preferable to an unbounded + * temporary allocation for generated patches and transcripts. */ +export function countJulesTokensInChunks( + text: string, + countTokens: (chunk: string) => number, +): number { + let total = 0; + let offset = 0; + while (offset < text.length) { + let end = Math.min(text.length, offset + JULES_TOKENIZER_CHUNK_CHARS); + if ( + end < text.length + && end > offset + && text.charCodeAt(end - 1) >= 0xD800 + && text.charCodeAt(end - 1) <= 0xDBFF + ) { + end -= 1; + } + total += countTokens(text.slice(offset, end)); + offset = end; + } + return total; +} + +function measureAddedDiffLines( + unidiffPatch: string, + countTokens: (text: string) => number, +): { chars: number; tokens: number } { + let batch = ""; + let chars = 0; + let tokens = 0; + let hasAddedLine = false; + let cursor = 0; + + const flush = () => { + if (!batch) { + return; + } + tokens += countJulesTokensInChunks(batch, countTokens); + batch = ""; + }; + + while (cursor <= unidiffPatch.length) { + const newline = unidiffPatch.indexOf("\n", cursor); + const end = newline >= 0 ? newline : unidiffPatch.length; + const line = unidiffPatch.slice(cursor, end); + if (line.startsWith("+") && !line.startsWith("+++")) { + const addedLine = line.slice(1); + const separator = hasAddedLine ? "\n" : ""; + chars += separator.length + addedLine.length; + if (batch.length + separator.length + addedLine.length > JULES_TOKENIZER_CHUNK_CHARS) { + flush(); + } + batch += separator + addedLine; + hasAddedLine = true; + } + if (newline < 0) { + break; + } + cursor = newline + 1; + } + flush(); + return { chars, tokens }; +} + function planToMarkdown(activity: JulesActivity): string { const steps = activity.planGenerated?.plan?.steps; if (!Array.isArray(steps) || steps.length === 0) { @@ -102,6 +173,7 @@ function sortByCreateTime(activities: JulesActivity[]): JulesActivity[] { */ export function estimateJulesUsage(input: JulesUsageEstimateInput): JulesUsageEstimate { const { countTokens, gitMetrics } = input; + const countBoundedTokens = (text: string) => countJulesTokensInChunks(text, countTokens); const activities = sortByCreateTime(input.activities || []); let inputTokens = 0; @@ -116,7 +188,7 @@ export function estimateJulesUsage(input: JulesUsageEstimateInput): JulesUsageEs const prompt = input.prompt || ""; if (prompt) { promptChars += prompt.length; - context = Math.min(JULES_CONTEXT_TOKEN_CAP, context + countTokens(prompt)); + context = Math.min(JULES_CONTEXT_TOKEN_CAP, context + countBoundedTokens(prompt)); } const addContext = (tokens: number) => { @@ -137,24 +209,24 @@ export function estimateJulesUsage(input: JulesUsageEstimateInput): JulesUsageEs if (activity.userMessaged?.userMessage) { const text = activity.userMessaged.userMessage; promptChars += text.length; - addContext(countTokens(text)); + addContext(countBoundedTokens(text)); } if (activity.planApproved?.planId) { const text = `Approved plan (ID: ${activity.planApproved.planId})`; promptChars += text.length; - addContext(countTokens(text)); + addContext(countBoundedTokens(text)); } // Agent-side model output turns. if (activity.agentMessaged?.agentMessage) { const text = activity.agentMessaged.agentMessage; transcriptChars += text.length; - billAgentTurn(countTokens(text)); + billAgentTurn(countBoundedTokens(text)); } if (activity.planGenerated?.plan?.steps) { const text = `Proposed plan:\n\n${planToMarkdown(activity)}`; transcriptChars += text.length; - billAgentTurn(countTokens(text)); + billAgentTurn(countBoundedTokens(text)); } if (activity.progressUpdated?.title || activity.progressUpdated?.description) { // Progress updates are short status lines the agent emits while driving @@ -164,13 +236,13 @@ export function estimateJulesUsage(input: JulesUsageEstimateInput): JulesUsageEs const text = `${title}\n${desc}`; transcriptChars += text.length; toolCallCount += 1; - billAgentTurn(countTokens(text)); + billAgentTurn(countBoundedTokens(text)); } if (activity.sessionCompleted !== undefined && activity.sessionCompleted !== null) { - billAgentTurn(countTokens("Jules session completed successfully.")); + billAgentTurn(countBoundedTokens("Jules session completed successfully.")); } if (activity.sessionFailed?.reason) { - billAgentTurn(countTokens(`Jules session failed: ${activity.sessionFailed.reason}`)); + billAgentTurn(countBoundedTokens(`Jules session failed: ${activity.sessionFailed.reason}`)); } // Code artifacts: the model produced a patch (a tool result). Count only the @@ -180,15 +252,14 @@ export function estimateJulesUsage(input: JulesUsageEstimateInput): JulesUsageEs if (unidiffPatch) { sawUnidiffPatch = true; toolCallCount += 1; - const addedCode = extractAddedDiffLines(unidiffPatch); - const codeTokens = countTokens(addedCode); - outputTokens += codeTokens; - transcriptChars += addedCode.length; - addContext(countTokens(unidiffPatch)); + const addedCode = measureAddedDiffLines(unidiffPatch, countTokens); + outputTokens += addedCode.tokens; + transcriptChars += addedCode.chars; + addContext(countBoundedTokens(unidiffPatch)); } const commitMessage = art.changeSet?.gitPatch?.suggestedCommitMessage; if (commitMessage) { - const msgTokens = countTokens(commitMessage); + const msgTokens = countBoundedTokens(commitMessage); outputTokens += msgTokens; transcriptChars += commitMessage.length; addContext(msgTokens); diff --git a/src/domain/jules/jules-usage-service.ts b/src/domain/jules/jules-usage-service.ts index 8998da4be3..c2bf9b5ef2 100644 --- a/src/domain/jules/jules-usage-service.ts +++ b/src/domain/jules/jules-usage-service.ts @@ -11,6 +11,7 @@ import type { JulesActivity, JulesSession } from "../../contracts/app-types.js"; import { estimateJulesUsage, type JulesUsageEstimate } from "./jules-usage-estimator.js"; import { MAX_MESSAGE_CONTENT_CHARS, + MAX_TOOL_PAYLOAD_CHARS, truncateForStorage, } from "../../services/invocation-message-limits.js"; import { isNotFoundError } from "../../integrations/jules-api-client.js"; @@ -35,6 +36,8 @@ const LIVE_SYNC_THROTTLE_MS = 8_000; export class JulesUsageService { private encoder: Tiktoken | null = null; private readonly lastLiveSyncMsBySession = new Map(); + private readonly liveSyncBySession = new Map>(); + private usageSyncTail: Promise = Promise.resolve(); constructor( private readonly julesClient: JulesClient, @@ -63,6 +66,24 @@ export class JulesUsageService { sessionId: string, passedPrompt?: string, gitMetrics?: GitMetrics + ): Promise { + await this.enqueueUsageSync(async () => { + await this.calculateAndSaveUsageForTaskSerial( + projectId, + taskId, + sessionId, + passedPrompt, + gitMetrics, + ); + }); + } + + private async calculateAndSaveUsageForTaskSerial( + projectId: string, + taskId: string, + sessionId: string, + passedPrompt?: string, + gitMetrics?: GitMetrics, ): Promise { try { const existingRecord = this.executionRepository.getLatestProviderInvocationUsageBySession(sessionId, "task_coding"); @@ -120,13 +141,42 @@ export class JulesUsageService { prompt?: string, gitMetrics?: GitMetrics ): Promise { + const inFlight = this.liveSyncBySession.get(sessionId); + if (inFlight) { + await inFlight; + return; + } + const now = Date.now(); const last = this.lastLiveSyncMsBySession.get(sessionId) ?? 0; if (now - last < LIVE_SYNC_THROTTLE_MS) { return; } - this.lastLiveSyncMsBySession.set(sessionId, now); + const pending = this.enqueueUsageSync(async () => { + try { + await this.syncLiveInvocationSerial(projectId, taskId, sessionId, prompt, gitMetrics); + } finally { + this.lastLiveSyncMsBySession.set(sessionId, Date.now()); + } + }); + this.liveSyncBySession.set(sessionId, pending); + try { + await pending; + } finally { + if (this.liveSyncBySession.get(sessionId) === pending) { + this.liveSyncBySession.delete(sessionId); + } + } + } + + private async syncLiveInvocationSerial( + projectId: string, + taskId: string, + sessionId: string, + prompt?: string, + gitMetrics?: GitMetrics, + ): Promise { try { const activities = await this.julesClient.getFullConversation(sessionId); if (activities.length === 0 && !prompt) { @@ -150,6 +200,12 @@ export class JulesUsageService { } } + private enqueueUsageSync(work: () => Promise): Promise { + const pending = this.usageSyncTail.then(work, work); + this.usageSyncTail = pending.catch(() => undefined); + return pending; + } + /** Resolves the session prompt and PR git stats, fetching the session only * when the caller did not already provide them. */ private async resolvePromptAndGitMetrics( @@ -197,12 +253,22 @@ export class JulesUsageService { }): void { const { projectId, taskId, sessionId, activities, prompt, gitMetrics, final } = args; - const estimate = estimateJulesUsage({ - prompt, - activities, - gitMetrics, - countTokens: (text) => this.countTokens(text), - }); + let estimate!: JulesUsageEstimate; + let conversationMessages!: AppendExecutionInvocationMessageInput[]; + try { + conversationMessages = this.buildConversationMessages(activities, prompt, ""); + estimate = estimateJulesUsage({ + prompt, + activities, + gitMetrics, + countTokens: (text) => this.countTokens(text), + }); + } finally { + // These arrays are freshly fetched for this sync. Release the raw remote + // messages, patches, and media before the SQLite reconciliation allocates + // its own bounded message view. + activities.length = 0; + } const status = final ? "completed" : "running"; @@ -280,9 +346,15 @@ export class JulesUsageService { // Activity history is cumulative. Reconcile by stable ordinal so an // eight-second live sync only inserts or updates the changed suffix. + if (prompt && conversationMessages.length > 0) { + conversationMessages[0] = { + ...conversationMessages[0], + createdAt: record.createdAt, + }; + } this.executionRepository.syncExecutionInvocationMessages( execInvocation.id, - this.buildConversationMessages(activities, prompt, record.createdAt), + conversationMessages, ); this.logger.info("Saved Jules usage telemetry and conversation transcript for task", { @@ -372,10 +444,14 @@ export class JulesUsageService { for (const art of activity.artifacts || []) { const patch = art.changeSet?.gitPatch?.unidiffPatch; if (patch) { + const contentPatch = truncateForStorage( + patch, + Math.max(0, MAX_MESSAGE_CONTENT_CHARS - 12), + ); push({ role: "tool", - contentMarkdown: cap(`\`\`\`diff\n${patch}\n\`\`\``), - toolCallsJson: { output: truncateForStorage(patch, MAX_MESSAGE_CONTENT_CHARS) }, + contentMarkdown: `\`\`\`diff\n${contentPatch}\n\`\`\``, + toolCallsJson: { output: truncateForStorage(patch, MAX_TOOL_PAYLOAD_CHARS) }, metadata: { ...base, kind: "tool_result", toolName: "apply_patch" }, }); } diff --git a/src/infrastructure/providers/cli/provider-logs/codex-log-parser.ts b/src/infrastructure/providers/cli/provider-logs/codex-log-parser.ts index 7c01144613..80f78cf7e0 100644 --- a/src/infrastructure/providers/cli/provider-logs/codex-log-parser.ts +++ b/src/infrastructure/providers/cli/provider-logs/codex-log-parser.ts @@ -18,10 +18,55 @@ export interface CodexLogResult extends ParsedProviderLogResult | null { return value && typeof value === "object" ? value as Record : null; } +function truncateRetainedField(value: string, maxChars: number): string { + if (!value || value.length <= maxChars) { + return value; + } + const omitted = value.length - maxChars; + const marker = `\n\n… [${omitted.toLocaleString("en-US")} characters truncated] …\n\n`; + const budget = Math.max(0, maxChars - marker.length); + const headChars = Math.ceil(budget * 0.7); + const tailChars = budget - headChars; + return `${value.slice(0, headChars)}${marker}${tailChars > 0 ? value.slice(-tailChars) : ""}`; +} + +function boundConversationTurn(turn: ParsedConversationTurn): ParsedConversationTurn { + return { + ...turn, + text: truncateRetainedField(turn.text, CODEX_MAX_RETAINED_TURN_TEXT_CHARS), + ...(turn.toolArguments !== undefined + ? { + toolArguments: truncateRetainedField( + turn.toolArguments, + CODEX_MAX_RETAINED_TOOL_PAYLOAD_CHARS, + ), + } + : {}), + ...(turn.toolOutput !== undefined + ? { + toolOutput: truncateRetainedField( + turn.toolOutput, + CODEX_MAX_RETAINED_TOOL_PAYLOAD_CHARS, + ), + } + : {}), + }; +} + /** Flattens a Codex message `content` array (input_text / output_text / text parts) to plain text. */ function flattenContent(content: unknown): string { if (typeof content === "string") { @@ -243,17 +288,43 @@ function upsertConversationGroup( if (turns.length === 0) { return null; } + const boundedTurns = turns.map(boundConversationTurn); const key = conversationItemKey(item); if (key) { const existingIndex = groupIndexes.get(key); if (existingIndex !== undefined) { - groups[existingIndex] = { key, turns }; + groups[existingIndex] = { key, turns: boundedTurns }; return existingIndex; } groupIndexes.set(key, groups.length); } - groups.push({ key, turns }); - return groups.length - 1; + groups.push({ key, turns: boundedTurns }); + if (groups.length <= CODEX_MAX_RETAINED_CONVERSATION_GROUPS) { + return groups.length - 1; + } + + groups.splice(0, groups.length - CODEX_MAX_RETAINED_CONVERSATION_GROUPS); + groupIndexes.clear(); + for (let index = 0; index < groups.length; index += 1) { + const retainedKey = groups[index]?.key; + if (retainedKey) { + groupIndexes.set(retainedKey, index); + } + } + return 0; +} + +function appendBoundedFallbackTurns( + conversation: ParsedConversationTurn[], + turns: ParsedConversationTurn[], +): number { + const changedFrom = conversation.length; + conversation.push(...turns.map(boundConversationTurn)); + if (conversation.length <= CODEX_MAX_RETAINED_CONVERSATION_GROUPS) { + return changedFrom; + } + conversation.splice(0, conversation.length - CODEX_MAX_RETAINED_CONVERSATION_GROUPS); + return 0; } function flattenConversationGroups(groups: ConversationTurnGroup[]): ParsedConversationTurn[] { @@ -579,8 +650,7 @@ function processCodexRolloutLine(state: CodexRolloutParserState, rawLine: string } const turns = eventMsgToTurns(payload, timestampMs); if (turns.length > 0 && isInWindow(timestampMs)) { - const changedFrom = state.fallbackEventConversation.length; - state.fallbackEventConversation.push(...turns); + const changedFrom = appendBoundedFallbackTurns(state.fallbackEventConversation, turns); state.conversationRevision += 1; state.conversationChangedFromIndex = state.conversationChangedFromIndex === null ? changedFrom @@ -643,20 +713,58 @@ function buildCodexRolloutResult(state: CodexRolloutParserState): CodexLogResult }; } +interface CodexRolloutChunkState { + pendingLine: string; + discardingOversizedLine: boolean; +} + function processCodexRolloutChunk( state: CodexRolloutParserState, chunk: string, pendingLine: string, -): string { - const lines = (pendingLine + chunk).split("\n"); - const finalLine = lines.pop() ?? ""; - for (const line of lines) { - processCodexRolloutLine(state, line); - } - if (!finalLine) { - return ""; + discardingOversizedLine: boolean, +): CodexRolloutChunkState { + let cursor = 0; + let retainedPendingLine = pendingLine; + let discarding = discardingOversizedLine; + + while (cursor < chunk.length) { + const newlineIndex = chunk.indexOf("\n", cursor); + const segmentEnd = newlineIndex >= 0 ? newlineIndex : chunk.length; + const segment = chunk.slice(cursor, segmentEnd); + + if (discarding) { + if (newlineIndex < 0) { + return { pendingLine: "", discardingOversizedLine: true }; + } + discarding = false; + cursor = newlineIndex + 1; + continue; + } + + if (retainedPendingLine.length + segment.length > CODEX_MAX_JSONL_RECORD_CHARS) { + retainedPendingLine = ""; + if (newlineIndex < 0) { + return { pendingLine: "", discardingOversizedLine: true }; + } + cursor = newlineIndex + 1; + continue; + } + + const record = retainedPendingLine ? retainedPendingLine + segment : segment; + retainedPendingLine = ""; + if (newlineIndex < 0) { + retainedPendingLine = processCodexRolloutLine(state, record) ? "" : record; + break; + } + processCodexRolloutLine(state, record); + cursor = newlineIndex + 1; } - return processCodexRolloutLine(state, finalLine) ? "" : finalLine; + + return { + pendingLine: retainedPendingLine, + discardingOversizedLine: discarding, + }; } /** Incremental parser for the append-only Codex rollout used by live telemetry. */ @@ -666,6 +774,7 @@ export class CodexRolloutAccumulator { private previousHead = ""; private previousBoundary = ""; private pendingLine = ""; + private discardingOversizedLine = false; private sourceId: string | null = null; private lastResult: CodexLogResult | null = null; @@ -684,7 +793,14 @@ export class CodexRolloutAccumulator { const chunk = canAppend ? jsonl.slice(this.previousLength) : jsonl; this.state.conversationChangedFromIndex = null; - this.pendingLine = processCodexRolloutChunk(this.state, chunk, this.pendingLine); + const chunkState = processCodexRolloutChunk( + this.state, + chunk, + this.pendingLine, + this.discardingOversizedLine, + ); + this.pendingLine = chunkState.pendingLine; + this.discardingOversizedLine = chunkState.discardingOversizedLine; this.previousLength = jsonl.length; this.previousHead = jsonl.slice(0, Math.min(4096, jsonl.length)); this.previousBoundary = jsonl.slice(Math.max(0, jsonl.length - 4096)); @@ -698,7 +814,14 @@ export class CodexRolloutAccumulator { this.reset(sourceId); } this.state.conversationChangedFromIndex = null; - this.pendingLine = processCodexRolloutChunk(this.state, text, this.pendingLine); + const chunkState = processCodexRolloutChunk( + this.state, + text, + this.pendingLine, + this.discardingOversizedLine, + ); + this.pendingLine = chunkState.pendingLine; + this.discardingOversizedLine = chunkState.discardingOversizedLine; this.sourceId = sourceId; // Full-snapshot prefix checks do not apply while consuming byte deltas. this.previousLength = 0; @@ -728,6 +851,7 @@ export class CodexRolloutAccumulator { this.previousHead = ""; this.previousBoundary = ""; this.pendingLine = ""; + this.discardingOversizedLine = false; this.sourceId = sourceId; this.lastResult = null; } @@ -802,7 +926,10 @@ export function parseCodexExecStdout(stdout: string): CodexLogResult { } if (type === "event_msg" && payload) { - fallbackEventConversation.push(...eventMsgToTurns(payload, timestampMs)); + appendBoundedFallbackTurns( + fallbackEventConversation, + eventMsgToTurns(payload, timestampMs), + ); continue; } diff --git a/tests/backend/domain/jules/jules-usage-memory.test.ts b/tests/backend/domain/jules/jules-usage-memory.test.ts new file mode 100644 index 0000000000..a816d171c0 --- /dev/null +++ b/tests/backend/domain/jules/jules-usage-memory.test.ts @@ -0,0 +1,83 @@ +import { describe, expect, it, vi } from "vitest"; +import type { JulesActivity } from "../../../../src/contracts/app-types.js"; +import type { JulesClient } from "../../../../src/domain/jules/jules-client.js"; +import { + countJulesTokensInChunks, + JULES_TOKENIZER_CHUNK_CHARS, +} from "../../../../src/domain/jules/jules-usage-estimator.js"; +import { JulesUsageService } from "../../../../src/domain/jules/jules-usage-service.js"; +import type { ExecutionRepository } from "../../../../src/repositories/execution-repository.js"; +import type { Logger } from "../../../../src/shared/logging/logger.js"; + +function createService(getFullConversation: (sessionId: string) => Promise): JulesUsageService { + const julesClient = { + getFullConversation, + } as unknown as JulesClient; + const executionRepository = { + getLatestProviderInvocationUsageBySession: vi.fn().mockReturnValue(null), + createProviderInvocationUsage: vi.fn().mockReturnValue({ + id: "usage-1", + createdAt: "2026-07-17T00:00:00.000Z", + sprintId: null, + taskId: "task-1", + sprintRunId: null, + dispatchId: null, + taskRunId: null, + attentionItemId: null, + }), + updateProviderInvocationUsage: vi.fn(), + listExecutionInvocationsByProviderInvocationId: vi.fn().mockReturnValue([]), + createExecutionInvocation: vi.fn().mockReturnValue({ id: "invocation-1" }), + syncExecutionInvocationMessages: vi.fn(), + } as unknown as ExecutionRepository; + const logger = { + info: vi.fn(), + error: vi.fn(), + warn: vi.fn(), + debug: vi.fn(), + } as unknown as Logger; + return new JulesUsageService(julesClient, executionRepository, logger); +} + +describe("Jules usage memory bounds", () => { + it("passes bounded slices to the tokenizer", () => { + const text = "x".repeat(JULES_TOKENIZER_CHUNK_CHARS * 3 + 17); + const chunkLengths: number[] = []; + + const tokens = countJulesTokensInChunks(text, (chunk) => { + chunkLengths.push(chunk.length); + return chunk.length; + }); + + expect(tokens).toBe(text.length); + expect(chunkLengths).toHaveLength(4); + expect(Math.max(...chunkLengths)).toBeLessThanOrEqual(JULES_TOKENIZER_CHUNK_CHARS); + }); + + it("keeps only one full-conversation fetch active across sessions", async () => { + const resolvers = new Map void>(); + let activeFetches = 0; + let maximumActiveFetches = 0; + const service = createService(async (sessionId) => { + activeFetches += 1; + maximumActiveFetches = Math.max(maximumActiveFetches, activeFetches); + const activities = await new Promise((resolve) => { + resolvers.set(sessionId, resolve); + }); + activeFetches -= 1; + return activities; + }); + + const first = service.syncLiveInvocation("project-1", "task-1", "session-a", "prompt"); + const second = service.syncLiveInvocation("project-1", "task-2", "session-b", "prompt"); + await vi.waitFor(() => expect(resolvers.has("session-a")).toBe(true)); + expect(resolvers.has("session-b")).toBe(false); + + resolvers.get("session-a")?.([]); + await vi.waitFor(() => expect(resolvers.has("session-b")).toBe(true)); + resolvers.get("session-b")?.([]); + await Promise.all([first, second]); + + expect(maximumActiveFetches).toBe(1); + }); +}); diff --git a/tests/backend/infrastructure/providers/cli/codex-log-parser.test.ts b/tests/backend/infrastructure/providers/cli/codex-log-parser.test.ts index 7f0d80c064..16df07f887 100644 --- a/tests/backend/infrastructure/providers/cli/codex-log-parser.test.ts +++ b/tests/backend/infrastructure/providers/cli/codex-log-parser.test.ts @@ -1,5 +1,8 @@ import { describe, it, expect } from "vitest"; import { + CODEX_MAX_JSONL_RECORD_CHARS, + CODEX_MAX_RETAINED_CONVERSATION_GROUPS, + CODEX_MAX_RETAINED_TOOL_PAYLOAD_CHARS, CodexRolloutAccumulator, parseCodexExecStdout, parseCodexRolloutJsonl, @@ -522,6 +525,70 @@ describe("CodexRolloutAccumulator", () => { expect.objectContaining({ kind: "user", text: "complete later" }), ]); }); + + it("discards an oversized split JSONL record and resumes at the next record", () => { + const accumulator = new CodexRolloutAccumulator(); + const oversizedPrefix = JSON.stringify({ + type: "response_item", + timestamp: "2026-06-01T10:00:00.000Z", + payload: { type: "function_call_output", call_id: "huge-output" }, + }).slice(0, -2) + ',"output":"'; + const oversized = oversizedPrefix + "x".repeat(CODEX_MAX_JSONL_RECORD_CHARS); + const splitAt = Math.floor(oversized.length / 2); + + accumulator.appendChunk(oversized.slice(0, splitAt), "rollout-large", true); + const result = accumulator.appendChunk([ + oversized.slice(splitAt), + userMessage("2026-06-01T10:00:01.000Z", "parser recovered"), + "", + ].join("\n"), "rollout-large"); + + expect(result.conversation).toEqual([ + expect.objectContaining({ kind: "user", text: "parser recovered" }), + ]); + }); + + it("bounds retained Codex tool payloads before live persistence", () => { + const largeOutput = `head-${"x".repeat(50_000)}-tail`; + const result = parseCodexRolloutJsonl(responseItem( + "2026-06-01T10:00:00.000Z", + { + type: "function_call_output", + call_id: "large-tool", + output: largeOutput, + }, + )); + + expect(result.conversation[0]?.toolOutput?.length) + .toBeLessThanOrEqual(CODEX_MAX_RETAINED_TOOL_PAYLOAD_CHARS); + expect(result.conversation[0]?.toolOutput).toContain("head-"); + expect(result.conversation[0]?.toolOutput).toContain("-tail"); + expect(result.conversation[0]?.toolOutput).toContain("characters truncated"); + }); + + it("retains only the newest bounded window of Codex conversation groups", () => { + const extraGroups = 12; + const jsonl = Array.from( + { length: CODEX_MAX_RETAINED_CONVERSATION_GROUPS + extraGroups }, + (_, index) => responseItem( + "2026-06-01T10:00:00.000Z", + { + id: `message-${index}`, + type: "message", + role: "assistant", + content: [{ type: "output_text", text: `answer-${index}` }], + }, + ), + ).join("\n"); + + const result = parseCodexRolloutJsonl(jsonl); + + expect(result.conversation).toHaveLength(CODEX_MAX_RETAINED_CONVERSATION_GROUPS); + expect(result.conversation[0]?.text).toBe(`answer-${extraGroups}`); + expect(result.conversation.at(-1)?.text) + .toBe(`answer-${CODEX_MAX_RETAINED_CONVERSATION_GROUPS + extraGroups - 1}`); + expect(result.conversationChangedFromIndex).toBe(0); + }); }); describe("parseCodexExecStdout", () => { diff --git a/tests/domain/jules/jules-usage-estimator.test.ts b/tests/domain/jules/jules-usage-estimator.test.ts index 292f9c4a6b..40309c5fcf 100644 --- a/tests/domain/jules/jules-usage-estimator.test.ts +++ b/tests/domain/jules/jules-usage-estimator.test.ts @@ -1,10 +1,12 @@ import { describe, it, expect } from "vitest"; import { getEncoding } from "js-tiktoken"; import { + countJulesTokensInChunks, estimateJulesUsage, extractAddedDiffLines, JULES_SYSTEM_PROMPT_TOKENS, JULES_CONTEXT_TOKEN_CAP, + JULES_TOKENIZER_CHUNK_CHARS, JULES_TOKENS_PER_ADDED_LINE, } from "../../../src/domain/jules/jules-usage-estimator.js"; import type { JulesActivity } from "../../../src/contracts/app-types.js"; @@ -28,6 +30,21 @@ describe("extractAddedDiffLines", () => { }); }); +describe("countJulesTokensInChunks", () => { + it("never passes an unbounded string to the tokenizer", () => { + const input = "x".repeat(JULES_TOKENIZER_CHUNK_CHARS * 3 + 17); + const seenChunkLengths: number[] = []; + const tokens = countJulesTokensInChunks(input, (chunk) => { + seenChunkLengths.push(chunk.length); + return chunk.length; + }); + + expect(tokens).toBe(input.length); + expect(seenChunkLengths.length).toBe(4); + expect(Math.max(...seenChunkLengths)).toBeLessThanOrEqual(JULES_TOKENIZER_CHUNK_CHARS); + }); +}); + describe("estimateJulesUsage", () => { it("bills an agent turn as input=context, output=generated", () => { const activities: JulesActivity[] = [ diff --git a/tests/domain/jules/jules-usage-service.test.ts b/tests/domain/jules/jules-usage-service.test.ts index 90caf77e1e..473d6895b9 100644 --- a/tests/domain/jules/jules-usage-service.test.ts +++ b/tests/domain/jules/jules-usage-service.test.ts @@ -4,6 +4,10 @@ import type { JulesClient } from "../../../src/domain/jules/jules-client.js"; import type { ExecutionRepository } from "../../../src/repositories/execution-repository.js"; import type { Logger } from "../../../src/shared/logging/logger.js"; import type { JulesActivity } from "../../../src/contracts/app-types.js"; +import { + MAX_MESSAGE_CONTENT_CHARS, + MAX_TOOL_PAYLOAD_CHARS, +} from "../../../src/services/invocation-message-limits.js"; describe("JulesUsageService", () => { let getFullConversationMock: ReturnType; @@ -163,6 +167,37 @@ describe("JulesUsageService", () => { expect(toolMsg!.metadata.toolName).toBe("apply_patch"); }); + it("bounds patch messages and releases raw activities before persistence", async () => { + const patch = `+${"generated asset data".repeat(20_000)}`; + const activities = [ + { + id: "1", + name: "1", + createTime: "2026-06-01T00:00:00Z", + artifacts: [{ changeSet: { gitPatch: { unidiffPatch: patch } } }], + }, + ] as JulesActivity[]; + getFullConversationMock.mockResolvedValue(activities); + syncMessagesMock.mockImplementation((_invocationId, messages) => { + expect(activities).toHaveLength(0); + const toolMessage = messages.find( + (message: { metadata?: { kind?: string } }) => message.metadata?.kind === "tool_result", + ); + expect(toolMessage.contentMarkdown.length).toBeLessThanOrEqual(MAX_MESSAGE_CONTENT_CHARS); + expect(toolMessage.toolCallsJson.output.length).toBeLessThanOrEqual(MAX_TOOL_PAYLOAD_CHARS); + return { inserted: 0, updated: 0, deleted: 0, unchanged: 0 }; + }); + + await service.calculateAndSaveUsageForTask( + "proj-1", + "task-1", + "session-1", + "Initial prompt for testing", + ); + + expect(syncMessagesMock).toHaveBeenCalledTimes(1); + }); + it("handles API failure gracefully and logs an error", async () => { getFullConversationMock.mockRejectedValue(new Error("API Error")); @@ -267,6 +302,41 @@ describe("JulesUsageService", () => { expect(getFullConversationMock).toHaveBeenCalledTimes(2); }); + it("deduplicates concurrent syncs for the same session", async () => { + let resolveConversation!: (activities: JulesActivity[]) => void; + getFullConversationMock.mockImplementation(() => new Promise((resolve) => { + resolveConversation = resolve; + })); + + const first = service.syncLiveInvocation("proj-1", "task-1", "session-a", "x"); + const second = service.syncLiveInvocation("proj-1", "task-1", "session-a", "x"); + await vi.waitFor(() => expect(getFullConversationMock).toHaveBeenCalledTimes(1)); + resolveConversation([]); + await Promise.all([first, second]); + + expect(getFullConversationMock).toHaveBeenCalledTimes(1); + }); + + it("serializes full-conversation fetches across distinct sessions", async () => { + const resolvers = new Map void>(); + getFullConversationMock.mockImplementation((sessionId: string) => ( + new Promise((resolve) => { + resolvers.set(sessionId, resolve); + }) + )); + + const first = service.syncLiveInvocation("proj-1", "task-1", "session-a", "x"); + const second = service.syncLiveInvocation("proj-1", "task-2", "session-b", "y"); + await vi.waitFor(() => expect(getFullConversationMock).toHaveBeenCalledTimes(1)); + expect(getFullConversationMock).toHaveBeenLastCalledWith("session-a"); + + resolvers.get("session-a")?.([]); + await vi.waitFor(() => expect(getFullConversationMock).toHaveBeenCalledTimes(2)); + expect(getFullConversationMock).toHaveBeenLastCalledWith("session-b"); + resolvers.get("session-b")?.([]); + await Promise.all([first, second]); + }); + it("handles 404 gracefully without logging a warning during live sync", async () => { const error404 = new Error("Request failed with status code 404"); (error404 as any).status = 404;