refactor: split query agent toolbox (#139)

apps/freya-backend/src/session/feed-source-provider.ts

@@ -1,7 +1,6 @@
 import type { FeedSource } from "@freya/core"
-import type { type } from "arktype"
 
-export type ConfigSchema = ReturnType<typeof type>
+export type ConfigSchema = (value: unknown) => unknown
 
 export interface FeedSourceProvider {
 	/** The source ID this provider is responsible for (e.g., "freya.location"). */

apps/freya-backend/src/session/user-session-manager.ts

@@ -14,13 +14,14 @@ import {
 	SourceNotFoundError,
 } from "../sources/errors.ts"
 import { sources } from "../sources/user-sources.ts"
-import { UserSession } from "./user-session.ts"
+import { UserSession, type UserSessionAgentConfig } from "./user-session.ts"
 
 export interface UserSessionManagerConfig {
 	db: Database
 	providers: FeedSourceProvider[]
 	feedEnhancer?: FeedEnhancer | null
 	credentialEncryptor?: CredentialEncryptor | null
+	queryAgent?: UserSessionAgentConfig
 }
 
 export class UserSessionManager {
@@ -30,6 +31,7 @@ export class UserSessionManager {
 	private readonly providers = new Map<string, FeedSourceProvider>()
 	private readonly feedEnhancer: FeedEnhancer | null
 	private readonly encryptor: CredentialEncryptor | null
+	private readonly queryAgentConfig: UserSessionAgentConfig | undefined
 
 	constructor(config: UserSessionManagerConfig) {
 		this.db = config.db
@@ -38,6 +40,7 @@ export class UserSessionManager {
 		}
 		this.feedEnhancer = config.feedEnhancer ?? null
 		this.encryptor = config.credentialEncryptor ?? null
+		this.queryAgentConfig = config.queryAgent
 	}
 
 	getProvider(sourceId: string): FeedSourceProvider | undefined {
@@ -99,6 +102,14 @@ export class UserSessionManager {
 		this.pending.delete(userId)
 	}
 
+	dispose(): void {
+		for (const session of this.sessions.values()) {
+			session.destroy()
+		}
+		this.sessions.clear()
+		this.pending.clear()
+	}
+
 	/**
 	 * Merges, validates, and persists a user's source config and/or enabled
 	 * state, then invalidates the cached session.
@@ -362,7 +373,7 @@ export class UserSessionManager {
 		}
 
 		if (promises.length === 0) {
-			return new UserSession(userId, [], this.feedEnhancer)
+			return new UserSession(userId, [], this.feedEnhancer, this.queryAgentConfig)
 		}
 
 		const results = await Promise.allSettled(promises)
@@ -386,7 +397,7 @@ export class UserSessionManager {
 			console.error("[UserSessionManager] Feed source provider failed:", error)
 		}
 
-		return new UserSession(userId, feedSources, this.feedEnhancer)
+		return new UserSession(userId, feedSources, this.feedEnhancer, this.queryAgentConfig)
 	}
 
 	/**

apps/freya-backend/src/session/user-session.test.ts

@@ -58,6 +58,15 @@ describe("UserSession", () => {
 		expect(session.getSource("test")).toBeUndefined()
 	})
 
+	test("destroy disposes query agent", () => {
+		const session = new UserSession("test-user", [createStubSource("test")])
+		const disposeSpy = spyOn(session.agent, "dispose")
+
+		session.destroy()
+
+		expect(disposeSpy).toHaveBeenCalled()
+	})
+
 	test("engine.executeAction routes to correct source", async () => {
 		const location = new LocationSource()
 		const session = new UserSession("test-user", [location])

apps/freya-backend/src/session/user-session.ts

@@ -6,11 +6,24 @@ import {
 	type FeedSource,
 } from "@freya/core"
 
+import type { QueryAgentToolbox } from "../agent/query-agent-toolbox.ts"
+import type { QueryAgent } from "../agent/query-agent.ts"
 import type { FeedEnhancer } from "../enhancement/enhance-feed.ts"
 
+import { PiQueryAgent } from "../agent/pi-query-agent.ts"
+import { UserSessionQueryAgentToolbox } from "../agent/user-session-query-agent-toolbox.ts"
+
+export interface UserSessionAgentConfig {
+	apiKey?: string
+	cwd?: string
+	systemPrompt?: string
+}
+
 export class UserSession {
 	readonly userId: string
 	readonly engine: FeedEngine
+	readonly toolbox: QueryAgentToolbox
+	readonly agent: QueryAgent
 	private sources = new Map<string, FeedSource>()
 	private readonly enhancer: FeedEnhancer | null
 	private enhancedItems: FeedItem[] | null = null
@@ -19,7 +32,12 @@ export class UserSession {
 	private enhancingPromise: Promise<void> | null = null
 	private unsubscribe: (() => void) | null = null
 
-	constructor(userId: string, sources: FeedSource[], enhancer?: FeedEnhancer | null) {
+	constructor(
+		userId: string,
+		sources: FeedSource[],
+		enhancer?: FeedEnhancer | null,
+		agentConfig?: UserSessionAgentConfig,
+	) {
 		this.userId = userId
 		this.engine = new FeedEngine()
 		this.enhancer = enhancer ?? null
@@ -35,6 +53,15 @@ export class UserSession {
 			})
 		}
 
+		this.toolbox = new UserSessionQueryAgentToolbox(this)
+		this.agent = new PiQueryAgent({
+			userId: this.userId,
+			toolbox: this.toolbox,
+			apiKey: agentConfig?.apiKey,
+			cwd: agentConfig?.cwd,
+			systemPrompt: agentConfig?.systemPrompt,
+		})
+
 		this.engine.start()
 	}
 
@@ -174,6 +201,7 @@ export class UserSession {
 	}
 
 	destroy(): void {
+		this.agent.dispose()
 		this.unsubscribe?.()
 		this.unsubscribe = null
 		this.engine.stop()