diff --git a/src/agent-loop.ts b/src/agent-loop.ts index 469ef11..5c01d03 100644 --- a/src/agent-loop.ts +++ b/src/agent-loop.ts @@ -376,6 +376,12 @@ export async function runAgentTurn(args: { result: Awaited> toolResult: PendingToolResult }> = [] + const neverPersistToolNames = new Set( + args.tools + .list() + .filter(tool => tool.maxResultSizeChars !== undefined && !Number.isFinite(tool.maxResultSizeChars)) + .map(tool => tool.name), + ) for (const call of next.calls) { args.onToolStart?.(call.toolName, call.input) @@ -390,13 +396,14 @@ export async function runAgentTurn(args: { } args.onToolResult?.(call.toolName, result.output, !result.ok) + const tool = args.tools.find(call.toolName) const toolResult = await replaceLargeToolResult({ role: 'tool_result', toolUseId: call.id, toolName: call.toolName, content: result.output, isError: !result.ok, - }, contentReplacementState) + }, contentReplacementState, tool?.maxResultSizeChars) executedToolResults.push({ call, @@ -405,12 +412,15 @@ export async function runAgentTurn(args: { }) } - const budgetedResults = await applyToolResultBudget( + const { results: budgetedToolResults } = await applyToolResultBudget( executedToolResults.map(entry => entry.toolResult), contentReplacementState, + undefined, + neverPersistToolNames, ) + const toolResultById = new Map( - budgetedResults.results.map(result => [result.toolUseId, result]), + budgetedToolResults.map(result => [result.toolUseId, result]), ) const toolCallMessages = executedToolResults.map((entry, i) => { diff --git a/src/index.ts b/src/index.ts index 7ab8ae2..cd0f71c 100644 --- a/src/index.ts +++ b/src/index.ts @@ -91,7 +91,7 @@ async function main(): Promise { }), }, ] - const contentReplacementState = createContentReplacementState() + let contentReplacementState = createContentReplacementState({ cwd }) const contextCollapseState = createContextCollapseState() async function refreshSystemPrompt(): Promise { @@ -119,6 +119,8 @@ async function main(): Promise { } } + contentReplacementState = createContentReplacementState({ cwd, sessionId }) + await runTtyApp({ runtime, tools, diff --git a/src/session-paths.ts b/src/session-paths.ts new file mode 100644 index 0000000..75b1793 --- /dev/null +++ b/src/session-paths.ts @@ -0,0 +1,89 @@ +import path from 'node:path' +import { MINI_CODE_PROJECTS_DIR } from './config.js' + +export const SESSION_FILE_EXTENSION = '.jsonl' +export const TOOL_RESULTS_SUBDIR = 'tool-results' + +/** + * Sanitizes a single path segment so it is safe to use as a file/directory name. + * + * - Replaces any character outside `[a-zA-Z0-9._-]` with an underscore. + * - Normalizes leading dots (hidden-file prefix / relative traversal) to underscores. + * - Falls back to `fallback` when the result would be empty. + */ +export function sanitizePathSegment(value: string, fallback = 'session'): string { + const trimmed = value.trim() + const sanitized = trimmed.replace(/[^a-zA-Z0-9._-]/g, '_') + const normalized = sanitized.replace(/^\.+/, match => '_'.repeat(match.length)) + return normalized.length > 0 ? normalized : fallback +} + +/** + * Throws if `filePath` is not contained within `rootDir`. + * Used as a defense-in-depth check after all sanitization has been applied. + */ +function assertPathContained(filePath: string, rootDir: string): void { + const relative = path.relative(rootDir, filePath) + if (relative.startsWith('..') || path.isAbsolute(relative)) { + throw new Error( + `[session-paths] Path escapes expected root.\n root: ${rootDir}\n path: ${filePath}`, + ) + } +} + +function sessionPathSegment(sessionId: string): string { + return sanitizePathSegment(sessionId, 'session') +} + +export function projectDirName(cwd: string): string { + return cwd.replace(/[/\\:]+/g, '-').replace(/^-+/, '') +} + +export function projectDir(cwd: string): string { + return path.join(MINI_CODE_PROJECTS_DIR, projectDirName(cwd)) +} + +export function sessionFilePath(cwd: string, sessionId: string): string { + return path.join(projectDir(cwd), `${sessionPathSegment(sessionId)}${SESSION_FILE_EXTENSION}`) +} + +/** + * Returns the artifact directory for a session. + * + * Layout under MINI_CODE_PROJECTS_DIR: + * / ← one directory per unique CWD (projectDirName) + * .jsonl ← session event log + * / ← sessionArtifactsDir: per-session artefacts + * tool-results/ ← sessionToolResultsDir: persisted large tool outputs + * .txt ← individual result file (string output) + * .json← individual result file (structured / block output) + */ +export function sessionArtifactsDir(cwd: string, sessionId: string): string { + return path.join(projectDir(cwd), sessionPathSegment(sessionId)) +} + +export function sessionToolResultsDir(cwd: string, sessionId: string): string { + return path.join(sessionArtifactsDir(cwd, sessionId), TOOL_RESULTS_SUBDIR) +} + +/** + * Returns the file path for persisting a single tool result. + * + * @param ext - `'txt'` for plain-string outputs (default), `'json'` for + * structured content-block arrays. + * + * The path is always contained within `sessionToolResultsDir`. An assertion + * is raised if sanitization somehow fails to prevent escape (defense-in-depth). + */ +export function sessionToolResultPath( + cwd: string, + sessionId: string, + toolUseId: string, + ext: 'txt' | 'json' = 'txt', +): string { + const dir = sessionToolResultsDir(cwd, sessionId) + const filename = `${sanitizePathSegment(toolUseId, 'tool-result')}.${ext}` + const filePath = path.join(dir, filename) + assertPathContained(filePath, dir) + return filePath +} diff --git a/src/session.ts b/src/session.ts index ade0e9a..7e5c924 100644 --- a/src/session.ts +++ b/src/session.ts @@ -10,12 +10,21 @@ import { import { randomUUID } from 'node:crypto' import path from 'node:path' import { MINI_CODE_PROJECTS_DIR } from './config.js' +import { + projectDir, + sessionArtifactsDir, + sessionFilePath, +} from './session-paths.js' import type { ChatMessage } from './types.js' import { createContextCollapseState, type CollapseSpan, type ContextCollapseState, } from './compact/context-collapse.js' +import { + reconstructContentReplacementState, + type ContentReplacementState, +} from './utils/tool-result-storage.js' const MAX_TITLE_LENGTH = 60 @@ -46,17 +55,7 @@ type SessionEvent = { title?: string } -function projectDirName(cwd: string): string { - return cwd.replace(/[/\\:]+/g, '-').replace(/^-+/, '') -} - -function projectDir(cwd: string): string { - return path.join(MINI_CODE_PROJECTS_DIR, projectDirName(cwd)) -} - -function sessionFilePath(cwd: string, sessionId: string): string { - return path.join(projectDir(cwd), `${sessionId}.jsonl`) -} +type SessionEventDraft = Omit function roleToType(role: string): EventType { switch (role) { @@ -79,17 +78,38 @@ function ensureMessageId(message: ChatMessage): string { return message.id } -function wrapEvent(message: ChatMessage, sessionId: string, cwd: string, parentUuid: string | null): string { +function buildSessionEvent( + draft: SessionEventDraft, + sessionId: string, + cwd: string, + timestamp: string, + parentUuid: string | null, + logicalParentUuid?: string | null, +): SessionEvent { + return { + ...draft, + timestamp, + sessionId, + cwd, + parentUuid, + logicalParentUuid, + } +} + +function buildMessageEvent( + message: ChatMessage, + sessionId: string, + cwd: string, + parentUuid: string | null, + timestamp = new Date().toISOString(), +): SessionEvent { const uuid = ensureMessageId(message) - const event: SessionEvent = { + const event = buildSessionEvent({ type: roleToType(message.role), message, uuid, - timestamp: new Date().toISOString(), - sessionId, - cwd, parentUuid, - } + }, sessionId, cwd, timestamp, parentUuid) if (message.role === 'snip_boundary') { event.snipMetadata = { type: 'snip_boundary', @@ -100,6 +120,10 @@ function wrapEvent(message: ChatMessage, sessionId: string, cwd: string, parentU createdAt: event.timestamp, } } + return event +} + +function serializeSessionEvent(event: SessionEvent): string { return JSON.stringify(event) } @@ -199,20 +223,114 @@ async function readLastEventUuid(filePath: string): Promise { } } -async function readExistingEventUuids(filePath: string): Promise> { +/** + * Reads the session file once and returns both the set of known event UUIDs + * and the UUID of the last event. Used by `saveSession` to avoid reading the + * file twice (once to check for duplicates, once to get the parent UUID). + */ +async function readSessionFileMeta(filePath: string): Promise<{ + existingIds: Set + lastUuid: string | null +}> { try { const content = await readFile(filePath, 'utf8') - const ids = new Set() - for (const line of content.trim().split('\n').filter(Boolean)) { + const lines = content.trim().split('\n').filter(Boolean) + const existingIds = new Set() + let lastUuid: string | null = null + for (const line of lines) { const event = parseEvent(line) if (event?.uuid) { - ids.add(event.uuid) + existingIds.add(event.uuid) + lastUuid = event.uuid } } - return ids + return { existingIds, lastUuid } + } catch { + return { existingIds: new Set(), lastUuid: null } + } +} + +/** + * Reads the raw JSONL lines for a session file. + * Returns null if the file does not exist or cannot be read. + */ +async function readSessionLines(cwd: string, sessionId: string): Promise { + try { + const content = await readFile(sessionFilePath(cwd, sessionId), 'utf8') + return content.trim().split('\n').filter(Boolean) } catch { - return new Set() + return null + } +} + +/** + * Scans backward to find the index of the last compact_boundary event. + * Returns -1 when no boundary exists (i.e., the full file is active). + */ +function findLastCompactBoundaryIndex(lines: string[]): number { + for (let i = lines.length - 1; i >= 0; i--) { + if (parseEvent(lines[i]!)?.type === 'compact_boundary') return i + } + return -1 +} + +/** + * Returns parsed events that are active after the last compact boundary. + * When no boundary exists the entire file is considered active. + */ +function activeEventsFromLines(lines: string[]): SessionEvent[] { + const startLine = findLastCompactBoundaryIndex(lines) + 1 // +1 on -1 gives 0 + const events: SessionEvent[] = [] + for (let i = startLine; i < lines.length; i++) { + const event = parseEvent(lines[i]!) + if (event) events.push(event) } + return events +} + +/** + * Ensures the project directory exists and returns the session file path. + * Used as the common preamble for all append operations. + */ +async function ensureSessionDir(cwd: string, sessionId: string): Promise { + const filePath = sessionFilePath(cwd, sessionId) + await mkdir(projectDir(cwd), { recursive: true }) + return filePath +} + +/** + * Ensures the session directory exists, then reads the UUID of the last + * persisted event. Returns the file path alongside the UUID so callers + * can append without a second round-trip. + */ +async function prepareAppend( + cwd: string, + sessionId: string, +): Promise<{ filePath: string; lastUuid: string | null }> { + const filePath = await ensureSessionDir(cwd, sessionId) + const lastUuid = await readLastEventUuid(filePath) + return { filePath, lastUuid } +} + +async function appendSessionEvents( + cwd: string, + sessionId: string, + drafts: SessionEventDraft[], + timestamp = new Date().toISOString(), +): Promise { + const { filePath, lastUuid } = await prepareAppend(cwd, sessionId) + let previousUuid = lastUuid + const lines: string[] = [] + + for (const [index, draft] of drafts.entries()) { + const parentUuid = draft.parentUuid ?? (index === 0 ? null : previousUuid) + const logicalParentUuid = draft.logicalParentUuid ?? (index === 0 ? lastUuid : undefined) + const event = buildSessionEvent(draft, sessionId, cwd, timestamp, parentUuid, logicalParentUuid) + previousUuid = event.uuid + lines.push(serializeSessionEvent(event)) + } + + await appendFile(filePath, lines.join('\n') + '\n', 'utf8') } export async function saveSession( @@ -221,11 +339,10 @@ export async function saveSession( messages: ChatMessage[], alreadySavedCount: number = 0, ): Promise { - const dir = projectDir(cwd) - const filePath = sessionFilePath(cwd, sessionId) - await mkdir(dir, { recursive: true }) + const filePath = await ensureSessionDir(cwd, sessionId) - const existingIds = await readExistingEventUuids(filePath) + // Single file read: get existing IDs (for dedup) and last UUID (for parent chain). + const { existingIds, lastUuid: initialLastUuid } = await readSessionFileMeta(filePath) const nonSystemMessages = messages.slice(1) const toSave = nonSystemMessages.filter((message, index) => { if (message.id && existingIds.has(message.id)) { @@ -238,13 +355,12 @@ export async function saveSession( }) if (toSave.length === 0) return - let parentUuid = await readLastEventUuid(filePath) + let parentUuid = initialLastUuid const lines: string[] = [] for (const m of toSave) { - const line = wrapEvent(m, sessionId, cwd, parentUuid) - const parsed = JSON.parse(line) as SessionEvent - parentUuid = parsed.uuid - lines.push(line) + const event = buildMessageEvent(m, sessionId, cwd, parentUuid) + parentUuid = event.uuid + lines.push(serializeSessionEvent(event)) } await appendFile(filePath, lines.join('\n') + '\n', 'utf8') } @@ -254,35 +370,23 @@ export async function appendSnipBoundary( sessionId: string, boundaryMessage: Extract, ): Promise { - const dir = projectDir(cwd) - const filePath = sessionFilePath(cwd, sessionId) - await mkdir(dir, { recursive: true }) - - const lastUuid = await readLastEventUuid(filePath) - const now = new Date().toISOString() const uuid = ensureMessageId(boundaryMessage) - - const event: SessionEvent = { + const timestamp = new Date().toISOString() + await appendSessionEvents(cwd, sessionId, [{ type: 'snip_boundary', subtype: 'snip_boundary', message: boundaryMessage, uuid, - timestamp: now, - sessionId, - cwd, parentUuid: null, - logicalParentUuid: lastUuid, snipMetadata: { type: 'snip_boundary', removedMessageIds: boundaryMessage.removedMessageIds, removedCount: boundaryMessage.removedCount, tokensFreed: boundaryMessage.tokensFreed, - timestamp: now, - createdAt: now, + timestamp, + createdAt: timestamp, }, - } - - await appendFile(filePath, JSON.stringify(event) + '\n', 'utf8') + }], timestamp) } export async function appendContextCollapseSpan( @@ -290,26 +394,13 @@ export async function appendContextCollapseSpan( sessionId: string, span: CollapseSpan, ): Promise { - const dir = projectDir(cwd) - const filePath = sessionFilePath(cwd, sessionId) - await mkdir(dir, { recursive: true }) - - const lastUuid = await readLastEventUuid(filePath) - const now = new Date().toISOString() - - const event: SessionEvent = { + await appendSessionEvents(cwd, sessionId, [{ type: 'context_collapse', subtype: 'context_collapse', uuid: span.id, - timestamp: now, - sessionId, - cwd, parentUuid: null, - logicalParentUuid: lastUuid, contextCollapseSpan: span, - } - - await appendFile(filePath, JSON.stringify(event) + '\n', 'utf8') + }]) } export async function appendCompactBoundary( @@ -321,48 +412,33 @@ export async function appendCompactBoundary( postTokens: number, retainedMessages: ChatMessage[] = [], ): Promise { - const dir = projectDir(cwd) - const filePath = sessionFilePath(cwd, sessionId) - await mkdir(dir, { recursive: true }) - - const lastUuid = await readLastEventUuid(filePath) - const now = new Date().toISOString() - - const boundary: SessionEvent = { - type: 'compact_boundary', - subtype: 'compact_boundary', - uuid: randomUUID(), - timestamp: now, - sessionId, - cwd, - parentUuid: null, - logicalParentUuid: lastUuid, - compactMetadata: { trigger, preTokens, postTokens }, - } - - const summary: SessionEvent = { - type: 'user', - message: { role: 'user', content: summaryText }, - uuid: randomUUID(), - timestamp: now, - sessionId, - cwd, - parentUuid: boundary.uuid, - } - - const lines = [ - JSON.stringify(boundary), - JSON.stringify(summary), + const boundaryUuid: string = randomUUID() + const summaryUuid: string = randomUUID() + + const drafts: SessionEventDraft[] = [ + { + type: 'compact_boundary', + subtype: 'compact_boundary', + uuid: boundaryUuid, + parentUuid: null, + compactMetadata: { trigger, preTokens, postTokens }, + }, + { + type: 'user', + message: { role: 'user', content: summaryText }, + uuid: summaryUuid, + parentUuid: boundaryUuid, + }, ] - let parentUuid = summary.uuid + + let parentUuid: string | null = summaryUuid for (const message of retainedMessages) { - const line = wrapEvent(message, sessionId, cwd, parentUuid) - const parsed = JSON.parse(line) as SessionEvent - parentUuid = parsed.uuid - lines.push(line) + const event = buildMessageEvent(message, sessionId, cwd, parentUuid) + parentUuid = event.uuid + drafts.push(event) } - await appendFile(filePath, lines.join('\n') + '\n', 'utf8') + await appendSessionEvents(cwd, sessionId, drafts) } export async function loadSession( @@ -370,33 +446,8 @@ export async function loadSession( sessionId: string, ): Promise { try { - const content = await readFile(sessionFilePath(cwd, sessionId), 'utf8') - const lines = content.trim().split('\n').filter(Boolean) - - // Find last compact_boundary - let lastBoundaryIndex = -1 - for (let i = lines.length - 1; i >= 0; i--) { - const event = parseEvent(lines[i]!) - if (event?.type === 'compact_boundary') { - lastBoundaryIndex = i - break - } - } - - const startLine = lastBoundaryIndex >= 0 ? lastBoundaryIndex + 1 : 0 - const activeEvents: SessionEvent[] = [] - for (let i = startLine; i < lines.length; i++) { - const event = parseEvent(lines[i]!) - if (event) activeEvents.push(event) - } - - const messages: ChatMessage[] = [] - for (const event of reconstructSnippedEvents(activeEvents)) { - const msg = unwrapMessage(event) - if (msg) messages.push(msg) - } - - return messages.length > 0 ? messages : null + const snapshot = await readActiveSessionSnapshot(cwd, sessionId) + return snapshot?.messages ?? null } catch { return null } @@ -407,31 +458,49 @@ export async function loadContextCollapseState( sessionId: string, ): Promise { try { - const content = await readFile(sessionFilePath(cwd, sessionId), 'utf8') - const lines = content.trim().split('\n').filter(Boolean) + const snapshot = await readActiveSessionSnapshot(cwd, sessionId) + return snapshot?.contextCollapseState ?? null + } catch { + return null + } +} - let lastBoundaryIndex = -1 - for (let i = lines.length - 1; i >= 0; i--) { - const event = parseEvent(lines[i]!) - if (event?.type === 'compact_boundary') { - lastBoundaryIndex = i - break - } +export type SessionRuntimeState = { + messages: ChatMessage[] | null + contentReplacementState: ContentReplacementState | null + contextCollapseState: ContextCollapseState | null +} + +export async function loadSessionRuntimeState( + cwd: string, + sessionId: string, +): Promise { + try { + const snapshot = await readActiveSessionSnapshot(cwd, sessionId) + if (!snapshot?.messages) return null + return { + messages: snapshot.messages, + contentReplacementState: reconstructContentReplacementState(snapshot.messages, { cwd, sessionId }), + contextCollapseState: snapshot.contextCollapseState, } + } catch { + return null + } +} - const state = createContextCollapseState() - for (let i = lastBoundaryIndex + 1; i < lines.length; i++) { - const event = parseEvent(lines[i]!) - if (event?.type !== 'context_collapse' || !event.contextCollapseSpan) { - continue - } - if (event.contextCollapseSpan.status !== 'committed') { - continue - } - state.spans.push(event.contextCollapseSpan) +export async function loadContentReplacementState( + cwd: string, + sessionId: string, + loadedMessages?: ChatMessage[], +): Promise { + try { + if (loadedMessages) { + return reconstructContentReplacementState(loadedMessages, { cwd, sessionId }) } - return state.spans.length > 0 ? state : null + const snapshot = await readActiveSessionSnapshot(cwd, sessionId) + if (!snapshot?.messages) return null + return reconstructContentReplacementState(snapshot.messages, { cwd, sessionId }) } catch { return null } @@ -447,6 +516,12 @@ export async function clearSession( // ignore } + try { + await rm(sessionArtifactsDir(cwd, sessionId), { recursive: true, force: true }) + } catch { + // ignore + } + try { const dir = projectDir(cwd) const files = await readdir(dir) @@ -520,8 +595,8 @@ export async function renameSession( sessionId, cwd, }) - await mkdir(projectDir(cwd), { recursive: true }) - await appendFile(sessionFilePath(cwd, sessionId), event + '\n', 'utf8') + const filePath = await ensureSessionDir(cwd, sessionId) + await appendFile(filePath, event + '\n', 'utf8') return true } @@ -569,10 +644,12 @@ export async function cleanupExpiredSessions( let removed = 0 for (const name of entries.filter(e => e.endsWith('.jsonl'))) { const filePath = path.join(dir, name) + const sessionId = name.slice(0, -'.jsonl'.length) try { const stats = await stat(filePath) if (now - stats.mtime.getTime() > maxAgeMs) { await unlink(filePath) + await rm(sessionArtifactsDir(cwd, sessionId), { recursive: true, force: true }) removed += 1 } } catch { @@ -648,12 +725,15 @@ export async function loadTranscript( sessionId: string, ): Promise { try { - const content = await readFile(sessionFilePath(cwd, sessionId), 'utf8') - const lines = content.trim().split('\n').filter(Boolean) + // Transcript shows full history (including pre-compact events), so we read + // all lines rather than just the active post-boundary segment. + const allLines = await readSessionLines(cwd, sessionId) + if (!allLines) return null + const entries: PersistedTranscriptEntry[] = [] const events = reconstructSnippedEvents( - lines + allLines .map(line => parseEvent(line)) .filter((event): event is SessionEvent => Boolean(event)), ) @@ -706,3 +786,35 @@ export async function loadTranscript( return null } } + +type ActiveSessionSnapshot = { + messages: ChatMessage[] | null + contextCollapseState: ContextCollapseState | null +} + +async function readActiveSessionSnapshot( + cwd: string, + sessionId: string, +): Promise { + const lines = await readSessionLines(cwd, sessionId) + if (!lines) return null + + const activeEvents = activeEventsFromLines(lines) + const messages: ChatMessage[] = [] + for (const event of reconstructSnippedEvents(activeEvents)) { + const msg = unwrapMessage(event) + if (msg) messages.push(msg) + } + + const state = createContextCollapseState() + for (const event of activeEvents) { + if (event.type !== 'context_collapse' || !event.contextCollapseSpan) continue + if (event.contextCollapseSpan.status !== 'committed') continue + state.spans.push(event.contextCollapseSpan) + } + + return { + messages: messages.length > 0 ? messages : null, + contextCollapseState: state.spans.length > 0 ? state : null, + } +} diff --git a/src/tool.ts b/src/tool.ts index f9974c1..91ca41f 100644 --- a/src/tool.ts +++ b/src/tool.ts @@ -29,6 +29,7 @@ export type ToolDefinition = { description: string inputSchema: Record schema: z.ZodType + maxResultSizeChars?: number run(input: TInput, context: ToolContext): Promise } diff --git a/src/tools/grep-files.ts b/src/tools/grep-files.ts index 453a519..b93f2cd 100644 --- a/src/tools/grep-files.ts +++ b/src/tools/grep-files.ts @@ -26,6 +26,7 @@ export const grepFilesTool: ToolDefinition = { pattern: z.string().min(1), path: z.string().optional(), }), + maxResultSizeChars: 20_000, async run(input, context) { const args = ['-n', '--no-heading', input.pattern] if (input.path) { diff --git a/src/tools/read-file.ts b/src/tools/read-file.ts index 3e81530..4678a88 100644 --- a/src/tools/read-file.ts +++ b/src/tools/read-file.ts @@ -30,6 +30,7 @@ export const readFileTool: ToolDefinition = { offset: z.number().int().min(0).optional(), limit: z.number().int().min(1).max(MAX_READ_LIMIT).optional(), }), + maxResultSizeChars: Number.POSITIVE_INFINITY, async run(input, context) { const target = await resolveToolPath(context, input.path, 'read') const content = await readFile(target, 'utf8') diff --git a/src/tools/run-command.ts b/src/tools/run-command.ts index 69ff5dc..46cd7fc 100644 --- a/src/tools/run-command.ts +++ b/src/tools/run-command.ts @@ -172,6 +172,7 @@ export const runCommandTool: ToolDefinition = { args: z.array(z.string()).optional(), cwd: z.string().optional(), }), + maxResultSizeChars: 30_000, async run(input, context) { const effectiveCwd = input.cwd ? await resolveToolPath(context, input.cwd, 'list') diff --git a/src/tty-app.ts b/src/tty-app.ts index 3300bd0..c26db3d 100644 --- a/src/tty-app.ts +++ b/src/tty-app.ts @@ -27,7 +27,7 @@ import { appendSnipBoundary, appendContextCollapseSpan, loadTranscript, - loadContextCollapseState, + loadSessionRuntimeState, forkSession, cleanupExpiredSessions, listAllProjects, @@ -1002,8 +1002,15 @@ async function resumeSession( body: `Session ${sessionId} resumed (${loaded.length} messages loaded).`, }) args.alreadySavedCount = loaded.length + const resumedState = await loadSessionRuntimeState(args.cwd, sessionId) + args.contentReplacementState = + resumedState?.contentReplacementState ?? + createContentReplacementState({ + cwd: args.cwd, + sessionId, + }) args.contextCollapseState = - await loadContextCollapseState(args.cwd, sessionId) ?? + resumedState?.contextCollapseState ?? createContextCollapseState() state.transcriptScrollOffset = 0 } @@ -1281,6 +1288,10 @@ async function handleInput( if (input === '/new') { args.sessionId = crypto.randomUUID().slice(0, 8) args.alreadySavedCount = 0 + args.contentReplacementState = createContentReplacementState({ + cwd: args.cwd, + sessionId: args.sessionId, + }) args.contextCollapseState = createContextCollapseState() state.transcript = [] args.messages.length = 0 @@ -1304,6 +1315,10 @@ async function handleInput( } args.sessionId = newId args.alreadySavedCount = args.messages.length - 1 + args.contentReplacementState = createContentReplacementState({ + cwd: args.cwd, + sessionId: newId, + }) args.contextCollapseState = createContextCollapseState() state.transcriptScrollOffset = 0 pushTranscriptEntry(state, { @@ -1682,7 +1697,10 @@ export async function runTtyApp(args: TtyAppArgs): Promise { const permissionArgs: TtyAppArgs = { ...args, contentReplacementState: - args.contentReplacementState ?? createContentReplacementState(), + args.contentReplacementState ?? createContentReplacementState({ + cwd: args.cwd, + sessionId: args.sessionId, + }), contextCollapseState: args.contextCollapseState ?? createContextCollapseState(), permissions: new PermissionManager( diff --git a/src/utils/tool-result-storage.ts b/src/utils/tool-result-storage.ts index 79eabbf..2dc5bdb 100644 --- a/src/utils/tool-result-storage.ts +++ b/src/utils/tool-result-storage.ts @@ -1,20 +1,47 @@ import { mkdir, writeFile } from 'node:fs/promises' import path from 'node:path' import { randomUUID } from 'node:crypto' -import { MINI_CODE_DIR } from '../config.js' +import { + TOOL_RESULTS_SUBDIR, + sessionToolResultPath, + sessionToolResultsDir, +} from '../session-paths.js' import type { ChatMessage } from '../types.js' -export const TOOL_RESULTS_SUBDIR = 'tool-results' +export { TOOL_RESULTS_SUBDIR } + export const PERSISTED_OUTPUT_TAG = '' export const PERSISTED_OUTPUT_CLOSING_TAG = '' export const DEFAULT_MAX_RESULT_SIZE_CHARS = 50_000 export const MAX_TOOL_RESULTS_PER_BATCH_CHARS = 200_000 +/** Default preview length for most tools. */ export const PREVIEW_SIZE_CHARS = 2_000 +/** + * Larger preview for shell/command tools whose output tends to be long and + * structured (e.g. build output, test runners, grep results). + */ +export const SHELL_PREVIEW_SIZE_CHARS = 5_000 + +/** + * Tool names whose output is expected to be long and command-like. + * These receive a larger in-context preview when their result is persisted. + */ +export const SHELL_TOOL_NAMES: ReadonlySet = new Set([ + 'run_command', + 'bash', + 'local_bash', +]) + +export type ToolResultStorageContext = { + cwd: string + sessionId: string +} export type ContentReplacementState = { seenIds: Set replacements: Map + storageContext: ToolResultStorageContext } export type ToolResultReplacementRecord = { @@ -31,48 +58,60 @@ type ReplacementCandidate = { size: number } -export function createContentReplacementState(): ContentReplacementState { +const fallbackSessionId = `ephemeral-${randomUUID()}` + +function resolveStorageContext( + storageContext?: Partial, +): ToolResultStorageContext { return { - seenIds: new Set(), - replacements: new Map(), + cwd: storageContext?.cwd ?? process.cwd(), + sessionId: storageContext?.sessionId?.trim() || fallbackSessionId, } } -function sanitizePathSegment(value: string): string { - const sanitized = value.replace(/[^a-zA-Z0-9._-]/g, '_') - return sanitized.length > 0 ? sanitized : randomUUID() +export function createContentReplacementState( + storageContext?: Partial, +): ContentReplacementState { + return { + seenIds: new Set(), + replacements: new Map(), + storageContext: resolveStorageContext(storageContext), + } } -const sessionId = sanitizePathSegment(randomUUID()) - -function getToolResultsDir(): string { - return path.join(MINI_CODE_DIR, TOOL_RESULTS_SUBDIR, sessionId) +export function getToolResultsDir( + storageContext?: Partial, +): string { + const resolved = resolveStorageContext(storageContext) + return sessionToolResultsDir(resolved.cwd, resolved.sessionId) } -function getToolResultPath(toolUseId: string): string { - const dir = getToolResultsDir() - const filepath = path.resolve(dir, `${sanitizePathSegment(toolUseId)}.txt`) - const relative = path.relative(dir, filepath) - if (relative.startsWith('..') || path.isAbsolute(relative)) { - return path.join(dir, `${randomUUID()}.txt`) - } - return filepath +function getToolResultPath( + toolUseId: string, + storageContext?: Partial, + ext: 'txt' | 'json' = 'txt', +): string { + const resolved = resolveStorageContext(storageContext) + return sessionToolResultPath(resolved.cwd, resolved.sessionId, toolUseId, ext) } function isAlreadyPersistedOutput(content: string): boolean { return content.startsWith(PERSISTED_OUTPUT_TAG) } -function generatePreview(content: string): { preview: string; hasMore: boolean } { - if (content.length <= PREVIEW_SIZE_CHARS) { +function generatePreview( + content: string, + previewSize = PREVIEW_SIZE_CHARS, +): { preview: string; hasMore: boolean } { + if (content.length <= previewSize) { return { preview: content, hasMore: false } } - const truncated = content.slice(0, PREVIEW_SIZE_CHARS) + const truncated = content.slice(0, previewSize) const lastNewline = truncated.lastIndexOf('\n') - const cutPoint = lastNewline > PREVIEW_SIZE_CHARS * 0.5 + const cutPoint = lastNewline > previewSize * 0.5 ? lastNewline - : PREVIEW_SIZE_CHARS + : previewSize return { preview: content.slice(0, cutPoint), @@ -86,16 +125,47 @@ function formatChars(chars: number): string { return `${chars} chars` } +/** + * A structured content block as returned by the Anthropic API (or similar). + * Blocks with `type === 'text'` contribute to the in-context preview; + * others are preserved in the JSON file but not shown in the preview. + */ +type ContentBlock = { type: string; text?: string; [key: string]: unknown } + +function isContentBlockArray(value: unknown): value is ContentBlock[] { + return ( + Array.isArray(value) && + value.length > 0 && + typeof (value as unknown[])[0] === 'object' && + (value as ContentBlock[])[0] !== null && + 'type' in (value as ContentBlock[])[0] + ) +} + +/** + * Extracts plain text from a content-block array for preview / normalization. + * Non-text blocks are represented as `[]` placeholders. + */ +function textFromContentBlocks(blocks: ContentBlock[]): string { + return blocks + .map(b => (typeof b.text === 'string' ? b.text : `[${b.type}]`)) + .join('\n') +} + export function normalizeToolResultContent(content: unknown): string { if (content == null) return '' - return typeof content === 'string' ? content : String(content) + if (typeof content === 'string') return content + if (isContentBlockArray(content)) return textFromContentBlocks(content) + return String(content) } async function persistToolResult( content: string, toolUseId: string, + storageContext?: Partial, + previewSize = PREVIEW_SIZE_CHARS, ): Promise<{ filepath: string; originalSize: number; preview: string; hasMore: boolean } | null> { - const filepath = getToolResultPath(toolUseId) + const filepath = getToolResultPath(toolUseId, storageContext, 'txt') try { await mkdir(path.dirname(filepath), { recursive: true }) await writeFile(filepath, content, { encoding: 'utf8', flag: 'wx' }) @@ -108,7 +178,7 @@ async function persistToolResult( } } - const { preview, hasMore } = generatePreview(content) + const { preview, hasMore } = generatePreview(content, previewSize) return { filepath, originalSize: content.length, @@ -117,6 +187,41 @@ async function persistToolResult( } } +/** + * Persists a structured content-block array as JSON. + * Returns the same shape as `persistToolResult` but uses the `.json` extension + * and serializes the raw blocks alongside a text-based preview. + */ +async function persistStructuredToolResult( + blocks: ContentBlock[], + toolUseId: string, + storageContext?: Partial, + previewSize = PREVIEW_SIZE_CHARS, +): Promise<{ filepath: string; originalSize: number; preview: string; hasMore: boolean } | null> { + const serialized = JSON.stringify(blocks, null, 2) + const filepath = getToolResultPath(toolUseId, storageContext, 'json') + try { + await mkdir(path.dirname(filepath), { recursive: true }) + await writeFile(filepath, serialized, { encoding: 'utf8', flag: 'wx' }) + } catch (error) { + const code = typeof error === 'object' && error !== null && 'code' in error + ? (error as { code?: unknown }).code + : undefined + if (code !== 'EEXIST') { + return null + } + } + + const textContent = textFromContentBlocks(blocks) + const { preview, hasMore } = generatePreview(textContent, previewSize) + return { + filepath, + originalSize: serialized.length, + preview, + hasMore, + } +} + function buildPersistedToolResultMessage(result: { filepath: string originalSize: number @@ -148,6 +253,9 @@ export async function replaceLargeToolResult( typeof stateOrThreshold === 'number' ? undefined : stateOrThreshold const threshold = typeof stateOrThreshold === 'number' ? stateOrThreshold : maybeThreshold + + // Detect structured (content-block array) before string normalization + const isStructured = isContentBlockArray(result.content) const content = normalizeToolResultContent(result.content) const normalizedResult: PendingToolResult = { ...result, @@ -176,11 +284,28 @@ export async function replaceLargeToolResult( return normalizedResult } - if (content.length <= threshold) { + if (!Number.isFinite(threshold) || content.length <= threshold) { return normalizedResult } - const persisted = await persistToolResult(content, result.toolUseId) + const previewSize = SHELL_TOOL_NAMES.has(result.toolName) + ? SHELL_PREVIEW_SIZE_CHARS + : PREVIEW_SIZE_CHARS + + const persisted = isStructured + ? await persistStructuredToolResult( + result.content as ContentBlock[], + result.toolUseId, + state?.storageContext, + previewSize, + ) + : await persistToolResult( + content, + result.toolUseId, + state?.storageContext, + previewSize, + ) + if (!persisted) { return normalizedResult } @@ -195,10 +320,28 @@ export async function replaceLargeToolResult( } } +export function reconstructContentReplacementState( + messages: ChatMessage[], + storageContext?: Partial, +): ContentReplacementState { + const state = createContentReplacementState(storageContext) + + for (const message of messages) { + if (message.role !== 'tool_result') continue + state.seenIds.add(message.toolUseId) + if (isAlreadyPersistedOutput(message.content)) { + state.replacements.set(message.toolUseId, message.content) + } + } + + return state +} + export async function applyToolResultBudget( results: PendingToolResult[], state: ContentReplacementState, limit = MAX_TOOL_RESULTS_PER_BATCH_CHARS, + skipToolNames: ReadonlySet = new Set(), ): Promise<{ results: PendingToolResult[] newlyReplaced: ToolResultReplacementRecord[] @@ -225,6 +368,12 @@ export async function applyToolResultBudget( continue } + if (skipToolNames.has(result.toolName)) { + state.seenIds.add(result.toolUseId) + visibleSize += content.length + continue + } + if (content.trim().length === 0) { state.seenIds.add(result.toolUseId) continue @@ -255,7 +404,11 @@ export async function applyToolResultBudget( for (const candidate of sortedFreshCandidates) { if (visibleSize <= limit) break - const persisted = await persistToolResult(candidate.content, candidate.toolUseId) + const persisted = await persistToolResult( + candidate.content, + candidate.toolUseId, + state.storageContext, + ) state.seenIds.add(candidate.toolUseId) if (!persisted) { continue diff --git a/test/run-tests.mjs b/test/run-tests.mjs index b0cd5a1..8769308 100644 --- a/test/run-tests.mjs +++ b/test/run-tests.mjs @@ -1,5 +1,6 @@ -import { readdir } from 'node:fs/promises' +import { mkdtemp, readdir, rm } from 'node:fs/promises' import { spawn } from 'node:child_process' +import os from 'node:os' import path from 'node:path' import { fileURLToPath } from 'node:url' @@ -16,17 +17,63 @@ if (testFiles.length === 0) { process.exit(1) } -const child = spawn( - process.execPath, - ['--import', 'tsx', '--test', ...testFiles], - { stdio: 'inherit' }, -) +let activeChild = null -child.on('exit', (code, signal) => { - if (signal) { - process.kill(process.pid, signal) - return +const forwardedSignals = ['SIGINT', 'SIGTERM'] +const forwardSignal = signal => { + if (activeChild && !activeChild.killed) { + activeChild.kill(signal) } +} + +for (const signal of forwardedSignals) { + process.on(signal, forwardSignal) +} + +async function runTestFile(filePath) { + const miniCodeHome = await mkdtemp(path.join(os.tmpdir(), 'minicode-test-home-')) + + try { + const exit = await new Promise(resolve => { + activeChild = spawn( + process.execPath, + ['--import', 'tsx', '--test', filePath], + { + stdio: 'inherit', + env: { + ...process.env, + MINI_CODE_HOME: miniCodeHome, + }, + }, + ) + + activeChild.once('close', (code, signal) => { + resolve({ code, signal }) + }) + }) + + if (exit.signal) { + process.kill(process.pid, exit.signal) + return false + } + + return (exit.code ?? 1) === 0 + } finally { + activeChild = null + await rm(miniCodeHome, { recursive: true, force: true }) + } +} + +let allPassed = true +for (const filePath of testFiles) { + const ok = await runTestFile(filePath) + if (!ok) { + allPassed = false + } +} + +for (const signal of forwardedSignals) { + process.off(signal, forwardSignal) +} - process.exit(code ?? 1) -}) +process.exit(allPassed ? 0 : 1) diff --git a/test/session.test.ts b/test/session.test.ts index 627c8ea..ba5bb72 100644 --- a/test/session.test.ts +++ b/test/session.test.ts @@ -1,6 +1,6 @@ import { describe, it, beforeEach, afterEach } from 'node:test' import assert from 'node:assert/strict' -import { mkdir, rm, readFile } from 'node:fs/promises' +import { mkdir, rm, readFile, writeFile } from 'node:fs/promises' import path from 'node:path' import os from 'node:os' import { @@ -13,18 +13,25 @@ import { appendSnipBoundary, appendContextCollapseSpan, loadContextCollapseState, + loadContentReplacementState, loadTranscript, + loadSessionRuntimeState, forkSession, cleanupExpiredSessions, listAllProjects, } from '../src/session.js' import { MINI_CODE_PROJECTS_DIR } from '../src/config.js' +import { projectDirName, sessionFilePath, sessionToolResultsDir } from '../src/session-paths.js' import type { AgentStep, ChatMessage, ModelAdapter } from '../src/types.js' import type { ContextStats } from '../src/utils/token-estimator.js' import { estimateMessagesTokens, tokenCountWithEstimation, } from '../src/utils/token-estimator.js' +import { + createContentReplacementState, + replaceLargeToolResult, +} from '../src/utils/tool-result-storage.js' import { snipCompactConversation } from '../src/compact/snipCompact.js' import { compactConversation } from '../src/compact/compact.js' import type { CollapseSpan } from '../src/compact/context-collapse.js' @@ -42,10 +49,6 @@ function makeMessages(count: number): ChatMessage[] { return messages } -function projectDirName(cwd: string): string { - return cwd.replace(/[/\\:]+/g, '-').replace(/^-+/, '') -} - function contextStats(messages: ChatMessage[], effectiveInput = 20_000): ContextStats { const accounting = tokenCountWithEstimation(messages) const utilization = accounting.totalTokens / effectiveInput @@ -135,13 +138,17 @@ describe('session persistence', () => { assert.equal(loaded, null) }) - it('clears an existing session', async () => { + it('clears an existing session and removes session artifacts', async () => { const cwd = path.join(testDir, 'project-b') + const toolResultsDir = sessionToolResultsDir(cwd, 'sess0001') await saveSession(cwd, 'sess0001', makeMessages(1)) + await mkdir(toolResultsDir, { recursive: true }) + await writeFile(path.join(toolResultsDir, 'artifact.txt'), 'temporary persisted output', 'utf8') assert.notEqual(await loadSession(cwd, 'sess0001'), null) await clearSession(cwd, 'sess0001') assert.equal(await loadSession(cwd, 'sess0001'), null) + await assert.rejects(readFile(path.join(toolResultsDir, 'artifact.txt'), 'utf8')) }) it('clearSession does not throw for nonexistent session', async () => { @@ -208,6 +215,24 @@ describe('session persistence', () => { assert.ok(content.length > 0) }) + it('keeps unsafe session ids inside the project directory', async () => { + const cwd = path.join(testDir, 'unsafe-session-id') + const sessionId = '../..\\nested/session' + await saveSession(cwd, sessionId, makeMessages(1)) + + const projectRoot = path.join(MINI_CODE_PROJECTS_DIR, projectDirName(cwd)) + const filePath = sessionFilePath(cwd, sessionId) + const fileRelative = path.relative(projectRoot, filePath) + const toolResultsDir = sessionToolResultsDir(cwd, sessionId) + const toolResultsRelative = path.relative(projectRoot, toolResultsDir) + + assert.ok(!fileRelative.startsWith('..')) + assert.ok(!path.isAbsolute(fileRelative)) + assert.ok(!toolResultsRelative.startsWith('..')) + assert.ok(!path.isAbsolute(toolResultsRelative)) + assert.ok((await readFile(filePath, 'utf8')).length > 0) + }) + it('appends only new messages with alreadySavedCount', async () => { const cwd = path.join(testDir, 'append-test') const msgs1 = makeMessages(1) @@ -553,6 +578,38 @@ describe('session persistence', () => { assert.equal(collapseState!.enabled, true) }) + it('loadSessionRuntimeState restores messages, replacement state, and collapse state together', async () => { + const cwd = path.join(testDir, 'runtime-state') + const sessionId = 'runtime001' + const messages = makeMessages(1) + + await saveSession(cwd, sessionId, messages) + const loadedBefore = await loadSession(cwd, sessionId) + assert.ok(loadedBefore) + + const span: CollapseSpan = { + id: 'runtime-collapse-span', + startMessageId: loadedBefore![0]!.id!, + endMessageId: loadedBefore![0]!.id!, + messageIds: [loadedBefore![0]!.id!], + summary: 'Runtime collapse state', + tokensBefore: 800, + tokensAfter: 120, + status: 'committed', + createdAt: 123, + reason: 'context_pressure', + } + + await appendContextCollapseSpan(cwd, sessionId, span) + + const state = await loadSessionRuntimeState(cwd, sessionId) + assert.ok(state) + assert.ok(state!.messages) + assert.ok(state!.contentReplacementState) + assert.ok(state!.contextCollapseState) + assert.deepEqual(state!.contextCollapseState!.spans, [span]) + }) + it('saveSession writes parentUuid chain linking consecutive events', async () => { const cwd = path.join(testDir, 'parent-chain') await saveSession(cwd, 'chain001', makeMessages(2), 0) @@ -748,6 +805,10 @@ describe('session persistence', () => { const pdir = path.join(MINI_CODE_PROJECTS_DIR, projectDirName(cwd)) const oldPath = path.join(pdir, 'old001.jsonl') await saveSession(cwd, 'old001', makeMessages(1)) + const oldToolResultsDir = sessionToolResultsDir(cwd, 'old001') + await mkdir(oldToolResultsDir, { recursive: true }) + await writeFile(path.join(oldToolResultsDir, 'artifact.txt'), 'temporary persisted output', 'utf8') + // Set mtime to 31 days ago const oldTime = Date.now() - 31 * 24 * 60 * 60 * 1000 const { utimes } = await import('node:fs/promises') @@ -757,10 +818,63 @@ describe('session persistence', () => { assert.equal(removed, 1) assert.equal(await loadSession(cwd, 'old001'), null) + await assert.rejects(readFile(path.join(oldToolResultsDir, 'artifact.txt'), 'utf8')) assert.notEqual(await loadSession(cwd, 'recent001'), null) assert.notEqual(await loadSession(cwd, 'recent002'), null) }) + it('loadContentReplacementState reconstructs persisted replacement decisions from saved tool_result content alone', async () => { + const cwd = path.join(testDir, 'replacement-state') + const sessionId = 'resume001' + const original = [ + '$ npm test', + 'src/example.test.ts:12: failing snapshot', + 'note: persisted output preview should remain stable on resume', + '', + ].join('\n') + 'L'.repeat(50_001) + const replacementState = createContentReplacementState({ cwd, sessionId }) + const replacementResult = await replaceLargeToolResult({ + role: 'tool_result', + toolUseId: 'big001', + toolName: 'run_command', + content: original, + isError: false, + }, replacementState, 50_000) + + const messages: ChatMessage[] = [ + { role: 'system', content: 'sys' }, + { role: 'assistant_tool_call', toolUseId: 'big001', toolName: 'run_command', input: { command: 'npm test' } }, + replacementResult, + { role: 'assistant_tool_call', toolUseId: 'small001', toolName: 'read_file', input: { path: 'README.md' } }, + { role: 'tool_result', toolUseId: 'small001', toolName: 'read_file', content: 'README preview', isError: false }, + ] + await saveSession(cwd, sessionId, messages) + + const rawSession = await readFile(sessionFilePath(cwd, sessionId), 'utf8') + assert.equal(rawSession.includes('"type":"content_replacement"'), false) + + const loaded = await loadSession(cwd, sessionId) + assert.ok(loaded) + const state = await loadContentReplacementState(cwd, sessionId) + assert.ok(state) + assert.ok(state!.seenIds.has('big001')) + assert.ok(state!.seenIds.has('small001')) + assert.equal(state!.replacements.get('big001'), replacementResult.content) + assert.equal(state!.replacements.has('small001'), false) + + const stateFromLoadedMessages = await loadContentReplacementState(cwd, sessionId, loaded!) + assert.deepEqual(stateFromLoadedMessages, state) + + const replayed = await replaceLargeToolResult({ + role: 'tool_result', + toolUseId: 'big001', + toolName: 'run_command', + content: original, + isError: false, + }, state!) + assert.equal(replayed.content, replacementResult.content) + }) + it('listAllProjects returns all projects with sessions', async () => { const cwdA = path.join(testDir, 'proj-a') const cwdB = path.join(testDir, 'proj-b') diff --git a/test/tool-result-storage.test.ts b/test/tool-result-storage.test.ts index ad8fd24..c06e429 100644 --- a/test/tool-result-storage.test.ts +++ b/test/tool-result-storage.test.ts @@ -3,14 +3,17 @@ import path from 'node:path' import { describe, it } from 'node:test' import assert from 'node:assert/strict' import type { ChatMessage } from '../src/types.js' -import { MINI_CODE_DIR } from '../src/config.js' +import { MINI_CODE_PROJECTS_DIR } from '../src/config.js' +import { projectDirName } from '../src/session-paths.js' import { MAX_TOOL_RESULTS_PER_BATCH_CHARS, PERSISTED_OUTPUT_TAG, - PREVIEW_SIZE_CHARS, + SHELL_PREVIEW_SIZE_CHARS, + SHELL_TOOL_NAMES, TOOL_RESULTS_SUBDIR, applyToolResultBudget, createContentReplacementState, + normalizeToolResultContent, replaceLargeToolResult, } from '../src/utils/tool-result-storage.js' @@ -27,6 +30,34 @@ function toolResult( } } +function createStorageState(sessionId = 'test-session') { + return createContentReplacementState({ + cwd: path.join(process.cwd(), 'tool-result-storage-fixture'), + sessionId, + }) +} + +function sizedOutput(label: string, size: number): string { + const header = [ + `$ ${label}`, + `src/${label}.ts:10: simulated tool output`, + 'note: this fixture keeps an exact character budget for threshold tests', + '', + ].join('\n') + const fillerLine = `[${label}] deterministic filler 0123456789 abcdefghijklmnopqrstuvwxyz\n` + let output = header + + while (output.length + fillerLine.length <= size) { + output += fillerLine + } + + if (output.length < size) { + output += '.'.repeat(size - output.length) + } + + return output +} + function extractSavedPath(content: string): string { const match = content.match(/Full output saved to: (.+)\n/) assert.ok(match, 'replacement should include saved path') @@ -41,8 +72,8 @@ function persistedMessages( describe('tool result replacement', () => { it('persists a single oversized tool result and preserves the full original output on disk', async () => { - const state = createContentReplacementState() - const original = 'x'.repeat(50_001) + const state = createStorageState() + const original = sizedOutput('single-large-output', 50_001) const result = await replaceLargeToolResult( toolResult('single-large', original), state, @@ -50,7 +81,7 @@ describe('tool result replacement', () => { assert.ok(result.content.startsWith(PERSISTED_OUTPUT_TAG)) assert.ok(result.content.endsWith('')) - assert.ok(result.content.includes(original.slice(0, PREVIEW_SIZE_CHARS))) + assert.ok(result.content.includes(original.slice(0, 200))) assert.equal(state.replacements.get('single-large'), result.content) const savedPath = extractSavedPath(result.content) @@ -59,8 +90,8 @@ describe('tool result replacement', () => { }) it('honors the 50_000 single-result boundary', async () => { - const exact = 'x'.repeat(50_000) - const over = 'y'.repeat(50_001) + const exact = sizedOutput('exact-single-boundary', 50_000) + const over = sizedOutput('over-single-boundary', 50_001) const exactResult = await replaceLargeToolResult(toolResult('exact-single', exact)) const overResult = await replaceLargeToolResult(toolResult('over-single', over)) @@ -87,19 +118,19 @@ describe('tool result replacement', () => { }) it('honors the 200_000 batch boundary', async () => { - const state = createContentReplacementState() + const state = createStorageState() const exact = await applyToolResultBudget([ - toolResult('batch-exact-a', 'a'.repeat(100_000)), - toolResult('batch-exact-b', 'b'.repeat(100_000)), + toolResult('batch-exact-a', sizedOutput('batch-exact-a', 100_000)), + toolResult('batch-exact-b', sizedOutput('batch-exact-b', 100_000)), ] as Extract[], state) assert.equal(exact.newlyReplaced.length, 0) assert.equal(persistedMessages(exact.results).length, 0) - const overState = createContentReplacementState() + const overState = createStorageState('batch-over') const over = await applyToolResultBudget([ - toolResult('batch-over-a', 'a'.repeat(100_001)), - toolResult('batch-over-b', 'b'.repeat(100_000)), + toolResult('batch-over-a', sizedOutput('batch-over-a', 100_001)), + toolResult('batch-over-b', sizedOutput('batch-over-b', 100_000)), ] as Extract[], overState) assert.equal(over.newlyReplaced.length, 1) @@ -107,11 +138,11 @@ describe('tool result replacement', () => { }) it('replaces the largest fresh batch results until visible content is under budget', async () => { - const state = createContentReplacementState() + const state = createStorageState() const result = await applyToolResultBudget([ - toolResult('largest-a', 'a'.repeat(100_000)), - toolResult('largest-b', 'b'.repeat(100_000)), - toolResult('largest-c', 'c'.repeat(100_000)), + toolResult('largest-a', sizedOutput('largest-a', 100_000)), + toolResult('largest-b', sizedOutput('largest-b', 100_000)), + toolResult('largest-c', sizedOutput('largest-c', 100_000)), ] as Extract[], state) const replacedIds = persistedMessages(result.results).map(message => message.toolUseId) @@ -126,11 +157,11 @@ describe('tool result replacement', () => { }) it('uses stable tie-breaking when same-size batch results need replacement', async () => { - const state = createContentReplacementState() + const state = createStorageState() const result = await applyToolResultBudget([ - toolResult('tie-c', 'c'.repeat(100_000)), - toolResult('tie-a', 'a'.repeat(100_000)), - toolResult('tie-b', 'b'.repeat(100_000)), + toolResult('tie-c', sizedOutput('tie-c', 100_000)), + toolResult('tie-a', sizedOutput('tie-a', 100_000)), + toolResult('tie-b', sizedOutput('tie-b', 100_000)), ] as Extract[], state) const replacedIds = persistedMessages(result.results).map(message => message.toolUseId) @@ -138,8 +169,8 @@ describe('tool result replacement', () => { }) it('replays single-result replacements byte-identically without regenerating them', async () => { - const state = createContentReplacementState() - const original = 'x'.repeat(50_001) + const state = createStorageState() + const original = sizedOutput('single-large-output', 50_001) const first = await replaceLargeToolResult( toolResult('single-replay', original), state, @@ -155,10 +186,10 @@ describe('tool result replacement', () => { }) it('replays batch replacements byte-identically without new replacement records', async () => { - const state = createContentReplacementState() + const state = createStorageState() const inputs = [ - toolResult('batch-replay-a', 'a'.repeat(100_001)), - toolResult('batch-replay-b', 'b'.repeat(100_000)), + toolResult('batch-replay-a', sizedOutput('batch-replay-a', 100_001)), + toolResult('batch-replay-b', sizedOutput('batch-replay-b', 100_000)), ] as Extract[] const first = await applyToolResultBudget(inputs, state) @@ -173,7 +204,7 @@ describe('tool result replacement', () => { }) it('does not re-persist content that is already a persisted-output replacement', async () => { - const state = createContentReplacementState() + const state = createStorageState() const replacement = [ PERSISTED_OUTPUT_TAG, 'Output too large. Full output saved to: /tmp/example.txt', @@ -190,17 +221,160 @@ describe('tool result replacement', () => { assert.equal(result.newlyReplaced.length, 0) }) + it('stores persisted outputs under the project/session tool-results directory', async () => { + const cwd = path.join(process.cwd(), 'tool-result-storage-fixture') + const sessionId = 'session-path-check' + const state = createContentReplacementState({ cwd, sessionId }) + const original = sizedOutput('project-session-path-check', 50_001) + const result = await replaceLargeToolResult(toolResult('path-check', original), state) + const savedPath = extractSavedPath(result.content) + const expectedRoot = path.join( + MINI_CODE_PROJECTS_DIR, + projectDirName(cwd), + sessionId, + TOOL_RESULTS_SUBDIR, + ) + const relative = path.relative(expectedRoot, savedPath) + + assert.ok(!relative.startsWith('..')) + assert.ok(!path.isAbsolute(relative)) + assert.equal(await readFile(savedPath, 'utf8'), original) + }) + it('keeps persisted paths under the tool-results directory for unsafe toolUseIds', async () => { - const original = 'z'.repeat(50_001) + const state = createStorageState('unsafe-id-session') + const original = sizedOutput('unsafe-tool-use-id', 50_001) const result = await replaceLargeToolResult( toolResult('../..\\evil/name', original), + state, ) const savedPath = extractSavedPath(result.content) - const root = path.join(MINI_CODE_DIR, TOOL_RESULTS_SUBDIR) + const root = path.join( + MINI_CODE_PROJECTS_DIR, + projectDirName(path.join(process.cwd(), 'tool-result-storage-fixture')), + 'unsafe-id-session', + TOOL_RESULTS_SUBDIR, + ) const relative = path.relative(root, savedPath) assert.ok(!relative.startsWith('..')) assert.ok(!path.isAbsolute(relative)) assert.equal(await readFile(savedPath, 'utf8'), original) }) + + it('respects never-persist thresholds and skip lists for bounded tools like read_file', async () => { + const state = createStorageState('never-persist') + const original = sizedOutput('read-file-never-persist', 120_000) + + const single = await replaceLargeToolResult( + { + role: 'tool_result', + toolUseId: 'read-large', + toolName: 'read_file', + content: original, + isError: false, + }, + state, + Number.POSITIVE_INFINITY, + ) + assert.equal(single.content, original) + + const batch = await applyToolResultBudget([ + { + role: 'tool_result', + toolUseId: 'read-large-a', + toolName: 'read_file', + content: original, + isError: false, + }, + { + role: 'tool_result', + toolUseId: 'read-large-b', + toolName: 'read_file', + content: original, + isError: false, + }, + ], state, 10_000, new Set(['read_file'])) + + assert.equal(batch.newlyReplaced.length, 0) + assert.equal(persistedMessages(batch.results).length, 0) + assert.equal(batch.results[0]?.content, original) + assert.equal(batch.results[1]?.content, original) + }) + + it('uses larger preview for shell tools and saves with .txt extension', async () => { + assert.ok(SHELL_TOOL_NAMES.has('run_command'), 'run_command should be a shell tool') + + const state = createStorageState('shell-preview') + // Build content larger than SHELL_PREVIEW_SIZE_CHARS to verify the larger preview is used + const lineCount = 500 + const lines = Array.from({ length: lineCount }, (_, i) => `line ${i}: build output entry`) + const original = lines.join('\n') // well over 50_001 chars + + // Pad to exceed the persist threshold + const padded = original + '\n' + 'x'.repeat(50_001 - original.length + 1) + + const result = await replaceLargeToolResult( + { + role: 'tool_result', + toolUseId: 'shell-preview-check', + toolName: 'run_command', + content: padded, + isError: false, + }, + state, + ) + + assert.ok(result.content.startsWith(PERSISTED_OUTPUT_TAG)) + // The preview for shell tools is SHELL_PREVIEW_SIZE_CHARS (5K) not the default 2K + const previewLines = result.content.split('\n') + const previewContent = previewLines.slice(4, -2).join('\n') // skip header + closing tag + assert.ok(previewContent.length >= SHELL_PREVIEW_SIZE_CHARS * 0.5, + `expected shell preview >= ${SHELL_PREVIEW_SIZE_CHARS * 0.5} chars, got ${previewContent.length}`) + + const savedPath = extractSavedPath(result.content) + assert.ok(savedPath.endsWith('.txt'), 'shell tool output should be saved as .txt') + assert.equal(await readFile(savedPath, 'utf8'), padded) + }) + + it('normalizes content-block arrays to text via normalizeToolResultContent', () => { + const blocks = [ + { type: 'text', text: 'first line' }, + { type: 'text', text: 'second line' }, + { type: 'image', source: { url: 'data:...' } }, + ] + const normalized = normalizeToolResultContent(blocks) + assert.ok(normalized.includes('first line')) + assert.ok(normalized.includes('second line')) + assert.ok(normalized.includes('[image]'), 'non-text blocks should become [type] placeholders') + }) + + it('persists structured content-block arrays as .json and saves raw blocks', async () => { + const state = createStorageState('structured-json') + const blocks = Array.from({ length: 3_000 }, (_, i) => ({ + type: 'text', + text: `line ${i}: ${'x'.repeat(20)}`, + })) + + const result = await replaceLargeToolResult( + { + role: 'tool_result', + toolUseId: 'structured-blocks', + toolName: 'web_fetch', + content: blocks, + isError: false, + }, + state, + ) + + assert.ok(result.content.startsWith(PERSISTED_OUTPUT_TAG), + 'large structured content should be persisted') + + const savedPath = extractSavedPath(result.content) + assert.ok(savedPath.endsWith('.json'), 'structured content should be saved as .json') + + const saved = JSON.parse(await readFile(savedPath, 'utf8')) as unknown[] + assert.equal(saved.length, blocks.length) + assert.deepEqual((saved[0] as { type: string; text: string }).type, 'text') + }) })