From d5588ea80da55392ce79865448cb86b10778b6be Mon Sep 17 00:00:00 2001 From: beardthelion <56458543+beardthelion@users.noreply.github.com> Date: Tue, 16 Jun 2026 22:02:54 -0500 Subject: [PATCH] feat(context-collapse): opt-in between-turns context collapse (span summarization) (#1619) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(context-collapse): implement context collapse for proactive context management * feat(context-collapse): add turn-boundary helpers for span selection * feat(context-collapse): deterministic turn-anchored span selection * feat(context-collapse): code-computed span risk score * feat(context-collapse): ctx-agent summarization instruction * feat(context-collapse): implement ctx-agent span summarization spawn * fix(context-collapse): make runtime activation opt-in (CLAUDE_CONTEXT_COLLAPSE) * fix(context-collapse): address review feedback on restore state and test rigor - restoreContextCollapseState now resets armed/lastSpawnTokens up front so a snapshot-less restore cannot carry stale spawn state across sessions. - projectView reuses a stable timestamp from the replaced span instead of new Date(), keeping the read-side projection deterministic. - Strengthen the disabled-state and turn-boundary assertions, drop an internal renderToolUseMessage assertion, and isolate the operations/persist/spawn tests from shared module and CLAUDE_CONTEXT_COLLAPSE env state. * test(context-collapse): re-init enablement in persist.test hooks resetContextCollapse() does not re-read CLAUDE_CONTEXT_COLLAPSE, so the afterEach env delete left enabled=true in module state, leaking to the next test file. Call initContextCollapse() in both hooks so module enablement stays synced to the env var. * test(context-collapse): stop spawnCtxAgent module stubs leaking across files spawnCtxAgent.test.ts stubs shared modules (tokens, forkedAgent, messages, analytics, log, spanSelection) via mock.module in beforeEach. bun's mock.restore() does not undo mock.module, so the tokens stub (() => 100000) bled into autoCompact/microCompact/runAgent tests run later in the full serial suite, making them see every conversation as over-threshold (4 spurious failures in test:full, all green in isolation). Restore each stub to its real implementation in afterEach. The reals are snapshotted into plain objects up front because 'import * as' yields a live namespace that mock.module mutates in place, so holding the namespace would restore the stub. autoCompact.js is deliberately not restored here since autoCompact.test.ts re-imports it fresh via a cache-busting nonce. Also reset+reinit the collapse module in afterEach so enabled state stays synced to the now-unset env var. * test(context-collapse): also restore autoCompact stub from spawnCtxAgent The getEffectiveContextWindowSize stub on ../compact/autoCompact.js was the one module the previous commit left unrestored, on the assumption that restoring it would clash with autoCompact.test.ts's nonce re-import. It doesn't: the nonce import uses a different specifier, and the snapshot restore is keyed by the plain specifier. compressToolHistory imports getEffectiveContextWindowSize and sizes tool-history truncation from it, so the leaked 20000-token window made it fully omit tool results ('chars omitted') instead of mid-truncating ('[…truncated') for large-context models, failing the openaiShim compression tests in the full serial suite. Restore all seven mocked modules. * fix(context-collapse): re-arm after reset and gate ctx_inspect on opt-in resetContextCollapse() left armed=false while enabled stayed true, so the first /compact, main-thread compaction cleanup, or rewind permanently disabled collapse for the rest of an opted-in session. Reset now mirrors restoreContextCollapseState and sets armed=enabled. CtxInspectTool.isEnabled() returned true unconditionally, advertising ctx_inspect to the model in every default session even when the runtime opt-in was off. It now returns isContextCollapseEnabled(). The opt-in is also exposed as the contextCollapseEnabled global config key, so it is reachable through /config instead of only the CLAUDE_CONTEXT_COLLAPSE env var. * refactor(context-collapse): drop no-op ternary in drainStaged persist call The (stagedQueue.length > 0 ? 0 : 0) subtrahend always evaluated to 0, so this is just persistCommits(processed.length). * fix(context-collapse): persist commits before advancing the snapshot drainStaged removed processed spans from the staged queue and then fired persistCommits and persistSnapshot in parallel. If the snapshot write (which no longer lists those spans as staged) landed while the commit write failed or the process died between them, restore would find the spans neither staged nor committed and the collapse would disappear on resume. Chain the snapshot write after the commit write so the commit log is durable first. * fix(context-collapse): project committed collapses on the query path, fix opt-in reach Three issues from review: - Committed collapses were never re-applied to the model input. The query path calls applyCollapsesIfNeeded but only drained staged spans; projectView (which replays the commit log) ran only in /context. Since messagesForQuery is rebuilt from full REPL history each turn and the commit log is repopulated on resume, the archived spans returned to the model on the next turn, undoing the collapse. applyCollapsesIfNeeded now runs projectView first (idempotent). Adds a regression that a committed collapse changes the next query input. - Cache-safe params were saved only for exact repl_main_thread/sdk sources, but the REPL tags non-default output styles as repl_main_thread:outputStyle:*, so those sessions left the ctx-agent without params (empty spawns). Matches repl_main_thread:* now, via a small tested helper. - contextCollapseEnabled had no settings control. Adds a /config toggle that refreshes runtime state (re-runs initContextCollapse) so it applies without a restart. * fix(context-collapse): clear already-committed staged spans; harden config toggle After projecting committed collapses before draining, a span present in both the commit log and the staged snapshot (a restore whose snapshot predates the matching commit write) could not be drained — projectView had already removed its messages — so it lingered in stagedQueue and distorted spawn/overflow checks. drainStaged now drops staged spans that are already committed and syncs the snapshot. Adds a regression covering the committed+staged overlap restore. Also wraps the /config context-collapse refresh in try/catch so a failed require/init can't crash the settings UI, and lists the toggle in the save-and-close change summary like the neighboring compaction settings. * fix(context-collapse): re-sync runtime state on config cancel The context-collapse toggle's onChange refreshes the module-level enabled/armed cache via initContextCollapse(). The revert path restored the config key on disk but left that cache untouched, so enabling the toggle and then pressing Escape kept collapse active for the rest of the session. Re-init context collapse after the global config snapshot is restored so cancel fully reverts runtime state. * fix(context-collapse): keep collapsed summaries visible to the model projectView and drainStaged replaced an archived span with a system informational placeholder, but normalizeMessagesForAPI filters out every system message that is not a local command. So once a collapse committed, the next model request lost both the archived messages and the summary meant to stand in for them, defeating the feature. Mark the placeholder with isCollapseSummary and let it take the same model-input path as local-command system messages (converted to a user message), so the summary survives normalization. Added a regression that runs the projected view through normalizeMessagesForAPI and asserts the summary is still present. * fix(context-collapse): avoid competing snapshot write after drain After an immediate post-spawn drain, drainStaged(messages, true) starts its own persistCommits().then(persistSnapshot) chain to guarantee commit durability before the snapshot stops listing the staged spans. The unconditional await persistSnapshot() that followed could win that race and persist a snapshot with no staged spans before the commits landed, reopening the crash window that drops collapses on restore. Only persist directly when nothing was drained. * fix(context-collapse): fall back, keep summaries non-snippable, gate /context Three review findings: - Suppress autocompact and the blocking preempt only when collapse holds a real committed/staged reduction, not on mere enablement. Adds hasActiveReduction(); a first over-threshold turn where spawnCtxAgent cannot produce a span (getLastCacheSafeParams() still null) now falls back to autocompact/blocking instead of sending an oversized transcript. - Preserve isMeta when converting a collapse-summary placeholder to a user message in normalizeMessagesForAPI, so the HISTORY_SNIP sweep cannot tag the only replacement for an archived span as snippable. - Gate the two /context projectView calls on isContextCollapseEnabled(), so a disabled session does not under-report token usage from a lingering commit log while the API receives the full transcript. Adds regressions for hasActiveReduction and for the summary surviving normalization as a non-snippable meta message. * fix(context-collapse): scope collapse to the main thread that owns the store The collapse store (commitLog/stagedQueue) is module-level and shared by in-process subagents (agent:*) and the ctx-agent (marble_origami), which run in the same process but do not own the main transcript. applyCollapsesIfNeeded only skipped marble_origami, so a subagent could stage or commit a span, flip the global hasActiveReduction(), and make the next main-thread turn suppress autocompact and the blocking prompt-too-long preempt while projectView() no-ops against the main messages, sending an oversized transcript to the API. Add isMainThreadSource() and gate both application (applyCollapsesIfNeeded, isWithheldPromptTooLong, recoverFromOverflow) and fallback suppression (autoCompact shouldAutoCompact, query collapseOwnsIt) to the owning thread. Subagents now autocompact and preempt their own oversized turns normally and never mutate the shared store. Also adds the staged-only hasActiveReduction regression CodeRabbit requested. * fix(context-collapse): persist archived count so resumed stats stay accurate restoreContextCollapseState rebuilt each commit with an empty archived list, and getStats summed that list, so after a resume /context, the context visualization, the token warning, and ctx_inspect reported 'N spans summarized (0 messages)' even though projectView was actively removing the archived spans. The persisted-entry docstring claimed projectView lazily refills the archive, but it only splices by boundary uuid and never does. The archived messages are never read back (only their count fed getStats), so replace the per-commit Message[] with a persisted archivedCount. It is written with each commit and restored on resume; pre-field sessions restore as 0. getStats now reports the same figure live and after resume. * fix(context-collapse): keep collapse summary non-snippable across user merge Preserving isMeta on the system->user conversion was not enough: when the collapsed span ends right before the next user turn, normalizeMessagesForAPI merges the summary into that real user message. Under HISTORY_SNIP mergeUserMessages clears isMeta whenever an operand is real user content and keeps the real turn's uuid, so the combined block — which carries the only replacement for the archived span — got a snip id and the model could queue it for removal. Carry an isCollapseSummary marker onto the converted user message and through mergeUserMessages (either operand), strip any snip id already baked into the real turn when the merge absorbs a summary, and skip such blocks in appendMessageTagToUserMessage. The merged block stays non-snippable regardless of merge direction or isMeta being cleared. * fix(context-collapse): preserve collapse marker on split, drop empty snip blocks normalizeMessages split path now forwards isCollapseSummary so an array-backed collapse summary keeps its non-snippable marker across API normalization. stripSnipTagsFromContent drops a text block whose only content was the snip marker, so the merge recovery path no longer emits an empty text block. --- scripts/build.ts | 2 +- scripts/feature-flags-source-guard.test.ts | 1 + .../context/context-noninteractive.ts | 9 +- src/commands/context/context.tsx | 11 +- src/components/Settings/Config.tsx | 41 +- src/query.ts | 10 +- src/query/stopHooks.test.ts | 24 + src/query/stopHooks.ts | 15 +- src/services/compact/autoCompact.ts | 16 +- .../contextCollapse/collapseUtils.test.ts | 163 +++++ src/services/contextCollapse/collapseUtils.ts | 70 ++ .../contextCollapse/ctxAgentPrompt.ts | 19 + src/services/contextCollapse/index.test.ts | 641 ++++++++++++++++++ src/services/contextCollapse/index.ts | 620 +++++++++++++++-- .../contextCollapse/operations.test.ts | 276 ++++++++ src/services/contextCollapse/operations.ts | 60 +- src/services/contextCollapse/persist.test.ts | 72 ++ src/services/contextCollapse/persist.ts | 20 +- .../contextCollapse/spanSelection.test.ts | 107 +++ src/services/contextCollapse/spanSelection.ts | 117 ++++ .../contextCollapse/spawnCtxAgent.test.ts | 201 ++++++ .../CtxInspectTool/CtxInspectTool.test.ts | 79 +++ src/tools/CtxInspectTool/CtxInspectTool.ts | 89 +++ src/tools/CtxInspectTool/prompt.ts | 9 + src/types/logs.ts | 16 +- src/types/message.ts | 15 + src/utils/config.ts | 3 + src/utils/messages.ts | 100 ++- src/utils/sessionStorage.ts | 1 + 29 files changed, 2713 insertions(+), 94 deletions(-) create mode 100644 src/query/stopHooks.test.ts create mode 100644 src/services/contextCollapse/collapseUtils.test.ts create mode 100644 src/services/contextCollapse/collapseUtils.ts create mode 100644 src/services/contextCollapse/ctxAgentPrompt.ts create mode 100644 src/services/contextCollapse/index.test.ts create mode 100644 src/services/contextCollapse/operations.test.ts create mode 100644 src/services/contextCollapse/persist.test.ts create mode 100644 src/services/contextCollapse/spanSelection.test.ts create mode 100644 src/services/contextCollapse/spanSelection.ts create mode 100644 src/services/contextCollapse/spawnCtxAgent.test.ts create mode 100644 src/tools/CtxInspectTool/CtxInspectTool.test.ts create mode 100644 src/tools/CtxInspectTool/CtxInspectTool.ts create mode 100644 src/tools/CtxInspectTool/prompt.ts diff --git a/scripts/build.ts b/scripts/build.ts index 156566a01..a64736e01 100644 --- a/scripts/build.ts +++ b/scripts/build.ts @@ -27,7 +27,7 @@ const featureFlags: Record = { DAEMON: false, // Background daemon process (stubbed in open build) AGENT_TRIGGERS: false, // Scheduled remote agent triggers ABLATION_BASELINE: false, // A/B testing harness for eval experiments - CONTEXT_COLLAPSE: false, // Context collapsing optimization (stubbed) + CONTEXT_COLLAPSE: true, // Context collapsing optimization COMMIT_ATTRIBUTION: false, // Co-Authored-By metadata in git commits HISTORY_SNIP: true, // Model-callable snip tool for context management UDS_INBOX: false, // Unix Domain Socket inter-session messaging diff --git a/scripts/feature-flags-source-guard.test.ts b/scripts/feature-flags-source-guard.test.ts index db7a85eb1..9a688807f 100644 --- a/scripts/feature-flags-source-guard.test.ts +++ b/scripts/feature-flags-source-guard.test.ts @@ -22,6 +22,7 @@ type FlagGuard = { const FLAG_REQUIRES_SOURCE: FlagGuard[] = [ { flag: 'MCP_SKILLS', source: 'src/skills/mcpSkills.ts' }, + { flag: 'CONTEXT_COLLAPSE', source: 'src/services/contextCollapse/index.ts' }, ] test('build feature flags are not enabled without their source files', () => { diff --git a/src/commands/context/context-noninteractive.ts b/src/commands/context/context-noninteractive.ts index b4b80b516..cafe6e4e9 100644 --- a/src/commands/context/context-noninteractive.ts +++ b/src/commands/context/context-noninteractive.ts @@ -51,8 +51,15 @@ export async function collectContextData( /* eslint-disable @typescript-eslint/no-require-imports */ const { projectView } = require('../../services/contextCollapse/operations.js') as typeof import('../../services/contextCollapse/operations.js') + const { isContextCollapseEnabled } = + require('../../services/contextCollapse/index.js') as typeof import('../../services/contextCollapse/index.js') /* eslint-enable @typescript-eslint/no-require-imports */ - apiView = projectView(apiView) + // Gate on runtime enablement, matching the query path. Once collapse is + // disabled the API receives the full transcript, so projecting a lingering + // commit log here would under-report real token usage. + if (isContextCollapseEnabled()) { + apiView = projectView(apiView) + } } const { messages: compactedMessages } = await microcompactMessages(apiView) diff --git a/src/commands/context/context.tsx b/src/commands/context/context.tsx index 595a3c594..4e4454e8f 100644 --- a/src/commands/context/context.tsx +++ b/src/commands/context/context.tsx @@ -22,8 +22,17 @@ function toApiView(messages: Message[]): Message[] { const { projectView } = require('../../services/contextCollapse/operations.js') as typeof import('../../services/contextCollapse/operations.js'); + const { + isContextCollapseEnabled + } = require('../../services/contextCollapse/index.js') as typeof import('../../services/contextCollapse/index.js'); /* eslint-enable @typescript-eslint/no-require-imports */ - view = projectView(view); + // Gate on runtime enablement, matching the query path (which only calls + // applyCollapsesIfNeeded while collapse is enabled). After the user turns + // collapse off, the API gets the full transcript again, so projecting a + // lingering commit log here would under-report the real token usage. + if (isContextCollapseEnabled()) { + view = projectView(view); + } } return view; } diff --git a/src/components/Settings/Config.tsx b/src/components/Settings/Config.tsx index 65ea72a24..14f47d134 100644 --- a/src/components/Settings/Config.tsx +++ b/src/components/Settings/Config.tsx @@ -329,7 +329,32 @@ export function Config({ enabled: toolHistoryCompressionEnabled }); } - }, { + }, ...(feature('CONTEXT_COLLAPSE') ? [{ + id: 'contextCollapseEnabled', + label: 'Context collapse (lossy)', + value: globalConfig.contextCollapseEnabled, + type: 'boolean' as const, + onChange(contextCollapseEnabled: boolean) { + saveGlobalConfig(current_cc => ({ + ...current_cc, + contextCollapseEnabled + })); + setGlobalConfig({ + ...getGlobalConfig(), + contextCollapseEnabled + }); + // Refresh runtime state so the toggle applies without a restart: + // initContextCollapse re-reads env + this config key. + try { + (require('../../services/contextCollapse/index.js') as typeof import('../../services/contextCollapse/index.js')).initContextCollapse(); + } catch (error) { + logError(`Failed to refresh context collapse state: ${error}`); + } + logEvent('tengu_context_collapse_setting_changed', { + enabled: contextCollapseEnabled + }); + } + }] : []), { id: 'showCacheStats', label: 'Cache stats display', value: globalConfig.showCacheStats, @@ -1223,6 +1248,9 @@ export function Config({ if (globalConfig.toolHistoryCompressionEnabled !== initialConfig.current.toolHistoryCompressionEnabled) { formattedChanges.push(`${globalConfig.toolHistoryCompressionEnabled ? 'Enabled' : 'Disabled'} tool history compression`); } + if (feature('CONTEXT_COLLAPSE') && globalConfig.contextCollapseEnabled !== initialConfig.current.contextCollapseEnabled) { + formattedChanges.push(`${globalConfig.contextCollapseEnabled ? 'Enabled' : 'Disabled'} context collapse (lossy)`); + } if (globalConfig.respectGitignore !== initialConfig.current.respectGitignore) { formattedChanges.push(`${globalConfig.respectGitignore ? 'Enabled' : 'Disabled'} respect .gitignore in file picker`); } @@ -1274,6 +1302,17 @@ export function Config({ // the returned ref equals current (test mode checks ref; prod writes to // disk but content is identical). saveGlobalConfig(() => initialConfig.current); + // Context collapse: the toggle's onChange calls initContextCollapse() to + // refresh the module-level enabled/armed cache. The global config restore + // above rewrites the key on disk but doesn't touch that cache, so re-init + // here to keep runtime state in sync with the reverted config. + if (feature('CONTEXT_COLLAPSE')) { + try { + (require('../../services/contextCollapse/index.js') as typeof import('../../services/contextCollapse/index.js')).initContextCollapse(); + } catch (error) { + logError(`Failed to refresh context collapse state on cancel: ${error}`); + } + } // Settings files: restore each key Config may have touched. undefined // deletes the key (updateSettingsForSource customizer at settings.ts:368). const il = initialLocalSettings; diff --git a/src/query.ts b/src/query.ts index 49f442db7..70ba660c1 100644 --- a/src/query.ts +++ b/src/query.ts @@ -778,10 +778,18 @@ async function* queryLoop( // API call and starve both recovery paths. The isAutoCompactEnabled() // conjunct preserves the user's explicit "no automatic anything" // config — if they set DISABLE_AUTO_COMPACT, they get the preempt. + // hasActiveReduction() (not mere enablement) means a turn where collapse + // could not reduce anything still hits the blocking preempt instead of + // sending an oversized request that only a real 413 could recover. let collapseOwnsIt = false if (feature('CONTEXT_COLLAPSE')) { + // Only the main thread that owns the reduction may skip the blocking + // preempt: the store is shared with in-process subagents (agent:*), and a + // subagent must still preempt its own oversized turn rather than defer to + // a reduction that does not apply to its messages. collapseOwnsIt = - (contextCollapse?.isContextCollapseEnabled() ?? false) && + (contextCollapse?.isMainThreadSource(querySource) ?? false) && + (contextCollapse?.hasActiveReduction() ?? false) && isAutoCompactEnabled() } // Hoist media-recovery gate once per turn. Withholding (inside the diff --git a/src/query/stopHooks.test.ts b/src/query/stopHooks.test.ts new file mode 100644 index 000000000..475949400 --- /dev/null +++ b/src/query/stopHooks.test.ts @@ -0,0 +1,24 @@ +import { describe, expect, test } from 'bun:test' +import type { QuerySource } from '../constants/querySource.js' +import { isMainThreadCacheParamSource } from './stopHooks.js' + +describe('isMainThreadCacheParamSource', () => { + test('matches the bare main-thread and sdk sources', () => { + expect(isMainThreadCacheParamSource('repl_main_thread' as QuerySource)).toBe(true) + expect(isMainThreadCacheParamSource('sdk' as QuerySource)).toBe(true) + }) + + test('matches output-style main-thread sources', () => { + expect( + isMainThreadCacheParamSource('repl_main_thread:outputStyle:explanatory' as QuerySource), + ).toBe(true) + expect( + isMainThreadCacheParamSource('repl_main_thread:outputStyle:custom' as QuerySource), + ).toBe(true) + }) + + test('does not match subagent sources', () => { + expect(isMainThreadCacheParamSource('marble_origami' as QuerySource)).toBe(false) + expect(isMainThreadCacheParamSource('user_prompt' as QuerySource)).toBe(false) + }) +}) diff --git a/src/query/stopHooks.ts b/src/query/stopHooks.ts index 8f4c885f9..65f8c7521 100644 --- a/src/query/stopHooks.ts +++ b/src/query/stopHooks.ts @@ -74,6 +74,15 @@ export type StopHookExecutionDeps = { isTeammate?: typeof isTeammate } +/** + * Whether a query source is the main session (vs a subagent) for the purpose of + * saving cache-safe params. Matches `repl_main_thread` and its suffixed forms + * (e.g. `repl_main_thread:outputStyle:custom`) plus `sdk`, but not subagents. + */ +export function isMainThreadCacheParamSource(querySource: QuerySource): boolean { + return querySource.startsWith('repl_main_thread') || querySource === 'sdk' +} + export async function* handleStopHooks( messagesForQuery: Message[], assistantMessages: AssistantMessage[], @@ -116,8 +125,10 @@ export async function* handleStopHooks( // Only save params for main session queries — subagents must not overwrite. // Outside the prompt-suggestion gate: the REPL /btw command and the // side_question SDK control_request both read this snapshot, and neither - // depends on prompt suggestions being enabled. - if (querySource === 'repl_main_thread' || querySource === 'sdk') { + // depends on prompt suggestions being enabled. Match repl_main_thread:* too, + // since the REPL tags non-default output styles as repl_main_thread:outputStyle:* + // and the context-collapse ctx-agent relies on this snapshot being present. + if (isMainThreadCacheParamSource(querySource)) { saveCacheSafeParams(createCacheSafeParams(stopHookContext)) } diff --git a/src/services/compact/autoCompact.ts b/src/services/compact/autoCompact.ts index 694978699..68ad02862 100644 --- a/src/services/compact/autoCompact.ts +++ b/src/services/compact/autoCompact.ts @@ -306,16 +306,22 @@ export async function shouldAutoCompact( // fallback (it consults isAutoCompactEnabled directly) and leaves // sessionMemory + manual /compact working. // - // Consult isContextCollapseEnabled (not the raw gate) so the - // CLAUDE_CONTEXT_COLLAPSE env override is honored here too. require() - // inside the block breaks the init-time cycle (this file exports + // hasActiveReduction() folds in the enablement check (so the + // CLAUDE_CONTEXT_COLLAPSE env override is honored here too) but also + // requires collapse to actually hold a committed/staged reduction. + // require() inside the block breaks the init-time cycle (this file exports // getEffectiveContextWindowSize which collapse's index imports). if (feature('CONTEXT_COLLAPSE')) { /* eslint-disable @typescript-eslint/no-require-imports */ - const { isContextCollapseEnabled } = + const { hasActiveReduction, isMainThreadSource } = require('../contextCollapse/index.js') as typeof import('../contextCollapse/index.js') /* eslint-enable @typescript-eslint/no-require-imports */ - if (isContextCollapseEnabled()) { + // Suppress only when collapse actually holds the headroom (a committed or + // staged reduction) AND this is the main thread that owns it. The store is + // shared across in-process subagents (agent:*); a subagent must still + // autocompact its own oversized transcript instead of being suppressed by a + // reduction that only applies to the main transcript. + if (isMainThreadSource(querySource) && hasActiveReduction()) { return false } } diff --git a/src/services/contextCollapse/collapseUtils.test.ts b/src/services/contextCollapse/collapseUtils.test.ts new file mode 100644 index 000000000..b1d4019f2 --- /dev/null +++ b/src/services/contextCollapse/collapseUtils.test.ts @@ -0,0 +1,163 @@ +import { randomUUID } from 'crypto' +import { describe, expect, test } from 'bun:test' +import type { UUID } from 'crypto' +import { + buildCollapsePlaceholder, + deriveCollapseId, + getSpanUuids, + isCollapsePlaceholder, + isWithinCollapsedSpan, + resetCollapseIdCounter, +} from './collapseUtils.js' +import type { Message } from '../../types/message.js' + +function uid(s: string): UUID { + // Replace underscores with dashes to produce valid UUID-looking strings + return `00000000-0000-4000-8000-${s.padStart(12, '0')}` as UUID +} + +function makeMsg(id: string): Message { + return { + type: 'user', + uuid: uid(id), + timestamp: new Date().toISOString(), + message: { content: 'hello', role: 'user' as const }, + } as unknown as Message +} + +function makeSysMsg(content: string): Message { + return { + type: 'system', + subtype: 'informational' as const, + content, + uuid: randomUUID(), + timestamp: new Date().toISOString(), + isMeta: true, + } as unknown as Message +} + +describe('deriveCollapseId', () => { + test('derives 16-digit sequential IDs', () => { + resetCollapseIdCounter() + expect(deriveCollapseId(uid('a'))).toBe('0000000000000001') + expect(deriveCollapseId(uid('b'))).toBe('0000000000000002') + }) + + test('resets and reseeds correctly', () => { + resetCollapseIdCounter(42) + expect(deriveCollapseId(uid('x'))).toBe('0000000000000043') + }) + + test('counter stability across calls', () => { + resetCollapseIdCounter() + expect(deriveCollapseId(uid('a'))).toBe('0000000000000001') + expect(deriveCollapseId(uid('b'))).toBe('0000000000000002') + expect(deriveCollapseId(uid('c'))).toBe('0000000000000003') + }) +}) + +describe('getSpanUuids', () => { + const msgs: Message[] = ['a', 'b', 'c', 'd', 'e'].map(makeMsg) + + test('extracts inclusive span', () => { + expect(getSpanUuids(msgs, uid('b'), uid('d'))).toEqual([uid('b'), uid('c'), uid('d')]) + }) + + test('single message span', () => { + expect(getSpanUuids(msgs, uid('c'), uid('c'))).toEqual([uid('c')]) + }) + + test('first boundary missing returns empty', () => { + expect(getSpanUuids(msgs, uid('x'), uid('d'))).toEqual([]) + }) + + test('last boundary missing returns empty', () => { + expect(getSpanUuids(msgs, uid('b'), uid('x'))).toEqual([]) + }) + + test('out-of-order boundaries return empty', () => { + expect(getSpanUuids(msgs, uid('d'), uid('b'))).toEqual([]) + }) + + test('full span from first to last', () => { + expect(getSpanUuids(msgs, uid('a'), uid('e'))).toEqual(['a', 'b', 'c', 'd', 'e'].map(uid)) + }) + + test('empty messages returns empty', () => { + expect(getSpanUuids([], uid('a'), uid('b'))).toEqual([]) + }) +}) + +describe('buildCollapsePlaceholder', () => { + test('builds valid placeholder string', () => { + const result = buildCollapsePlaceholder('0000000000000001', 'summarized content') + expect(result).toBe('summarized content') + }) + + test('escapes nothing - raw summary is embedded', () => { + const result = buildCollapsePlaceholder('42', 'text with and "quotes"') + expect(result).toContain('') + expect(result).toContain('"quotes"') + }) +}) + +describe('isCollapsePlaceholder', () => { + test('detects collapse placeholder system messages', () => { + const msg = makeSysMsg('summary') + expect(isCollapsePlaceholder(msg)).toBe(true) + }) + + test('rejects non-system messages', () => { + const msg = makeMsg('uuid-1') + expect(isCollapsePlaceholder(msg)).toBe(false) + }) + + test('rejects system messages without collapsed prefix', () => { + const msg = makeSysMsg('regular system message') + expect(isCollapsePlaceholder(msg)).toBe(false) + }) + + test('rejects system messages that only happen to contain collapsed word', () => { + const msg = makeSysMsg('something about a collapsed bridge') + expect(isCollapsePlaceholder(msg)).toBe(false) + }) +}) + +describe('isWithinCollapsedSpan', () => { + const msgs: Message[] = ['a', 'b', 'c', 'd', 'e'].map(makeMsg) + + test('message inside a single span returns true', () => { + const commits = [{ firstArchivedUuid: uid('b'), lastArchivedUuid: uid('d') }] + expect(isWithinCollapsedSpan(makeMsg('c'), msgs, commits)).toBe(true) + }) + + test('message at span boundary returns true', () => { + const commits = [{ firstArchivedUuid: uid('b'), lastArchivedUuid: uid('d') }] + expect(isWithinCollapsedSpan(makeMsg('b'), msgs, commits)).toBe(true) + expect(isWithinCollapsedSpan(makeMsg('d'), msgs, commits)).toBe(true) + }) + + test('message outside span returns false', () => { + const commits = [{ firstArchivedUuid: uid('b'), lastArchivedUuid: uid('d') }] + expect(isWithinCollapsedSpan(makeMsg('a'), msgs, commits)).toBe(false) + expect(isWithinCollapsedSpan(makeMsg('e'), msgs, commits)).toBe(false) + }) + + test('message not in array at all returns false', () => { + const commits = [{ firstArchivedUuid: uid('a'), lastArchivedUuid: uid('e') }] + expect(isWithinCollapsedSpan(makeMsg('z'), msgs, commits)).toBe(false) + }) + + test('no commits returns false', () => { + expect(isWithinCollapsedSpan(makeMsg('c'), msgs, [])).toBe(false) + }) + + test('multiple commits - message in second span', () => { + const commits = [ + { firstArchivedUuid: uid('a'), lastArchivedUuid: uid('a') }, + { firstArchivedUuid: uid('c'), lastArchivedUuid: uid('e') }, + ] + expect(isWithinCollapsedSpan(makeMsg('d'), msgs, commits)).toBe(true) + expect(isWithinCollapsedSpan(makeMsg('b'), msgs, commits)).toBe(false) + }) +}) diff --git a/src/services/contextCollapse/collapseUtils.ts b/src/services/contextCollapse/collapseUtils.ts new file mode 100644 index 000000000..1e2057fc8 --- /dev/null +++ b/src/services/contextCollapse/collapseUtils.ts @@ -0,0 +1,70 @@ +import type { Message } from '../../types/message.js' + +let idCounter = 0 + +export function resetCollapseIdCounter(seed?: number): void { + idCounter = seed ?? 0 +} + +export function deriveCollapseId(_uuid: string): string { + idCounter++ + return String(idCounter).padStart(16, '0') +} + +export function getSpanUuids( + messages: Message[], + firstUuid: string, + lastUuid: string, +): string[] { + const firstIdx = messages.findIndex(m => m.uuid === firstUuid) + const lastIdx = findLastIndex(messages, m => m.uuid === lastUuid) + + if (firstIdx === -1 || lastIdx === -1 || firstIdx > lastIdx) { + return [] + } + + return messages.slice(firstIdx, lastIdx + 1).map(m => m.uuid as string) +} + +export function buildCollapsePlaceholder( + collapseId: string, + summary: string, +): string { + return `${summary}` +} + +export function isCollapsePlaceholder(message: Message): boolean { + if (message.type !== 'system') return false + const content = 'content' in message ? (message as Record).content : '' + return typeof content === 'string' && content.startsWith(' m.uuid === msgUuid) + if (msgIdx === -1) return false + + return commits.some(c => { + const firstIdx = messages.findIndex(m => m.uuid === c.firstArchivedUuid) + const lastIdx = findLastIndex(messages, m => m.uuid === c.lastArchivedUuid) + return firstIdx !== -1 && lastIdx !== -1 && firstIdx <= msgIdx && msgIdx <= lastIdx + }) +} + +function findLastIndex( + arr: T[], + predicate: (item: T) => boolean, +): number { + for (let i = arr.length - 1; i >= 0; i--) { + if (predicate(arr[i]!)) return i + } + return -1 +} + +export { findLastIndex } diff --git a/src/services/contextCollapse/ctxAgentPrompt.ts b/src/services/contextCollapse/ctxAgentPrompt.ts new file mode 100644 index 000000000..799a96f68 --- /dev/null +++ b/src/services/contextCollapse/ctxAgentPrompt.ts @@ -0,0 +1,19 @@ +/** + * Instruction handed to the forked ctx-agent. The span's actual messages are + * supplied as fork context; this asks for a faithful, compact replacement summary. + * Provider-neutral: it runs on whatever model/provider the session is using. + */ +export const CTX_AGENT_INSTRUCTION = [ + 'You are compacting an older portion of this conversation to save context.', + 'Write a single compact summary of the conversation messages above so that', + 'work can continue without re-reading them. Preserve, concisely:', + '- decisions made and the reasoning behind them', + '- the current state, configuration, and any values of record', + '- file paths, identifiers, commands, and other concrete references touched', + '- unresolved threads, open questions, and TODOs', + '- any facts later steps depend on', + '', + 'Rules: include only information present in the messages above; do not invent', + 'or speculate. No preamble, no headings, no closing remarks — output only the', + 'summary prose. Be substantially shorter than the original.', +].join('\n') diff --git a/src/services/contextCollapse/index.test.ts b/src/services/contextCollapse/index.test.ts new file mode 100644 index 000000000..be2871f4e --- /dev/null +++ b/src/services/contextCollapse/index.test.ts @@ -0,0 +1,641 @@ +import { randomUUID } from 'crypto' +import { afterEach, beforeEach, describe, expect, test } from 'bun:test' +import type { UUID } from 'crypto' +import type { Message } from '../../types/message.js' +import type { ContextCollapseCommitEntry, ContextCollapseSnapshotEntry } from '../../types/logs.js' + +beforeEach(() => { + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' +}) +afterEach(() => { + delete process.env.CLAUDE_CONTEXT_COLLAPSE +}) + +function uid(s: string): UUID { + return `00000000-0000-4000-8000-${s.padStart(12, '0')}` as UUID +} + +function makeUserMsg(id: string, content = 'hello'): Message { + return { + type: 'user', + uuid: uid(id), + timestamp: new Date().toISOString(), + message: { content, role: 'user' as const }, + } as unknown as Message +} + +function makeAssistantMsg(id: string, content = 'response'): Message { + return { + type: 'assistant', + uuid: uid(id), + timestamp: new Date().toISOString(), + message: { + id: randomUUID(), + model: 'claude-sonnet-4', + role: 'assistant', + stop_reason: 'end_turn', + stop_sequence: '', + type: 'message', + usage: { input_tokens: 10, output_tokens: 5 }, + content: [{ type: 'text' as const, text: content }], + context_management: null, + }, + } as unknown as Message +} + +function makeFakeToolUseContext() { + return { + options: { + commands: [], + debug: false, + mainLoopModel: 'claude-sonnet-4', + tools: [] as any, + verbose: false, + thinkingConfig: {}, + mcpClients: [], + mcpResources: {}, + isNonInteractiveSession: false, + agentDefinitions: { maxTurns: 10 }, + }, + abortController: new AbortController(), + readFileState: {} as any, + getAppState: () => ({} as any), + setAppState: (_f: any) => {}, + messages: [], + } as any +} + +// Module state is shared across ALL test files. We clean between groups. +async function cleanState() { + const idx = await import('./index.js') + idx.resetContextCollapse() +} + +describe('init and enable', () => { + test('initContextCollapse enables when CLAUDE_CONTEXT_COLLAPSE=1', async () => { + await cleanState() + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + const idx = await import('./index.js') + idx.initContextCollapse() + expect(idx.isContextCollapseEnabled()).toBe(true) + }) + + test('initContextCollapse defaults to OFF without the env opt-in', async () => { + await cleanState() + delete process.env.CLAUDE_CONTEXT_COLLAPSE + const idx = await import('./index.js') + idx.initContextCollapse() + expect(idx.isContextCollapseEnabled()).toBe(false) + // restore for subsequent tests (beforeEach also sets it, but be explicit) + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + }) + + test('resetContextCollapse re-arms an enabled session', async () => { + await cleanState() + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + const idx = await import('./index.js') + idx.initContextCollapse() + expect(idx.getContextCollapseState()!.armed).toBe(true) + // Compaction cleanup / rewind call reset; the session stays opted-in and + // must remain able to spawn again, not silently disable for the session. + idx.resetContextCollapse() + expect(idx.getContextCollapseState()!.armed).toBe(true) + }) + + test('resetContextCollapse leaves a disabled session disarmed', async () => { + await cleanState() + delete process.env.CLAUDE_CONTEXT_COLLAPSE + const idx = await import('./index.js') + idx.initContextCollapse() + idx.resetContextCollapse() + expect(idx.getContextCollapseState()).toBeNull() + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + }) + + test('getContextCollapseState returns valid shape when enabled', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + const state = idx.getContextCollapseState()! + expect(state).not.toBeNull() + expect(typeof state.committedSpans).toBe('number') + expect(typeof state.stagedSpans).toBe('number') + expect(typeof state.armed).toBe('boolean') + }) + + test('getContextCollapseState returns null when not enabled', async () => { + // Deterministic disabled state: clear the opt-in env and re-init so + // enabled=false, then assert the null contract directly. + await cleanState() + delete process.env.CLAUDE_CONTEXT_COLLAPSE + const idx = await import('./index.js') + idx.initContextCollapse() + expect(idx.getContextCollapseState()).toBeNull() + // Restore the opt-in for subsequent tests (beforeEach also sets it). + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + }) +}) + +describe('stats and subscribe', () => { + test('getStats returns zero stats on fresh state', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + const stats = idx.getStats() + expect(stats.collapsedSpans).toBe(0) + expect(stats.collapsedMessages).toBe(0) + expect(stats.stagedSpans).toBe(0) + }) + + test('subscribe returns unsubscribe function', async () => { + await cleanState() + const idx = await import('./index.js') + const unsub = idx.subscribe(() => {}) + expect(typeof unsub).toBe('function') + unsub() + }) + + test('subscribe listener fires on resetContextCollapse', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + let called = false + idx.subscribe(() => { called = true }) + idx.resetContextCollapse() + expect(called).toBe(true) + }) +}) + +describe('hasActiveReduction', () => { + // Gates autocompact/blocking suppression: enablement alone is not enough, or + // a turn where collapse could not produce a span would suppress the fallback + // and send an oversized transcript to the API. + test('false when collapse is disabled', async () => { + await cleanState() + delete process.env.CLAUDE_CONTEXT_COLLAPSE + const idx = await import('./index.js') + idx.initContextCollapse() + expect(idx.hasActiveReduction()).toBe(false) + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + }) + + test('false when enabled but nothing committed or staged', async () => { + await cleanState() + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + const idx = await import('./index.js') + idx.initContextCollapse() + expect(idx.isContextCollapseEnabled()).toBe(true) + expect(idx.hasActiveReduction()).toBe(false) + }) + + test('true once a span is committed', async () => { + await cleanState() + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + const idx = await import('./index.js') + idx.initContextCollapse() + idx.restoreContextCollapseState( + [ + { + type: 'marble-origami-commit' as const, + sessionId: uid('s1'), + collapseId: '0000000000000001', + summaryUuid: uid('sum1'), + summaryContent: 'test summary', + summary: 'test summary', + firstArchivedUuid: uid('a'), + lastArchivedUuid: uid('b'), + }, + ], + undefined, + ) + expect(idx.hasActiveReduction()).toBe(true) + }) + + test('true when only staged (no committed spans)', async () => { + // The pre-commit staged state is what decides whether autocompact and the + // blocking preempt are suppressed; a regression here recreates the + // fallback bug without any committed-state test failing. + await cleanState() + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + const idx = await import('./index.js') + idx.initContextCollapse() + idx.restoreContextCollapseState([], { + type: 'marble-origami-snapshot' as const, + sessionId: uid('s1'), + staged: [ + { startUuid: uid('a'), endUuid: uid('b'), summary: 'pending', risk: 0.7, stagedAt: Date.now() }, + ], + armed: true, + lastSpawnTokens: 1000, + }) + expect(idx.getStats().collapsedSpans).toBe(0) + expect(idx.getStats().stagedSpans).toBe(1) + expect(idx.hasActiveReduction()).toBe(true) + }) +}) + +describe('isMainThreadSource', () => { + test('true for main-thread sources, false for subagent/fork sources', async () => { + const idx = await import('./index.js') + expect(idx.isMainThreadSource('repl_main_thread' as any)).toBe(true) + expect(idx.isMainThreadSource('repl_main_thread:resume' as any)).toBe(true) + expect(idx.isMainThreadSource('sdk' as any)).toBe(true) + expect(idx.isMainThreadSource(undefined)).toBe(true) + expect(idx.isMainThreadSource('agent:explore' as any)).toBe(false) + expect(idx.isMainThreadSource('marble_origami' as any)).toBe(false) + expect(idx.isMainThreadSource('compact' as any)).toBe(false) + expect(idx.isMainThreadSource('session_memory' as any)).toBe(false) + }) +}) + +describe('subagent sources do not touch the shared store', () => { + // Regression for the fallback bug: an in-process subagent (agent:*) shares the + // module-level collapse store but does not own the main transcript. It must + // not apply/stage/commit (which would flip hasActiveReduction() globally and + // suppress the main thread's autocompact/blocking fallback while projectView() + // no-ops on main messages). + function committedSpan() { + return [ + { + type: 'marble-origami-commit' as const, + sessionId: uid('s1'), + collapseId: '0000000000000001', + summaryUuid: uid('sum'), + summaryContent: 'summary', + summary: 'summary', + firstArchivedUuid: uid('a1'), + lastArchivedUuid: uid('a3'), + }, + ] + } + const fullHistory: Message[] = [ + makeUserMsg('u0'), + makeUserMsg('a1'), + makeAssistantMsg('a2'), + makeUserMsg('a3'), + makeUserMsg('u4'), + ] + + test('applyCollapsesIfNeeded does not project a committed collapse for an agent:* source', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + idx.restoreContextCollapseState(committedSpan(), undefined) + + const { messages } = await idx.applyCollapsesIfNeeded( + fullHistory, + makeFakeToolUseContext(), + 'agent:explore' as any, + ) + // Untouched: the archived span is still present, no summary injected. + expect(messages).toEqual(fullHistory) + + // Sanity: the same state DOES project on the main thread. + const main = await idx.applyCollapsesIfNeeded( + fullHistory, + makeFakeToolUseContext(), + 'repl_main_thread' as any, + ) + expect(main.messages.map(m => m.uuid)).toContain(uid('sum')) + expect(main.messages.map(m => m.uuid)).not.toContain(uid('a2')) + }) + + test('isWithheldPromptTooLong is false for an agent:* source even with staged spans', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + idx.restoreContextCollapseState([], { + type: 'marble-origami-snapshot' as const, + sessionId: uid('s1'), + staged: [ + { startUuid: uid('a'), endUuid: uid('b'), summary: 'pending', risk: 0.7, stagedAt: Date.now() }, + ], + armed: true, + lastSpawnTokens: 1000, + }) + const msg = makeAssistantMsg('a', 'prompt too long') + expect(idx.isWithheldPromptTooLong(msg, () => true, 'agent:explore' as any)).toBe(false) + // Main thread with the same staged span does withhold. + expect(idx.isWithheldPromptTooLong(msg, () => true, 'repl_main_thread' as any)).toBe(true) + }) + + test('recoverFromOverflow does not drain staged spans for an agent:* source', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + idx.restoreContextCollapseState([], { + type: 'marble-origami-snapshot' as const, + sessionId: uid('s1'), + staged: [ + { startUuid: uid('a1'), endUuid: uid('a3'), summary: 'pending', risk: 0.7, stagedAt: Date.now() }, + ], + armed: true, + lastSpawnTokens: 1000, + }) + const result = idx.recoverFromOverflow(fullHistory, 'agent:explore' as any) + expect(result.committed).toBe(0) + expect(result.messages).toEqual(fullHistory) + expect(idx.getStats().stagedSpans).toBe(1) + }) +}) + +describe('core API (no staged spans)', () => { + test('applyCollapsesIfNeeded: identity when nothing staged', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + const msgs = [makeUserMsg('a'), makeAssistantMsg('b')] + const result = await idx.applyCollapsesIfNeeded(msgs, makeFakeToolUseContext(), 'user_prompt' as any) + expect(result.messages).toEqual(msgs) + }) + + test('applyCollapsesIfNeeded: skips marble_origami source', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + const msgs = [makeUserMsg('a'), makeAssistantMsg('b')] + const result = await idx.applyCollapsesIfNeeded(msgs, makeFakeToolUseContext(), 'marble_origami' as any) + expect(result.messages).toEqual(msgs) + }) + + test('isWithheldPromptTooLong: false with no staged', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + const msg = makeAssistantMsg('a', 'prompt too long') + expect(idx.isWithheldPromptTooLong(msg, () => true, 'user_prompt' as any)).toBe(false) + }) + + test('isWithheldPromptTooLong: false for non-assistant', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + const user = makeUserMsg('a') + expect(idx.isWithheldPromptTooLong(user, () => true, 'user_prompt' as any)).toBe(false) + }) + + test('isWithheldPromptTooLong: false for undefined', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + expect(idx.isWithheldPromptTooLong(undefined, () => true, 'user_prompt' as any)).toBe(false) + }) + + test('recoverFromOverflow: zero committed on clean state', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + const msgs = [makeUserMsg('a'), makeAssistantMsg('b')] + const result = idx.recoverFromOverflow(msgs, 'user_prompt' as any) + expect(result.committed).toBe(0) + expect(result.messages).toEqual(msgs) + }) +}) + +describe('restoreContextCollapseState', () => { + test('rebuilds from commits and snapshot', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + + const commits: ContextCollapseCommitEntry[] = [ + { + type: 'marble-origami-commit' as const, + sessionId: uid('s1'), + collapseId: '0000000000000007', + summaryUuid: uid('s1'), + summaryContent: 'summary', + summary: 'summary', + firstArchivedUuid: uid('a'), + lastArchivedUuid: uid('b'), + }, + ] + + const snapshot: ContextCollapseSnapshotEntry = { + type: 'marble-origami-snapshot' as const, + sessionId: uid('s1'), + staged: [ + { startUuid: uid('c'), endUuid: uid('d'), summary: 'pending', risk: 0.7, stagedAt: Date.now() }, + ], + armed: true, + lastSpawnTokens: 10000, + } + + idx.restoreContextCollapseState(commits, snapshot) + expect(idx.getStats().collapsedSpans).toBe(1) + expect(idx.getStats().stagedSpans).toBe(1) + }) + + test('restored commit reports its archivedCount in collapsedMessages', async () => { + // After a resume the archived messages are not held per-commit; getStats + // must use the persisted count so /context, ctx_inspect, and the token + // warning report the same figure as a live session, not "0 messages". + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + idx.restoreContextCollapseState( + [ + { + type: 'marble-origami-commit' as const, + sessionId: uid('s1'), + collapseId: '0000000000000007', + summaryUuid: uid('sum'), + summaryContent: 'summary', + summary: 'summary', + firstArchivedUuid: uid('a'), + lastArchivedUuid: uid('b'), + archivedCount: 5, + }, + ], + undefined, + ) + expect(idx.getStats().collapsedSpans).toBe(1) + expect(idx.getStats().collapsedMessages).toBe(5) + }) + + test('pre-field commit (no archivedCount) restores as 0 messages', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + idx.restoreContextCollapseState( + [ + { + type: 'marble-origami-commit' as const, + sessionId: uid('s1'), + collapseId: '0000000000000008', + summaryUuid: uid('sum2'), + summaryContent: 'summary', + summary: 'summary', + firstArchivedUuid: uid('a'), + lastArchivedUuid: uid('b'), + }, + ], + undefined, + ) + expect(idx.getStats().collapsedSpans).toBe(1) + expect(idx.getStats().collapsedMessages).toBe(0) + }) + + test('ID counter reseeded from max collapseId', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + + idx.restoreContextCollapseState([ + { + type: 'marble-origami-commit' as const, + sessionId: uid('s1'), + collapseId: '0000000000000042', + summaryUuid: uid('s1'), + summaryContent: 'x', + summary: 'x', + firstArchivedUuid: uid('a'), + lastArchivedUuid: uid('b'), + }, + ], undefined) + + expect(idx.getStats().collapsedSpans).toBe(1) + }) + + test('snapshot last-wins', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + + const s1: ContextCollapseSnapshotEntry = { + type: 'marble-origami-snapshot' as const, sessionId: uid('s1'), + staged: [{ startUuid: uid('a'), endUuid: uid('b'), summary: 'first', risk: 0.3, stagedAt: Date.now() }], + armed: true, lastSpawnTokens: 1000, + } + idx.restoreContextCollapseState([], s1) + expect(idx.getStats().stagedSpans).toBe(1) + + const s2: ContextCollapseSnapshotEntry = { + type: 'marble-origami-snapshot' as const, sessionId: uid('s1'), + staged: [ + { startUuid: uid('c'), endUuid: uid('d'), summary: 'second-a', risk: 0.9, stagedAt: Date.now() }, + { startUuid: uid('e'), endUuid: uid('f'), summary: 'second-b', risk: 0.5, stagedAt: Date.now() }, + ], + armed: false, lastSpawnTokens: 2000, + } + idx.restoreContextCollapseState([], s2) + expect(idx.getStats().stagedSpans).toBe(2) + }) +}) + +describe('applyCollapsesIfNeeded projection', () => { + test('re-applies a committed collapse to the next turn input, not just /context', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + + // A collapse committed on a prior turn (e.g. restored on resume). + idx.restoreContextCollapseState( + [ + { + type: 'marble-origami-commit' as const, + sessionId: uid('s1'), + collapseId: '0000000000000001', + summaryUuid: uid('sum'), + summaryContent: 'summary', + summary: 'summary', + firstArchivedUuid: uid('a1'), + lastArchivedUuid: uid('a3'), + }, + ], + undefined, + ) + + // The REPL rebuilds messagesForQuery from the full history every turn, so + // the archived span is present again on entry. + const fullHistory: Message[] = [ + makeUserMsg('u0'), + makeUserMsg('a1'), + makeAssistantMsg('a2'), + makeUserMsg('a3'), + makeUserMsg('u4'), + ] + + const { messages } = await idx.applyCollapsesIfNeeded( + fullHistory, + makeFakeToolUseContext(), + 'repl_main_thread' as any, + ) + + const uuids = messages.map(m => m.uuid) + expect(uuids).not.toContain(uid('a1')) + expect(uuids).not.toContain(uid('a2')) + expect(uuids).not.toContain(uid('a3')) + expect(uuids).toContain(uid('sum')) + expect(uuids).toContain(uid('u0')) + expect(uuids).toContain(uid('u4')) + }) + + test('clears a staged span that is already committed (crash-window restore)', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + + // Restore where the same span is in BOTH the commit log and the staged + // snapshot — possible if a crash landed the commit write but not the + // snapshot that drops it from staged. + idx.restoreContextCollapseState( + [ + { + type: 'marble-origami-commit' as const, + sessionId: uid('s1'), + collapseId: '0000000000000001', + summaryUuid: uid('sum'), + summaryContent: 'summary', + summary: 'summary', + firstArchivedUuid: uid('a1'), + lastArchivedUuid: uid('a3'), + }, + ], + { + type: 'marble-origami-snapshot' as const, + sessionId: uid('s1'), + staged: [ + { startUuid: uid('a1'), endUuid: uid('a3'), summary: 'summary', risk: 0.8, stagedAt: Date.now() }, + ], + armed: true, + lastSpawnTokens: 0, + }, + ) + expect(idx.getStats().stagedSpans).toBe(1) + + const fullHistory: Message[] = [ + makeUserMsg('u0'), + makeUserMsg('a1'), + makeAssistantMsg('a2'), + makeUserMsg('a3'), + makeUserMsg('u4'), + ] + + const { messages } = await idx.applyCollapsesIfNeeded( + fullHistory, + makeFakeToolUseContext(), + 'repl_main_thread' as any, + ) + + // Collapse applied, and the stale staged span is gone (not stuck forever). + const uuids = messages.map(m => m.uuid) + expect(uuids).toContain(uid('sum')) + expect(uuids).not.toContain(uid('a2')) + expect(idx.getStats().stagedSpans).toBe(0) + }) +}) + +describe('health', () => { + test('health fields initialized correctly', async () => { + await cleanState() + const idx = await import('./index.js') + idx.initContextCollapse() + const s = idx.getStats() + expect(s.health.totalSpawns).toBe(0) + expect(s.health.totalErrors).toBe(0) + expect(s.health.totalEmptySpawns).toBe(0) + expect(s.health.lastError).toBeNull() + expect(s.health.emptySpawnWarningEmitted).toBe(false) + }) +}) diff --git a/src/services/contextCollapse/index.ts b/src/services/contextCollapse/index.ts index 692875483..9aef71017 100644 --- a/src/services/contextCollapse/index.ts +++ b/src/services/contextCollapse/index.ts @@ -1,6 +1,4 @@ -// Stub — contextCollapse not included in source snapshot (feature-gated). -// All paths are behind feature('CONTEXT_COLLAPSE') + isContextCollapseEnabled(), -// which returns false here, so these inert implementations preserve behavior. +import { randomUUID } from 'crypto' import type { QuerySource } from '../../constants/querySource.js' import type { ToolUseContext } from '../../Tool.js' import type { @@ -8,16 +6,46 @@ import type { Message, StreamEvent, } from '../../types/message.js' +import { logError } from '../../utils/log.js' +import type { CanUseToolFn } from '../../hooks/useCanUseTool.js' +import { selectCollapseSpan, computeRisk } from './spanSelection.js' +import { CTX_AGENT_INSTRUCTION } from './ctxAgentPrompt.js' +import { + buildCollapsePlaceholder, + deriveCollapseId, + resetCollapseIdCounter, +} from './collapseUtils.js' -export function isContextCollapseEnabled(): boolean { - return false +import type { + ContextCollapseCommitEntry, + ContextCollapseSnapshotEntry, +} from '../../types/logs.js' + +// ── Internal types ────────────────────────────────────────────────────────── + +type CommittedCollapse = { + collapseId: string + summaryUuid: string + summaryContent: string + summary: string + firstArchivedUuid: string + lastArchivedUuid: string + // Count only: the archived messages are never read back (projectView splices + // by boundary uuid), so we keep the count for getStats and persist it so a + // resumed session reports the same figure as a live one. + archivedCount: number } -export function getContextCollapseState() { - return null +type StagedSpan = { + startUuid: string + endUuid: string + summary: string + risk: number + stagedAt: number } -/** Spawn/error health counters surfaced in /context and the token warning. */ +// ── Public types ──────────────────────────────────────────────────────────── + export type ContextCollapseHealth = { totalSpawns: number totalErrors: number @@ -33,68 +61,558 @@ export type ContextCollapseStats = { health: ContextCollapseHealth } -const INERT_STATS: ContextCollapseStats = { - collapsedSpans: 0, - collapsedMessages: 0, - stagedSpans: 0, - health: { +// ── Module-level state ───────────────────────────────────────────────────── + +let commitLog: CommittedCollapse[] = [] +let stagedQueue: StagedSpan[] = [] +let uuidToCollapseId: Map = new Map() +let collapseIdToUuid: Map = new Map() +let enabled = false +let armed = false +let lastSpawnTokens = 0 +let health: ContextCollapseHealth = { + totalSpawns: 0, + totalErrors: 0, + totalEmptySpawns: 0, + lastError: null, + emptySpawnWarningEmitted: false, +} +let listeners: Set<() => void> = new Set() +let spawnInProgress = false + +// ── Public API ────────────────────────────────────────────────────────────── + +export function isContextCollapseEnabled(): boolean { + return enabled +} + +/** + * Whether collapse currently has a real reduction in effect (committed or + * staged). Autocompact/blocking suppression keys off this rather than mere + * enablement: when collapse is on but has not reduced anything yet (e.g. the + * first over-threshold turn, before getLastCacheSafeParams() is populated, so + * spawnCtxAgent produces no span), the oversized transcript must still fall + * back to autocompact/blocking instead of reaching the API unguarded. + */ +export function hasActiveReduction(): boolean { + return enabled && (commitLog.length > 0 || stagedQueue.length > 0) +} + +/** + * Whether a query source owns the main-thread transcript that the collapse + * store reduces. In-process subagents (agent:*) and the ctx-agent + * (marble_origami) run in the SAME process and share this module-level store, + * but their message arrays are not the main transcript. Collapse must apply and + * suppress fallbacks only for the owning thread: a subagent staging/committing + * here would flip hasActiveReduction() globally and make the next main-thread + * turn suppress autocompact/blocking while projectView() no-ops on main + * messages, sending an oversized transcript to the API. Mirrors the + * main-thread classification in postCompactCleanup.ts. + */ +export function isMainThreadSource(querySource?: QuerySource): boolean { + return ( + querySource === undefined || + querySource.startsWith('repl_main_thread') || + querySource === 'sdk' + ) +} + +export function getContextCollapseState(): { + committedSpans: number + stagedSpans: number + armed: boolean + lastSpawnTokens: number + health: ContextCollapseHealth +} | null { + if (!enabled) return null + return { + committedSpans: commitLog.length, + stagedSpans: stagedQueue.length, + armed, + lastSpawnTokens, + health: { ...health }, + } +} + +export function initContextCollapse(): void { + const v = + typeof process !== 'undefined' ? process.env.CLAUDE_CONTEXT_COLLAPSE : undefined + const envOptIn = v === '1' || v === 'true' + let configOptIn = false + try { + /* eslint-disable @typescript-eslint/no-require-imports */ + const { getGlobalConfig } = + require('../../utils/config.js') as typeof import('../../utils/config.js') + /* eslint-enable @typescript-eslint/no-require-imports */ + configOptIn = getGlobalConfig().contextCollapseEnabled === true + } catch { + configOptIn = false + } + const optedIn = envOptIn || configOptIn + enabled = optedIn + armed = optedIn +} + +export function getStats(): ContextCollapseStats { + let collapsedMessages = 0 + for (const c of commitLog) { + collapsedMessages += c.archivedCount + } + return { + collapsedSpans: commitLog.length, + collapsedMessages, + stagedSpans: stagedQueue.length, + health: { ...health }, + } +} + +export function subscribe(listener: () => void): () => void { + listeners.add(listener) + return () => { + listeners.delete(listener) + } +} + +function notifyListeners(): void { + for (const listener of listeners) { + listener() + } +} + +export function resetContextCollapse(): void { + commitLog = [] + stagedQueue = [] + uuidToCollapseId = new Map() + collapseIdToUuid = new Map() + health = { totalSpawns: 0, totalErrors: 0, totalEmptySpawns: 0, lastError: null, emptySpawnWarningEmitted: false, - }, + } + lastSpawnTokens = 0 + spawnInProgress = false + resetCollapseIdCounter() + // Re-arm enabled sessions: reset is called by main-thread compaction cleanup + // and conversation rewind, which clear stale spans but should not permanently + // disable a session the user opted into. Mirrors restoreContextCollapseState. + armed = enabled + notifyListeners() } -/** One-time startup hook (setup.ts). No-op in this snapshot. */ -export function initContextCollapse(): void {} +// ── Core collapse algorithm ───────────────────────────────────────────────── -export function getStats(): ContextCollapseStats { - return INERT_STATS -} - -/** - * useSyncExternalStore-compatible subscription. The store never mutates in - * this snapshot, so the listener is never invoked. - */ -export function subscribe(_listener: () => void): () => void { - return () => {} -} - -/** Reset collapse state after a full compaction or rewind. No-op here. */ -export function resetContextCollapse(): void {} - -/** - * Apply any staged collapses before the API call. Inert: returns the - * input messages unchanged. - */ export async function applyCollapsesIfNeeded( messages: Message[], - _toolUseContext: ToolUseContext, - _querySource: QuerySource, + toolUseContext: ToolUseContext, + querySource: QuerySource, ): Promise<{ messages: Message[] }> { + if (!enabled) return { messages } + // Main-thread only: subagents (agent:*) and the ctx-agent (marble_origami) + // share this store but do not own the main transcript. Applying/staging here + // for them would mutate the shared store against the wrong messages. This + // subsumes the marble_origami skip. + if (!isMainThreadSource(querySource)) return { messages } + + // Re-apply committed collapses. messagesForQuery is rebuilt from the REPL's + // full history every turn (and the commit log is repopulated on resume), so + // without replaying the log here the archived spans would return to the model + // on the next turn. projectView is idempotent: already-collapsed spans no-op. + /* eslint-disable @typescript-eslint/no-require-imports */ + const { projectView } = + require('./operations.js') as typeof import('./operations.js') + /* eslint-enable @typescript-eslint/no-require-imports */ + messages = projectView(messages) + + // Commit drain: process all staged spans + if (stagedQueue.length > 0) { + messages = drainStaged(messages, true) + } + + // Spawn check: at ~95% of effective window, block and spawn ctx-agent + messages = await maybeSpawnCtxAgent(messages, toolUseContext, querySource) + return { messages } } -/** - * Whether a prompt-too-long error should be withheld pending a collapse - * drain. Always false — nothing is ever staged in this snapshot. - */ export function isWithheldPromptTooLong( - _message: Message | StreamEvent | undefined, - _isPromptTooLongMessage: (msg: AssistantMessage) => boolean, - _querySource: QuerySource, + message: Message | StreamEvent | undefined, + isPromptTooLongMessage: (msg: AssistantMessage) => boolean, + querySource: QuerySource, ): boolean { - return false + if (!enabled) return false + // The staged queue belongs to the main transcript; a subagent's PTL must not + // be withheld against it (see isMainThreadSource). + if (!isMainThreadSource(querySource)) return false + if (stagedQueue.length === 0) return false + if (!message || message.type !== 'assistant') return false + return isPromptTooLongMessage(message as AssistantMessage) } -/** - * Drain staged collapses to recover from a context overflow. Inert: nothing - * staged, so nothing is committed and messages pass through unchanged. - */ export function recoverFromOverflow( messages: Message[], - _querySource: QuerySource, + querySource: QuerySource, ): { messages: Message[]; committed: number } { - return { messages, committed: 0 } + if (!enabled) return { messages, committed: 0 } + // Draining the shared staged queue against a subagent's messages would corrupt + // the main-thread reduction; only the owning thread recovers here. + if (!isMainThreadSource(querySource)) return { messages, committed: 0 } + + const beforeCount = commitLog.length + messages = drainStaged(messages, true) + const committed = commitLog.length - beforeCount + + return { messages, committed } +} + +// ── Drain staged collapses ────────────────────────────────────────────────── + +function drainStaged( + messages: Message[], + persist: boolean, +): Message[] { + // A staged span can also already be in the commit log if a restore's snapshot + // predates the matching commit write (crash between the two persists). After + // projectView those messages are gone, so the span can't be drained normally; + // drop it here so it doesn't linger in stagedQueue and distort spawn/overflow. + const stale = stagedQueue.filter(s => + commitLog.some( + c => c.firstArchivedUuid === s.startUuid && c.lastArchivedUuid === s.endUuid, + ), + ) + if (stale.length > 0) { + stagedQueue = stagedQueue.filter(s => !stale.includes(s)) + } + + const processed: StagedSpan[] = [] + + for (const span of stagedQueue.sort((a, b) => a.stagedAt - b.stagedAt)) { + const firstIdx = messages.findIndex(m => m.uuid === span.startUuid) + const lastIdx = findLastIndex(messages, m => m.uuid === span.endUuid) + + if (firstIdx === -1 || lastIdx === -1 || firstIdx > lastIdx) continue + + const archivedCount = lastIdx - firstIdx + 1 + const collapseId = deriveCollapseId(span.startUuid) + const summaryUuid = randomUUID() + const summaryContent = buildCollapsePlaceholder(collapseId, span.summary) + + const placeholder: Message = { + type: 'system', + subtype: 'informational', + content: summaryContent, + uuid: summaryUuid, + timestamp: new Date().toISOString(), + isMeta: true, + // Survives normalizeMessagesForAPI as a user message so the summary + // reaches the model after the archived span is removed. + isCollapseSummary: true, + } as Message + + const committed: CommittedCollapse = { + collapseId, + summaryUuid, + summaryContent, + summary: span.summary, + firstArchivedUuid: span.startUuid, + lastArchivedUuid: span.endUuid, + archivedCount, + } + + commitLog.push(committed) + uuidToCollapseId.set(summaryUuid, collapseId) + collapseIdToUuid.set(collapseId, summaryUuid) + + messages = [ + ...messages.slice(0, firstIdx), + placeholder, + ...messages.slice(lastIdx + 1), + ] + + processed.push(span) + } + + if (processed.length > 0) { + stagedQueue = stagedQueue.filter(s => !processed.includes(s)) + + if (persist) { + // Persist commits before advancing the snapshot: if the snapshot (which no + // longer lists these spans as staged) landed first and the commit write + // failed, restore would find the spans neither staged nor committed and + // the collapse would vanish on resume. + void persistCommits(processed.length) + .then(() => persistSnapshot()) + .catch(() => {}) + } + + notifyListeners() + } else if (stale.length > 0) { + // Only cleared already-committed staged spans; sync the snapshot so the + // dropped entries don't reappear on the next restore. + if (persist) void persistSnapshot().catch(() => {}) + notifyListeners() + } + + return messages +} + +// ── Spawn mechanism ───────────────────────────────────────────────────────── + +const SPAWN_THRESHOLD_RATIO = 0.95 // 95% of effective window triggers spawn + +async function maybeSpawnCtxAgent( + messages: Message[], + toolUseContext: ToolUseContext, + querySource: QuerySource, +): Promise { + if (!armed) return messages + if (spawnInProgress) return messages + if (querySource === 'marble_origami') return messages + + /* eslint-disable @typescript-eslint/no-require-imports */ + const { getEffectiveContextWindowSize } = + require('../compact/autoCompact.js') as typeof import('../compact/autoCompact.js') + const { tokenCountWithEstimation } = + require('../../utils/tokens.js') as typeof import('../../utils/tokens.js') + /* eslint-enable @typescript-eslint/no-require-imports */ + + const model = toolUseContext.options.mainLoopModel ?? 'claude-sonnet-4' + const effectiveWindow = getEffectiveContextWindowSize(model) + const currentTokens = tokenCountWithEstimation(messages) + const threshold = Math.floor(effectiveWindow * SPAWN_THRESHOLD_RATIO) + + if (currentTokens < threshold) return messages + + // Don't re-spawn within the same token band unless we have new staged results + if ( + lastSpawnTokens > 0 && + currentTokens <= lastSpawnTokens + 500 && + stagedQueue.length === 0 + ) { + return messages + } + + spawnInProgress = true + health.totalSpawns++ + + try { + const spawnedSpans = await spawnCtxAgent(messages, toolUseContext, effectiveWindow) + + if (!spawnedSpans || spawnedSpans.length === 0) { + health.totalEmptySpawns++ + if (!health.emptySpawnWarningEmitted) { + health.emptySpawnWarningEmitted = true + } + } else { + for (const span of spawnedSpans) { + stagedQueue.push(span) + } + // Commit immediately after spawn at 95% + messages = drainStaged(messages, true) + } + + lastSpawnTokens = currentTokens + // When spans were drained, drainStaged(messages, true) already enforces + // commit -> snapshot ordering on its own async chain. A competing snapshot + // write here could land before those commits are durable, reopening the + // crash window that drops collapses on restore. Only persist directly when + // nothing was drained. + if (!spawnedSpans || spawnedSpans.length === 0) { + await persistSnapshot() + } + } catch (err) { + health.totalErrors++ + health.lastError = String(err) + logError(`contextCollapse spawn failed: ${String(err)}`) + } finally { + spawnInProgress = false + notifyListeners() + } + + return messages +} + +const denyCtxAgentTools: CanUseToolFn = async () => ({ + behavior: 'deny' as const, + message: 'Tool use is not allowed during context collapse', + decisionReason: { + type: 'other' as const, + reason: 'ctx-agent should only produce a text summary', + }, +}) + +async function spawnCtxAgent( + messages: Message[], + _toolUseContext: ToolUseContext, + effectiveWindow: number, +): Promise { + const span = selectCollapseSpan(messages, effectiveWindow) + if (!span) return [] + + const startIdx = messages.findIndex(m => (m.uuid as string) === span.startUuid) + const endIdx = messages.findIndex(m => (m.uuid as string) === span.endUuid) + if (startIdx === -1 || endIdx === -1 || startIdx > endIdx) return [] + const spanMessages = messages.slice(startIdx, endIdx + 1) + + /* eslint-disable @typescript-eslint/no-require-imports */ + const { runForkedAgent, getLastCacheSafeParams } = + require('../../utils/forkedAgent.js') as typeof import('../../utils/forkedAgent.js') + const { createUserMessage, getLastAssistantMessage, getAssistantMessageText } = + require('../../utils/messages.js') as typeof import('../../utils/messages.js') + /* eslint-enable @typescript-eslint/no-require-imports */ + + const base = getLastCacheSafeParams() + if (!base) return [] + const cacheSafeParams = { ...base, forkContextMessages: spanMessages } + + const result = await runForkedAgent({ + promptMessages: [createUserMessage({ content: CTX_AGENT_INSTRUCTION })], + cacheSafeParams, + canUseTool: denyCtxAgentTools, + querySource: 'marble_origami', + forkLabel: 'ctx-collapse', + maxTurns: 1, + skipCacheWrite: true, + }) + + const assistantMsg = getLastAssistantMessage(result.messages) + const summary = assistantMsg ? getAssistantMessageText(assistantMsg) : null + if (!assistantMsg || !summary || assistantMsg.isApiErrorMessage) return [] + const trimmed = summary.trim() + if (!trimmed) return [] + + const risk = computeRisk(startIdx, messages.length, span.tokenEstimate, effectiveWindow) + return [ + { + startUuid: span.startUuid, + endUuid: span.endUuid, + summary: trimmed, + risk, + stagedAt: Date.now(), + }, + ] +} + +// ── Persistence helpers ───────────────────────────────────────────────────── + +async function persistCommits(count: number): Promise { + /* eslint-disable @typescript-eslint/no-require-imports */ + const { recordContextCollapseCommit } = + require('../../utils/sessionStorage.js') as typeof import('../../utils/sessionStorage.js') + /* eslint-enable @typescript-eslint/no-require-imports */ + + const start = Math.max(0, commitLog.length - count) + for (let i = start; i < commitLog.length; i++) { + const c = commitLog[i]! + await recordContextCollapseCommit({ + collapseId: c.collapseId, + summaryUuid: c.summaryUuid, + summaryContent: c.summaryContent, + summary: c.summary, + firstArchivedUuid: c.firstArchivedUuid, + lastArchivedUuid: c.lastArchivedUuid, + archivedCount: c.archivedCount, + }) + } +} + +async function persistSnapshot(): Promise { + /* eslint-disable @typescript-eslint/no-require-imports */ + const { recordContextCollapseSnapshot } = + require('../../utils/sessionStorage.js') as typeof import('../../utils/sessionStorage.js') + /* eslint-enable @typescript-eslint/no-require-imports */ + + await recordContextCollapseSnapshot({ + staged: stagedQueue.map(s => ({ + startUuid: s.startUuid, + endUuid: s.endUuid, + summary: s.summary, + risk: s.risk, + stagedAt: s.stagedAt, + })), + armed, + lastSpawnTokens, + }) +} + +// ── Read-side projection (called by operations.ts) ───────────────────────── + +export function getCommitLogForProjection(): ReadonlyArray<{ + collapseId: string + summaryUuid: string + summaryContent: string + summary: string + firstArchivedUuid: string + lastArchivedUuid: string +}> { + return commitLog +} + +// ── Persist restore (called by persist.ts) ────────────────────────────────── + +export function restoreContextCollapseState( + commits: ContextCollapseCommitEntry[], + snapshot: ContextCollapseSnapshotEntry | undefined, +): void { + commitLog = [] + stagedQueue = [] + uuidToCollapseId = new Map() + collapseIdToUuid = new Map() + // Reset transient spawn state up front so a snapshot-less restore doesn't + // carry stale arming/last-spawn values from a previous session. + armed = enabled + lastSpawnTokens = 0 + + let maxCollapseId = 0 + + for (const entry of commits) { + const collapseIdNum = parseInt(entry.collapseId, 10) + if (collapseIdNum > maxCollapseId) { + maxCollapseId = collapseIdNum + } + + const c: CommittedCollapse = { + collapseId: entry.collapseId, + summaryUuid: entry.summaryUuid, + summaryContent: entry.summaryContent, + summary: entry.summary, + firstArchivedUuid: entry.firstArchivedUuid, + lastArchivedUuid: entry.lastArchivedUuid, + // Pre-field sessions persisted no count; report 0 rather than guess. + archivedCount: entry.archivedCount ?? 0, + } + + commitLog.push(c) + uuidToCollapseId.set(entry.summaryUuid, entry.collapseId) + collapseIdToUuid.set(entry.collapseId, entry.summaryUuid) + } + + resetCollapseIdCounter(maxCollapseId) + + if (snapshot) { + stagedQueue = snapshot.staged.map(s => ({ + startUuid: s.startUuid, + endUuid: s.endUuid, + summary: s.summary, + risk: s.risk, + stagedAt: s.stagedAt, + })) + armed = snapshot.armed + lastSpawnTokens = snapshot.lastSpawnTokens + } + + notifyListeners() +} + +// ── Utility ───────────────────────────────────────────────────────────────── + +function findLastIndex( + arr: T[], + predicate: (item: T) => boolean, +): number { + for (let i = arr.length - 1; i >= 0; i--) { + if (predicate(arr[i]!)) return i + } + return -1 } diff --git a/src/services/contextCollapse/operations.test.ts b/src/services/contextCollapse/operations.test.ts new file mode 100644 index 000000000..4a5d15c26 --- /dev/null +++ b/src/services/contextCollapse/operations.test.ts @@ -0,0 +1,276 @@ +import { randomUUID } from 'crypto' +import { afterEach, beforeEach, describe, expect, test } from 'bun:test' +import type { UUID } from 'crypto' +import type { Message } from '../../types/message.js' + +function uid(s: string): UUID { + return `00000000-0000-4000-8000-${s.padStart(12, '0')}` as UUID +} + +function makeUserMsg(id: string): Message { + return { + type: 'user', + uuid: uid(id), + timestamp: new Date().toISOString(), + message: { content: 'hello', role: 'user' as const }, + } as unknown as Message +} + +function makeAssistantMsg(id: string): Message { + return { + type: 'assistant', + uuid: uid(id), + timestamp: new Date().toISOString(), + message: { + id: randomUUID(), + model: 'claude-sonnet-4', + role: 'assistant', + stop_reason: 'end_turn', + stop_sequence: '', + type: 'message', + usage: { input_tokens: 10, output_tokens: 5 }, + content: [{ type: 'text' as const, text: 'ok' }], + context_management: null, + }, + } as unknown as Message +} + +describe('projectView', () => { + // Reset and rebuild a known commit before EACH test so outcomes never depend + // on shared module state or test execution order. + beforeEach(async () => { + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + const idx = await import('./index.js') + idx.resetContextCollapse() + idx.initContextCollapse() + idx.restoreContextCollapseState( + [ + { + type: 'marble-origami-commit' as const, + sessionId: uid('s1'), + collapseId: '0000000000000001', + summaryUuid: uid('sum1'), + summaryContent: 'test summary', + summary: 'test summary', + firstArchivedUuid: uid('a'), + lastArchivedUuid: uid('b'), + }, + ], + undefined, + ) + }) + + afterEach(async () => { + delete process.env.CLAUDE_CONTEXT_COLLAPSE + const idx = await import('./index.js') + idx.resetContextCollapse() + }) + + test('replays commit, replacing span with placeholder', async () => { + const mod = await import('./operations.js') + const msgs: Message[] = [makeUserMsg('a'), makeAssistantMsg('b'), makeUserMsg('c')] + const result = mod.projectView(msgs) + + expect(result.length).toBe(2) + expect(result[0]!.type).toBe('system') + expect((result[0]! as any).content).toContain('test summary') + expect(result[1]!.uuid).toBe(uid('c')) + }) + + test('collapsed summary survives normalizeMessagesForAPI as a user message', async () => { + // Regression: the projected placeholder is a system message, and + // normalizeMessagesForAPI drops system messages that are not local commands. + // Without the isCollapseSummary carve-out the summary (and the + // archived span it replaced) would vanish from the model's input. + const mod = await import('./operations.js') + const { normalizeMessagesForAPI } = await import('../../utils/messages.js') + const msgs: Message[] = [makeUserMsg('a'), makeAssistantMsg('b'), makeUserMsg('c')] + const projected = mod.projectView(msgs) + + const normalized = normalizeMessagesForAPI(projected) + const serialized = JSON.stringify(normalized) + expect(serialized).toContain('test summary') + expect(serialized).toContain(' { + // Regression: the system->user conversion in normalizeMessagesForAPI must + // preserve isMeta on the collapse placeholder. The snip-tag sweep + // (appendMessageTagToUserMessage) skips isMeta messages; if the conversion + // drops the flag, HISTORY_SNIP tags the summary with a snip_id and SnipTool + // can remove the only replacement for the archived span. + const mod = await import('./operations.js') + const { normalizeMessagesForAPI, appendMessageTagToUserMessage } = + await import('../../utils/messages.js') + const msgs: Message[] = [makeUserMsg('a'), makeAssistantMsg('b'), makeUserMsg('c')] + const normalized = normalizeMessagesForAPI(mod.projectView(msgs)) + + const summaryMsg = normalized.find(m => + JSON.stringify(m.message.content).includes('test summary'), + ) + expect(summaryMsg).toBeDefined() + expect(summaryMsg!.type).toBe('user') + expect((summaryMsg as { isMeta?: boolean }).isMeta).toBe(true) + + // With snip injection enabled the sweep leaves the meta summary untagged... + const swept = appendMessageTagToUserMessage(summaryMsg as never) + expect(JSON.stringify(swept.message.content)).not.toContain('snip_id=') + // ...while a normal (non-meta) user message still gets a snip id, so the + // exemption above is meaningful rather than a no-op. + const plain = appendMessageTagToUserMessage(makeUserMsg('c') as never) + expect(JSON.stringify(plain.message.content)).toContain('snip_id=') + }) + + test('collapse summary merged with the next user turn stays non-snippable', async () => { + // Regression: when the collapsed span ends right before the next user turn, + // normalizeMessagesForAPI merges the summary (now a user message) with that + // real user turn (Bedrock can't take consecutive user messages). Under + // HISTORY_SNIP that merge clears isMeta, so the combined block — which holds + // the only replacement for the archived span — must keep the + // isCollapseSummary marker or it gets a snip id and the model can drop it. + const mod = await import('./operations.js') + const { normalizeMessagesForAPI, appendMessageTagToUserMessage } = + await import('../../utils/messages.js') + // span [a..b] is committed in beforeEach; c is the adjacent real user turn. + const msgs: Message[] = [makeUserMsg('a'), makeAssistantMsg('b'), makeUserMsg('c')] + const normalized = normalizeMessagesForAPI(mod.projectView(msgs)) + + const merged = normalized.find(m => + JSON.stringify(m.message.content).includes('test summary'), + ) + expect(merged).toBeDefined() + expect(merged!.type).toBe('user') + // The real user turn was folded in, and the marker survived the merge. + expect(JSON.stringify(merged!.message.content)).toContain('hello') + expect((merged as { isCollapseSummary?: boolean }).isCollapseSummary).toBe(true) + // So the snip sweep leaves it untagged. + const swept = appendMessageTagToUserMessage(merged as never) + expect(JSON.stringify(swept.message.content)).not.toContain('snip_id=') + }) + + test('appendMessageTagToUserMessage skips a merged collapse block even when isMeta was cleared', async () => { + // The production HISTORY_SNIP merge sets isMeta=undefined; isCollapseSummary + // alone must keep the block non-snippable. + const { appendMessageTagToUserMessage } = await import('../../utils/messages.js') + const merged = { + type: 'user' as const, + uuid: uid('c'), + timestamp: new Date().toISOString(), + isMeta: undefined, + isCollapseSummary: true, + message: { role: 'user' as const, content: 's\nhello' }, + } + const swept = appendMessageTagToUserMessage(merged as never) + expect(JSON.stringify(swept.message.content)).not.toContain('snip_id=') + }) + + test('merging a collapse summary strips a snip id already baked into the user turn', async () => { + // The real user turn may be tagged before the merge; the combined block must + // shed that id so no resolvable snip id points at the summary. + const { mergeUserMessages } = await import('../../utils/messages.js') + const summary = { + type: 'user' as const, + uuid: uid('sum1'), + timestamp: new Date().toISOString(), + isMeta: true, + isCollapseSummary: true, + message: { role: 'user' as const, content: 's' }, + } + const realUser = { + type: 'user' as const, + uuid: uid('c'), + timestamp: new Date().toISOString(), + message: { + role: 'user' as const, + content: + 'hello\nsnip_id=abc123; system-generated; for snip tool use only; do not discuss in thinking or responses.', + }, + } + const merged = mergeUserMessages(summary as never, realUser as never) + expect((merged as { isCollapseSummary?: boolean }).isCollapseSummary).toBe(true) + expect(JSON.stringify(merged.message.content)).not.toContain('snip_id=') + expect(JSON.stringify(merged.message.content)).toContain('hello') + expect(JSON.stringify(merged.message.content)).toContain('collapsed') + }) + + test('normalizeMessages preserves isCollapseSummary when splitting array content', async () => { + // Regression: an already-merged collapse summary can carry array content. + // normalizeMessages splits multi-block user messages into single-block ones; + // if it forwards isMeta but drops isCollapseSummary, a later API-normalization + // pass tags the summary with a snip id and the model can drop the only + // replacement for the archived span. + const { normalizeMessages } = await import('../../utils/messages.js') + const summary = { + type: 'user' as const, + uuid: uid('sum1'), + timestamp: new Date().toISOString(), + isCollapseSummary: true, + message: { + role: 'user' as const, + content: [ + { type: 'text' as const, text: 's' }, + { type: 'text' as const, text: 'hello' }, + ], + }, + } + const normalized = normalizeMessages([summary as never]) + expect(normalized.length).toBe(2) + for (const m of normalized) { + expect((m as { isCollapseSummary?: boolean }).isCollapseSummary).toBe(true) + } + }) + + test('merging a collapse summary drops a text block that was only a snip marker', async () => { + // Regression: stripSnipTagsFromContent removed the marker text but kept the + // now-empty text block. Merging a collapse summary with a tool-result user + // turn whose trailing text block is solely the snip marker must not leave an + // empty { type:"text", text:"" } block in the API-bound content. + const { mergeUserMessages } = await import('../../utils/messages.js') + const summary = { + type: 'user' as const, + uuid: uid('sum1'), + timestamp: new Date().toISOString(), + isMeta: true, + isCollapseSummary: true, + message: { + role: 'user' as const, + content: [{ type: 'text' as const, text: 's' }], + }, + } + const realUser = { + type: 'user' as const, + uuid: uid('c'), + timestamp: new Date().toISOString(), + message: { + role: 'user' as const, + content: [ + { type: 'tool_result' as const, tool_use_id: 'tu1', content: 'result' }, + { + type: 'text' as const, + text: '\nsnip_id=abc123; system-generated; for snip tool use only; do not discuss in thinking or responses.', + }, + ], + }, + } + const merged = mergeUserMessages(summary as never, realUser as never) + const content = merged.message.content + expect(Array.isArray(content)).toBe(true) + const blocks = content as Array<{ type: string; text?: string }> + expect(blocks.some(b => b.type === 'text' && b.text === '')).toBe(false) + expect(JSON.stringify(content)).not.toContain('snip_id=') + expect(JSON.stringify(content)).toContain('collapsed') + }) + + test('silently skips missing boundaries', async () => { + const mod = await import('./operations.js') + const msgs: Message[] = [makeUserMsg('x'), makeAssistantMsg('y')] + const result = mod.projectView(msgs) + expect(result.length).toBe(2) + }) + + test('handles empty messages', async () => { + const mod = await import('./operations.js') + const result = mod.projectView([]) + expect(Array.isArray(result)).toBe(true) + }) +}) diff --git a/src/services/contextCollapse/operations.ts b/src/services/contextCollapse/operations.ts index 1aa1e23d5..8b49353ae 100644 --- a/src/services/contextCollapse/operations.ts +++ b/src/services/contextCollapse/operations.ts @@ -1,13 +1,55 @@ -// Stub — contextCollapse not included in source snapshot (feature-gated). -// projectView is the read-side projection of committed collapses onto the -// message history; with no collapses it is the identity function. import type { Message } from '../../types/message.js' -/** - * Project committed collapses onto the API view of the conversation. - * Inert: no collapses ever exist in this snapshot, so this returns the - * input unchanged. - */ export function projectView(messages: Message[]): Message[] { - return messages + /* eslint-disable @typescript-eslint/no-require-imports */ + const { getCommitLogForProjection } = + require('./index.js') as typeof import('./index.js') + /* eslint-enable @typescript-eslint/no-require-imports */ + + const commits = getCommitLogForProjection() + if (commits.length === 0) return messages + + let result = messages + + for (const c of commits) { + const firstIdx = result.findIndex(m => m.uuid === c.firstArchivedUuid) + const lastIdx = findLastIndex(result, m => m.uuid === c.lastArchivedUuid) + + if (firstIdx === -1 || lastIdx === -1 || firstIdx > lastIdx) continue + + // Reuse a stable timestamp from the replaced span so projectView stays a + // pure, deterministic projection (identical input -> identical output). + const placeholder: Message = { + type: 'system', + subtype: 'informational', + content: c.summaryContent, + uuid: c.summaryUuid, + timestamp: + result[firstIdx]?.timestamp ?? + result[lastIdx]?.timestamp ?? + new Date(0).toISOString(), + isMeta: true, + // Survives normalizeMessagesForAPI as a user message so the summary + // reaches the model after the archived span is removed. + isCollapseSummary: true, + } as Message + + result = [ + ...result.slice(0, firstIdx), + placeholder, + ...result.slice(lastIdx + 1), + ] + } + + return result +} + +function findLastIndex( + arr: T[], + predicate: (item: T) => boolean, +): number { + for (let i = arr.length - 1; i >= 0; i--) { + if (predicate(arr[i]!)) return i + } + return -1 } diff --git a/src/services/contextCollapse/persist.test.ts b/src/services/contextCollapse/persist.test.ts new file mode 100644 index 000000000..0df7d3f1e --- /dev/null +++ b/src/services/contextCollapse/persist.test.ts @@ -0,0 +1,72 @@ +import { afterEach, beforeEach, describe, expect, test } from 'bun:test' +import type { UUID } from 'crypto' + +function uid(s: string): UUID { + return `00000000-0000-4000-8000-${s.padStart(12, '0')}` as UUID +} + +describe('restoreFromEntries', () => { + // Isolate each test from shared module/env state. + beforeEach(async () => { + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + const idx = await import('./index.js') + idx.initContextCollapse() + idx.resetContextCollapse() + }) + + afterEach(async () => { + delete process.env.CLAUDE_CONTEXT_COLLAPSE + const idx = await import('./index.js') + idx.initContextCollapse() + idx.resetContextCollapse() + }) + + test('module loads and restoreFromEntries is callable', async () => { + const mod = await import('./persist.js') + expect(typeof mod.restoreFromEntries).toBe('function') + }) + + test('restores with empty inputs without error', async () => { + const mod = await import('./persist.js') + // First make sure context collapse is initialized + const idx = await import('./index.js') + idx.initContextCollapse() + + // restore should not throw + expect(() => mod.restoreFromEntries([], undefined)).not.toThrow() + }) + + test('restores with commits and snapshot', async () => { + const mod = await import('./persist.js') + const idx = await import('./index.js') + idx.initContextCollapse() + + mod.restoreFromEntries( + [ + { + type: 'marble-origami-commit' as const, + sessionId: uid('s1'), + collapseId: '0000000000000001', + summaryUuid: uid('sum1'), + summaryContent: 'test', + summary: 'test', + firstArchivedUuid: uid('a'), + lastArchivedUuid: uid('b'), + }, + ], + { + type: 'marble-origami-snapshot' as const, + sessionId: uid('s1'), + staged: [ + { startUuid: uid('c'), endUuid: uid('d'), summary: 'pending', risk: 0.5, stagedAt: Date.now() }, + ], + armed: true, + lastSpawnTokens: 10000, + }, + ) + + const stats = idx.getStats() + expect(stats.collapsedSpans).toBe(1) + expect(stats.stagedSpans).toBe(1) + }) +}) diff --git a/src/services/contextCollapse/persist.ts b/src/services/contextCollapse/persist.ts index b714e35e9..415648779 100644 --- a/src/services/contextCollapse/persist.ts +++ b/src/services/contextCollapse/persist.ts @@ -1,16 +1,16 @@ -// Stub — contextCollapse not included in source snapshot (feature-gated). -// Persistence/restore of collapse state across --resume. Inert no-op. import type { ContextCollapseCommitEntry, ContextCollapseSnapshotEntry, } from '../../types/logs.js' -/** - * Restore collapse state from transcript log entries on session resume. - * Inert: collapse is disabled in this snapshot, so restored entries are - * intentionally dropped. - */ export function restoreFromEntries( - _commits: ContextCollapseCommitEntry[], - _snapshot: ContextCollapseSnapshotEntry | undefined, -): void {} + commits: ContextCollapseCommitEntry[], + snapshot: ContextCollapseSnapshotEntry | undefined, +): void { + /* eslint-disable @typescript-eslint/no-require-imports */ + const { restoreContextCollapseState } = + require('./index.js') as typeof import('./index.js') + /* eslint-enable @typescript-eslint/no-require-imports */ + + restoreContextCollapseState(commits, snapshot) +} diff --git a/src/services/contextCollapse/spanSelection.test.ts b/src/services/contextCollapse/spanSelection.test.ts new file mode 100644 index 000000000..e07bd7863 --- /dev/null +++ b/src/services/contextCollapse/spanSelection.test.ts @@ -0,0 +1,107 @@ +import { describe, expect, test } from 'bun:test' +import type { Message } from '../../types/message.js' +import { isToolResultMessage, isTurnStart, MIN_COLLAPSE_TOKENS, selectCollapseSpan, computeRisk } from './spanSelection.js' + +function userMsg(content: any = 'hi', extra: Record = {}): Message { + return { type: 'user', uuid: 'u', timestamp: '', message: { role: 'user', content }, ...extra } as unknown as Message +} +function assistantMsg(): Message { + return { type: 'assistant', uuid: 'a', timestamp: '', message: { role: 'assistant', content: 'ok' } } as unknown as Message +} + +describe('isToolResultMessage', () => { + test('true for user message with a tool_result block', () => { + expect(isToolResultMessage(userMsg([{ type: 'tool_result', tool_use_id: 'x', content: 'r' }]))).toBe(true) + }) + test('false for plain user text', () => { + expect(isToolResultMessage(userMsg('hello'))).toBe(false) + }) + test('false for assistant message', () => { + expect(isToolResultMessage(assistantMsg())).toBe(false) + }) +}) + +describe('isTurnStart', () => { + test('true for a plain user message', () => { + expect(isTurnStart(userMsg('hello'))).toBe(true) + }) + test('false for a tool_result user message', () => { + expect(isTurnStart(userMsg([{ type: 'tool_result', tool_use_id: 'x', content: 'r' }]))).toBe(false) + }) + test('false for a meta user message', () => { + expect(isTurnStart(userMsg('hello', { isMeta: true }))).toBe(false) + }) + test('false for an assistant message', () => { + expect(isTurnStart(assistantMsg())).toBe(false) + }) +}) + +// Build an alternating user/assistant transcript with stable uuids. +function transcript(n: number): Message[] { + const out: Message[] = [] + for (let i = 0; i < n; i++) { + const isUser = i % 2 === 0 + out.push({ + type: isUser ? 'user' : 'assistant', + uuid: `m${i}`, + timestamp: '', + message: { role: isUser ? 'user' : 'assistant', content: isUser ? `q${i}` : `a${i}` }, + } as unknown as Message) + } + return out +} + +// Fake estimator: every message counts as 100 tokens. Deterministic + easy to reason about. +const flat100 = (msgs: Message[]) => msgs.length * 100 + +describe('selectCollapseSpan', () => { + test('returns null when there is no collapsible region', () => { + // 4 messages, window 1000: protected tail 30% = 300 => protects last 3; head protects m0. + expect(selectCollapseSpan(transcript(4), 1000, flat100)).toBeNull() + }) + + test('selects an oldest span anchored on turn boundaries', () => { + // 40 messages * 100 = 4000 tokens. window 1000. + // protected tail = 300 tokens => last 3 msgs protected (tail starts ~m37, snapped to a turn-start). + // head protects m0. candidate starts at the next turn-start (m2). + const span = selectCollapseSpan(transcript(40), 1000, flat100) + expect(span).not.toBeNull() + expect(span!.startUuid).toBe('m2') + expect(span!.startIndex).toBe(2) + expect(span!.tokenEstimate).toBeGreaterThanOrEqual(MIN_COLLAPSE_TOKENS) + }) + + test('returns null when the candidate span is below MIN_COLLAPSE_TOKENS', () => { + // 8 messages * 100 = 800 tokens total; any candidate span is well under 2000. + expect(selectCollapseSpan(transcript(8), 1000, flat100)).toBeNull() + }) + + test('span ends on a turn boundary (never mid-turn)', () => { + const msgs = transcript(40) + const span = selectCollapseSpan(msgs, 1000, flat100) + expect(span).not.toBeNull() + const endIdx = msgs.findIndex(m => (m.uuid as string) === span!.endUuid) + expect(endIdx).toBeGreaterThanOrEqual(0) + // The message immediately after the span must begin a new turn (or the span + // ends at the last message) — proving the span never cuts mid-turn. + const endsOnBoundary = endIdx === msgs.length - 1 || isTurnStart(msgs[endIdx + 1]!) + expect(endsOnBoundary).toBe(true) + }) +}) + +describe('computeRisk', () => { + test('older span (lower startIndex) scores higher than newer', () => { + const older = computeRisk(0, 100, 1000, 10000) + const newer = computeRisk(80, 100, 1000, 10000) + expect(older).toBeGreaterThan(newer) + }) + test('bigger span scores higher than smaller', () => { + const big = computeRisk(10, 100, 5000, 10000) + const small = computeRisk(10, 100, 500, 10000) + expect(big).toBeGreaterThan(small) + }) + test('always clamped to [0,1]', () => { + expect(computeRisk(0, 100, 999999, 10000)).toBeLessThanOrEqual(1) + expect(computeRisk(100, 100, 0, 10000)).toBeGreaterThanOrEqual(0) + }) +}) diff --git a/src/services/contextCollapse/spanSelection.ts b/src/services/contextCollapse/spanSelection.ts new file mode 100644 index 000000000..830c6e0fb --- /dev/null +++ b/src/services/contextCollapse/spanSelection.ts @@ -0,0 +1,117 @@ +import type { Message } from '../../types/message.js' +import { tokenCountWithEstimation } from '../../utils/tokens.js' + +/** Collapse when projected context exceeds this fraction; size the span to drop under it. */ +export const COLLAPSE_TARGET_RATIO = 0.7 +/** Most-recent fraction of the window that is never collapsed (the working set). */ +export const PROTECTED_TAIL_RATIO = 0.3 +/** Below this many estimated tokens, a span is not worth a model call. */ +export const MIN_COLLAPSE_TOKENS = 2000 + +/** A user message carrying a tool_result block (the back half of a tool exchange). */ +export function isToolResultMessage(msg: Message): boolean { + if (msg.type !== 'user') return false + const content = (msg as { message?: { content?: unknown } }).message?.content + if (!Array.isArray(content)) return false + return content.some( + (block: unknown) => + typeof block === 'object' && block !== null && (block as { type?: string }).type === 'tool_result', + ) +} + +/** Start of a real conversational turn: a non-meta user message that is not a tool_result. */ +export function isTurnStart(msg: Message): boolean { + return msg.type === 'user' && !(msg as { isMeta?: boolean }).isMeta && !isToolResultMessage(msg) +} + +export type SelectedSpan = { + startUuid: string + endUuid: string + startIndex: number + tokenEstimate: number +} + +/** + * Pick ONE oldest collapsible span of whole turns. + * + * Protects the first turn (task framing) and the most-recent PROTECTED_TAIL_RATIO + * of the window (the working set). Anchors both boundaries on turn-starts so a + * tool_use/tool_result pair is never split. Grows the span turn-by-turn until the + * projected post-collapse total drops under COLLAPSE_TARGET_RATIO of the window. + * Returns null when there is no candidate or the span is below MIN_COLLAPSE_TOKENS. + * + * `estimateTokens` is injectable for deterministic tests; defaults to the real estimator. + */ +export function selectCollapseSpan( + messages: Message[], + effectiveWindow: number, + estimateTokens: (msgs: Message[]) => number = tokenCountWithEstimation, +): SelectedSpan | null { + if (messages.length === 0 || effectiveWindow <= 0) return null + + // Protected head: index of the first real turn-start. + const headIdx = messages.findIndex(isTurnStart) + if (headIdx === -1) return null + + // Protected tail: walk from the end accumulating tokens until we've reserved + // PROTECTED_TAIL_RATIO of the window, then snap forward to a turn-start so the + // span ends cleanly just before a turn. + const tailBudget = effectiveWindow * PROTECTED_TAIL_RATIO + let tailTokens = 0 + let tailStartIdx = messages.length + for (let i = messages.length - 1; i >= 0; i--) { + tailTokens += estimateTokens([messages[i]!]) + tailStartIdx = i + if (tailTokens >= tailBudget) break + } + while (tailStartIdx < messages.length && !isTurnStart(messages[tailStartIdx]!)) { + tailStartIdx++ + } + + // Candidate region begins at the first turn-start AFTER the protected head. + let candidateStart = -1 + for (let i = headIdx + 1; i < tailStartIdx; i++) { + if (isTurnStart(messages[i]!)) { + candidateStart = i + break + } + } + if (candidateStart === -1) return null + + // Grow the span turn-by-turn until projected total < target. + const total = estimateTokens(messages) + const target = effectiveWindow * COLLAPSE_TARGET_RATIO + let endIdx = candidateStart + let spanTokens = 0 + let i = candidateStart + while (i < tailStartIdx) { + let next = i + 1 + while (next < tailStartIdx && !isTurnStart(messages[next]!)) next++ + endIdx = next - 1 + spanTokens = estimateTokens(messages.slice(candidateStart, endIdx + 1)) + if (total - spanTokens < target) break + i = next + } + + if (spanTokens < MIN_COLLAPSE_TOKENS) return null + + return { + startUuid: messages[candidateStart]!.uuid as string, + endUuid: messages[endIdx]!.uuid as string, + startIndex: candidateStart, + tokenEstimate: spanTokens, + } +} + +/** Drain-priority score in [0,1]: blends span age (older = higher) and size (bigger = higher). */ +export function computeRisk( + startIndex: number, + totalMessages: number, + spanTokens: number, + effectiveWindow: number, +): number { + const ageFactor = totalMessages > 0 ? 1 - startIndex / totalMessages : 0 + const sizeFactor = effectiveWindow > 0 ? Math.min(spanTokens / effectiveWindow, 1) : 0 + const risk = 0.5 * ageFactor + 0.5 * sizeFactor + return Math.max(0, Math.min(1, risk)) +} diff --git a/src/services/contextCollapse/spawnCtxAgent.test.ts b/src/services/contextCollapse/spawnCtxAgent.test.ts new file mode 100644 index 000000000..a7af0e852 --- /dev/null +++ b/src/services/contextCollapse/spawnCtxAgent.test.ts @@ -0,0 +1,201 @@ +import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test' +import type { Message } from '../../types/message.js' + +// Capture the real modules up front. mock.module() is global and mock.restore() +// does NOT undo it (see effort.codex.test.ts), so the beforeEach stubs below +// would otherwise bleed into later test files (e.g. the tokens stub makes +// autoCompact see every conversation as over-threshold). We re-mock each back +// to its real implementation in afterEach. +// +// `import * as` yields a LIVE namespace that bun's mock.module mutates in +// place, so we must snapshot the real exports into a plain object now (before +// any mock runs) rather than hold the namespace. The autoCompact.js stub +// (getEffectiveContextWindowSize) in particular leaks into compressToolHistory, +// which imports that function and uses it to size truncation. Restoring by +// specifier is safe: autoCompact.test.ts re-imports autoCompact fresh via a +// cache-busting nonce (a different specifier), so it is unaffected. +import * as spanSelectionNs from './spanSelection.js' +import * as forkedAgentNs from '../../utils/forkedAgent.js' +import * as tokensNs from '../../utils/tokens.js' +import * as autoCompactNs from '../compact/autoCompact.js' +import * as analyticsNs from '../../services/analytics/index.js' +import * as logNs from '../../utils/log.js' +import * as messagesNs from '../../utils/messages.js' + +const realSpanSelection = { ...spanSelectionNs } +const realForkedAgent = { ...forkedAgentNs } +const realTokens = { ...tokensNs } +const realAutoCompact = { ...autoCompactNs } +const realAnalytics = { ...analyticsNs } +const realLog = { ...logNs } +const realMessages = { ...messagesNs } + +// Build a transcript big enough to yield a collapsible span. +function bigTranscript(): Message[] { + const out: Message[] = [] + for (let i = 0; i < 200; i++) { + const isUser = i % 2 === 0 + out.push({ + type: isUser ? 'user' : 'assistant', + uuid: `m${i}`, + timestamp: '', + message: { role: isUser ? 'user' : 'assistant', content: 'x'.repeat(500) }, + } as unknown as Message) + } + return out +} + +let forkReturn: { messages: Message[]; totalUsage: any } | Error = { + messages: [ + { type: 'assistant', uuid: 'sum', timestamp: '', message: { role: 'assistant', content: [{ type: 'text', text: 'A concise summary.' }] } } as unknown as Message, + ], + totalUsage: {}, +} + +beforeEach(() => { + // Mock spanSelection to return a deterministic span so we can test the fork path. + mock.module('./spanSelection.js', () => ({ + isToolResultMessage: (msg: any) => false, + isTurnStart: (msg: any) => msg.type === 'user', + selectCollapseSpan: () => ({ + startUuid: 'm0', + endUuid: 'm19', + startIndex: 0, + tokenEstimate: 5000, + }), + computeRisk: () => 0.5, + COLLAPSE_TARGET_RATIO: 0.7, + PROTECTED_TAIL_RATIO: 0.3, + MIN_COLLAPSE_TOKENS: 2000, + })) + + // Mock forkedAgent to control the fork path. + mock.module('../../utils/forkedAgent.js', () => ({ + runForkedAgent: async () => { + if (forkReturn instanceof Error) throw forkReturn + return forkReturn + }, + getLastCacheSafeParams: () => ({ + systemPrompt: {} as any, + userContext: {}, + systemContext: {}, + toolUseContext: {} as any, + forkContextMessages: [], + }), + saveCacheSafeParams: () => {}, + createCacheSafeParams: () => ({ + systemPrompt: {} as any, + userContext: {}, + systemContext: {}, + toolUseContext: {} as any, + forkContextMessages: [], + }), + createSubagentContext: () => ({}), + createGetAppStateWithAllowedTools: () => () => ({}), + prepareForkedCommandContext: async () => ({}), + extractResultText: () => '', + cloneFileStateCache: () => ({}), + createChildAbortController: () => new AbortController(), + cloneContentReplacementState: () => ({}), + accumulateUsage: () => ({}), + updateUsage: () => ({}), + parseToolListFromCLI: () => ({}), + createDenialTrackingState: () => ({}), + recordSidechainTranscript: async () => {}, + })) + + // Mock tokens so the 95% threshold is easily hit. + mock.module('../../utils/tokens.js', () => ({ + tokenCountWithEstimation: () => 100000, + getTokenUsage: () => undefined, + tokenCountFromLastAPIResponse: () => 0, + getIncrementalTokenCounter: () => ({}), + getTokenCountFromUsage: () => 0, + roughTokenCountEstimation: () => 100, + roughTokenCountEstimationForMessages: () => 500, + })) + + // Mock autoCompact since maybeSpawnCtxAgent requires it. + mock.module('../compact/autoCompact.js', () => ({ + getEffectiveContextWindowSize: () => 20000, + })) + + // Silence analytics/log noise. + mock.module('../../services/analytics/index.js', () => ({ + logEvent: () => {}, + })) + mock.module('../../utils/log.js', () => ({ + logError: () => {}, + })) + + // Mock messages.js since spawnCtxAgent uses require() for it. + mock.module('../../utils/messages.js', () => ({ + createUserMessage: ({ content }: { content: string }) => ({ + type: 'user', + uuid: 'prompt-uuid', + timestamp: '', + message: { role: 'user', content }, + }), + getLastAssistantMessage: (msgs: Message[]) => + msgs.findLast((m: Message) => m.type === 'assistant'), + getAssistantMessageText: (msg: Message) => { + if (msg.type !== 'assistant') return null + if (typeof msg.message.content === 'string') return msg.message.content + return null + }, + })) +}) + +afterEach(async () => { + mock.restore() + // Restore module stubs to their real implementations (mock.restore() does + // not undo mock.module) so they do not bleed into other test files. + mock.module('./spanSelection.js', () => realSpanSelection) + mock.module('../../utils/forkedAgent.js', () => realForkedAgent) + mock.module('../../utils/tokens.js', () => realTokens) + mock.module('../compact/autoCompact.js', () => realAutoCompact) + mock.module('../../services/analytics/index.js', () => realAnalytics) + mock.module('../../utils/log.js', () => realLog) + mock.module('../../utils/messages.js', () => realMessages) + delete process.env.CLAUDE_CONTEXT_COLLAPSE + // Re-sync enablement to the now-unset env so enabled=true does not leak + // into later test files. + const idx = await import('./index.js') + idx.resetContextCollapse() + idx.initContextCollapse() +}) + +function ctx(): any { + return { options: { mainLoopModel: 'claude-sonnet-4', tools: [] }, abortController: new AbortController() } +} + +describe('spawnCtxAgent (via maybeSpawnCtxAgent)', () => { + test('stages exactly one span when the fork returns summary text', async () => { + forkReturn = { + messages: [ + { type: 'assistant', uuid: 'sum', timestamp: '', message: { role: 'assistant', content: 'A concise summary.' } } as unknown as Message, + ], + totalUsage: {}, + } + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + const idx = await import('./index.js') + idx.resetContextCollapse() + idx.initContextCollapse() + const before = idx.getStats() + await idx.applyCollapsesIfNeeded(bigTranscript(), ctx(), 'repl_main_thread') + const stats = idx.getStats() + // At least one span was collapsed or staged — real work happened. + expect(stats.collapsedSpans + stats.stagedSpans).toBeGreaterThan(before.collapsedSpans + before.stagedSpans) + expect(stats.health.totalSpawns).toBeGreaterThan(0) + }) + + test('records an error when the fork throws', async () => { + forkReturn = new Error('boom') + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + const idx = await import('./index.js') + idx.resetContextCollapse() + idx.initContextCollapse() + await idx.applyCollapsesIfNeeded(bigTranscript(), ctx(), 'repl_main_thread') + expect(idx.getStats().health.totalErrors).toBeGreaterThan(0) + }) +}) diff --git a/src/tools/CtxInspectTool/CtxInspectTool.test.ts b/src/tools/CtxInspectTool/CtxInspectTool.test.ts new file mode 100644 index 000000000..6e1cb6158 --- /dev/null +++ b/src/tools/CtxInspectTool/CtxInspectTool.test.ts @@ -0,0 +1,79 @@ +import { describe, expect, test } from 'bun:test' +import { CtxInspectTool } from './CtxInspectTool.js' +import { CTX_INSPECT_TOOL_NAME, getPrompt } from './prompt.js' + +describe('CtxInspectTool', () => { + test('tool name matches constant', () => { + expect(CtxInspectTool.name).toBe(CTX_INSPECT_TOOL_NAME) + expect(CTX_INSPECT_TOOL_NAME).toBe('ctx_inspect') + }) + + test('tool is read-only', () => { + expect(CtxInspectTool.isReadOnly()).toBe(true) + }) + + test('tool is gated on context-collapse opt-in', () => { + const mod = require('../../services/contextCollapse/index.js') + const prev = process.env.CLAUDE_CONTEXT_COLLAPSE + + delete process.env.CLAUDE_CONTEXT_COLLAPSE + mod.initContextCollapse() + expect(CtxInspectTool.isEnabled()).toBe(false) + + process.env.CLAUDE_CONTEXT_COLLAPSE = '1' + mod.initContextCollapse() + expect(CtxInspectTool.isEnabled()).toBe(true) + + if (prev === undefined) delete process.env.CLAUDE_CONTEXT_COLLAPSE + else process.env.CLAUDE_CONTEXT_COLLAPSE = prev + mod.initContextCollapse() + }) + + test('tool is concurrency safe', () => { + expect(CtxInspectTool.isConcurrencySafe()).toBe(true) + }) + + test('prompt returns non-empty string', () => { + const prompt = getPrompt() + expect(typeof prompt).toBe('string') + expect(prompt.length).toBeGreaterThan(0) + expect(prompt).toContain('context collapse') + }) + + test('description returns prompt', async () => { + const desc = await CtxInspectTool.description() + expect(desc).toBe(getPrompt()) + }) + + test('mapToolResultToToolResultBlockParam formats JSON output', () => { + const output = { + committedSpans: 2, + collapsedMessages: 10, + stagedSpans: 1, + armed: true, + health: { + totalSpawns: 3, + totalErrors: 0, + totalEmptySpawns: 1, + lastError: null, + emptySpawnWarningEmitted: false, + }, + } + const result = CtxInspectTool.mapToolResultToToolResultBlockParam( + output, + 'toolu_abc', + ) + expect(result.type).toBe('tool_result') + expect(result.tool_use_id).toBe('toolu_abc') + const content = JSON.parse(result.content as string) + expect(content.committedSpans).toBe(2) + expect(content.collapsedMessages).toBe(10) + }) + + test('tool loads and has expected shape', () => { + expect(typeof CtxInspectTool.name).toBe('string') + expect(typeof CtxInspectTool.description).toBe('function') + expect(typeof CtxInspectTool.prompt).toBe('function') + expect(typeof CtxInspectTool.call).toBe('function') + }) +}) diff --git a/src/tools/CtxInspectTool/CtxInspectTool.ts b/src/tools/CtxInspectTool/CtxInspectTool.ts new file mode 100644 index 000000000..4127de721 --- /dev/null +++ b/src/tools/CtxInspectTool/CtxInspectTool.ts @@ -0,0 +1,89 @@ +import type { ToolResultBlockParam } from '@anthropic-ai/sdk/resources/index.mjs' +import { z } from 'zod/v4' +import { buildTool, type ToolDef } from '../../Tool.js' +import { lazySchema } from '../../utils/lazySchema.js' +import { getPrompt, CTX_INSPECT_TOOL_NAME } from './prompt.js' + +const inputSchema = lazySchema(() => + z.strictObject({}), +) +type InputSchema = ReturnType + +type CtxInspectOutput = { + committedSpans: number + collapsedMessages: number + stagedSpans: number + armed: boolean + health: { + totalSpawns: number + totalErrors: number + totalEmptySpawns: number + lastError: string | null + emptySpawnWarningEmitted: boolean + } +} + +export const CtxInspectTool = buildTool({ + name: CTX_INSPECT_TOOL_NAME, + isEnabled() { + /* eslint-disable @typescript-eslint/no-require-imports */ + const { isContextCollapseEnabled } = + require('../../services/contextCollapse/index.js') as typeof import('../../services/contextCollapse/index.js') + /* eslint-enable @typescript-eslint/no-require-imports */ + return isContextCollapseEnabled() + }, + isConcurrencySafe() { + return true + }, + isReadOnly() { + return true + }, + async description() { + return getPrompt() + }, + async prompt() { + return getPrompt() + }, + get inputSchema(): InputSchema { + return inputSchema() + }, + async call(_input, _context) { + /* eslint-disable @typescript-eslint/no-require-imports */ + const { getStats, getContextCollapseState } = + require('../../services/contextCollapse/index.js') as typeof import('../../services/contextCollapse/index.js') + /* eslint-enable @typescript-eslint/no-require-imports */ + + const stats = getStats() + const state = getContextCollapseState() + + const result: CtxInspectOutput = { + committedSpans: stats.collapsedSpans, + collapsedMessages: stats.collapsedMessages, + stagedSpans: stats.stagedSpans, + armed: state?.armed ?? false, + health: { + totalSpawns: stats.health.totalSpawns, + totalErrors: stats.health.totalErrors, + totalEmptySpawns: stats.health.totalEmptySpawns, + lastError: stats.health.lastError, + emptySpawnWarningEmitted: stats.health.emptySpawnWarningEmitted, + }, + } + + return { data: result } + }, + renderToolUseMessage() { + return null + }, + maxResultSizeChars: 4096, + mapToolResultToToolResultBlockParam( + output: CtxInspectOutput, + toolUseID: string, + ): ToolResultBlockParam { + return { + type: 'tool_result', + tool_use_id: toolUseID, + content: JSON.stringify(output, null, 2), + } + }, +} satisfies ToolDef) diff --git a/src/tools/CtxInspectTool/prompt.ts b/src/tools/CtxInspectTool/prompt.ts new file mode 100644 index 000000000..2b518b8f5 --- /dev/null +++ b/src/tools/CtxInspectTool/prompt.ts @@ -0,0 +1,9 @@ +export const CTX_INSPECT_TOOL_NAME = 'ctx_inspect' + +export function getPrompt(): string { + return `Inspect context collapse state to understand what spans have been collapsed and what spans are staged for collapse. + +Use this tool when you need to understand the current state of context management - what's been summarized away and what's queued for summarization. + +This is a read-only introspection tool. It has no side effects.` +} diff --git a/src/types/logs.ts b/src/types/logs.ts index 5fdba1dc0..cb2f7165e 100644 --- a/src/types/logs.ts +++ b/src/types/logs.ts @@ -271,11 +271,11 @@ export type SpeculationAcceptMessage = { * Persisted context-collapse commit. The archived messages themselves are * NOT persisted — they're already in the transcript as ordinary user/ * assistant messages. We only persist enough to reconstruct the splice - * instruction (boundary uuids) and the summary placeholder (which is NOT - * in the transcript because it's never yielded to the REPL). - * - * On restore, the store reconstructs CommittedCollapse with archived=[]; - * projectView lazily fills the archive the first time it finds the span. + * instruction (boundary uuids), the summary placeholder (which is NOT + * in the transcript because it's never yielded to the REPL), and the count + * of archived messages so getStats reports the same figure after a resume as + * it did live (projectView removes the span but does not refill any per-commit + * message list). * * Discriminator is obfuscated to match the gate name. sessionStorage.ts * isn't feature-gated (it's the generic transcript plumbing used by every @@ -297,6 +297,12 @@ export type ContextCollapseCommitEntry = { /** Span boundaries — projectView finds these in the resumed Message[]. */ firstArchivedUuid: string lastArchivedUuid: string + /** + * Number of messages the span archived. Optional for back-compat with + * sessions persisted before this field existed (those restore as 0). Lets + * getStats report accurate collapsedMessages after a resume. + */ + archivedCount?: number } /** diff --git a/src/types/message.ts b/src/types/message.ts index 12905aa1c..7b7f7eb72 100644 --- a/src/types/message.ts +++ b/src/types/message.ts @@ -100,6 +100,14 @@ export interface UserMessage { isVisibleInTranscriptOnly?: boolean isVirtual?: boolean isCompactSummary?: boolean + /** + * Set when a context-collapse summary placeholder is converted to (or merged + * into) a user message. Keeps the `` summary non-snippable: the + * snip-tag sweep skips these, and merges that absorb a summary inherit the + * flag and drop any pre-baked snip id, so the only replacement for an archived + * span can never be queued for removal. + */ + isCollapseSummary?: boolean summarizeMetadata?: { messagesSummarized: number userContext?: string @@ -202,6 +210,13 @@ interface SystemMessageBase { export interface SystemInformationalMessage extends SystemMessageBase { subtype: 'informational' content: string + /** + * Marks a context-collapse summary placeholder. The transcript renders it + * like any informational notice, but normalizeMessagesForAPI converts it into + * a user message so the `` summary still reaches the model after + * the archived span is removed. + */ + isCollapseSummary?: boolean } export interface SystemPermissionRetryMessage extends SystemMessageBase { diff --git a/src/utils/config.ts b/src/utils/config.ts index bada21859..a50996ea2 100644 --- a/src/utils/config.ts +++ b/src/utils/config.ts @@ -283,6 +283,7 @@ export type GlobalConfig = { bypassPermissionsModeAccepted?: boolean hasUsedBackslashReturn?: boolean autoCompactEnabled: boolean // Controls whether auto-compact is enabled + contextCollapseEnabled: boolean // Opt-in: collapse old transcript spans into summaries (lossy; off by default) toolHistoryCompressionEnabled: boolean // Compress old tool_result content for small-context providers showTurnDuration: boolean // Controls whether to show turn duration message (e.g., "Cooked for 1m 6s") // Controls whether to show per-query cache hit/miss stats at the end of each turn. @@ -696,6 +697,7 @@ function createDefaultGlobalConfig(): GlobalConfig { verbose: false, editorMode: 'normal', autoCompactEnabled: true, + contextCollapseEnabled: false, toolHistoryCompressionEnabled: true, showTurnDuration: true, showCacheStats: 'compact', @@ -748,6 +750,7 @@ export const GLOBAL_CONFIG_KEYS = [ 'editorMode', 'hasUsedBackslashReturn', 'autoCompactEnabled', + 'contextCollapseEnabled', 'toolHistoryCompressionEnabled', 'showTurnDuration', 'showCacheStats', diff --git a/src/utils/messages.ts b/src/utils/messages.ts index 0a684c84d..3457bd2b7 100644 --- a/src/utils/messages.ts +++ b/src/utils/messages.ts @@ -468,6 +468,7 @@ export function createUserMessage({ isVisibleInTranscriptOnly, isVirtual, isCompactSummary, + isCollapseSummary, summarizeMetadata, toolUseResult, mcpMeta, @@ -483,6 +484,7 @@ export function createUserMessage({ isVisibleInTranscriptOnly?: boolean isVirtual?: boolean isCompactSummary?: boolean + isCollapseSummary?: boolean toolUseResult?: unknown // Matches tool's `Output` type /** MCP protocol metadata to pass through to SDK consumers (never sent to model) */ mcpMeta?: { @@ -519,6 +521,7 @@ export function createUserMessage({ isVisibleInTranscriptOnly, isVirtual, isCompactSummary, + isCollapseSummary, summarizeMetadata, uuid: (uuid as UUID | undefined) || randomUUID(), timestamp: timestamp ?? new Date().toISOString(), @@ -818,6 +821,7 @@ export function normalizeMessages(messages: Message[]): NormalizedMessage[] { toolUseResult: message.toolUseResult, mcpMeta: message.mcpMeta, isMeta: message.isMeta, + isCollapseSummary: message.isCollapseSummary, isVisibleInTranscriptOnly: message.isVisibleInTranscriptOnly, isVirtual: message.isVirtual, timestamp: message.timestamp, @@ -1534,6 +1538,19 @@ export function isSystemLocalCommandMessage( return message.type === 'system' && message.subtype === 'local_command' } +/** + * A context-collapse summary placeholder. Like local-command system messages, + * its content must survive model-input normalization (converted to a user + * message) so the collapsed-span summary stays visible to the model. + */ +export function isCollapseSummaryMessage(message: Message): boolean { + return ( + message.type === 'system' && + message.subtype === 'informational' && + (message as { isCollapseSummary?: boolean }).isCollapseSummary === true + ) +} + /** * Strips tool_reference blocks for tools that no longer exist from tool_result content. * This handles the case where a session was saved with MCP tools that are no longer @@ -1622,7 +1639,10 @@ function stripUnavailableToolReferencesFromUserMessage( export function appendMessageTagToUserMessage( message: UserMessage, ): UserMessage { - if (message.isMeta) { + // isCollapseSummary blocks must never carry a snip id: the model could queue + // the only replacement for an archived span for removal. A merge can clear + // isMeta while keeping isCollapseSummary, so both are checked here. + if (message.isMeta || message.isCollapseSummary) { return message } @@ -1713,6 +1733,42 @@ export function appendMessageTagToUserMessage( } } +// Matches the exact internal snip marker appended by appendMessageTagToUserMessage +// (with or without the leading newline used for the no-text-block variant). The +// body has no '<' chars, so [^<]* terminates cleanly at the closing tag. +const SNIP_TAG_PATTERN = + /\n?snip_id=[^<]*<\/system-reminder>/g + +/** + * Remove any internal snip marker from user content. Used when a merge folds a + * collapse summary into a real user turn: the real turn may have been tagged + * before the merge, and the merged block must not present a snip id (it carries + * the only replacement for an archived span). + */ +function stripSnipTagsFromContent( + content: string | ContentBlockParam[], +): string | ContentBlockParam[] { + if (typeof content === 'string') { + return content.replace(SNIP_TAG_PATTERN, '') + } + if (!Array.isArray(content)) return content + const result: ContentBlockParam[] = [] + for (const block of content) { + if (block?.type === 'text') { + const original = (block as TextBlockParam).text + const text = original.replace(SNIP_TAG_PATTERN, '') + // Drop a text block whose only content was the snip marker; sending an + // empty text block alongside the collapse summary is invalid. Pre-existing + // empty blocks are left untouched so this stays scoped to the merge path. + if (text === '' && original !== '') continue + result.push({ ...block, text }) + } else { + result.push(block) + } + } + return result +} + /** * Strips tool_reference blocks from tool_result content in a user message. * tool_reference blocks are only valid when the tool search beta is enabled. @@ -2136,10 +2192,13 @@ export function normalizeMessagesForAPI( | UserMessage | AssistantMessage | AttachmentMessage - | SystemLocalCommandMessage => { + | SystemLocalCommandMessage + | SystemInformationalMessage => { if ( _.type === 'progress' || - (_.type === 'system' && !isSystemLocalCommandMessage(_)) || + (_.type === 'system' && + !isSystemLocalCommandMessage(_) && + !isCollapseSummaryMessage(_)) || isSyntheticApiErrorMessage(_) ) { return false @@ -2151,11 +2210,25 @@ export function normalizeMessagesForAPI( switch (message.type) { case 'system': { // local_command system messages need to be included as user messages - // so the model can reference previous command output in later turns + // so the model can reference previous command output in later turns. + // Context-collapse summaries take the same path so the + // summary stays visible after its archived span is removed. + // + // Preserve isMeta: collapse-summary placeholders are created isMeta so + // the snip-tag sweep (appendMessageTagToUserMessage skips isMeta) does + // not mark the only replacement for an archived span as snippable, + // which would let the model remove the summary collapse relies on. + // local_command messages carry no isMeta and stay snippable as before. const userMsg = createUserMessage({ content: message.content, uuid: message.uuid, timestamp: message.timestamp, + isMeta: message.isMeta, + // Carry the collapse-summary marker onto the user message so it + // stays non-snippable even after a merge clears isMeta (a merge + // with an adjacent real user turn would otherwise expose the + // summary under a snippable id). + isCollapseSummary: isCollapseSummaryMessage(message), }) const lastMessage = last(result) if (lastMessage?.type === 'user') { @@ -2499,6 +2572,15 @@ function isToolResultMessage(msg: Message): boolean { export function mergeUserMessages(a: UserMessage, b: UserMessage): UserMessage { const lastContent = normalizeUserTextContent(a.message.content) const currentContent = normalizeUserTextContent(b.message.content) + // A merge that absorbs a collapse summary stays non-snippable: the combined + // block holds the only replacement for an archived span, so it must keep the + // marker and shed any snip id a real-user operand was tagged with pre-merge. + const isCollapseSummary = + a.isCollapseSummary || b.isCollapseSummary ? (true as const) : undefined + const finalize = ( + content: string | ContentBlockParam[], + ): string | ContentBlockParam[] => + isCollapseSummary ? stripSnipTagsFromContent(content) : content if (feature('HISTORY_SNIP')) { // A merged message is only meta if ALL merged messages are meta. If any // operand is real user content, the result must not be flagged isMeta @@ -2514,11 +2596,12 @@ export function mergeUserMessages(a: UserMessage, b: UserMessage): UserMessage { return { ...a, isMeta: a.isMeta && b.isMeta ? (true as const) : undefined, + isCollapseSummary, uuid: a.isMeta ? b.uuid : a.uuid, message: { ...a.message, - content: hoistToolResults( - joinTextAtSeam(lastContent, currentContent), + content: finalize( + hoistToolResults(joinTextAtSeam(lastContent, currentContent)), ), }, } @@ -2526,12 +2609,15 @@ export function mergeUserMessages(a: UserMessage, b: UserMessage): UserMessage { } return { ...a, + isCollapseSummary, // Preserve the non-meta message's uuid so snip ids (derived from uuid) // stay stable across API calls (meta messages like system context get fresh uuids each call) uuid: a.isMeta ? b.uuid : a.uuid, message: { ...a.message, - content: hoistToolResults(joinTextAtSeam(lastContent, currentContent)), + content: finalize( + hoistToolResults(joinTextAtSeam(lastContent, currentContent)), + ), }, } } diff --git a/src/utils/sessionStorage.ts b/src/utils/sessionStorage.ts index 9500bb94a..066533b5b 100644 --- a/src/utils/sessionStorage.ts +++ b/src/utils/sessionStorage.ts @@ -1590,6 +1590,7 @@ export async function recordContextCollapseCommit(commit: { summary: string firstArchivedUuid: string lastArchivedUuid: string + archivedCount: number }): Promise { const sessionId = getSessionId() as UUID if (!sessionId) return