feat(context-collapse): opt-in between-turns context collapse (span summarization) (#1619)

* 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
<collapsed> 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
<collapsed> 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.
This commit is contained in:
beardthelion
2026-06-17 11:02:54 +08:00
committed by GitHub
parent 650fae952d
commit d5588ea80d
29 changed files with 2713 additions and 94 deletions
+1 -1
View File
@@ -27,7 +27,7 @@ const featureFlags: Record<string, boolean> = {
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
@@ -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', () => {
@@ -51,9 +51,16 @@ 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 */
// 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)
const appState = getAppState()
+9
View File
@@ -22,9 +22,18 @@ 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 */
// 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;
}
export async function call(onDone: LocalJSXCommandOnDone, context: LocalJSXCommandContext): Promise<React.ReactNode> {
+40 -1
View File
@@ -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;
+9 -1
View File
@@ -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
+24
View File
@@ -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)
})
})
+13 -2
View File
@@ -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))
}
+11 -5
View File
@@ -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
}
}
@@ -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('<collapsed id="0000000000000001">summarized content</collapsed>')
})
test('escapes nothing - raw summary is embedded', () => {
const result = buildCollapsePlaceholder('42', 'text with <tags> and "quotes"')
expect(result).toContain('<tags>')
expect(result).toContain('"quotes"')
})
})
describe('isCollapsePlaceholder', () => {
test('detects collapse placeholder system messages', () => {
const msg = makeSysMsg('<collapsed id="1">summary</collapsed>')
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)
})
})
@@ -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 `<collapsed id="${collapseId}">${summary}</collapsed>`
}
export function isCollapsePlaceholder(message: Message): boolean {
if (message.type !== 'system') return false
const content = 'content' in message ? (message as Record<string, unknown>).content : ''
return typeof content === 'string' && content.startsWith('<collapsed')
}
export function isWithinCollapsedSpan(
message: Message,
messages: Message[],
commits: { firstArchivedUuid: string; lastArchivedUuid: string }[],
): boolean {
const msgUuid = message.uuid as string | undefined
if (!msgUuid) return false
const msgIdx = messages.findIndex(m => 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<T>(
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 }
@@ -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')
+641
View File
@@ -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: '<collapsed id="0000000000000001">test summary</collapsed>',
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: '<collapsed id="0000000000000001">summary</collapsed>',
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: '<collapsed id="0000000000000007">summary</collapsed>',
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: '<collapsed id="0000000000000007">summary</collapsed>',
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: '<collapsed id="0000000000000008">summary</collapsed>',
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: '<collapsed id="0000000000000042">x</collapsed>',
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: '<collapsed id="0000000000000001">summary</collapsed>',
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: '<collapsed id="0000000000000001">summary</collapsed>',
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)
})
})
+563 -45
View File
@@ -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<string, string> = new Map()
let collapseIdToUuid: Map<string, string> = 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
}
/** One-time startup hook (setup.ts). No-op in this snapshot. */
export function initContextCollapse(): void {}
/**
* 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 {
return INERT_STATS
let collapsedMessages = 0
for (const c of commitLog) {
collapsedMessages += c.archivedCount
}
return {
collapsedSpans: commitLog.length,
collapsedMessages,
stagedSpans: stagedQueue.length,
health: { ...health },
}
}
/**
* useSyncExternalStore-compatible subscription. The store never mutates in
* this snapshot, so the listener is never invoked.
*/
export function subscribe(_listener: () => void): () => void {
return () => {}
export function subscribe(listener: () => void): () => void {
listeners.add(listener)
return () => {
listeners.delete(listener)
}
}
/** Reset collapse state after a full compaction or rewind. No-op here. */
export function resetContextCollapse(): void {}
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()
}
// ── Core collapse algorithm ─────────────────────────────────────────────────
/**
* 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<Message[]> {
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<StagedSpan[]> {
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<void> {
/* 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<void> {
/* 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<T>(
arr: T[],
predicate: (item: T) => boolean,
): number {
for (let i = arr.length - 1; i >= 0; i--) {
if (predicate(arr[i]!)) return i
}
return -1
}
@@ -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: '<collapsed id="0000000000000001">test summary</collapsed>',
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 <collapsed> 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('<collapsed')
})
test('collapsed summary stays meta so the snip sweep cannot remove it', async () => {
// 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 <collapsed> 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: '<collapsed id="1">s</collapsed>\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: '<collapsed id="1">s</collapsed>' },
}
const realUser = {
type: 'user' as const,
uuid: uid('c'),
timestamp: new Date().toISOString(),
message: {
role: 'user' as const,
content:
'hello\n<system-reminder>snip_id=abc123; system-generated; for snip tool use only; do not discuss in thinking or responses.</system-reminder>',
},
}
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
// <collapsed> 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: '<collapsed id="1">s</collapsed>' },
{ 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: '<collapsed id="1">s</collapsed>' }],
},
}
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: '\n<system-reminder>snip_id=abc123; system-generated; for snip tool use only; do not discuss in thinking or responses.</system-reminder>',
},
],
},
}
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)
})
})
+51 -9
View File
@@ -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<T>(
arr: T[],
predicate: (item: T) => boolean,
): number {
for (let i = arr.length - 1; i >= 0; i--) {
if (predicate(arr[i]!)) return i
}
return -1
}
@@ -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: '<collapsed id="0000000000000001">test</collapsed>',
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)
})
})
+10 -10
View File
@@ -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)
}
@@ -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<string, unknown> = {}): 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)
})
})
@@ -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))
}
@@ -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)
})
})
@@ -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')
})
})
@@ -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<typeof inputSchema>
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<InputSchema, CtxInspectOutput>)
+9
View File
@@ -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.`
}
+11 -5
View File
@@ -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
}
/**
+15
View File
@@ -100,6 +100,14 @@ export interface UserMessage<C = string | ContentBlockParam[]> {
isVisibleInTranscriptOnly?: boolean
isVirtual?: boolean
isCompactSummary?: boolean
/**
* Set when a context-collapse summary placeholder is converted to (or merged
* into) a user message. Keeps the `<collapsed>` 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 `<collapsed>` summary still reaches the model after
* the archived span is removed.
*/
isCollapseSummary?: boolean
}
export interface SystemPermissionRetryMessage extends SystemMessageBase {
+3
View File
@@ -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',
+93 -7
View File
@@ -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?<system-reminder>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 <collapsed>
// 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
// <collapsed> 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)),
),
},
}
}
+1
View File
@@ -1590,6 +1590,7 @@ export async function recordContextCollapseCommit(commit: {
summary: string
firstArchivedUuid: string
lastArchivedUuid: string
archivedCount: number
}): Promise<void> {
const sessionId = getSessionId() as UUID
if (!sessionId) return