From 95f6c99f19b5dd3fb3a6687ad3ab345b1d956b43 Mon Sep 17 00:00:00 2001 From: Kenneth Date: Mon, 15 Jun 2026 20:58:07 +0100 Subject: [PATCH] refactor: split query agent toolbox (#139) --- apps/freya-backend/.env.example | 2 - apps/freya-backend/src/agent/http.test.ts | 11 +- apps/freya-backend/src/agent/http.ts | 16 +- .../src/agent/pi-query-agent.test.ts | 59 ++- .../freya-backend/src/agent/pi-query-agent.ts | 100 ++--- .../src/agent/query-agent-toolbox.ts | 93 +++++ apps/freya-backend/src/agent/query-agent.ts | 2 - apps/freya-backend/src/agent/tools.test.ts | 116 ++++++ apps/freya-backend/src/agent/tools.ts | 355 +++++++----------- .../agent/user-session-query-agent-toolbox.ts | 253 +++++++++++++ apps/freya-backend/src/lib/env.test.ts | 21 -- apps/freya-backend/src/lib/env.ts | 2 - apps/freya-backend/src/server.ts | 17 +- .../src/session/feed-source-provider.ts | 3 +- .../src/session/user-session-manager.ts | 17 +- .../src/session/user-session.test.ts | 9 + .../freya-backend/src/session/user-session.ts | 30 +- 17 files changed, 771 insertions(+), 335 deletions(-) create mode 100644 apps/freya-backend/src/agent/query-agent-toolbox.ts create mode 100644 apps/freya-backend/src/agent/tools.test.ts create mode 100644 apps/freya-backend/src/agent/user-session-query-agent-toolbox.ts diff --git a/apps/freya-backend/.env.example b/apps/freya-backend/.env.example index ffd672e..bf9e5d0 100644 --- a/apps/freya-backend/.env.example +++ b/apps/freya-backend/.env.example @@ -12,8 +12,6 @@ BETTER_AUTH_URL=http://localhost:3000 # OpenRouter (LLM feed enhancement) OPENROUTER_API_KEY= -# Optional: override the default model (default: openai/gpt-4.1-mini) -# OPENROUTER_MODEL=openai/gpt-4.1-mini # Apple WeatherKit credentials WEATHERKIT_PRIVATE_KEY= diff --git a/apps/freya-backend/src/agent/http.test.ts b/apps/freya-backend/src/agent/http.test.ts index d5c3aa3..4dfdf6d 100644 --- a/apps/freya-backend/src/agent/http.test.ts +++ b/apps/freya-backend/src/agent/http.test.ts @@ -1,6 +1,7 @@ import { describe, expect, test } from "bun:test" import { Hono } from "hono" +import type { UserSessionManager } from "../session/index.ts" import type { QueryDebugTools, QueryDebugToolDefinition } from "./debug-tools.ts" import type { QueryAgent, QueryAgentAsk, QueryAgentEvent } from "./query-agent.ts" @@ -24,8 +25,6 @@ class FakeQueryAgent implements QueryAgent { } } - disposeUser(): void {} - dispose(): void {} } @@ -52,8 +51,14 @@ class FakeDebugTools implements QueryDebugTools { function buildTestApp(queryAgent: QueryAgent, userId?: string) { const app = new Hono() + const sessionManager = { + async getOrCreate() { + return { agent: queryAgent } + }, + } as unknown as UserSessionManager + registerAgentHttpHandlers(app, { - queryAgent, + sessionManager, authSessionMiddleware: mockAuthSessionMiddleware(userId), }) return app diff --git a/apps/freya-backend/src/agent/http.ts b/apps/freya-backend/src/agent/http.ts index 9d212fd..77bbd01 100644 --- a/apps/freya-backend/src/agent/http.ts +++ b/apps/freya-backend/src/agent/http.ts @@ -4,14 +4,14 @@ import { type } from "arktype" import { createMiddleware } from "hono/factory" import type { AuthSessionMiddleware } from "../auth/session-middleware.ts" +import type { UserSessionManager } from "../session/index.ts" import type { QueryDebugTools } from "./debug-tools.ts" -import type { QueryAgent } from "./query-agent.ts" import { collectQueryAgentResponse, QueryAgentError } from "./query-agent.ts" type Env = { Variables: { - queryAgent: QueryAgent + sessionManager: UserSessionManager } } @@ -22,7 +22,7 @@ type DebugEnv = { } interface AgentHttpHandlersDeps { - queryAgent: QueryAgent + sessionManager: UserSessionManager authSessionMiddleware: AuthSessionMiddleware } @@ -39,10 +39,10 @@ const AgentAskRequestBody = type({ export function registerAgentHttpHandlers( app: Hono, - { queryAgent, authSessionMiddleware }: AgentHttpHandlersDeps, + { sessionManager, authSessionMiddleware }: AgentHttpHandlersDeps, ) { const inject = createMiddleware(async (c, next) => { - c.set("queryAgent", queryAgent) + c.set("sessionManager", sessionManager) await next() }) @@ -76,11 +76,11 @@ async function handleAgentAsk(c: Context) { } const user = c.get("user")! - const queryAgent = c.get("queryAgent") + const sessionManager = c.get("sessionManager") try { - const response = await collectQueryAgentResponse(queryAgent, { - userId: user.id, + const session = await sessionManager.getOrCreate(user.id) + const response = await collectQueryAgentResponse(session.agent, { message: parsed.message, }) return c.json(response) diff --git a/apps/freya-backend/src/agent/pi-query-agent.test.ts b/apps/freya-backend/src/agent/pi-query-agent.test.ts index 78398cf..9c9fc98 100644 --- a/apps/freya-backend/src/agent/pi-query-agent.test.ts +++ b/apps/freya-backend/src/agent/pi-query-agent.test.ts @@ -1,6 +1,6 @@ import { beforeEach, describe, expect, mock, test } from "bun:test" -import type { UserSessionManager } from "../session/index.ts" +import type { QueryAgentToolbox } from "./query-agent-toolbox.ts" import type { QueryAgentEvent } from "./query-agent.ts" interface FakePiSession { @@ -11,6 +11,8 @@ interface FakePiSession { let createAgentSessionCalls = 0 let createAgentSessionOptions: unknown +let runtimeApiKeyCalls: Array<{ provider: string; apiKey: string }> = [] +let modelFindCalls: Array<{ provider: string; modelId: string }> = [] let promptCalls = 0 let unsubscribeCalls = 0 let sessionListeners: Array<(event: unknown) => void> = [] @@ -53,7 +55,9 @@ mock.module("@earendil-works/pi-coding-agent", () => ({ AuthStorage: { inMemory() { return { - setRuntimeApiKey(_provider: string, _apiKey: string): void {}, + setRuntimeApiKey(provider: string, apiKey: string): void { + runtimeApiKeyCalls.push({ provider, apiKey }) + }, } }, }, @@ -73,7 +77,8 @@ mock.module("@earendil-works/pi-coding-agent", () => ({ ModelRegistry: { inMemory(_authStorage: unknown) { return { - find(_provider: string, _modelId: string): unknown { + find(provider: string, modelId: string): unknown { + modelFindCalls.push({ provider, modelId }) return { id: "mock-model" } }, } @@ -94,6 +99,8 @@ mock.module("@earendil-works/pi-coding-agent", () => ({ beforeEach(() => { createAgentSessionCalls = 0 createAgentSessionOptions = undefined + runtimeApiKeyCalls = [] + modelFindCalls = [] promptCalls = 0 unsubscribeCalls = 0 sessionListeners = [] @@ -124,16 +131,15 @@ 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({ - sessionManager: createStubSessionManager(), - modelProvider: "mock", - modelId: "mock-model", + userId: "user-1", + toolbox: createStubToolbox(), + apiKey: "test-api-key", cwd: "/tmp/freya-pi-query-agent-test", systemPrompt: "test", }) const firstEvents = collectEvents( agent.ask({ - userId: "user-1", message: "first", }), ) @@ -142,7 +148,6 @@ describe("PiQueryAgent", () => { const secondEvents = await collectEvents( agent.ask({ - userId: "user-1", message: "second", }), ) @@ -154,6 +159,8 @@ describe("PiQueryAgent", () => { }, ]) expect(createAgentSessionCalls).toBe(1) + expect(runtimeApiKeyCalls).toEqual([{ provider: "openrouter", apiKey: "test-api-key" }]) + expect(modelFindCalls).toEqual([{ provider: "openrouter", modelId: "z-ai/glm-4.7-flash" }]) expect(promptCalls).toBe(0) releaseSessionCreation() @@ -175,9 +182,8 @@ 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({ - sessionManager: createStubSessionManager(), - modelProvider: "mock", - modelId: "mock-model", + userId: "user-1", + toolbox: createStubToolbox(), cwd: "/tmp/freya-pi-query-agent-test", systemPrompt: "test", }) @@ -195,7 +201,6 @@ describe("PiQueryAgent", () => { const events = collectEvents( agent.ask({ - userId: "user-1", message: "hello", }), ) @@ -214,9 +219,8 @@ 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({ - sessionManager: createStubSessionManager(), - modelProvider: "mock", - modelId: "mock-model", + userId: "user-1", + toolbox: createStubToolbox(), cwd: "/tmp/freya-pi-query-agent-test", systemPrompt: "test", }) @@ -236,7 +240,6 @@ describe("PiQueryAgent", () => { const events = collectEvents( agent.ask({ - userId: "user-1", message: "hello", }), ) @@ -261,12 +264,30 @@ async function collectEvents(events: AsyncIterable): Promise { + async listSources(): Promise { throw new Error("not used") }, - } as unknown as UserSessionManager + async getContext(): Promise { + throw new Error("not used") + }, + async getFeedItem(): Promise { + throw new Error("not used") + }, + async queryContext(): Promise { + throw new Error("not used") + }, + async listContext(): Promise { + throw new Error("not used") + }, + async getSourceData(): Promise { + throw new Error("not used") + }, + async executeAction(): Promise { + throw new Error("not used") + }, + } } function isRecord(value: unknown): value is Record { diff --git a/apps/freya-backend/src/agent/pi-query-agent.ts b/apps/freya-backend/src/agent/pi-query-agent.ts index 989e64e..56a3e1a 100644 --- a/apps/freya-backend/src/agent/pi-query-agent.ts +++ b/apps/freya-backend/src/agent/pi-query-agent.ts @@ -9,7 +9,7 @@ import { } from "@earendil-works/pi-coding-agent" import { tmpdir } from "node:os" -import type { UserSessionManager } from "../session/index.ts" +import type { QueryAgentToolbox } from "./query-agent-toolbox.ts" import type { QueryAgent, QueryAgentAsk, QueryAgentEvent } from "./query-agent.ts" import { InMemoryResourceLoader } from "./in-memory-resource-loader.ts" @@ -22,36 +22,37 @@ type PiAgentMessage = PiMessageEndEvent["message"] type PiAgentEndEvent = Extract export interface PiQueryAgentConfig { - sessionManager: UserSessionManager - modelProvider: string - modelId: string + userId: string + toolbox: QueryAgentToolbox apiKey?: string cwd?: string systemPrompt?: string } +const MODEL_PROVIDER = "openrouter" +const MODEL_ID = "z-ai/glm-4.7-flash" + export class PiQueryAgent implements QueryAgent { - private readonly sessionManager: UserSessionManager + private readonly userId: string + private readonly toolbox: QueryAgentToolbox private readonly cwd: string private readonly systemPrompt: string - private readonly modelProvider: string - private readonly modelId: string private readonly apiKey: string | undefined - private readonly sessions = new Map() - private readonly pendingSessions = new Map>() - private readonly activeRuns = new Map() + private session: PiSession | null = null + private pendingSession: Promise | null = null + private activeRun: symbol | null = null + private disposed = false constructor(config: PiQueryAgentConfig) { - this.sessionManager = config.sessionManager - this.modelProvider = config.modelProvider - this.modelId = config.modelId + this.userId = config.userId + this.toolbox = config.toolbox this.apiKey = config.apiKey this.cwd = config.cwd ?? tmpdir() this.systemPrompt = config.systemPrompt ?? defaultSystemPrompt } async *ask(input: QueryAgentAsk): AsyncIterable { - if (this.activeRuns.has(input.userId)) { + if (this.activeRun) { yield { type: "error", message: "A query is already running for this user", @@ -59,14 +60,14 @@ export class PiQueryAgent implements QueryAgent { return } - const run = Symbol(input.userId) - this.activeRuns.set(input.userId, run) + const run = Symbol(this.userId) + this.activeRun = run let session: PiSession try { - session = await this.getOrCreateSession(input.userId) + session = await this.getOrCreateSession() } catch (err) { - this.clearActiveRun(input.userId, run) + this.clearActiveRun(run) yield { type: "error", message: `Failed to create query session: ${errorMessage(err)}`, @@ -117,7 +118,7 @@ export class PiQueryAgent implements QueryAgent { }) .finally(() => { unsubscribe() - this.clearActiveRun(input.userId, run) + this.clearActiveRun(run) close() }) @@ -134,62 +135,62 @@ export class PiQueryAgent implements QueryAgent { } } - disposeUser(userId: string): void { - const session = this.sessions.get(userId) - session?.dispose() - this.sessions.delete(userId) - this.pendingSessions.delete(userId) - this.activeRuns.delete(userId) - } - dispose(): void { - for (const session of this.sessions.values()) { - session.dispose() - } - this.sessions.clear() - this.pendingSessions.clear() - this.activeRuns.clear() + this.disposed = true + this.session?.dispose() + this.session = null + this.pendingSession = null + this.activeRun = null } - private clearActiveRun(userId: string, run: symbol): void { - if (this.activeRuns.get(userId) === run) { - this.activeRuns.delete(userId) + private clearActiveRun(run: symbol): void { + if (this.activeRun === run) { + this.activeRun = null } } - private async getOrCreateSession(userId: string): Promise { - const existing = this.sessions.get(userId) - if (existing) return existing + private async getOrCreateSession(): Promise { + if (this.disposed) { + throw new Error("Query agent is disposed") + } - const pending = this.pendingSessions.get(userId) + if (this.session) return this.session + + const pending = this.pendingSession if (pending) return pending - const promise = this.createSession(userId) - this.pendingSessions.set(userId, promise) + const promise = this.createSession() + this.pendingSession = promise try { const session = await promise - this.sessions.set(userId, session) + if (this.disposed) { + session.dispose() + throw new Error("Query agent is disposed") + } + this.session = session return session } finally { - this.pendingSessions.delete(userId) + if (this.pendingSession === promise) { + this.pendingSession = null + } } } - private async createSession(userId: string): Promise { + private async createSession(): Promise { const settingsManager = SettingsManager.inMemory({ compaction: { enabled: true }, retry: { enabled: true, maxRetries: 2 }, }) const authStorage = AuthStorage.inMemory() if (this.apiKey) { - authStorage.setRuntimeApiKey(this.modelProvider, this.apiKey) + authStorage.setRuntimeApiKey(MODEL_PROVIDER, this.apiKey) } const modelRegistry = ModelRegistry.inMemory(authStorage) - const model = modelRegistry.find(this.modelProvider, this.modelId) + const model = modelRegistry.find(MODEL_PROVIDER, MODEL_ID) if (!model) { - throw new Error(`Pi model not found: ${this.modelProvider}/${this.modelId}`) + throw new Error(`Pi model not found: ${MODEL_PROVIDER}/${MODEL_ID}`) } const { session } = await createAgentSession({ @@ -202,8 +203,7 @@ export class PiQueryAgent implements QueryAgent { sessionManager: SessionManager.inMemory(this.cwd), noTools: "builtin", customTools: createFreyaAgentTools({ - userId, - sessionManager: this.sessionManager, + toolbox: this.toolbox, }), tools: [...FREYA_AGENT_TOOL_NAMES], }) diff --git a/apps/freya-backend/src/agent/query-agent-toolbox.ts b/apps/freya-backend/src/agent/query-agent-toolbox.ts new file mode 100644 index 0000000..9367c17 --- /dev/null +++ b/apps/freya-backend/src/agent/query-agent-toolbox.ts @@ -0,0 +1,93 @@ +import type { ContextKeyPart } from "@freya/core" + +export interface QueryAgentToolResult { + content: Array<{ type: "text"; text: string }> + details: Record +} + +/** + * Implementation boundary for FREYA query-agent tools. + * + * The Pi-facing tool definitions in `tools.ts` should stay thin: they declare + * schemas, validate and narrow raw model-provided parameters, then delegate to + * this toolbox. Concrete implementations own the actual data gathering, + * source/action lookups, result shaping, and any session-specific behavior. + */ +export interface QueryAgentToolbox { + /** + * Summarizes every source currently visible to the user's session. + * + * Implementations should refresh or read the current feed as needed, then + * return a compact source inventory including feed item counts, context + * entry counts, available action IDs/descriptions, and source errors. This + * is the broad discovery tool an agent can use before deciding which more + * targeted tool call to make. + */ + listSources(): Promise + + /** + * Reads context entries from the current FREYA context graph. + * + * `key` is a tuple-style context key. With `match: "exact"`, the implementation + * should return only the value at that exact key and indicate whether it was + * found. With `match: "prefix"`, it should return all entries whose keys + * begin with the provided key parts, plus a count. Implementations may refresh + * the feed first so the context reflects the latest source data. + */ + getContext(key: ContextKeyPart[], match: "exact" | "prefix"): Promise + + /** + * Reads one feed item by ID and includes source-local diagnostics. + * + * Implementations should search the current feed for `feedItemId`. When found, + * the result should include the item plus related context entries, source + * action summaries, and source errors. When missing, the result should clearly + * report `found: false` and return `item: null`. + */ + getFeedItem(feedItemId: string): Promise + + /** + * Returns the broad context bundle needed to answer a natural-language query. + * + * `question` is included in the result for traceability. If `feedItemId` is + * provided, implementations should also include the matching selected item + * when present. The result should expose the current feed items, context graph + * entries, available source actions, and source errors so the agent can + * synthesize an answer from the user's personal data. + */ + queryContext(question: string, feedItemId?: string): Promise + + /** + * Lists every current context graph entry. + * + * This is a lower-level inspection tool than `queryContext`: it should return + * all context entries and a count, without feed items or action summaries. + * Implementations may refresh the feed first to ensure source-provided + * context has been materialized. + */ + listContext(): Promise + + /** + * Returns all currently available data for one source. + * + * Implementations should include whether the source is enabled, all feed + * items from `sourceId`, context entries owned by that source, available + * action summaries, and errors from that source. If `feedItemId` is provided, + * the result should also include the matching selected item from that source + * when present. + */ + getSourceData(sourceId: string, feedItemId?: string): Promise + + /** + * Executes a source action and returns a serializable execution result. + * + * `sourceId` identifies the source, `actionId` identifies the action within + * that source, and `params` is the source-specific action payload. Tool + * wrappers validate the action envelope, while the source action schema owns + * payload validation. Implementations should let source/action validation + * errors propagate, and on success should return an `ok: true` result plus + * `details.actionExecution` for callers that need a structured record of + * what ran. + */ + executeAction(sourceId: string, actionId: string, params?: unknown): Promise +} diff --git a/apps/freya-backend/src/agent/query-agent.ts b/apps/freya-backend/src/agent/query-agent.ts index 20d16d4..ad76e66 100644 --- a/apps/freya-backend/src/agent/query-agent.ts +++ b/apps/freya-backend/src/agent/query-agent.ts @@ -1,5 +1,4 @@ export interface QueryAgentAsk { - userId: string message: string } @@ -12,7 +11,6 @@ export type QueryAgentEvent = export interface QueryAgent { ask(input: QueryAgentAsk): AsyncIterable - disposeUser(userId: string): void dispose(): void } diff --git a/apps/freya-backend/src/agent/tools.test.ts b/apps/freya-backend/src/agent/tools.test.ts new file mode 100644 index 0000000..42d9416 --- /dev/null +++ b/apps/freya-backend/src/agent/tools.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, mock, test } from "bun:test" + +import type { QueryAgentToolResult, QueryAgentToolbox } from "./query-agent-toolbox.ts" + +mock.module("@earendil-works/pi-coding-agent", () => ({ + defineTool(tool: unknown): unknown { + return tool + }, +})) + +interface TestTool { + name: string + parameters: unknown + execute(toolCallId: string, params: unknown): Promise +} + +describe("FREYA agent tools", () => { + test("rejects unknown top-level params", async () => { + const { createFreyaAgentTools, FREYA_GET_CONTEXT_TOOL } = await import("./tools.ts") + const tool = expectTool( + createFreyaAgentTools({ toolbox: createStubToolbox() }), + FREYA_GET_CONTEXT_TOOL, + ) + + await expect( + tool.execute("tool-call-1", { + key: ["freya.location"], + extra: true, + }), + ).rejects.toThrow("extra") + }) + + test("rejects invalid context keys", async () => { + const { createFreyaAgentTools, FREYA_GET_CONTEXT_TOOL } = await import("./tools.ts") + const tool = expectTool( + createFreyaAgentTools({ toolbox: createStubToolbox() }), + FREYA_GET_CONTEXT_TOOL, + ) + + await expect(tool.execute("tool-call-1", { key: [] })).rejects.toThrow("key") + await expect(tool.execute("tool-call-1", { key: [["freya.location"]] })).rejects.toThrow("key") + await expect( + tool.execute("tool-call-1", { key: [{ nested: { invalid: true } }] }), + ).rejects.toThrow("nested") + }) + + test("marks tool schemas as closed objects", async () => { + const { createFreyaAgentTools } = await import("./tools.ts") + const tools = createFreyaAgentTools({ toolbox: createStubToolbox() }) + + for (const tool of tools.map(expectTestTool)) { + expect(expectRecord(tool.parameters).additionalProperties).toBe(false) + } + }) +}) + +function createStubToolbox(): QueryAgentToolbox { + return { + async listSources() { + return toolResult({ sources: [] }) + }, + async getContext(key, match) { + return toolResult({ key, match }) + }, + async getFeedItem(feedItemId) { + return toolResult({ feedItemId }) + }, + async queryContext(question, feedItemId) { + return toolResult({ question, feedItemId }) + }, + async listContext() { + return toolResult({ entries: [] }) + }, + async getSourceData(sourceId, feedItemId) { + return toolResult({ sourceId, feedItemId }) + }, + async executeAction(sourceId, actionId, params) { + return toolResult({ sourceId, actionId, params }) + }, + } +} + +function toolResult(result: unknown): QueryAgentToolResult { + return { + content: [{ type: "text", text: JSON.stringify(result) }], + details: {}, + } +} + +function expectTool(tools: unknown[], name: string): TestTool { + const tool = tools.map(expectTestTool).find((candidate) => candidate.name === name) + if (!tool) { + throw new Error(`Missing test tool: ${name}`) + } + return tool +} + +function expectTestTool(value: unknown): TestTool { + const record = expectRecord(value) + const execute = record.execute + if (typeof record.name !== "string" || typeof execute !== "function") { + throw new Error("Expected test tool") + } + return { + name: record.name, + parameters: record.parameters, + execute: execute as TestTool["execute"], + } +} + +function expectRecord(value: unknown): Record { + expect(typeof value).toBe("object") + expect(value).not.toBeNull() + expect(Array.isArray(value)).toBe(false) + return value as Record +} diff --git a/apps/freya-backend/src/agent/tools.ts b/apps/freya-backend/src/agent/tools.ts index 6a2a5b1..42a815b 100644 --- a/apps/freya-backend/src/agent/tools.ts +++ b/apps/freya-backend/src/agent/tools.ts @@ -1,14 +1,11 @@ import { defineTool } from "@earendil-works/pi-coding-agent" +import { type } from "arktype" import { Type } from "typebox" -import type { UserSessionManager } from "../session/index.ts" -import type { QueryDebugTools } from "./debug-tools.ts" - -import { createQueryDebugTools } from "./debug-tools.ts" +import type { QueryAgentToolbox } from "./query-agent-toolbox.ts" interface CreateFreyaAgentToolsConfig { - userId: string - sessionManager: UserSessionManager + toolbox: QueryAgentToolbox } export const FREYA_QUERY_CONTEXT_TOOL = "freya_query_context" @@ -19,6 +16,41 @@ export const FREYA_GET_SOURCE_DATA_TOOL = "freya_get_source_data" export const FREYA_GET_FEED_ITEM_TOOL = "freya_get_feed_item" export const FREYA_EXECUTE_ACTION_TOOL = "freya_execute_action" +const ContextKeyObjectPart = type("Record").narrow( + (value) => !Array.isArray(value), +) +const ContextKeyPart = type("string | number").or(ContextKeyObjectPart) + +const GetContextToolParams = type({ + "+": "reject", + key: ContextKeyPart.array().atLeastLength(1), + "match?": "'exact' | 'prefix'", +}) + +const GetFeedItemToolParams = type({ + "+": "reject", + feedItemId: type.string.atLeastLength(1), +}) + +const QueryContextToolParams = type({ + "+": "reject", + question: type.string.atLeastLength(1), + "feedItemId?": "string", +}) + +const GetSourceDataToolParams = type({ + "+": "reject", + sourceId: type.string.atLeastLength(1), + "feedItemId?": "string", +}) + +const ExecuteActionToolParams = type({ + "+": "reject", + sourceId: type.string.atLeastLength(1), + actionId: type.string.atLeastLength(1), + "params?": "unknown", +}) + export const FREYA_AGENT_TOOL_NAMES = [ FREYA_LIST_SOURCES_TOOL, FREYA_GET_CONTEXT_TOOL, @@ -30,16 +62,13 @@ export const FREYA_AGENT_TOOL_NAMES = [ ] export function createFreyaAgentTools(config: CreateFreyaAgentToolsConfig) { - const { userId } = config - const debugTools = createQueryDebugTools(config.sessionManager) - const listSourcesTool = defineTool({ name: FREYA_LIST_SOURCES_TOOL, label: "List FREYA Sources", description: "List enabled FREYA source IDs and summarize available feed items, context entries, actions, and errors.", - parameters: Type.Object({}), - execute: async () => executeDebugTool(debugTools, userId, FREYA_LIST_SOURCES_TOOL, {}), + parameters: Type.Object({}, { additionalProperties: false }), + execute: async () => executeListSourcesTool(config.toolbox), }) const getContextTool = defineTool({ @@ -47,30 +76,34 @@ export function createFreyaAgentTools(config: CreateFreyaAgentToolsConfig) { label: "Get FREYA Context", description: "Read specific FREYA context entries by key. Use prefix matching to discover entries under a source ID, or exact matching when you know the full key.", - parameters: Type.Object({ - key: Type.Array(Type.Unknown(), { - description: - 'Context key array, for example ["freya.location"] or ["freya.location", "location"].', - }), - match: Type.Optional( - Type.Union([Type.Literal("exact"), Type.Literal("prefix")], { - description: "Match mode. Defaults to prefix.", + parameters: Type.Object( + { + key: Type.Array(Type.Unknown(), { + description: + 'Context key array, for example ["freya.location"] or ["freya.location", "location"].', }), - ), - }), - execute: async (_toolCallId, params) => - executeDebugTool(debugTools, userId, FREYA_GET_CONTEXT_TOOL, params), + match: Type.Optional( + Type.Union([Type.Literal("exact"), Type.Literal("prefix")], { + description: "Match mode. Defaults to prefix.", + }), + ), + }, + { additionalProperties: false }, + ), + execute: async (_toolCallId, params) => executeGetContextTool(config.toolbox, params), }) const getFeedItemTool = defineTool({ name: FREYA_GET_FEED_ITEM_TOOL, label: "Get FREYA Feed Item", description: "Read one feed item by ID, including related source context, actions, and errors.", - parameters: Type.Object({ - feedItemId: Type.String({ description: "Feed item ID to inspect." }), - }), - execute: async (_toolCallId, params) => - executeDebugTool(debugTools, userId, FREYA_GET_FEED_ITEM_TOOL, params), + parameters: Type.Object( + { + feedItemId: Type.String({ description: "Feed item ID to inspect." }), + }, + { additionalProperties: false }, + ), + execute: async (_toolCallId, params) => executeGetFeedItemTool(config.toolbox, params), }) const queryContextTool = defineTool({ @@ -78,17 +111,20 @@ export function createFreyaAgentTools(config: CreateFreyaAgentToolsConfig) { label: "Query FREYA Context", description: "Read the user's current FREYA feed, source graph context, source errors, and available actions.", - parameters: Type.Object({ - question: Type.String({ - description: "The specific personal-context question to answer.", - }), - feedItemId: Type.Optional( - Type.String({ - description: "Optional feed item ID when the user is asking about a specific card.", + parameters: Type.Object( + { + question: Type.String({ + description: "The specific personal-context question to answer.", }), - ), - }), - execute: async (_toolCallId, params) => executeQueryContextTool(config, params), + feedItemId: Type.Optional( + Type.String({ + description: "Optional feed item ID when the user is asking about a specific card.", + }), + ), + }, + { additionalProperties: false }, + ), + execute: async (_toolCallId, params) => executeQueryContextTool(config.toolbox, params), }) const listContextTool = defineTool({ @@ -96,8 +132,8 @@ export function createFreyaAgentTools(config: CreateFreyaAgentToolsConfig) { label: "List FREYA Context", description: "List all current FREYA context graph entries for the user. Use this to inspect what personal context is available.", - parameters: Type.Object({}), - execute: async () => executeListContextTool(config), + parameters: Type.Object({}, { additionalProperties: false }), + execute: async () => executeListContextTool(config.toolbox), }) const getSourceDataTool = defineTool({ @@ -105,17 +141,20 @@ export function createFreyaAgentTools(config: CreateFreyaAgentToolsConfig) { label: "Get FREYA Source Data", description: "Get current feed items, context entries, actions, and errors for a specific FREYA source ID.", - parameters: Type.Object({ - sourceId: Type.String({ - description: "Source ID, for example freya.location, freya.tfl, or freya.weather.", - }), - feedItemId: Type.Optional( - Type.String({ - description: "Optional feed item ID to select one item from the source.", + parameters: Type.Object( + { + sourceId: Type.String({ + description: "Source ID, for example freya.location, freya.tfl, or freya.weather.", }), - ), - }), - execute: async (_toolCallId, params) => executeGetSourceDataTool(config, params), + feedItemId: Type.Optional( + Type.String({ + description: "Optional feed item ID to select one item from the source.", + }), + ), + }, + { additionalProperties: false }, + ), + execute: async (_toolCallId, params) => executeGetSourceDataTool(config.toolbox, params), }) const executeActionTool = defineTool({ @@ -123,16 +162,19 @@ export function createFreyaAgentTools(config: CreateFreyaAgentToolsConfig) { label: "Execute FREYA Action", description: "Execute an available FREYA source action immediately without creating a proposal.", - parameters: Type.Object({ - sourceId: Type.String({ description: "Source ID that should execute the action." }), - actionId: Type.String({ description: "Source action ID to execute." }), - params: Type.Optional( - Type.Unknown({ - description: "Parameters to pass to the source action.", - }), - ), - }), - execute: async (_toolCallId, params) => executeActionToolCall(config, params), + parameters: Type.Object( + { + sourceId: Type.String({ description: "Source ID that should execute the action." }), + actionId: Type.String({ description: "Source action ID to execute." }), + params: Type.Optional( + Type.Unknown({ + description: "Parameters to pass to the source action.", + }), + ), + }, + { additionalProperties: false }, + ), + execute: async (_toolCallId, params) => executeActionToolCall(config.toolbox, params), }) return [ @@ -146,166 +188,57 @@ export function createFreyaAgentTools(config: CreateFreyaAgentToolsConfig) { ] } -async function executeDebugTool( - debugTools: QueryDebugTools, - userId: string, - toolName: string, - params: unknown, -) { - const result = await debugTools.execute(userId, toolName, params) - - return { - content: [ - { - type: "text" as const, - text: JSON.stringify(result), - }, - ], - details: {}, - } +async function executeListSourcesTool(toolbox: QueryAgentToolbox) { + return toolbox.listSources() } -async function executeQueryContextTool( - config: CreateFreyaAgentToolsConfig, - params: { question: string; feedItemId?: string }, -) { - const userSession = await config.sessionManager.getOrCreate(config.userId) - const feed = await userSession.feed() - const context = userSession.engine.currentContext() - const feedItemId = params.feedItemId - const selectedItem = - typeof feedItemId === "string" ? feed.items.find((item) => item.id === feedItemId) : undefined - const actions = await userSession.listActions() - - return { - content: [ - { - type: "text" as const, - text: JSON.stringify({ - time: context.time.toISOString(), - question: params.question, - feedItemId: feedItemId ?? null, - selectedItem: selectedItem ?? null, - items: feed.items, - context: context.entries(), - availableActions: actions.map((entry) => ({ - sourceId: entry.sourceId, - actions: Object.values(entry.actions).map((action) => ({ - id: action.id, - description: action.description ?? null, - })), - })), - errors: feed.errors.map((error) => ({ - sourceId: error.sourceId, - message: error.error.message, - })), - }), - }, - ], - details: {}, +async function executeGetContextTool(toolbox: QueryAgentToolbox, rawParams: unknown) { + const params = GetContextToolParams(rawParams) + if (params instanceof type.errors) { + throw new Error(params.summary) } + + const match = params.match ?? "prefix" + + return toolbox.getContext(params.key, match) } -async function executeListContextTool(config: CreateFreyaAgentToolsConfig) { - const userSession = await config.sessionManager.getOrCreate(config.userId) - await userSession.feed() - const context = userSession.engine.currentContext() - const entries = context.entries() - - return { - content: [ - { - type: "text" as const, - text: JSON.stringify({ - time: context.time.toISOString(), - count: entries.length, - entries, - }), - }, - ], - details: {}, +async function executeGetFeedItemTool(toolbox: QueryAgentToolbox, rawParams: unknown) { + const params = GetFeedItemToolParams(rawParams) + if (params instanceof type.errors) { + throw new Error(params.summary) } + + return toolbox.getFeedItem(params.feedItemId) } -async function executeGetSourceDataTool( - config: CreateFreyaAgentToolsConfig, - params: { sourceId: string; feedItemId?: string }, -) { - const userSession = await config.sessionManager.getOrCreate(config.userId) - const feed = await userSession.feed() - const context = userSession.engine.currentContext() - const sourceActions = userSession.hasSource(params.sourceId) - ? await userSession.engine.listActions(params.sourceId) - : {} - - const items = feed.items.filter((item) => item.sourceId === params.sourceId) - const selectedItem = - params.feedItemId !== undefined - ? items.find((item) => item.id === params.feedItemId) - : undefined - const contextEntries = context.entries().filter((entry) => entry.key[0] === params.sourceId) - const errors = feed.errors - .filter((error) => error.sourceId === params.sourceId) - .map((error) => ({ - sourceId: error.sourceId, - message: error.error.message, - })) - - return { - content: [ - { - type: "text" as const, - text: JSON.stringify({ - time: context.time.toISOString(), - sourceId: params.sourceId, - hasSource: userSession.hasSource(params.sourceId), - feedItemId: params.feedItemId ?? null, - selectedItem: selectedItem ?? null, - items, - context: contextEntries, - actions: Object.values(sourceActions).map((action) => ({ - id: action.id, - description: action.description ?? null, - })), - errors, - }), - }, - ], - details: {}, +async function executeQueryContextTool(toolbox: QueryAgentToolbox, rawParams: unknown) { + const params = QueryContextToolParams(rawParams) + if (params instanceof type.errors) { + throw new Error(params.summary) } + + return toolbox.queryContext(params.question, params.feedItemId) } -async function executeActionToolCall( - config: CreateFreyaAgentToolsConfig, - params: { - sourceId: string - actionId: string - params?: unknown - }, -) { - const userSession = await config.sessionManager.getOrCreate(config.userId) - const result = await userSession.engine.executeAction( - params.sourceId, - params.actionId, - params.params, - ) - - const actionExecution = { - sourceId: params.sourceId, - actionId: params.actionId, - result: result ?? null, - } - - return { - content: [ - { - type: "text" as const, - text: JSON.stringify({ - ok: true, - ...actionExecution, - }), - }, - ], - details: { actionExecution }, - } +async function executeListContextTool(toolbox: QueryAgentToolbox) { + return toolbox.listContext() +} + +async function executeGetSourceDataTool(toolbox: QueryAgentToolbox, rawParams: unknown) { + const params = GetSourceDataToolParams(rawParams) + if (params instanceof type.errors) { + throw new Error(params.summary) + } + + return toolbox.getSourceData(params.sourceId, params.feedItemId) +} + +async function executeActionToolCall(toolbox: QueryAgentToolbox, rawParams: unknown) { + const params = ExecuteActionToolParams(rawParams) + if (params instanceof type.errors) { + throw new Error(params.summary) + } + + return toolbox.executeAction(params.sourceId, params.actionId, params.params) } diff --git a/apps/freya-backend/src/agent/user-session-query-agent-toolbox.ts b/apps/freya-backend/src/agent/user-session-query-agent-toolbox.ts new file mode 100644 index 0000000..1bc3305 --- /dev/null +++ b/apps/freya-backend/src/agent/user-session-query-agent-toolbox.ts @@ -0,0 +1,253 @@ +import { contextKey, type ContextKeyPart } from "@freya/core" + +import type { UserSession } from "../session/user-session.ts" +import type { QueryAgentToolResult, QueryAgentToolbox } from "./query-agent-toolbox.ts" + +export class UserSessionQueryAgentToolbox implements QueryAgentToolbox { + constructor(private readonly session: UserSession) {} + + async listSources(): Promise { + const feed = await this.session.feed() + const context = this.session.engine.currentContext() + const contextEntries = context.entries() + const actions = await this.session.listActions() + + const feedCounts = countBy(feed.items.map((item) => item.sourceId)) + const contextCounts = countBy( + contextEntries + .map((entry) => entry.key[0]) + .filter((part): part is string => typeof part === "string"), + ) + const errors = groupErrorsBySource( + feed.errors.map((error) => ({ + sourceId: error.sourceId, + message: error.error.message, + })), + ) + const actionEntries = new Map(actions.map((entry) => [entry.sourceId, entry.actions])) + const sourceIds = new Set([ + ...actionEntries.keys(), + ...feedCounts.keys(), + ...contextCounts.keys(), + ...errors.keys(), + ]) + + return toolResult({ + time: context.time.toISOString(), + sources: [...sourceIds].sort().map((sourceId) => { + const sourceActions = actionEntries.get(sourceId) ?? {} + const feedItemCount = feedCounts.get(sourceId) ?? 0 + const contextEntryCount = contextCounts.get(sourceId) ?? 0 + + return { + sourceId, + hasFeedItems: feedItemCount > 0, + feedItemCount, + hasContext: contextEntryCount > 0, + contextEntryCount, + actions: Object.values(sourceActions).map((action) => ({ + id: action.id, + description: action.description ?? null, + })), + errors: errors.get(sourceId) ?? [], + } + }), + }) + } + + async getContext( + key: ContextKeyPart[], + match: "exact" | "prefix", + ): Promise { + await this.session.feed() + const context = this.session.engine.currentContext() + const keyObject = contextKey(...key) + + if (match === "exact") { + const value = context.get(keyObject) + return toolResult({ + time: context.time.toISOString(), + match, + key, + found: value !== undefined, + value: value ?? null, + }) + } + + const entries = context.find(keyObject) + return toolResult({ + time: context.time.toISOString(), + match, + key, + count: entries.length, + entries, + }) + } + + async getFeedItem(feedItemId: string): Promise { + const feed = await this.session.feed() + const context = this.session.engine.currentContext() + const item = feed.items.find((candidate) => candidate.id === feedItemId) + + if (!item) { + return toolResult({ + time: context.time.toISOString(), + feedItemId, + found: false, + item: null, + }) + } + + const sourceActions = this.session.hasSource(item.sourceId) + ? await this.session.engine.listActions(item.sourceId) + : {} + const errors = feed.errors + .filter((error) => error.sourceId === item.sourceId) + .map((error) => ({ + sourceId: error.sourceId, + message: error.error.message, + })) + + return toolResult({ + time: context.time.toISOString(), + feedItemId, + found: true, + item, + source: { + sourceId: item.sourceId, + hasSource: this.session.hasSource(item.sourceId), + context: context.entries().filter((entry) => entry.key[0] === item.sourceId), + actions: Object.values(sourceActions).map((action) => ({ + id: action.id, + description: action.description ?? null, + })), + errors, + }, + }) + } + + async queryContext(question: string, feedItemId?: string): Promise { + const feed = await this.session.feed() + const context = this.session.engine.currentContext() + const selectedItem = feedItemId ? feed.items.find((item) => item.id === feedItemId) : undefined + const actions = await this.session.listActions() + + return toolResult({ + time: context.time.toISOString(), + question, + feedItemId: feedItemId ?? null, + selectedItem: selectedItem ?? null, + items: feed.items, + context: context.entries(), + availableActions: actions.map((entry) => ({ + sourceId: entry.sourceId, + actions: Object.values(entry.actions).map((action) => ({ + id: action.id, + description: action.description ?? null, + })), + })), + errors: feed.errors.map((error) => ({ + sourceId: error.sourceId, + message: error.error.message, + })), + }) + } + + async listContext(): Promise { + await this.session.feed() + const context = this.session.engine.currentContext() + const entries = context.entries() + + return toolResult({ + time: context.time.toISOString(), + count: entries.length, + entries, + }) + } + + async getSourceData(sourceId: string, feedItemId?: string): Promise { + const feed = await this.session.feed() + const context = this.session.engine.currentContext() + const sourceActions = this.session.hasSource(sourceId) + ? await this.session.engine.listActions(sourceId) + : {} + + const items = feed.items.filter((item) => item.sourceId === sourceId) + const selectedItem = feedItemId ? items.find((item) => item.id === feedItemId) : undefined + const contextEntries = context.entries().filter((entry) => entry.key[0] === sourceId) + const errors = feed.errors + .filter((error) => error.sourceId === sourceId) + .map((error) => ({ + sourceId: error.sourceId, + message: error.error.message, + })) + + return toolResult({ + time: context.time.toISOString(), + sourceId, + hasSource: this.session.hasSource(sourceId), + feedItemId: feedItemId ?? null, + selectedItem: selectedItem ?? null, + items, + context: contextEntries, + actions: Object.values(sourceActions).map((action) => ({ + id: action.id, + description: action.description ?? null, + })), + errors, + }) + } + + async executeAction( + sourceId: string, + actionId: string, + params?: unknown, + ): Promise { + const result = await this.session.engine.executeAction(sourceId, actionId, params) + const actionExecution = { + sourceId, + actionId, + result: result ?? null, + } + + return toolResult( + { + ok: true, + ...actionExecution, + }, + { actionExecution }, + ) + } +} + +function toolResult(result: unknown, details: Record = {}): QueryAgentToolResult { + return { + content: [ + { + type: "text" as const, + text: JSON.stringify(result), + }, + ], + details, + } +} + +function countBy(values: string[]): Map { + const result = new Map() + for (const value of values) { + result.set(value, (result.get(value) ?? 0) + 1) + } + return result +} + +function groupErrorsBySource( + errors: Array<{ sourceId: string; message: string }>, +): Map> { + const result = new Map>() + for (const error of errors) { + const group = result.get(error.sourceId) ?? [] + group.push(error) + result.set(error.sourceId, group) + } + return result +} diff --git a/apps/freya-backend/src/lib/env.test.ts b/apps/freya-backend/src/lib/env.test.ts index ded8ac5..d59f7fe 100644 --- a/apps/freya-backend/src/lib/env.test.ts +++ b/apps/freya-backend/src/lib/env.test.ts @@ -11,7 +11,6 @@ describe("ensureEnv", () => { EXA_API_KEY: " exa-key ", GOOGLE_MAPS_API_KEY: " google-maps-key ", OPENROUTER_API_KEY: " openrouter-key ", - OPENROUTER_MODEL: " model-name ", TFL_API_KEY: " tfl-key ", WEATHERKIT_KEY_ID: " weather-key-id ", WEATHERKIT_PRIVATE_KEY: " weather-private-key ", @@ -26,7 +25,6 @@ describe("ensureEnv", () => { exaApiKey: "exa-key", googleMapsApiKey: "google-maps-key", openrouterApiKey: "openrouter-key", - openrouterModel: "model-name", tflApiKey: "tfl-key", weatherkitKeyId: "weather-key-id", weatherkitPrivateKey: "weather-private-key", @@ -53,25 +51,6 @@ describe("ensureEnv", () => { ).toThrow("Missing required environment variables: GOOGLE_MAPS_API_KEY") }) - test("allows openrouter model to be omitted", () => { - const env = ensureEnv({ - BETTER_AUTH_SECRET: "auth-secret", - CREDENTIAL_ENCRYPTION_KEY: "credential-key", - DATABASE_URL: "postgres://example", - EXA_API_KEY: "exa-key", - GOOGLE_MAPS_API_KEY: "google-maps-key", - OPENROUTER_API_KEY: "openrouter-key", - TFL_API_KEY: "tfl-key", - WEATHERKIT_KEY_ID: "weather-key-id", - WEATHERKIT_PRIVATE_KEY: "weather-private-key", - WEATHERKIT_SERVICE_ID: "weather-service-id", - WEATHERKIT_TEAM_ID: "weather-team-id", - }) - - expect(env.googleMapsApiKey).toBe("google-maps-key") - expect(env.openrouterModel).toBeUndefined() - }) - test("throws with all missing required env names", () => { expect(() => ensureEnv({})).toThrow( "Missing required environment variables: BETTER_AUTH_SECRET, CREDENTIAL_ENCRYPTION_KEY, DATABASE_URL, EXA_API_KEY, OPENROUTER_API_KEY, TFL_API_KEY, WEATHERKIT_PRIVATE_KEY, WEATHERKIT_KEY_ID, WEATHERKIT_TEAM_ID, WEATHERKIT_SERVICE_ID, GOOGLE_MAPS_API_KEY", diff --git a/apps/freya-backend/src/lib/env.ts b/apps/freya-backend/src/lib/env.ts index 59e485f..196f2c4 100644 --- a/apps/freya-backend/src/lib/env.ts +++ b/apps/freya-backend/src/lib/env.ts @@ -5,7 +5,6 @@ export interface ServerEnv { exaApiKey: string googleMapsApiKey: string openrouterApiKey: string - openrouterModel: string | undefined tflApiKey: string weatherkitKeyId: string weatherkitPrivateKey: string @@ -39,7 +38,6 @@ export function ensureEnv(env: Record): ServerEnv { exaApiKey, googleMapsApiKey, openrouterApiKey, - openrouterModel: readOptionalEnv(env, "OPENROUTER_MODEL"), tflApiKey, weatherkitKeyId, weatherkitPrivateKey, diff --git a/apps/freya-backend/src/server.ts b/apps/freya-backend/src/server.ts index 06e759f..3d43def 100644 --- a/apps/freya-backend/src/server.ts +++ b/apps/freya-backend/src/server.ts @@ -4,7 +4,6 @@ import { cors } from "hono/cors" import { registerAdminHttpHandlers } from "./admin/http.ts" import { createQueryDebugTools } from "./agent/debug-tools.ts" import { registerAgentHttpHandlers, registerDebugAgentHttpHandlers } from "./agent/http.ts" -import { PiQueryAgent } from "./agent/pi-query-agent.ts" import { createRequireAdmin } from "./auth/admin-middleware.ts" import { registerAuthHandlers } from "./auth/http.ts" import { createAuth } from "./auth/index.ts" @@ -35,11 +34,11 @@ function main() { const feedEnhancer = createFeedEnhancer({ client: createLlmClient({ apiKey: env.openrouterApiKey, - model: env.openrouterModel, }), }) const credentialEncryptor = new CredentialEncryptor(env.credentialEncryptionKey) + const piApiKey = process.env.PI_API_KEY ?? env.openrouterApiKey const sessionManager = new UserSessionManager({ db, @@ -63,13 +62,9 @@ function main() { ], feedEnhancer, credentialEncryptor, - }) - const piApiKey = process.env.PI_API_KEY ?? env.openrouterApiKey - const queryAgent = new PiQueryAgent({ - sessionManager, - modelProvider: process.env.PI_MODEL_PROVIDER ?? "openrouter", - modelId: process.env.PI_MODEL ?? env.openrouterModel ?? "z-ai/glm-4.7-flash", - apiKey: piApiKey, + queryAgent: { + apiKey: piApiKey, + }, }) if (!piApiKey) { console.warn("[query] PI_API_KEY or OPENROUTER_API_KEY not set — query agent unavailable") @@ -120,7 +115,7 @@ function main() { registerLocationHttpHandlers(app, { sessionManager, authSessionMiddleware }) registerSourcesHttpHandlers(app, { sessionManager, authSessionMiddleware }) registerAgentHttpHandlers(app, { - queryAgent, + sessionManager, authSessionMiddleware, }) if (isDebugMode) { @@ -133,7 +128,7 @@ function main() { registerAdminHttpHandlers(app, { sessionManager, adminMiddleware, db }) process.on("SIGTERM", async () => { - queryAgent.dispose() + sessionManager.dispose() await closeDb() process.exit(0) }) diff --git a/apps/freya-backend/src/session/feed-source-provider.ts b/apps/freya-backend/src/session/feed-source-provider.ts index da1d95b..5fca270 100644 --- a/apps/freya-backend/src/session/feed-source-provider.ts +++ b/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 +export type ConfigSchema = (value: unknown) => unknown export interface FeedSourceProvider { /** The source ID this provider is responsible for (e.g., "freya.location"). */ diff --git a/apps/freya-backend/src/session/user-session-manager.ts b/apps/freya-backend/src/session/user-session-manager.ts index ab2f875..b4f6adf 100644 --- a/apps/freya-backend/src/session/user-session-manager.ts +++ b/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() 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) } /** diff --git a/apps/freya-backend/src/session/user-session.test.ts b/apps/freya-backend/src/session/user-session.test.ts index e7a75a4..9dd4b46 100644 --- a/apps/freya-backend/src/session/user-session.test.ts +++ b/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]) diff --git a/apps/freya-backend/src/session/user-session.ts b/apps/freya-backend/src/session/user-session.ts index 4d316a8..da5c8d9 100644 --- a/apps/freya-backend/src/session/user-session.ts +++ b/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() private readonly enhancer: FeedEnhancer | null private enhancedItems: FeedItem[] | null = null @@ -19,7 +32,12 @@ export class UserSession { private enhancingPromise: Promise | 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()