import { agentJournalItemKey } from '../../shared/agent-session-journal-item-key' import { createAgentSessionDeltaCoalescer } from '../native-chat/agent-session-wire/agent-session-delta-coalescer' import { CodexItemStreamRetention } from './codex-item-stream-retention' import { appendCodexItemAndPublish } from './codex-structured-journal-sink' import { codexJournalItem, codexStreamingJournalItem, type CodexThreadItem } from './codex-structured-item-translation' import { codexStructuredItemKey, MAX_CODEX_ITEM_STREAM_PENDING_PATCHES, MAX_CODEX_ITEM_STREAM_PENDING_PATCH_BYTES, MAX_CODEX_ITEM_STREAM_RETAINED_BYTES, boundStreamItem, pendingPatchBytes } from './codex-structured-item-stream-bounds' import { CODEX_ITEM_STREAM_TYPES, codexPatchChangeBytes, PATCH_UPDATED_METHOD, readCodexItemStreamRecord, readCodexItemStreamString, REASONING_PART_METHOD, TERMINAL_INTERACTION_METHOD } from './codex-structured-item-stream-events' import type { CodexItemStreamDeps, CodexItemStreamState, CodexPendingItemPatch, CodexStructuredItemStreamAdmission, CodexStructuredItemStreams } from './codex-structured-item-stream-contracts' export type { CodexStructuredItemStreamAdmission, CodexStructuredItemStreamHandleResult, CodexStructuredItemStreams } from './codex-structured-item-stream-contracts' export { codexStructuredItemKey } from './codex-structured-item-stream-bounds' export { MAX_CODEX_ITEM_STREAM_PENDING_PATCH_BYTES, MAX_CODEX_ITEM_STREAM_RETAINED_BYTES, MAX_CODEX_ITEM_STREAM_PENDING_PATCHES, MAX_CODEX_ITEM_STREAM_STATES } from './codex-structured-item-stream-bounds' /** Delta-only item ids are provider input; retain only a deterministic recent window. */ export function createCodexStructuredItemStreams( deps: CodexItemStreamDeps ): CodexStructuredItemStreams { const states = new CodexItemStreamRetention(deps.maxMetadataBytes) const checkpointLengths = new Map() const pendingCheckpoints = new Set() // Which thread and turn produced each stream, resolved to linkage per append // so a parent learned after the first checkpoint still reaches the row. const producers = new Map() const linkageOf = (key: string) => { const producer = producers.get(key) return producer ? deps.linkageFor(producer.threadId, producer.turnId) : {} } // Patch updates are authoritative item snapshots. Keep the latest rejected // snapshot until the journal admits it; unlike streamed deltas, there is no // coalescer timer to retry these events for us. const pendingPatches = new Map() let retainedPatchBytes = 0 const forgetState = (key: string): void => { coalescer.forget(key) states.forget(key) checkpointLengths.delete(key) pendingCheckpoints.delete(key) producers.delete(key) const pending = pendingPatches.get(key) if (pending) { retainedPatchBytes = Math.max(0, retainedPatchBytes - pendingPatchBytes(pending)) pendingPatches.delete(key) } } const trimStates = (): void => { while (states.overCapacity) { const oldest = states.oldestEvictable() if (typeof oldest !== 'string') { break } const pending = coalescer.snapshot(oldest) if (pending && pending.text.length > 0 && !persist(oldest, pending.text, true)) { // Keep the state (and its buffered text) until the sink recovers. A // bounded map is preferable to silently losing streamed output. break } forgetState(oldest) } } const trimPendingPatches = (): void => { while (pendingPatches.size > MAX_CODEX_ITEM_STREAM_PENDING_PATCHES) { const oldest = pendingPatches.keys().next().value if (typeof oldest !== 'string') { break } const pending = pendingPatches.get(oldest) if (pending) { retainedPatchBytes = Math.max(0, retainedPatchBytes - pendingPatchBytes(pending)) } pendingPatches.delete(oldest) } } const append = (key: string, state: CodexItemStreamState, text: string): boolean => { const translated = codexStreamingJournalItem(state.item, text) if (!translated.body) { return true } return appendCodexItemAndPublish(deps.sink, state.identity, translated.body, { coalescingKey: `checkpoint:${agentJournalItemKey(state.identity)}`, ...linkageOf(key) }).accepted } const persist = (key: string, text: string, force: boolean): boolean => { const checkpointLength = checkpointLengths.get(key) ?? 0 const nextLength = Math.max(checkpointLength + 32, Math.ceil(checkpointLength * 1.125)) if (!force && checkpointLength > 0 && text.length < nextLength) { return true } const state = states.get(key) if (state && append(key, state, text)) { checkpointLengths.set(key, text.length) pendingCheckpoints.delete(key) return true } return false } const coalescer = createAgentSessionDeltaCoalescer({ windowMs: deps.coalesceMs, maxRetainedBytes: deps.maxRetainedBytes, maxTotalRetainedBytes: deps.maxTotalRetainedBytes, isProtected: (key) => states.isPersistent(key), schedule: deps.schedule, emit: (key, text) => { return persist(key, text, false) } }) const ensureState = ( threadId: string, itemId: string, type: string, params: unknown ): CodexItemStreamState | null => { const key = codexStructuredItemKey(threadId, itemId) const existing = states.get(key) if (existing) { return existing } const item = { type, id: itemId } const turnId = deps.turnIdFor(threadId, params) const state = { item, identity: deps.identityFor(threadId, turnId, item) } if (!states.retain(key, state)) { return null } producers.set(key, { threadId, turnId }) trimStates() return state } const flush = (): boolean => { let flushed = coalescer.flushAll() for (const key of pendingCheckpoints) { const snapshot = coalescer.snapshot(key) if (snapshot && checkpointLengths.get(key) !== snapshot.text.length) { flushed = persist(key, snapshot.text, true) && flushed } else { pendingCheckpoints.delete(key) } } for (const key of pendingPatches.keys()) { flushed = flushPatch(key).accepted && flushed } return flushed } const flushPatch = (key: string): CodexStructuredItemStreamAdmission => { const pending = pendingPatches.get(key) if (!pending) { return { accepted: true } } const admission = appendCodexItemAndPublish( deps.sink, pending.identity, pending.body, linkageOf(key) ) if (!admission.accepted) { return admission } retainedPatchBytes = Math.max(0, retainedPatchBytes - pendingPatchBytes(pending)) pendingPatches.delete(key) return { accepted: true } } return { get persistentCount() { return states.persistentSize }, canTrack: (threadId, item, identity) => states.canRetain(codexStructuredItemKey(threadId, item.id), { item: boundStreamItem(item) as CodexThreadItem, identity }), track: (threadId, turnId, item, identity) => { const key = codexStructuredItemKey(threadId, item.id) if (!states.retain(key, { item: boundStreamItem(item) as CodexThreadItem, identity })) { return false } producers.set(key, { threadId, turnId }) trimStates() return true }, handle: (threadId, method, params) => { const paramsRecord = readCodexItemStreamRecord(params) const itemId = readCodexItemStreamString(paramsRecord, 'itemId') if (method === PATCH_UPDATED_METHOD) { if (!itemId || !Array.isArray(paramsRecord.changes)) { return { handled: true, admission: { accepted: true } } } if ( codexPatchChangeBytes(paramsRecord.changes) > MAX_CODEX_ITEM_STREAM_PENDING_PATCH_BYTES ) { return { handled: true, admission: { accepted: false, reason: 'backpressure' } } } const key = codexStructuredItemKey(threadId, itemId) const streamFlushed = coalescer.flush(key) const state = ensureState(threadId, itemId, 'fileChange', params) if (!state) { return { handled: true, admission: { accepted: false, reason: 'failed' } } } state.item = { ...state.item, changes: paramsRecord.changes } const translated = codexJournalItem(state.item) if (translated.body) { const nextPending: CodexPendingItemPatch = { identity: state.identity, body: translated.body } const previous = pendingPatches.get(key) const previousBytes = previous ? pendingPatchBytes(previous) : 0 const nextBytes = pendingPatchBytes(nextPending) if ( nextBytes > MAX_CODEX_ITEM_STREAM_PENDING_PATCH_BYTES || retainedPatchBytes - previousBytes + nextBytes > MAX_CODEX_ITEM_STREAM_RETAINED_BYTES ) { return { handled: true, admission: { accepted: false, reason: 'backpressure' } } } retainedPatchBytes = Math.max(0, retainedPatchBytes - previousBytes) + nextBytes pendingPatches.set(key, nextPending) trimPendingPatches() if (streamFlushed) { const admission = flushPatch(key) if (!admission.accepted) { return { handled: true, admission } } } } return { handled: true, admission: { accepted: true } } } if (method === TERMINAL_INTERACTION_METHOD) { return { handled: true, admission: { accepted: true } } } const type = CODEX_ITEM_STREAM_TYPES[method as keyof typeof CODEX_ITEM_STREAM_TYPES] if (!type && method !== REASONING_PART_METHOD) { return { handled: false, admission: { accepted: true } } } if (!itemId) { return { handled: true, admission: { accepted: true } } } const state = ensureState(threadId, itemId, type ?? 'reasoning', params) if (!state) { return { handled: true, admission: { accepted: false, reason: 'failed' } } } const delta = method === REASONING_PART_METHOD ? '\n' : paramsRecord.delta if (typeof delta === 'string') { const key = codexStructuredItemKey(threadId, state.item.id) pendingCheckpoints.add(key) const accepted = coalescer.append(key, delta) if (!accepted) { return { handled: true, admission: { accepted: false, reason: 'backpressure' } } } } return { handled: true, admission: { accepted: true } } }, forget: (threadId, itemId) => { const key = codexStructuredItemKey(threadId, itemId) forgetState(key) }, flush, dispose: () => { coalescer.dispose() states.clear() checkpointLengths.clear() pendingCheckpoints.clear() producers.clear() pendingPatches.clear() retainedPatchBytes = 0 }, snapshot: (threadId, itemId) => { const key = codexStructuredItemKey(threadId, itemId) return coalescer.snapshot(key) } } }