feat: add conversation storage (#140)

apps/freya-backend/src/agent/pi-query-agent.test.ts

@@ -1,7 +1,10 @@
 import { beforeEach, describe, expect, mock, test } from "bun:test"
 
 import type { QueryAgentToolbox } from "./query-agent-toolbox.ts"
-import type { QueryAgentEvent } from "./query-agent.ts"
+import type { QueryAgentStreamEvent } from "./query-agent.ts"
+
+import { ConversationEntryKind } from "../conversations/types.ts"
+import { QueryAgentEvent } from "./query-agent.ts"
 
 interface FakePiSession {
 	subscribe(listener: (event: unknown) => void): () => void
@@ -9,6 +12,61 @@ interface FakePiSession {
 	dispose(): void
 }
 
+type CapturedExtensionHandler = (event: unknown) => Promise<unknown> | unknown
+
+interface CapturedExtensionApi {
+	on(event: string, handler: CapturedExtensionHandler): void
+}
+
+type CapturedExtensionFactory = (pi: CapturedExtensionApi) => Promise<void> | void
+
+interface CapturedExtension {
+	handlers: Map<string, CapturedExtensionHandler[]>
+}
+
+interface CapturedResourceLoader {
+	getExtensions(): unknown
+}
+
+interface CapturedDefaultResourceLoaderOptions {
+	extensionFactories?: CapturedExtensionFactory[]
+}
+
+class FakeDefaultResourceLoader implements CapturedResourceLoader {
+	private readonly extensionFactories: CapturedExtensionFactory[]
+	private extensionsResult: { extensions: CapturedExtension[] }
+
+	constructor(options: unknown) {
+		this.extensionFactories = isDefaultResourceLoaderOptions(options)
+			? (options.extensionFactories ?? [])
+			: []
+		this.extensionsResult = { extensions: [] }
+	}
+
+	async reload(): Promise<void> {
+		const handlers: CapturedExtension["handlers"] = new Map()
+		const api: CapturedExtensionApi = {
+			on(event: string, handler: CapturedExtensionHandler): void {
+				const existing = handlers.get(event) ?? []
+				existing.push(handler)
+				handlers.set(event, existing)
+			},
+		}
+
+		for (const factory of this.extensionFactories) {
+			await factory(api)
+		}
+
+		this.extensionsResult = {
+			extensions: [{ handlers }],
+		}
+	}
+
+	getExtensions(): unknown {
+		return this.extensionsResult
+	}
+}
+
 let createAgentSessionCalls = 0
 let createAgentSessionOptions: unknown
 let runtimeApiKeyCalls: Array<{ provider: string; apiKey: string }> = []
@@ -51,6 +109,44 @@ const fakeSession: FakePiSession = {
 	dispose(): void {},
 }
 
+class FakeSessionManager {
+	private messages: unknown[] = []
+	private compaction: { summary: string; tokensBefore: number; timestamp: number } | null = null
+
+	appendMessage(message: unknown): string {
+		this.messages.push(message)
+		return `message-${this.messages.length}`
+	}
+
+	appendCompaction(summary: string, _firstKeptEntryId: string, tokensBefore: number): string {
+		this.compaction = {
+			summary,
+			tokensBefore,
+			timestamp: Date.now(),
+		}
+		this.messages = []
+		return "compaction-1"
+	}
+
+	buildSessionContext(): unknown {
+		const messages = [...this.messages]
+		if (this.compaction) {
+			messages.unshift({
+				role: "compactionSummary",
+				summary: this.compaction.summary,
+				tokensBefore: this.compaction.tokensBefore,
+				timestamp: this.compaction.timestamp,
+			})
+		}
+
+		return {
+			messages,
+			thinkingLevel: "off",
+			model: modelFromMessages(messages),
+		}
+	}
+}
+
 mock.module("@earendil-works/pi-coding-agent", () => ({
 	AuthStorage: {
 		inMemory() {
@@ -71,6 +167,7 @@ mock.module("@earendil-works/pi-coding-agent", () => ({
 	createExtensionRuntime() {
 		return {}
 	},
+	DefaultResourceLoader: FakeDefaultResourceLoader,
 	defineTool(tool: unknown): unknown {
 		return tool
 	},
@@ -86,7 +183,7 @@ mock.module("@earendil-works/pi-coding-agent", () => ({
 	},
 	SessionManager: {
 		inMemory(_cwd: string): unknown {
-			return {}
+			return new FakeSessionManager()
 		},
 	},
 	SettingsManager: {
@@ -131,7 +228,6 @@ describe("PiQueryAgent", () => {
 	test("rejects a concurrent first query while the Pi session is being created", async () => {
 		const { PiQueryAgent } = await import("./pi-query-agent.ts")
 		const agent = new PiQueryAgent({
-			userId: "user-1",
 			toolbox: createStubToolbox(),
 			apiKey: "test-api-key",
 			cwd: "/tmp/freya-pi-query-agent-test",
@@ -155,7 +251,7 @@ describe("PiQueryAgent", () => {
 		expect(secondEvents).toEqual([
 			{
 				type: "error",
-				message: "A query is already running for this user",
+				message: "A query is already running",
 			},
 		])
 		expect(createAgentSessionCalls).toBe(1)
@@ -175,6 +271,228 @@ describe("PiQueryAgent", () => {
 		}
 		expect("agentDir" in createAgentSessionOptions).toBe(false)
 		expect(createAgentSessionOptions.resourceLoader).toBeDefined()
+		expect(typeof sessionCompactHandlerFromCapturedOptions()).toBe("function")
+
+		agent.dispose()
+	})
+
+	test("hydrates initial entries into the Pi session manager", async () => {
+		const { PiQueryAgent } = await import("./pi-query-agent.ts")
+		const agent = new PiQueryAgent({
+			toolbox: createStubToolbox(),
+			cwd: "/tmp/freya-pi-query-agent-test",
+			systemPrompt: "test",
+			initialEntries: [
+				{
+					id: "entry-1",
+					sequence: 1,
+					kind: ConversationEntryKind.UserMessage,
+					payload: {
+						role: "user",
+						parts: [{ type: "text", text: "stored hello" }],
+					},
+					metadata: {},
+					createdAt: new Date("2026-06-15T00:00:00.000Z"),
+				},
+				{
+					id: "entry-2",
+					sequence: 2,
+					kind: ConversationEntryKind.AssistantMessage,
+					payload: {
+						role: "assistant",
+						parts: [{ type: "text", text: "stored reply" }],
+					},
+					metadata: {},
+					createdAt: new Date("2026-06-15T00:00:01.000Z"),
+				},
+			],
+		})
+
+		const events = collectEvents(
+			agent.ask({
+				message: "hello",
+			}),
+		)
+
+		await sessionCreationStarted
+		if (!isRecord(createAgentSessionOptions)) {
+			throw new Error("createAgentSession options were not captured")
+		}
+		const sessionManager = createAgentSessionOptions.sessionManager
+		if (!(sessionManager instanceof FakeSessionManager)) {
+			throw new Error("session manager was not hydrated by PiQueryAgent")
+		}
+		const context = sessionManager.buildSessionContext()
+		if (!isRecord(context) || !Array.isArray(context.messages)) {
+			throw new Error("session context messages were not captured")
+		}
+		expect(context.messages[0]).toEqual({
+			role: "user",
+			content: "stored hello",
+			timestamp: new Date("2026-06-15T00:00:00.000Z").getTime(),
+		})
+		expect(context.messages[1]).toMatchObject({
+			role: "assistant",
+			provider: "openrouter",
+			model: "z-ai/glm-4.7-flash",
+			stopReason: "stop",
+			timestamp: new Date("2026-06-15T00:00:01.000Z").getTime(),
+		})
+
+		releaseSessionCreation()
+		await promptStarted
+		releasePrompt()
+
+		expect(await events).toEqual([{ type: "done" }])
+
+		agent.dispose()
+	})
+
+	test("emits Pi compaction events for the active conversation", async () => {
+		const recordedCompactions: unknown[] = []
+		const { PiQueryAgent } = await import("./pi-query-agent.ts")
+		const agent = new PiQueryAgent({
+			toolbox: createStubToolbox(),
+			cwd: "/tmp/freya-pi-query-agent-test",
+			systemPrompt: "test",
+		})
+		agent.addEventListener(QueryAgentEvent.Compaction, (event) => {
+			recordedCompactions.push(event)
+		})
+
+		const events = collectEvents(
+			agent.ask({
+				conversationId: "conversation-1",
+				message: "hello",
+			}),
+		)
+
+		await sessionCreationStarted
+		releaseSessionCreation()
+		await promptStarted
+
+		const handler = sessionCompactHandlerFromCapturedOptions()
+		await handler({
+			type: "session_compact",
+			fromExtension: false,
+			compactionEntry: {
+				type: "compaction",
+				id: "pi-compaction-1",
+				timestamp: "2026-06-15T00:00:00.000Z",
+				summary: "The user prefers concise updates.",
+				firstKeptEntryId: "pi-entry-7",
+				tokensBefore: 1234,
+				details: { reason: "threshold" },
+			},
+		})
+
+		expect(recordedCompactions).toEqual([
+			{
+				type: QueryAgentEvent.Compaction,
+				conversationId: "conversation-1",
+				summary: "The user prefers concise updates.",
+				firstKeptEntryId: "pi-entry-7",
+				compactedEntryRange: null,
+				tokensBefore: 1234,
+				details: { reason: "threshold" },
+				fromExtension: false,
+			},
+		])
+
+		releasePrompt()
+
+		expect(await events).toEqual([{ type: "done" }])
+		expect(unsubscribeCalls).toBe(1)
+
+		agent.dispose()
+	})
+
+	test("emits Freya coverage through the entry before Pi's kept boundary", async () => {
+		const recordedCompactions: unknown[] = []
+		const { PiQueryAgent } = await import("./pi-query-agent.ts")
+		const agent = new PiQueryAgent({
+			toolbox: createStubToolbox(),
+			cwd: "/tmp/freya-pi-query-agent-test",
+			systemPrompt: "test",
+			initialEntries: [
+				{
+					id: "entry-1",
+					sequence: 1,
+					kind: ConversationEntryKind.UserMessage,
+					payload: {
+						role: "user",
+						parts: [{ type: "text", text: "old hello" }],
+					},
+					metadata: {},
+					createdAt: new Date("2026-06-15T00:00:00.000Z"),
+				},
+				{
+					id: "entry-2",
+					sequence: 2,
+					kind: ConversationEntryKind.AssistantMessage,
+					payload: {
+						role: "assistant",
+						parts: [{ type: "text", text: "kept reply" }],
+					},
+					metadata: {},
+					createdAt: new Date("2026-06-15T00:00:01.000Z"),
+				},
+			],
+		})
+		agent.addEventListener(QueryAgentEvent.Compaction, (event) => {
+			recordedCompactions.push(event)
+		})
+
+		const events = collectEvents(
+			agent.ask({
+				conversationId: "conversation-1",
+				message: "hello",
+			}),
+		)
+
+		await sessionCreationStarted
+
+		await extensionHandlerFromCapturedOptions("session_before_compact")({
+			type: "session_before_compact",
+			preparation: {
+				firstKeptEntryId: "message-2",
+			},
+			branchEntries: [{ id: "message-1" }, { id: "message-2" }],
+		})
+		await extensionHandlerFromCapturedOptions("session_compact")({
+			type: "session_compact",
+			fromExtension: false,
+			compactionEntry: {
+				type: "compaction",
+				id: "pi-compaction-1",
+				timestamp: "2026-06-15T00:00:00.000Z",
+				summary: "Old hello was discussed.",
+				firstKeptEntryId: "message-2",
+				tokensBefore: 1234,
+			},
+		})
+
+		expect(recordedCompactions).toEqual([
+			{
+				type: QueryAgentEvent.Compaction,
+				conversationId: "conversation-1",
+				summary: "Old hello was discussed.",
+				firstKeptEntryId: "message-2",
+				compactedEntryRange: {
+					startSequence: 1,
+					endSequence: 1,
+				},
+				tokensBefore: 1234,
+				details: undefined,
+				fromExtension: false,
+			},
+		])
+
+		releaseSessionCreation()
+		await promptStarted
+		releasePrompt()
+
+		expect(await events).toEqual([{ type: "done" }])
 
 		agent.dispose()
 	})
@@ -182,7 +500,6 @@ describe("PiQueryAgent", () => {
 	test("surfaces Pi message_end provider errors instead of done", async () => {
 		const { PiQueryAgent } = await import("./pi-query-agent.ts")
 		const agent = new PiQueryAgent({
-			userId: "user-1",
 			toolbox: createStubToolbox(),
 			cwd: "/tmp/freya-pi-query-agent-test",
 			systemPrompt: "test",
@@ -219,7 +536,6 @@ describe("PiQueryAgent", () => {
 	test("surfaces Pi agent_end provider errors instead of done", async () => {
 		const { PiQueryAgent } = await import("./pi-query-agent.ts")
 		const agent = new PiQueryAgent({
-			userId: "user-1",
 			toolbox: createStubToolbox(),
 			cwd: "/tmp/freya-pi-query-agent-test",
 			systemPrompt: "test",
@@ -256,8 +572,10 @@ describe("PiQueryAgent", () => {
 	})
 })
 
-async function collectEvents(events: AsyncIterable<QueryAgentEvent>): Promise<QueryAgentEvent[]> {
-	const result: QueryAgentEvent[] = []
+async function collectEvents(
+	events: AsyncIterable<QueryAgentStreamEvent>,
+): Promise<QueryAgentStreamEvent[]> {
+	const result: QueryAgentStreamEvent[] = []
 	for await (const event of events) {
 		result.push(event)
 	}
@@ -290,6 +608,79 @@ function createStubToolbox(): QueryAgentToolbox {
 	}
 }
 
+function sessionCompactHandlerFromCapturedOptions(): CapturedExtensionHandler {
+	return extensionHandlerFromCapturedOptions("session_compact")
+}
+
+function extensionHandlerFromCapturedOptions(eventName: string): CapturedExtensionHandler {
+	if (!isRecord(createAgentSessionOptions)) {
+		throw new Error("createAgentSession options were not captured")
+	}
+
+	const resourceLoader = createAgentSessionOptions.resourceLoader
+	if (!isCapturedResourceLoader(resourceLoader)) {
+		throw new Error("resourceLoader was not captured")
+	}
+
+	const extensionsResult = resourceLoader.getExtensions()
+	if (!isRecord(extensionsResult) || !Array.isArray(extensionsResult.extensions)) {
+		throw new Error("extensions were not captured")
+	}
+
+	const extension = extensionsResult.extensions[0]
+	if (!isCapturedExtension(extension)) {
+		throw new Error("compaction extension was not captured")
+	}
+
+	const handlers = extension.handlers.get(eventName)
+	const handler = handlers?.[0]
+	if (!handler) {
+		throw new Error(`${eventName} handler was not captured`)
+	}
+
+	return handler
+}
+
+function isCapturedResourceLoader(value: unknown): value is CapturedResourceLoader {
+	return isRecord(value) && typeof value.getExtensions === "function"
+}
+
+function isCapturedExtension(value: unknown): value is CapturedExtension {
+	return isRecord(value) && value.handlers instanceof Map
+}
+
+function isDefaultResourceLoaderOptions(
+	value: unknown,
+): value is CapturedDefaultResourceLoaderOptions {
+	return (
+		isRecord(value) &&
+		(value.extensionFactories === undefined ||
+			(Array.isArray(value.extensionFactories) &&
+				value.extensionFactories.every(isCapturedExtensionFactory)))
+	)
+}
+
+function isCapturedExtensionFactory(value: unknown): value is CapturedExtensionFactory {
+	return typeof value === "function"
+}
+
+function modelFromMessages(messages: unknown[]): { provider: string; modelId: string } | null {
+	let model: { provider: string; modelId: string } | null = null
+
+	for (const message of messages) {
+		if (!isRecord(message)) continue
+		if (message.role !== "assistant") continue
+		if (typeof message.provider !== "string" || typeof message.model !== "string") continue
+
+		model = {
+			provider: message.provider,
+			modelId: message.model,
+		}
+	}
+
+	return model
+}
+
 function isRecord(value: unknown): value is Record<string, unknown> {
 	return typeof value === "object" && value !== null
 }