diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index f1b78228e18..a0b85830212 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -76,10 +76,29 @@ import { createTranscriptShadow, defaultStorage, diffTranscript, + parseTranscriptRuntimeState, + prefixFingerprint, + restoreModelLane, + type TranscriptChange, type TranscriptChangeReason, + type TranscriptRuntimeState, type TranscriptShadow, + type TranscriptStorage, type TranscriptStorageContext, } from "./transcriptStorage.js"; + +let transcriptStorageOverride: TranscriptStorage | undefined; + +/** + * Test-only override for the storage `chat.agent` persists through, so a + * test can capture the exact changesets the runtime produces. + * @internal + */ +export function __setTranscriptStorageForTests( + storage: TranscriptStorage | undefined +): void { + transcriptStorageOverride = storage; +} import { type ChatInputChunk, type ChatTaskWirePayload, @@ -6870,8 +6889,18 @@ function chatAgent< // collectively cost ~600ms on every first-message TTFC. Both reads // swallow errors internally; the agent stays available either way. const sessionIdForSnapshot = payload.sessionId ?? payload.chatId; - const transcriptStorage = defaultStorage; + const transcriptStorage = transcriptStorageOverride ?? defaultStorage; let transcriptShadow: TranscriptShadow = createTranscriptShadow([]); + let bootTranscriptState: unknown = null; + /** + * True while the model lane holds a compaction summary, so it cannot be + * rebuilt from the transcript and has to be persisted as state. Reset + * wherever the lane is reconverted from the UI lane. + */ + let laneCompacted = false; + /** Conversational `chat.inject` messages in the lane, anchored to the transcript. */ + let laneInjections: NonNullable = []; + let persistedStateSet = false; let bootSnapshot: | { messages: TUIMessage[]; lastOutEventId?: string; lastInEventId?: string } | undefined; @@ -6919,6 +6948,23 @@ function chatAgent< const { changes, shadow } = diffTranscript(transcriptShadow, opts.messages, { nonFinalIds: opts.nonFinalIds, }); + const lastId = opts.messages.at(-1)?.id; + const runtimeState: TranscriptRuntimeState | null = + laneCompacted && lastId !== undefined + ? { + v: 1, + compaction: { + modelMessages: accumulatedMessages, + throughId: lastId, + fingerprint: prefixFingerprint(shadow, lastId), + }, + } + : laneInjections.length > 0 + ? { v: 1, injections: laneInjections } + : null; + if (runtimeState !== null || persistedStateSet) { + changes.push({ op: "state", value: runtimeState } satisfies TranscriptChange); + } const inCursor = chatInputRouter().resumeFloor(); await transcriptStorage.save( { @@ -6939,6 +6985,7 @@ function chatAgent< } ); transcriptShadow = shadow; + persistedStateSet = runtimeState !== null; }; /** @@ -7023,6 +7070,8 @@ function chatAgent< clientData: bootClientData, }); transcriptShadow = createTranscriptShadow(loaded.messages); + bootTranscriptState = loaded.state; + persistedStateSet = loaded.state !== null && loaded.state !== undefined; bootSnapshot = { messages: loaded.messages, lastOutEventId: loaded.cursors?.lastOutEventId, @@ -7367,7 +7416,14 @@ function chatAgent< } } try { - accumulatedMessages = await toModelMessages(accumulatedUIMessages); + const restored = await restoreModelLane( + accumulatedUIMessages, + parseTranscriptRuntimeState(bootTranscriptState), + (messages) => toModelMessages(messages) + ); + accumulatedMessages = restored.messages; + laneCompacted = restored.compacted; + laneInjections = restored.injections; } catch (error) { logger.warn("chat.agent: toModelMessages failed at boot; starting empty", { error: error instanceof Error ? error.message : String(error), @@ -8039,6 +8095,8 @@ function chatAgent< ); accumulatedUIMessages = [...hydrated] as TUIMessage[]; accumulatedMessages = await toModelMessages(hydrated); + laneCompacted = false; + laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); } @@ -8076,6 +8134,8 @@ function chatAgent< locals.set(chatOverrideMessagesKey, undefined); accumulatedUIMessages = [...actionOverride] as TUIMessage[]; accumulatedMessages = await toModelMessages(actionOverride); + laneCompacted = false; + laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); actionChangedHistory = true; @@ -8208,6 +8268,8 @@ function chatAgent< accumulatedUIMessages = merged; accumulatedMessages = await toModelMessages(merged); + laneCompacted = false; + laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); // Track new messages for onTurnComplete.newUIMessages. @@ -8257,6 +8319,8 @@ function chatAgent< accumulatedUIMessages.pop(); } accumulatedMessages = await toModelMessages(accumulatedUIMessages); + laneCompacted = false; + laneInjections = []; } else if (cleanedUIMessages.length > 0) { // Submit-message (and the special-cased // handover-prepare → submit-message rewrite earlier in @@ -8310,6 +8374,8 @@ function chatAgent< "chat.agent: replaced message not found at the model lane tail; reconverting the lane" ); accumulatedMessages = await toModelMessages(accumulatedUIMessages); + laneCompacted = false; + laneInjections = []; } } else { const incomingModelMessages = await toModelMessages(cleanedUIMessages); @@ -8491,6 +8557,8 @@ function chatAgent< locals.set(chatOverrideMessagesKey, undefined); accumulatedUIMessages = [...turnStartOverride] as TUIMessage[]; accumulatedMessages = await toModelMessages(turnStartOverride); + laneCompacted = false; + laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); } }, @@ -8556,7 +8624,12 @@ function chatAgent< const lastAccumulated = accumulatedMessages[accumulatedMessages.length - 1]; const bgQueue = locals.get(chatBackgroundQueueKey); if (bgQueue && bgQueue.length > 0 && lastAccumulated?.role !== "tool") { - accumulatedMessages.push(...bgQueue.splice(0)); + const injected = bgQueue.splice(0); + accumulatedMessages.push(...injected); + laneInjections.push({ + afterId: accumulatedUIMessages.at(-1)?.id ?? "", + messages: injected, + }); } if (isHeadStartFinalTurn) { @@ -8749,6 +8822,8 @@ function chatAgent< accumulatedMessages = await toModelMessages( runOverride.filter((m) => !pendingIds.has(m.id)) ); + laneCompacted = false; + laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); } @@ -8774,6 +8849,8 @@ function chatAgent< accumulatedMessages = taskCompactionConfig?.compactModelMessages ? await taskCompactionConfig.compactModelMessages(compactEvent) : modelOnlyOverride; + laneCompacted = true; + laneInjections = []; // Apply UI messages: callback or default (preserve all) if (taskCompactionConfig?.compactUIMessages) { @@ -8866,6 +8943,8 @@ function chatAgent< "chat.agent: replaced response not found at the model lane tail; reconverting the lane" ); accumulatedMessages = await toModelMessages(accumulatedUIMessages); + laneCompacted = false; + laneInjections = []; } } else { accumulatedMessages.push(...responseModelMessages); @@ -8985,6 +9064,9 @@ function chatAgent< }, ]; + laneCompacted = true; + laneInjections = []; + // UI messages: callback or default (preserve all) if (outerCompaction.compactUIMessages) { accumulatedUIMessages = (await outerCompaction.compactUIMessages( @@ -9089,6 +9171,8 @@ function chatAgent< locals.set(chatOverrideMessagesKey, undefined); accumulatedUIMessages = [...override] as TUIMessage[]; accumulatedMessages = await toModelMessages(override); + laneCompacted = false; + laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); // Update event so onTurnComplete sees compacted messages turnCompleteEvent.messages = accumulatedMessages; @@ -9148,6 +9232,8 @@ function chatAgent< locals.set(chatOverrideMessagesKey, undefined); accumulatedUIMessages = [...turnCompleteOverride] as TUIMessage[]; accumulatedMessages = await toModelMessages(turnCompleteOverride); + laneCompacted = false; + laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); } }, @@ -9512,6 +9598,8 @@ function chatAgent< "chat.agent: replaced partial not found at the model lane tail; reconverting the lane" ); accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial); + laneCompacted = false; + laneInjections = []; } } accumulatedUIMessages = erroredUIMessagesWithPartial; diff --git a/packages/trigger-sdk/src/v3/transcriptStorage.ts b/packages/trigger-sdk/src/v3/transcriptStorage.ts index 6596fc8dc41..6aca0564290 100644 --- a/packages/trigger-sdk/src/v3/transcriptStorage.ts +++ b/packages/trigger-sdk/src/v3/transcriptStorage.ts @@ -1,5 +1,5 @@ import type { TaskRunContext, TranscriptSnapshotEntry } from "@trigger.dev/core/v3"; -import type { UIMessage } from "ai"; +import type { ModelMessage, UIMessage } from "ai"; import { readChatSnapshot, writeChatSnapshot } from "./chatSnapshotIo.js"; /** @@ -233,23 +233,12 @@ export function snapshotTranscriptStorage(): TranscriptStorage { : emptyTranscriptState(); transcripts.set(scope.chatId, full as TranscriptState); - let entries = full.entries; - if (opts?.before !== undefined) { - const idx = entries.findIndex((e) => e.id === opts.before); - if (idx !== -1) entries = entries.slice(0, idx); - } - let nextCursor: string | undefined; - if (opts?.limit !== undefined && entries.length > opts.limit) { - entries = entries.slice(entries.length - opts.limit); - nextCursor = entries[0]?.id; - } return { - messages: entries.map((e) => e.message), + ...pageEntries(full.entries, opts), state: full.state, cursors: snapshot ? { lastOutEventId: snapshot.lastOutEventId, lastInEventId: snapshot.lastInEventId } : undefined, - nextCursor, }; }, @@ -271,3 +260,188 @@ export function snapshotTranscriptStorage(): TranscriptStorage { /** The storage `chat.agent` uses when none is configured: {@link snapshotTranscriptStorage}. */ export const defaultStorage: TranscriptStorage = snapshotTranscriptStorage(); + +function pageEntries( + all: TranscriptSnapshotEntry[], + opts: TranscriptLoadOptions | undefined +): { messages: TUIMessage[]; nextCursor: string | undefined } { + let entries = all; + if (opts?.before !== undefined) { + const idx = entries.findIndex((e) => e.id === opts.before); + if (idx !== -1) entries = entries.slice(0, idx); + } + let nextCursor: string | undefined; + if (opts?.limit !== undefined && entries.length > opts.limit) { + entries = entries.slice(entries.length - opts.limit); + nextCursor = entries[0]?.id; + } + return { messages: entries.map((e) => e.message), nextCursor }; +} + +export type MemoryTranscriptStorage = TranscriptStorage & { + /** The stored transcript for a chat, or `undefined` when nothing has been saved. */ + transcript(chatId: string): (TranscriptState & { cursors?: TranscriptCursors }) | undefined; + /** Every changeset `save` received, in order. */ + readonly changesets: Array<{ + ctx: TranscriptStorageContext; + changeset: TranscriptChangeset; + }>; +}; + +/** + * A storage that keeps every transcript in process memory. The reference + * implementation for the conformance tests, and a way to inspect exactly + * what the runtime hands a storage. + */ +export function memoryTranscriptStorage(): MemoryTranscriptStorage { + const transcripts = new Map(); + const changesets: MemoryTranscriptStorage["changesets"] = []; + + return { + changesets, + transcript(chatId) { + return transcripts.get(chatId); + }, + async load( + scope: TranscriptScope, + opts?: TranscriptLoadOptions + ): Promise> { + const stored = transcripts.get(scope.chatId); + if (!stored) return { messages: [], state: null, cursors: undefined, nextCursor: undefined }; + return { + ...pageEntries(stored.entries as TranscriptSnapshotEntry[], opts), + state: stored.state, + cursors: stored.cursors, + }; + }, + async save(ctx, changeset) { + changesets.push({ ctx, changeset }); + const prev = transcripts.get(ctx.chatId); + const next = reduceTranscriptChanges(prev ?? emptyTranscriptState(), changeset.changes); + transcripts.set(ctx.chatId, { ...next, cursors: changeset.cursors ?? prev?.cursors }); + }, + }; +} + +type ModelLaneInjection = { afterId: string; messages: ModelMessage[] }; + +/** + * What the runtime keeps in the storage's `state` slot: the parts of the + * model's context that cannot be rebuilt from the transcript. Opaque to a + * storage; only the runtime reads it. + * + * `compaction` is the whole model lane after a compaction, valid for the + * transcript prefix ending at `throughId` whose fingerprint matches, so a + * rollback or edit of that prefix makes it unusable and the next save + * clears it. `injections` are conversational messages `chat.inject` added, + * anchored after the transcript message they followed. + */ +export type TranscriptRuntimeState = { + v: 1; + compaction?: { modelMessages: ModelMessage[]; throughId: string; fingerprint: string }; + injections?: ModelLaneInjection[]; +}; + +export function parseTranscriptRuntimeState(value: unknown): TranscriptRuntimeState | undefined { + if (!value || typeof value !== "object") return undefined; + const record = value as Record; + if (record.v !== 1) return undefined; + const out: TranscriptRuntimeState = { v: 1 }; + const compaction = record.compaction as Record | undefined; + if ( + compaction && + typeof compaction === "object" && + Array.isArray(compaction.modelMessages) && + typeof compaction.throughId === "string" && + typeof compaction.fingerprint === "string" + ) { + out.compaction = { + modelMessages: compaction.modelMessages as ModelMessage[], + throughId: compaction.throughId, + fingerprint: compaction.fingerprint, + }; + } + if (Array.isArray(record.injections)) { + out.injections = (record.injections as unknown[]).flatMap((entry) => { + const inj = entry as Record | null; + return inj && typeof inj.afterId === "string" && Array.isArray(inj.messages) + ? [{ afterId: inj.afterId, messages: inj.messages as ModelMessage[] }] + : []; + }); + } + return out; +} + +/** + * A 32-bit FNV-1a hash over the fingerprints of the messages up to and + * including `throughId`, in order. Cheap enough to compute on every save + * because the per-message fingerprints already exist in the shadow. + */ +export function prefixFingerprint(shadow: TranscriptShadow, throughId: string): string { + let hash = 0x811c9dc5; + const mix = (s: string) => { + for (let i = 0; i < s.length; i++) { + hash ^= s.charCodeAt(i); + hash = Math.imul(hash, 0x01000193) >>> 0; + } + }; + for (const id of shadow.ids) { + mix(id); + mix(""); + mix(shadow.fingerprints.get(id) ?? ""); + mix(""); + if (id === throughId) return hash.toString(16).padStart(8, "0"); + } + return ""; +} + +/** + * Rebuild the model lane for a transcript at boot. Uses the persisted + * compacted lane when the transcript prefix it covers is unchanged, then + * converts the rest of the transcript, re-inserting persisted injections + * after the messages they followed. + */ +export async function restoreModelLane( + messages: TUIMessage[], + state: TranscriptRuntimeState | undefined, + convert: (messages: TUIMessage[]) => Promise +): Promise<{ messages: ModelMessage[]; compacted: boolean; injections: ModelLaneInjection[] }> { + const lane: ModelMessage[] = []; + let start = 0; + let compacted = false; + + if (state?.compaction) { + const idx = messages.findIndex((m) => m.id === state.compaction!.throughId); + if ( + idx !== -1 && + prefixFingerprint(createTranscriptShadow(messages), state.compaction.throughId) === + state.compaction.fingerprint + ) { + lane.push(...state.compaction.modelMessages); + start = idx + 1; + compacted = true; + } + } + + const anchored = (state?.injections ?? []) + .map((inj) => ({ + inj, + idx: inj.afterId === "" ? -1 : messages.findIndex((m) => m.id === inj.afterId), + })) + .filter(({ inj, idx }) => (inj.afterId === "" ? start === 0 : idx >= start)) + .sort((a, b) => a.idx - b.idx); + + let cursor = start; + for (const { inj, idx } of anchored) { + if (idx + 1 > cursor) { + lane.push(...(await convert(messages.slice(cursor, idx + 1)))); + cursor = idx + 1; + } + lane.push(...inj.messages); + } + if (cursor < messages.length) { + lane.push(...(await convert(messages.slice(cursor)))); + } + + return { messages: lane, compacted, injections: anchored.map(({ inj }) => inj) }; +} diff --git a/packages/trigger-sdk/test/transcript-changesets.test.ts b/packages/trigger-sdk/test/transcript-changesets.test.ts new file mode 100644 index 00000000000..a746f91723b --- /dev/null +++ b/packages/trigger-sdk/test/transcript-changesets.test.ts @@ -0,0 +1,333 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import type { ModelMessage, UIMessage } from "ai"; +import { simulateReadableStream, stepCountIs, streamText, tool } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { z } from "zod"; +import { __setTranscriptStorageForTests, chat } from "../src/v3/ai.js"; +import { + memoryTranscriptStorage, + type MemoryTranscriptStorage, + type TranscriptChange, + type TranscriptRuntimeState, +} from "../src/v3/transcriptStorage.js"; + +const usage = { + inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 10, reasoning: undefined }, +}; + +function userMessage(text: string, id: string): UIMessage { + return { id, role: "user", parts: [{ type: "text", text }] }; +} + +function textChunks(text: string): LanguageModelV3StreamPart[] { + return [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: text }, + { type: "text-end", id: "t1" }, + { type: "finish", finishReason: { unified: "stop", raw: "stop" }, usage }, + ]; +} + +function promptText(prompt: unknown): string { + return JSON.stringify(prompt); +} + +function recordingModel(prompts: unknown[], reply = "ack") { + return new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + prompts.push(prompt); + return { stream: simulateReadableStream({ chunks: textChunks(reply) }) }; + }, + }); +} + +async function waitFor(check: () => boolean, label: string, timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} + +const ops = (changes: TranscriptChange[]) => changes.map((c) => c.op); +const putIds = (changes: TranscriptChange[]) => + changes.flatMap((c) => (c.op === "put" ? [c.message.id] : [])); +const stateOf = (changes: TranscriptChange[]) => + changes.find((c) => c.op === "state")?.value as TranscriptRuntimeState | null | undefined; + +let storage: MemoryTranscriptStorage; + +beforeEach(() => { + storage = memoryTranscriptStorage(); + __setTranscriptStorageForTests(storage); +}); + +afterEach(() => { + __setTranscriptStorageForTests(undefined); +}); + +describe("chat.agent transcript changesets", () => { + it("saves a turn as puts for the new user and assistant messages with cursors", async () => { + const prompts: unknown[] = []; + const agent = chat.agent({ + id: "changeset-turn", + run: async ({ messages, signal }) => + streamText({ model: recordingModel(prompts), messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { chatId: "changeset-turn" }); + try { + await harness.sendMessage(userMessage("hello", "u1")); + await waitFor(() => storage.changesets.length === 1, "first save"); + + const { ctx, changeset } = storage.changesets[0]!; + expect(ctx.chatId).toBe("changeset-turn"); + expect(ctx.trigger).toBe("submit-message"); + expect(ctx.turn).toBe(0); + expect(changeset.reason).toBe("turn-complete"); + expect(ops(changeset.changes)).toEqual(["put", "put"]); + expect(putIds(changeset.changes)[0]).toBe("u1"); + expect(changeset.cursors?.lastOutEventId).toBeDefined(); + + await harness.sendMessage(userMessage("again", "u2")); + await waitFor(() => storage.changesets.length === 2, "second save"); + expect(ops(storage.changesets[1]!.changeset.changes)).toEqual(["put", "put"]); + expect(putIds(storage.changesets[1]!.changeset.changes)[0]).toBe("u2"); + expect(storage.transcript("changeset-turn")!.entries.map((e) => e.message.role)).toEqual([ + "user", + "assistant", + "user", + "assistant", + ]); + } finally { + await harness.close(); + } + }); + + it("puts a steering message the drain consumed into the turn's changeset", async () => { + const send = { fn: async () => {} }; + let call = 0; + const model = new MockLanguageModelV3({ + doStream: async () => { + call += 1; + if (call === 1) { + await send.fn(); + return { + stream: simulateReadableStream({ + chunks: [ + { type: "tool-input-start", id: "c1", toolName: "lookup" }, + { type: "tool-input-delta", id: "c1", delta: "{}" }, + { type: "tool-input-end", id: "c1" }, + { type: "tool-call", toolCallId: "c1", toolName: "lookup", input: "{}" }, + { + type: "finish", + finishReason: { unified: "tool-calls", raw: "tool-calls" }, + usage, + }, + ] satisfies LanguageModelV3StreamPart[], + }), + }; + } + return { stream: simulateReadableStream({ chunks: textChunks("done") }) }; + }, + }); + + const agent = chat.agent({ + id: "changeset-steer", + tools: { + lookup: tool({ + description: "look something up", + inputSchema: z.object({}), + execute: async () => ({ ok: true }), + }), + }, + pendingMessages: { shouldInject: ({ steps }) => steps.length > 0 }, + run: async ({ messages, tools, signal }) => + streamText({ + ...chat.toStreamTextOptions({ tools }), + model, + messages, + abortSignal: signal, + stopWhen: stepCountIs(5), + }), + }); + const harness = mockChatAgent(agent, { chatId: "changeset-steer" }); + send.fn = async () => { + await harness.sendPendingMessage(userMessage("only the platform one", "steer-1")); + }; + try { + await harness.sendMessage(userMessage("summarise every project", "u1")); + await waitFor(() => storage.changesets.length === 1, "save"); + + const ids = putIds(storage.changesets[0]!.changeset.changes); + expect(ids).toContain("steer-1"); + expect(ids.indexOf("steer-1")).toBeGreaterThan(ids.indexOf("u1")); + expect(storage.transcript("changeset-steer")!.entries.map((e) => e.id)).toEqual(ids); + } finally { + await harness.close(); + } + }); + + it("persists a compaction as state and boots a continuation from the summary", async () => { + const chatId = "changeset-compaction"; + let compactions = 0; + const makeAgent = (prompts: unknown[]) => + chat.agent({ + id: "changeset-compaction", + compaction: { + shouldCompact: ({ source }) => source === "outer" && compactions === 0, + summarize: async () => { + compactions += 1; + return "SUMMARY-OF-EVERYTHING"; + }, + }, + run: async ({ messages, signal }) => + streamText({ model: recordingModel(prompts), messages, abortSignal: signal }), + }); + + const firstPrompts: unknown[] = []; + const first = mockChatAgent(makeAgent(firstPrompts), { chatId }); + try { + await first.sendMessage(userMessage("the early message", "u1")); + await waitFor(() => storage.changesets.length === 1, "turn 1 save"); + expect(compactions).toBe(1); + + const state = stateOf(storage.changesets[0]!.changeset.changes); + expect(state?.compaction).toBeDefined(); + expect(state!.compaction!.throughId).toBe( + putIds(storage.changesets[0]!.changeset.changes).at(-1) + ); + expect(JSON.stringify(state!.compaction!.modelMessages)).toContain("SUMMARY-OF-EVERYTHING"); + expect(JSON.stringify(state!.compaction!.modelMessages)).not.toContain("the early message"); + + await first.sendMessage(userMessage("a follow-up", "u2")); + await waitFor(() => storage.changesets.length === 2, "turn 2 save"); + expect(promptText(firstPrompts[1])).toContain("SUMMARY-OF-EVERYTHING"); + expect(promptText(firstPrompts[1])).not.toContain("the early message"); + expect(stateOf(storage.changesets[1]!.changeset.changes)?.compaction).toBeDefined(); + } finally { + await first.close(); + } + + expect(storage.transcript(chatId)!.entries.map((e) => e.id)).toHaveLength(4); + expect(storage.transcript(chatId)!.state).not.toBeNull(); + + const secondPrompts: unknown[] = []; + const second = mockChatAgent(makeAgent(secondPrompts), { + chatId, + continuation: true, + previousRunId: "run_first", + }); + try { + await second.sendMessage(userMessage("after the continuation", "u3")); + await waitFor(() => secondPrompts.length === 1, "continuation turn"); + + const prompt = promptText(secondPrompts[0]); + expect(prompt).toContain("SUMMARY-OF-EVERYTHING"); + expect(prompt).toContain("a follow-up"); + expect(prompt).toContain("after the continuation"); + expect(prompt).not.toContain("the early message"); + expect(compactions).toBe(1); + } finally { + await second.close(); + } + }); + + it("clears the compaction state in the same changeset as a rollback", async () => { + const chatId = "changeset-rollback"; + let compactions = 0; + const prompts: unknown[] = []; + const agent = chat.agent({ + id: "changeset-rollback", + actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("undo") })]), + compaction: { + shouldCompact: ({ source }) => source === "outer" && compactions === 0, + summarize: async () => { + compactions += 1; + return "SUMMARY"; + }, + }, + onAction: async ({ action }) => { + if (action.type === "undo") chat.history.slice(0, -2); + }, + run: async ({ messages, signal }) => + streamText({ model: recordingModel(prompts), messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { chatId }); + try { + await harness.sendMessage(userMessage("one", "u1")); + await harness.sendMessage(userMessage("two", "u2")); + await waitFor(() => storage.changesets.length === 2, "two turns"); + expect(stateOf(storage.changesets[1]!.changeset.changes)?.compaction).toBeDefined(); + + await harness.sendAction({ type: "undo" }); + await waitFor(() => storage.changesets.length === 3, "action save"); + + const { ctx, changeset } = storage.changesets[2]!; + expect(ctx.trigger).toBe("action"); + expect(changeset.reason).toBe("action"); + expect(ops(changeset.changes)).toEqual(["truncateAfter", "state"]); + expect(stateOf(changeset.changes)).toBeNull(); + expect(storage.transcript(chatId)!.entries.map((e) => e.id)).toHaveLength(2); + expect(storage.transcript(chatId)!.state).toBeNull(); + } finally { + await harness.close(); + } + }); + + it("persists conversational injections anchored to the transcript and restores them at boot", async () => { + const chatId = "changeset-inject"; + const makeAgent = (prompts: unknown[]) => + chat.agent({ + id: "changeset-inject", + onTurnComplete: async ({ turn }) => { + if (turn === 0) { + chat.inject([{ role: "user", content: "[note] inventory is low" } as ModelMessage]); + } + }, + run: async ({ messages, signal }) => + streamText({ model: recordingModel(prompts), messages, abortSignal: signal }), + }); + + const firstPrompts: unknown[] = []; + const first = mockChatAgent(makeAgent(firstPrompts), { chatId }); + try { + await first.sendMessage(userMessage("one", "u1")); + await first.sendMessage(userMessage("two", "u2")); + await waitFor(() => storage.changesets.length === 2, "two turns"); + + expect(promptText(firstPrompts[1])).toContain("[note] inventory is low"); + const state = stateOf(storage.changesets[1]!.changeset.changes); + expect(state?.injections).toHaveLength(1); + expect(state!.injections![0]!.afterId).toBe("u2"); + expect(stateOf(storage.changesets[0]!.changeset.changes)).toBeUndefined(); + } finally { + await first.close(); + } + + const secondPrompts: unknown[] = []; + const second = mockChatAgent(makeAgent(secondPrompts), { + chatId, + continuation: true, + previousRunId: "run_first", + }); + try { + await second.sendMessage(userMessage("three", "u3")); + await waitFor(() => secondPrompts.length === 1, "continuation turn"); + const prompt = secondPrompts[0] as { role: string; content: unknown }[]; + const text = promptText(prompt); + expect(text).toContain("[note] inventory is low"); + const noteIdx = prompt.findIndex((m) => promptText(m).includes("[note] inventory is low")); + const u2Idx = prompt.findIndex((m) => promptText(m).includes('"two"')); + const u3Idx = prompt.findIndex((m) => promptText(m).includes('"three"')); + expect(noteIdx).toBeGreaterThan(u2Idx); + expect(noteIdx).toBeLessThan(u3Idx); + } finally { + await second.close(); + } + }); +});