diff --git a/.gitignore b/.gitignore index 835aa9ef3..8dd699bec 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,6 @@ coverage internal .isaac/ + +# local scratch (not committed) +scratch/ diff --git a/apps/dev-playground/server/index.ts b/apps/dev-playground/server/index.ts index eb88c2a97..800dd21f7 100644 --- a/apps/dev-playground/server/index.ts +++ b/apps/dev-playground/server/index.ts @@ -12,7 +12,7 @@ import { serving, WRITE_ACTIONS, } from "@databricks/appkit"; -import { agents, aiSearch } from "@databricks/appkit/beta"; +import { agents, aiSearch, LakebaseThreadStore } from "@databricks/appkit/beta"; import { lakebaseExamples } from "./lakebase-examples-plugin"; import { reconnect } from "./reconnect-plugin"; @@ -115,6 +115,13 @@ createApp({ // conversational default for the bare `/agent` route (the markdown agents // are dispatchers or ephemeral and don't make sense as the landing agent). defaultAgent: "helper", + // Persist threads across restarts when a Lakebase is bound (same + // LAKEBASE_ENDPOINT signal the lakebase plugin uses above). The store + // self-bootstraps its tables on setup; without Lakebase we fall back to + // the in-memory store, so local dev needs no database. + ...(process.env.LAKEBASE_ENDPOINT + ? { threadStore: new LakebaseThreadStore() } + : {}), }), aiSearch({ indexes: { diff --git a/docs/docs/api/appkit/Class.LakebaseThreadStore.md b/docs/docs/api/appkit/Class.LakebaseThreadStore.md new file mode 100644 index 000000000..f81d8fd38 --- /dev/null +++ b/docs/docs/api/appkit/Class.LakebaseThreadStore.md @@ -0,0 +1,249 @@ +# Class: LakebaseThreadStore + +Persistent [ThreadStore](Interface.ThreadStore.md) backed by Databricks Lakebase (Postgres). + +Threads and messages live in two `user_id`-scoped tables (`agent_threads`, +`agent_messages`, FK cascade). The app service principal owns the tables; +**every** query filters `WHERE user_id = $` — that is the isolation +boundary, so a user can never read or mutate another user's threads. + +The schema is self-bootstrapping: [init](#init) issues idempotent +`CREATE TABLE IF NOT EXISTS` (once-guarded) and verifies connectivity, so +a fresh Lakebase database works with no migration step. + +Pass it to the agents plugin for a deployment that survives restarts: +```ts +agents({ threadStore: new LakebaseThreadStore() }) +``` + +## Implements + +- [`ThreadStore`](Interface.ThreadStore.md) + +## Constructors + +### Constructor + +```ts +new LakebaseThreadStore(__namedParameters: LakebaseThreadStoreOptions): LakebaseThreadStore; +``` + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `__namedParameters` | [`LakebaseThreadStoreOptions`](Interface.LakebaseThreadStoreOptions.md) | + +#### Returns + +`LakebaseThreadStore` + +## Methods + +### addMessage() + +```ts +addMessage( + threadId: string, + userId: string, +message: Message): Promise; +``` + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `threadId` | `string` | +| `userId` | `string` | +| `message` | [`Message`](Interface.Message.md) | + +#### Returns + +`Promise`\<`void`\> + +#### Implementation of + +[`ThreadStore`](Interface.ThreadStore.md).[`addMessage`](Interface.ThreadStore.md#addmessage) + +*** + +### close() + +```ts +close(): Promise; +``` + +Close the pool only when this store created it. + +#### Returns + +`Promise`\<`void`\> + +#### Implementation of + +[`ThreadStore`](Interface.ThreadStore.md).[`close`](Interface.ThreadStore.md#close) + +*** + +### create() + +```ts +create(userId: string): Promise; +``` + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `userId` | `string` | + +#### Returns + +`Promise`\<[`Thread`](Interface.Thread.md)\> + +#### Implementation of + +[`ThreadStore`](Interface.ThreadStore.md).[`create`](Interface.ThreadStore.md#create) + +*** + +### delete() + +```ts +delete(threadId: string, userId: string): Promise; +``` + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `threadId` | `string` | +| `userId` | `string` | + +#### Returns + +`Promise`\<`boolean`\> + +#### Implementation of + +[`ThreadStore`](Interface.ThreadStore.md).[`delete`](Interface.ThreadStore.md#delete) + +*** + +### get() + +```ts +get(threadId: string, userId: string): Promise; +``` + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `threadId` | `string` | +| `userId` | `string` | + +#### Returns + +`Promise`\<[`Thread`](Interface.Thread.md) \| `null`\> + +#### Implementation of + +[`ThreadStore`](Interface.ThreadStore.md).[`get`](Interface.ThreadStore.md#get) + +*** + +### init() + +```ts +init(): Promise; +``` + +Verify connectivity and create the tables once (idempotent). + +#### Returns + +`Promise`\<`void`\> + +#### Implementation of + +[`ThreadStore`](Interface.ThreadStore.md).[`init`](Interface.ThreadStore.md#init) + +*** + +### list() + +```ts +list(userId: string): Promise; +``` + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `userId` | `string` | + +#### Returns + +`Promise`\<[`Thread`](Interface.Thread.md)[]\> + +#### Implementation of + +[`ThreadStore`](Interface.ThreadStore.md).[`list`](Interface.ThreadStore.md#list) + +*** + +### listSummaries() + +```ts +listSummaries(userId: string): Promise; +``` + +Optional cheap list projection for a history sidebar — summaries only, no +message bodies. When a store omits it, the agents plugin falls back to +deriving summaries from [list](Interface.ThreadStore.md#list) (correct, just heavier). + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `userId` | `string` | + +#### Returns + +`Promise`\<[`ThreadSummary`](Interface.ThreadSummary.md)[]\> + +#### Implementation of + +[`ThreadStore`](Interface.ThreadStore.md).[`listSummaries`](Interface.ThreadStore.md#listsummaries) + +*** + +### rename() + +```ts +rename( + threadId: string, + userId: string, +title: string): Promise; +``` + +Optional rename of a thread's title (user-scoped). Returns `false` when no +matching thread exists for the user. When a store omits it, the rename +route reports the operation as unsupported. + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `threadId` | `string` | +| `userId` | `string` | +| `title` | `string` | + +#### Returns + +`Promise`\<`boolean`\> + +#### Implementation of + +[`ThreadStore`](Interface.ThreadStore.md).[`rename`](Interface.ThreadStore.md#rename) diff --git a/docs/docs/api/appkit/Interface.LakebaseThreadStoreOptions.md b/docs/docs/api/appkit/Interface.LakebaseThreadStoreOptions.md new file mode 100644 index 000000000..3c6415cee --- /dev/null +++ b/docs/docs/api/appkit/Interface.LakebaseThreadStoreOptions.md @@ -0,0 +1,27 @@ +# Interface: LakebaseThreadStoreOptions + +## Properties + +### pool? + +```ts +optional pool: Pool; +``` + +An existing `pg.Pool` to run on. When omitted, the store creates its own +pool via `createLakebasePool()` (OAuth token refresh handled inside) and +closes it on [LakebaseThreadStore.close](Class.LakebaseThreadStore.md#close). An injected pool is never +closed — the caller owns its lifecycle. + +*** + +### tableSchema? + +```ts +optional tableSchema: string; +``` + +Optional Postgres schema to hold the two tables. Created on init if it +does not exist. Defaults to the connection's search_path (usually +`public`). Validated as a plain lowercase identifier — it is interpolated +into DDL, not parameterizable, so anything else is rejected. diff --git a/docs/docs/api/appkit/Interface.Thread.md b/docs/docs/api/appkit/Interface.Thread.md index e9f15fee0..22de32cfe 100644 --- a/docs/docs/api/appkit/Interface.Thread.md +++ b/docs/docs/api/appkit/Interface.Thread.md @@ -26,6 +26,18 @@ messages: Message[]; *** +### title? + +```ts +optional title: string; +``` + +Optional human title. Defaults to a value derived from the first user +message (see [ThreadStore.listSummaries](Interface.ThreadStore.md#listsummaries)); an explicit rename via +[ThreadStore.rename](Interface.ThreadStore.md#rename) takes precedence. Undefined until renamed. + +*** + ### updatedAt ```ts diff --git a/docs/docs/api/appkit/Interface.ThreadStore.md b/docs/docs/api/appkit/Interface.ThreadStore.md index 215b76a2c..6e67364f1 100644 --- a/docs/docs/api/appkit/Interface.ThreadStore.md +++ b/docs/docs/api/appkit/Interface.ThreadStore.md @@ -25,6 +25,21 @@ message: Message): Promise; *** +### close()? + +```ts +optional close(): Promise; +``` + +Optional teardown — e.g. close an owned connection pool. Called during +agents-plugin shutdown. In-memory stores omit it. + +#### Returns + +`Promise`\<`void`\> + +*** + ### create() ```ts @@ -81,6 +96,22 @@ get(threadId: string, userId: string): Promise; *** +### init()? + +```ts +optional init(): Promise; +``` + +Optional one-time initialization — e.g. verify connectivity and bootstrap +a backing schema. Called once during agents-plugin setup, so a failure +here fails boot fast. In-memory stores omit it. + +#### Returns + +`Promise`\<`void`\> + +*** + ### list() ```ts @@ -96,3 +127,52 @@ list(userId: string): Promise; #### Returns `Promise`\<[`Thread`](Interface.Thread.md)[]\> + +*** + +### listSummaries()? + +```ts +optional listSummaries(userId: string): Promise; +``` + +Optional cheap list projection for a history sidebar — summaries only, no +message bodies. When a store omits it, the agents plugin falls back to +deriving summaries from [list](#list) (correct, just heavier). + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `userId` | `string` | + +#### Returns + +`Promise`\<[`ThreadSummary`](Interface.ThreadSummary.md)[]\> + +*** + +### rename()? + +```ts +optional rename( + threadId: string, + userId: string, +title: string): Promise; +``` + +Optional rename of a thread's title (user-scoped). Returns `false` when no +matching thread exists for the user. When a store omits it, the rename +route reports the operation as unsupported. + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `threadId` | `string` | +| `userId` | `string` | +| `title` | `string` | + +#### Returns + +`Promise`\<`boolean`\> diff --git a/docs/docs/api/appkit/Interface.ThreadSummary.md b/docs/docs/api/appkit/Interface.ThreadSummary.md new file mode 100644 index 000000000..61211aab9 --- /dev/null +++ b/docs/docs/api/appkit/Interface.ThreadSummary.md @@ -0,0 +1,46 @@ +# Interface: ThreadSummary + +Lightweight thread projection for a history list — no message bodies, so a +sidebar of many threads stays cheap. `title` is already resolved (explicit +rename, else derived from the first user message; may be empty when neither +exists). Returned by [ThreadStore.listSummaries](Interface.ThreadStore.md#listsummaries). + +## Properties + +### createdAt + +```ts +createdAt: Date; +``` + +*** + +### id + +```ts +id: string; +``` + +*** + +### messageCount + +```ts +messageCount: number; +``` + +*** + +### title + +```ts +title: string; +``` + +*** + +### updatedAt + +```ts +updatedAt: Date; +``` diff --git a/docs/docs/api/appkit/index.md b/docs/docs/api/appkit/index.md index 42afc4a50..02dcac4ee 100644 --- a/docs/docs/api/appkit/index.md +++ b/docs/docs/api/appkit/index.md @@ -22,6 +22,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [DatabricksAdapter](Class.DatabricksAdapter.md) | Adapter that talks directly to Databricks Model Serving `/invocations` endpoint. | | [ExecutionError](Class.ExecutionError.md) | Error thrown when an operation execution fails. Use for statement failures, canceled operations, or unexpected states. | | [InitializationError](Class.InitializationError.md) | Error thrown when a service or component is not properly initialized. Use when accessing services before they are ready. | +| [LakebaseThreadStore](Class.LakebaseThreadStore.md) | Persistent [ThreadStore](Interface.ThreadStore.md) backed by Databricks Lakebase (Postgres). | | [MlflowClient](Class.MlflowClient.md) | A thin client over the Databricks workspace REST API, owning the host + bearer token so callers (eval-run creation, assessment writes, the judge's serving endpoint) don't each re-derive URLs or re-attach auth. The host is normalized once at construction. | | [Plugin](Class.Plugin.md) | Base abstract class for creating AppKit plugins. | | [PolicyDeniedError](Class.PolicyDeniedError.md) | Thrown when a policy denies an action. | @@ -78,6 +79,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [LakebasePool](Interface.LakebasePool.md) | Subset of `pg.Pool` exposed by the Lakebase plugin. | | [LakebasePoolConfig](Interface.LakebasePoolConfig.md) | Configuration for creating a Lakebase connection pool | | [LakebasePoolManager](Interface.LakebasePoolManager.md) | Manages multiple Lakebase connection pools keyed by an identifier (e.g. userId). | +| [LakebaseThreadStoreOptions](Interface.LakebaseThreadStoreOptions.md) | - | | [MatchResult](Interface.MatchResult.md) | Result of a deterministic matcher run against a value. | | [McpConnectAllResult](Interface.McpConnectAllResult.md) | Per-endpoint outcome of [AppKitMcpClient.connectAll](Class.AppKitMcpClient.md#connectall). Callers (the agents plugin in particular) use the split to warn at startup when some MCP servers are unreachable without aborting boot for the rest. | | [Message](Interface.Message.md) | - | @@ -110,6 +112,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [TestContext](Interface.TestContext.md) | The `t` context passed to an eval's `test` function. | | [Thread](Interface.Thread.md) | - | | [ThreadStore](Interface.ThreadStore.md) | - | +| [ThreadSummary](Interface.ThreadSummary.md) | Lightweight thread projection for a history list — no message bodies, so a sidebar of many threads stays cheap. `title` is already resolved (explicit rename, else derived from the first user message; may be empty when neither exists). Returned by [ThreadStore.listSummaries](Interface.ThreadStore.md#listsummaries). | | [ToolAnnotations](Interface.ToolAnnotations.md) | - | | [ToolConfig](Interface.ToolConfig.md) | - | | [ToolEntry](Interface.ToolEntry.md) | Single-tool entry for a plugin's internal tool registry. | diff --git a/docs/docs/api/appkit/typedoc-sidebar.ts b/docs/docs/api/appkit/typedoc-sidebar.ts index 6dd11de11..1d2cb23a9 100644 --- a/docs/docs/api/appkit/typedoc-sidebar.ts +++ b/docs/docs/api/appkit/typedoc-sidebar.ts @@ -61,6 +61,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Class.InitializationError", label: "InitializationError" }, + { + type: "doc", + id: "api/appkit/Class.LakebaseThreadStore", + label: "LakebaseThreadStore" + }, { type: "doc", id: "api/appkit/Class.MlflowClient", @@ -322,6 +327,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Interface.LakebasePoolManager", label: "LakebasePoolManager" }, + { + type: "doc", + id: "api/appkit/Interface.LakebaseThreadStoreOptions", + label: "LakebaseThreadStoreOptions" + }, { type: "doc", id: "api/appkit/Interface.MatchResult", @@ -482,6 +492,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Interface.ThreadStore", label: "ThreadStore" }, + { + type: "doc", + id: "api/appkit/Interface.ThreadSummary", + label: "ThreadSummary" + }, { type: "doc", id: "api/appkit/Interface.ToolAnnotations", diff --git a/docs/docs/plugins/agents.md b/docs/docs/plugins/agents.md index 19b0f6da2..37b36c0e0 100644 --- a/docs/docs/plugins/agents.md +++ b/docs/docs/plugins/agents.md @@ -465,6 +465,69 @@ Supervisor and chat-completions adapters can both appear in the same `agents({ a Some hosted tool kinds return their final assistant text without incremental `output_text.delta` events. The adapter has a recovery path that pulls the text out of `response.completed.output[]` so the turn is not silently empty. Set `DEBUG=appkit:agents:supervisor-api` to log the per-turn event-type histogram if you want to verify which path a turn took. ::: +## Thread persistence + +Conversation threads are held by a `ThreadStore`. The default is +`InMemoryThreadStore` — fine for local dev and single-process demos, but it +**loses every thread on restart** and grows without bound. For any real +deployment, pass a persistent store. + +AppKit ships `LakebaseThreadStore`, backed by Databricks Lakebase (Postgres): + +```ts +import { agents, LakebaseThreadStore } from "@databricks/appkit/beta"; + +agents({ + threadStore: new LakebaseThreadStore(), +}); +``` + +With no arguments it creates its own connection pool via `createLakebasePool()` +(OAuth token refresh handled internally) and closes it on plugin shutdown. It +**self-bootstraps** its schema on setup — two tables, `agent_threads` and +`agent_messages` (FK `ON DELETE CASCADE`) — with `CREATE TABLE IF NOT EXISTS`, +so a fresh Lakebase database needs no migration step. If setup can't reach the +database, the app fails boot fast rather than silently degrading. + +**Per-user isolation.** The app service principal owns the tables, and **every** +query filters `WHERE user_id = $` — a user can never read or mutate another +user's threads. This is the security boundary; there is no cross-user read path. + +**Deploying with Lakebase.** The agents manifest declares an **optional** +`postgres` resource. Bind it at deploy time so the app's Lakebase host, +database, and endpoint are injected as `PGHOST` / `PGDATABASE` / +`LAKEBASE_ENDPOINT`. Apps that don't bind it keep the in-memory default. The +service principal needs `CAN_CONNECT_AND_CREATE` (the store issues `CREATE +TABLE`). + +**Options.** + +```ts +new LakebaseThreadStore({ + pool, // reuse an existing pg.Pool (not closed on shutdown — you own it) + tableSchema, // optional Postgres schema to hold the tables (default: search_path) +}); +``` + +**Custom stores.** Any object implementing the `ThreadStore` contract works. +The two lifecycle hooks are optional — implement them when your backing store +needs setup or teardown: + +```ts +interface ThreadStore { + create(userId: string): Promise; + get(threadId: string, userId: string): Promise; + list(userId: string): Promise; + addMessage(threadId: string, userId: string, message: Message): Promise; + delete(threadId: string, userId: string): Promise; + init?(): Promise; // called once in setup (fail-fast connectivity/bootstrap) + close?(): Promise; // called in shutdown (release resources) +} +``` + +For the exact exported symbols, run `npx @databricks/appkit docs` and open the +`appkit` API reference. + ## Configuration reference ```ts diff --git a/packages/appkit/src/beta.ts b/packages/appkit/src/beta.ts index b7c457cae..516449a42 100644 --- a/packages/appkit/src/beta.ts +++ b/packages/appkit/src/beta.ts @@ -15,6 +15,7 @@ export type { Message, Thread, ThreadStore, + ThreadSummary, ToolAnnotations, ToolProvider, } from "shared"; @@ -105,6 +106,8 @@ export type { export { agentIdFromMarkdownPath, isToolkitEntry, + LakebaseThreadStore, + type LakebaseThreadStoreOptions, loadAgentFromFile, loadAgentsFromDir, } from "./plugins/agents"; diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index adcf2e077..002522fd0 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -13,6 +13,7 @@ import type { ResponseOutputMessage, ResponseStreamEvent, Thread, + ThreadStore, ToolProvider, } from "shared"; @@ -79,6 +80,7 @@ import { cancelRequestSchema, chatRequestSchema, invocationsRequestSchema, + renameThreadRequestSchema, } from "./schemas"; import { dispatchSkillTool, @@ -87,7 +89,7 @@ import { renderForcedSkill, resolveAgentSkills, } from "./skill-loader"; -import { InMemoryThreadStore } from "./thread-store"; +import { deriveThreadSummary, InMemoryThreadStore } from "./thread-store"; import { ToolApprovalGate } from "./tool-approval-gate"; import { dispatchToolCall, @@ -114,7 +116,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { */ private streams = new ActiveStreamTracker(); private mcpClient: AppKitMcpClient | null = null; - private threadStore; + private threadStore: ThreadStore; private approvalGate = new ToolApprovalGate(); /** Guards the `agents({ agents })` deprecation warning to once per instance. */ private agentsMapDeprecationWarned = false; @@ -201,6 +203,9 @@ export class AgentsPlugin extends Plugin implements ToolProvider { } async setup() { + // Fail boot fast if the thread store can't initialize (e.g. Lakebase + // unreachable or schema bootstrap denied). No-op for in-memory stores. + await this.threadStore.init?.(); await initAgentTracing(); // Seed mlflow's config right after TelemetryManager.start() (before the // server serves), so the first turn's request-root span is forwarded and @@ -946,6 +951,12 @@ export class AgentsPlugin extends Plugin implements ToolProvider { path: "/threads/:threadId", handler: async (req, res) => this._handleGetThread(req, res), }); + this.route(router, { + name: "renameThread", + method: "patch", + path: "/threads/:threadId", + handler: async (req, res) => this._handleRenameThread(req, res), + }); this.route(router, { name: "deleteThread", method: "delete", @@ -1686,10 +1697,45 @@ export class AgentsPlugin extends Plugin implements ToolProvider { res: express.Response, ) { const userId = this.resolveUserId(req); - const threads = await this.threadStore.list(userId); + // Prefer the store's cheap summary projection; fall back to deriving from + // full threads for custom stores that don't implement listSummaries. + const threads = this.threadStore.listSummaries + ? await this.threadStore.listSummaries(userId) + : (await this.threadStore.list(userId)).map(deriveThreadSummary); res.json({ threads }); } + private async _handleRenameThread( + req: express.Request, + res: express.Response, + ) { + if (!this.threadStore.rename) { + res + .status(501) + .json({ error: "This thread store does not support renaming" }); + return; + } + const parsed = renameThreadRequestSchema.safeParse(req.body); + if (!parsed.success) { + res.status(400).json({ + error: "Invalid request", + details: parsed.error.flatten().fieldErrors, + }); + return; + } + const userId = this.resolveUserId(req); + const renamed = await this.threadStore.rename( + req.params.threadId, + userId, + parsed.data.title, + ); + if (!renamed) { + res.status(404).json({ error: "Thread not found" }); + return; + } + res.json({ renamed: true, title: parsed.data.title }); + } + private async _handleGetThread(req: express.Request, res: express.Response) { const userId = this.resolveUserId(req); const thread = await this.threadStore.get(req.params.threadId, userId); @@ -1744,6 +1790,8 @@ export class AgentsPlugin extends Plugin implements ToolProvider { await this.mcpClient.close(); this.mcpClient = null; } + // Release any pool the thread store owns. No-op for in-memory stores. + await this.threadStore.close?.(); } exports() { diff --git a/packages/appkit/src/plugins/agents/index.ts b/packages/appkit/src/plugins/agents/index.ts index 869333afb..bcb35c038 100644 --- a/packages/appkit/src/plugins/agents/index.ts +++ b/packages/appkit/src/plugins/agents/index.ts @@ -28,3 +28,7 @@ export { type ToolkitOptions, } from "../../core/agent/types"; export { AgentsPlugin, agents } from "./agents"; +export { + LakebaseThreadStore, + type LakebaseThreadStoreOptions, +} from "./lakebase-thread-store"; diff --git a/packages/appkit/src/plugins/agents/lakebase-thread-store.ts b/packages/appkit/src/plugins/agents/lakebase-thread-store.ts new file mode 100644 index 000000000..4f87f0187 --- /dev/null +++ b/packages/appkit/src/plugins/agents/lakebase-thread-store.ts @@ -0,0 +1,324 @@ +import { randomUUID } from "node:crypto"; + +import type { Pool } from "pg"; +import type { + Message, + Thread, + ThreadStore, + ThreadSummary, + ToolCall, +} from "shared"; + +import { createLakebasePool } from "../../connectors/lakebase"; + +/** Postgres identifier: lowercase, digits, underscore, not leading with a digit. */ +const IDENTIFIER = /^[a-z_][a-z0-9_]*$/; + +export interface LakebaseThreadStoreOptions { + /** + * An existing `pg.Pool` to run on. When omitted, the store creates its own + * pool via `createLakebasePool()` (OAuth token refresh handled inside) and + * closes it on {@link LakebaseThreadStore.close}. An injected pool is never + * closed — the caller owns its lifecycle. + */ + pool?: Pool; + /** + * Optional Postgres schema to hold the two tables. Created on init if it + * does not exist. Defaults to the connection's search_path (usually + * `public`). Validated as a plain lowercase identifier — it is interpolated + * into DDL, not parameterizable, so anything else is rejected. + */ + tableSchema?: string; +} + +interface ThreadRow { + id: string; + user_id: string; + title: string | null; + created_at: string | Date; + updated_at: string | Date; +} + +interface SummaryRow { + id: string; + title: string | null; + message_count: number; + created_at: string | Date; + updated_at: string | Date; +} + +interface MessageRow { + thread_id: string; + id: string; + role: string; + content: string; + tool_call_id: string | null; + tool_calls: ToolCall[] | null; + created_at: string | Date; +} + +/** + * Persistent {@link ThreadStore} backed by Databricks Lakebase (Postgres). + * + * Threads and messages live in two `user_id`-scoped tables (`agent_threads`, + * `agent_messages`, FK cascade). The app service principal owns the tables; + * **every** query filters `WHERE user_id = $` — that is the isolation + * boundary, so a user can never read or mutate another user's threads. + * + * The schema is self-bootstrapping: {@link init} issues idempotent + * `CREATE TABLE IF NOT EXISTS` (once-guarded) and verifies connectivity, so + * a fresh Lakebase database works with no migration step. + * + * Pass it to the agents plugin for a deployment that survives restarts: + * ```ts + * agents({ threadStore: new LakebaseThreadStore() }) + * ``` + */ +export class LakebaseThreadStore implements ThreadStore { + private readonly pool: Pool; + private readonly ownsPool: boolean; + private readonly threads: string; + private readonly messages: string; + private readonly schema?: string; + private initPromise: Promise | null = null; + + constructor({ pool, tableSchema }: LakebaseThreadStoreOptions = {}) { + this.ownsPool = !pool; + this.pool = pool ?? createLakebasePool(); + if (tableSchema !== undefined && !IDENTIFIER.test(tableSchema)) { + throw new Error( + `LakebaseThreadStore: invalid tableSchema "${tableSchema}" (expected a lowercase identifier)`, + ); + } + this.schema = tableSchema; + const prefix = tableSchema ? `${tableSchema}.` : ""; + this.threads = `${prefix}agent_threads`; + this.messages = `${prefix}agent_messages`; + } + + /** Verify connectivity and create the tables once (idempotent). */ + init(): Promise { + if (!this.initPromise) this.initPromise = this.bootstrap(); + return this.initPromise; + } + + /** Close the pool only when this store created it. */ + async close(): Promise { + if (this.ownsPool) await this.pool.end(); + } + + private async bootstrap(): Promise { + // Fail fast on an unauthenticated/unreachable pool before the DDL, mirroring + // the DatabasePlugin readiness probe. + await this.pool.query("select 1"); + if (this.schema) { + await this.pool.query(`CREATE SCHEMA IF NOT EXISTS ${this.schema}`); + } + await this.pool.query(` + CREATE TABLE IF NOT EXISTS ${this.threads} ( + id uuid PRIMARY KEY, + user_id text NOT NULL, + title text, + created_at timestamptz NOT NULL DEFAULT now(), + updated_at timestamptz NOT NULL DEFAULT now() + ) + `); + // Bring already-created tables (from an earlier version) up to schema. + await this.pool.query( + `ALTER TABLE ${this.threads} ADD COLUMN IF NOT EXISTS title text`, + ); + await this.pool.query( + `CREATE INDEX IF NOT EXISTS agent_threads_user_updated_idx ON ${this.threads} (user_id, updated_at DESC)`, + ); + await this.pool.query(` + CREATE TABLE IF NOT EXISTS ${this.messages} ( + seq bigserial PRIMARY KEY, + id text NOT NULL, + thread_id uuid NOT NULL REFERENCES ${this.threads}(id) ON DELETE CASCADE, + user_id text NOT NULL, + role text NOT NULL, + content text NOT NULL, + tool_call_id text, + tool_calls jsonb, + created_at timestamptz NOT NULL DEFAULT now() + ) + `); + await this.pool.query( + `CREATE INDEX IF NOT EXISTS agent_messages_thread_seq_idx ON ${this.messages} (thread_id, seq)`, + ); + } + + async create(userId: string): Promise { + const id = randomUUID(); + const { rows } = await this.pool.query( + `INSERT INTO ${this.threads} (id, user_id) + VALUES ($1, $2) + RETURNING id, user_id, title, created_at, updated_at`, + [id, userId], + ); + return this.toThread(rows[0], []); + } + + async get(threadId: string, userId: string): Promise { + const threadResult = await this.pool.query( + `SELECT id, user_id, title, created_at, updated_at + FROM ${this.threads} + WHERE id = $1 AND user_id = $2`, + [threadId, userId], + ); + const row = threadResult.rows[0]; + if (!row) return null; + + const messageResult = await this.pool.query( + `SELECT thread_id, id, role, content, tool_call_id, tool_calls, created_at + FROM ${this.messages} + WHERE thread_id = $1 AND user_id = $2 + ORDER BY seq`, + [threadId, userId], + ); + return this.toThread(row, messageResult.rows.map(toMessage)); + } + + async list(userId: string): Promise { + const threadResult = await this.pool.query( + `SELECT id, user_id, title, created_at, updated_at + FROM ${this.threads} + WHERE user_id = $1 + ORDER BY updated_at DESC`, + [userId], + ); + if (threadResult.rows.length === 0) return []; + + // One query for every message this user owns, grouped in-app by thread — + // avoids an N+1 across their threads. + const messageResult = await this.pool.query( + `SELECT thread_id, id, role, content, tool_call_id, tool_calls, created_at + FROM ${this.messages} + WHERE user_id = $1 + ORDER BY thread_id, seq`, + [userId], + ); + const byThread = new Map(); + for (const row of messageResult.rows) { + const list = byThread.get(row.thread_id) ?? []; + list.push(toMessage(row)); + byThread.set(row.thread_id, list); + } + return threadResult.rows.map((row) => + this.toThread(row, byThread.get(row.id) ?? []), + ); + } + + async addMessage( + threadId: string, + userId: string, + message: Message, + ): Promise { + // Single atomic statement: bump the thread's updated_at only if it exists + // for this user, then insert the message off that CTE. When the thread row + // is absent (unknown id, or owned by another user) `updated` is empty, the + // INSERT ... SELECT writes zero rows, and we throw — matching InMemory. + const result = await this.pool.query( + `WITH updated AS ( + UPDATE ${this.threads} SET updated_at = now() + WHERE id = $1 AND user_id = $2 + RETURNING id + ) + INSERT INTO ${this.messages} + (id, thread_id, user_id, role, content, tool_call_id, tool_calls) + SELECT $3::text, $1::uuid, $2::text, $4::text, $5::text, $6::text, $7::jsonb + FROM updated`, + [ + threadId, + userId, + message.id, + message.role, + message.content, + message.toolCallId ?? null, + message.toolCalls ? JSON.stringify(message.toolCalls) : null, + ], + ); + if ((result.rowCount ?? 0) === 0) { + throw new Error(`Thread ${threadId} not found`); + } + } + + async delete(threadId: string, userId: string): Promise { + // Messages cascade via the FK. + const result = await this.pool.query( + `DELETE FROM ${this.threads} WHERE id = $1 AND user_id = $2`, + [threadId, userId], + ); + return (result.rowCount ?? 0) > 0; + } + + async listSummaries(userId: string): Promise { + // Summary-only projection — no message bodies. The title resolves to the + // explicit rename, else the first user message truncated, else empty; the + // count is a scalar subquery. A LATERAL join fetches just the first user + // message per thread rather than all messages. + const { rows } = await this.pool.query( + `SELECT t.id, + COALESCE(t.title, left(fm.content, 80), '') AS title, + (SELECT count(*)::int FROM ${this.messages} m + WHERE m.thread_id = t.id) AS message_count, + t.created_at, t.updated_at + FROM ${this.threads} t + LEFT JOIN LATERAL ( + SELECT content FROM ${this.messages} m + WHERE m.thread_id = t.id AND m.role = 'user' + ORDER BY seq LIMIT 1 + ) fm ON true + WHERE t.user_id = $1 + ORDER BY t.updated_at DESC`, + [userId], + ); + return rows.map((row) => ({ + id: row.id, + title: row.title ?? "", + messageCount: row.message_count, + createdAt: new Date(row.created_at), + updatedAt: new Date(row.updated_at), + })); + } + + async rename( + threadId: string, + userId: string, + title: string, + ): Promise { + // Does NOT bump updated_at — recency ordering reflects activity, not renames. + const result = await this.pool.query( + `UPDATE ${this.threads} SET title = $3 WHERE id = $1 AND user_id = $2`, + [threadId, userId, title], + ); + return (result.rowCount ?? 0) > 0; + } + + private toThread(row: ThreadRow, messages: Message[]): Thread { + const thread: Thread = { + id: row.id, + userId: row.user_id, + messages, + createdAt: new Date(row.created_at), + updatedAt: new Date(row.updated_at), + }; + if (row.title != null) thread.title = row.title; + return thread; + } +} + +/** Reconstruct a `Message`, reviving its Date and passing tool_calls verbatim. */ +function toMessage(row: MessageRow): Message { + const message: Message = { + id: row.id, + role: row.role as Message["role"], + content: row.content, + createdAt: new Date(row.created_at), + }; + if (row.tool_call_id != null) message.toolCallId = row.tool_call_id; + // jsonb is parsed to a JS value by node-pg, so `thoughtSignature` (and every + // other ToolCall field) survives the round trip unchanged. + if (row.tool_calls != null) message.toolCalls = row.tool_calls; + return message; +} diff --git a/packages/appkit/src/plugins/agents/manifest.json b/packages/appkit/src/plugins/agents/manifest.json index 529c640bc..96ce13e98 100644 --- a/packages/appkit/src/plugins/agents/manifest.json +++ b/packages/appkit/src/plugins/agents/manifest.json @@ -50,6 +50,76 @@ } } } + }, + { + "type": "postgres", + "alias": "Agent thread storage", + "resourceKey": "postgres", + "description": "Optional Lakebase Postgres for persistent agent threads (LakebaseThreadStore). Unbound, agents fall back to the in-memory thread store.", + "permission": "CAN_CONNECT_AND_CREATE", + "fields": { + "project": { + "description": "Lakebase project resource name", + "examples": ["projects/{project-id}"], + "discovery": { + "type": "kind", + "resourceKind": "postgres_project", + "select": "name" + } + }, + "branch": { + "description": "Lakebase branch resource name", + "examples": ["projects/{project-id}/branches/{branch-id}"], + "discovery": { + "type": "kind", + "resourceKind": "postgres_branch", + "select": "name", + "dependsOn": "project" + } + }, + "database": { + "description": "Lakebase database resource name", + "examples": [ + "projects/{project-id}/branches/{branch-id}/databases/{database-id}" + ], + "discovery": { + "type": "kind", + "resourceKind": "postgres_database", + "select": "name", + "dependsOn": "branch" + } + }, + "host": { + "env": "PGHOST", + "localOnly": true, + "resolve": "postgres:host", + "description": "Postgres host" + }, + "databaseName": { + "env": "PGDATABASE", + "localOnly": true, + "resolve": "postgres:databaseName", + "description": "Postgres database name" + }, + "endpointPath": { + "env": "LAKEBASE_ENDPOINT", + "bundleIgnore": true, + "resolve": "postgres:endpointPath", + "description": "Lakebase endpoint resource name" + }, + "port": { + "env": "PGPORT", + "localOnly": true, + "value": "5432", + "description": "Postgres port" + }, + "sslmode": { + "env": "PGSSLMODE", + "localOnly": true, + "value": "require", + "description": "Postgres SSL mode" + } + } } ] } diff --git a/packages/appkit/src/plugins/agents/schemas.ts b/packages/appkit/src/plugins/agents/schemas.ts index 24e38f1d5..a691ea917 100644 --- a/packages/appkit/src/plugins/agents/schemas.ts +++ b/packages/appkit/src/plugins/agents/schemas.ts @@ -79,6 +79,20 @@ export const invocationsRequestSchema = z.object({ mlflowRunId: z.string().max(64).optional(), }); +/** Max characters for a thread title on `PATCH /threads/:id`. */ +const MAX_TITLE_CHARS = 200; + +export const renameThreadRequestSchema = z.object({ + title: z + .string() + .trim() + .min(1, "title must not be empty") + .max( + MAX_TITLE_CHARS, + `title exceeds the ${MAX_TITLE_CHARS}-character limit`, + ), +}); + export const approvalRequestSchema = z.object({ streamId: z.string().min(1, "streamId is required"), approvalId: z.string().min(1, "approvalId is required"), diff --git a/packages/appkit/src/plugins/agents/tests/lakebase-thread-store.test.ts b/packages/appkit/src/plugins/agents/tests/lakebase-thread-store.test.ts new file mode 100644 index 000000000..5dd2bc144 --- /dev/null +++ b/packages/appkit/src/plugins/agents/tests/lakebase-thread-store.test.ts @@ -0,0 +1,438 @@ +import type { Pool, QueryResult } from "pg"; +import type { Message, ToolCall } from "shared"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ createLakebasePool: vi.fn() })); + +vi.mock("../../../connectors/lakebase", () => ({ + createLakebasePool: mocks.createLakebasePool, +})); + +import { LakebaseThreadStore } from "../lakebase-thread-store"; + +const USER = "user-a"; +const OTHER = "user-b"; + +/** Build a full-shape QueryResult so mock returns satisfy pg's type. */ +function qr( + rows: Record[] = [], + rowCount = rows.length, +): QueryResult { + return { rows, rowCount, command: "", oid: 0, fields: [] } as QueryResult; +} + +/** A pool whose `query` returns empty by default; program per-call as needed. */ +function makePool() { + const query = vi.fn(async () => qr()); + const end = vi.fn(async () => undefined); + const pool = { query, end } as unknown as Pool; + return { pool, query, end }; +} + +/** Collapse whitespace so SQL assertions aren't formatting-coupled. */ +const sql = (call: unknown[]) => String(call[0]).replace(/\s+/g, " ").trim(); +const params = (call: unknown[]) => call[1] as unknown[]; + +beforeEach(() => { + mocks.createLakebasePool.mockReset(); +}); +afterEach(() => vi.restoreAllMocks()); + +describe("LakebaseThreadStore construction & pool ownership", () => { + test("creates its own pool when none is injected", () => { + const { pool } = makePool(); + mocks.createLakebasePool.mockReturnValue(pool); + new LakebaseThreadStore(); + expect(mocks.createLakebasePool).toHaveBeenCalledTimes(1); + }); + + test("uses an injected pool and never calls the factory", () => { + const { pool } = makePool(); + new LakebaseThreadStore({ pool }); + expect(mocks.createLakebasePool).not.toHaveBeenCalled(); + }); + + test("close() ends an owned pool", async () => { + const { pool, end } = makePool(); + mocks.createLakebasePool.mockReturnValue(pool); + await new LakebaseThreadStore().close(); + expect(end).toHaveBeenCalledTimes(1); + }); + + test("close() does NOT end an injected pool", async () => { + const { pool, end } = makePool(); + await new LakebaseThreadStore({ pool }).close(); + expect(end).not.toHaveBeenCalled(); + }); + + test("rejects a non-identifier tableSchema before touching the pool", () => { + const { pool } = makePool(); + expect( + () => new LakebaseThreadStore({ pool, tableSchema: "a; drop" }), + ).toThrow(/invalid tableSchema/); + }); +}); + +describe("LakebaseThreadStore.init (bootstrap)", () => { + test("verifies connectivity and creates both tables + indexes, once", async () => { + const { pool, query } = makePool(); + const store = new LakebaseThreadStore({ pool }); + + await store.init(); + const statements = query.mock.calls.map(sql); + expect(statements[0]).toBe("select 1"); + expect( + statements.some((s) => + s.includes("CREATE TABLE IF NOT EXISTS agent_threads"), + ), + ).toBe(true); + expect( + statements.some((s) => + s.includes("CREATE TABLE IF NOT EXISTS agent_messages"), + ), + ).toBe(true); + expect(statements.some((s) => s.includes("ON DELETE CASCADE"))).toBe(true); + // title column present in the fresh CREATE and back-filled via ALTER. + expect( + statements.some((s) => + /CREATE TABLE IF NOT EXISTS agent_threads .*title text/.test(s), + ), + ).toBe(true); + expect( + statements.some((s) => + s.includes("ALTER TABLE agent_threads ADD COLUMN IF NOT EXISTS title"), + ), + ).toBe(true); + expect( + statements.some((s) => + s.includes("CREATE INDEX IF NOT EXISTS agent_threads_user_updated_idx"), + ), + ).toBe(true); + expect( + statements.some((s) => + s.includes("CREATE INDEX IF NOT EXISTS agent_messages_thread_seq_idx"), + ), + ).toBe(true); + + const afterFirst = query.mock.calls.length; + await store.init(); // once-guarded: no new DDL + expect(query.mock.calls.length).toBe(afterFirst); + }); + + test("qualifies tables with a valid tableSchema and creates the schema", async () => { + const { pool, query } = makePool(); + await new LakebaseThreadStore({ pool, tableSchema: "appkit" }).init(); + const statements = query.mock.calls.map(sql); + expect( + statements.some((s) => s.includes("CREATE SCHEMA IF NOT EXISTS appkit")), + ).toBe(true); + expect( + statements.some((s) => + s.includes("CREATE TABLE IF NOT EXISTS appkit.agent_threads"), + ), + ).toBe(true); + expect(statements.some((s) => s.includes("appkit.agent_messages"))).toBe( + true, + ); + }); +}); + +describe("LakebaseThreadStore.create", () => { + test("inserts a thread scoped to the user and revives dates", async () => { + const { pool, query } = makePool(); + const created = "2026-01-02T03:04:05.000Z"; + query.mockResolvedValueOnce( + qr([ + { id: "t1", user_id: USER, created_at: created, updated_at: created }, + ]), + ); + + const thread = await new LakebaseThreadStore({ pool }).create(USER); + + expect(sql(query.mock.calls[0])).toContain("INSERT INTO agent_threads"); + // A generated uuid, then the user id. + expect(params(query.mock.calls[0])[1]).toBe(USER); + expect(thread.messages).toEqual([]); + expect(thread.createdAt).toBeInstanceOf(Date); + expect(thread.createdAt.toISOString()).toBe(created); + expect(thread.userId).toBe(USER); + }); +}); + +describe("LakebaseThreadStore.get", () => { + test("returns null and skips the message query when the thread is absent", async () => { + const { pool, query } = makePool(); + query.mockResolvedValueOnce(qr([])); + + const result = await new LakebaseThreadStore({ pool }).get("missing", USER); + + expect(result).toBeNull(); + expect(query).toHaveBeenCalledTimes(1); // no second (messages) query + }); + + test("scopes both queries to the user, orders messages, and preserves tool_calls verbatim", async () => { + const { pool, query } = makePool(); + const ts = "2026-02-02T00:00:00.000Z"; + const toolCalls: ToolCall[] = [ + { id: "c1", name: "search", args: { q: "x" }, thoughtSignature: "SIG==" }, + ]; + query + .mockResolvedValueOnce( + qr([{ id: "t1", user_id: USER, created_at: ts, updated_at: ts }]), + ) + .mockResolvedValueOnce( + qr([ + { + thread_id: "t1", + id: "m1", + role: "assistant", + content: "hi", + tool_call_id: "c1", + tool_calls: toolCalls, + created_at: ts, + }, + ]), + ); + + const thread = await new LakebaseThreadStore({ pool }).get("t1", USER); + + // Both queries carry the user_id as the isolation boundary. + expect(sql(query.mock.calls[0])).toContain( + "WHERE id = $1 AND user_id = $2", + ); + expect(params(query.mock.calls[0])).toEqual(["t1", USER]); + expect(sql(query.mock.calls[1])).toContain( + "WHERE thread_id = $1 AND user_id = $2", + ); + expect(sql(query.mock.calls[1])).toContain("ORDER BY seq"); + expect(params(query.mock.calls[1])).toEqual(["t1", USER]); + + expect(thread).not.toBeNull(); + const msg = thread?.messages[0]; + expect(msg?.createdAt).toBeInstanceOf(Date); + expect(msg?.toolCallId).toBe("c1"); + // thoughtSignature survives the jsonb round trip untouched. + expect(msg?.toolCalls).toEqual(toolCalls); + expect(msg?.toolCalls?.[0].thoughtSignature).toBe("SIG=="); + }); +}); + +describe("LakebaseThreadStore.list", () => { + test("returns [] without a message query when the user has no threads", async () => { + const { pool, query } = makePool(); + query.mockResolvedValueOnce(qr([])); + + const result = await new LakebaseThreadStore({ pool }).list(USER); + + expect(result).toEqual([]); + expect(query).toHaveBeenCalledTimes(1); + }); + + test("groups messages by thread in one query and leaves empty threads empty", async () => { + const { pool, query } = makePool(); + const ts = "2026-03-03T00:00:00.000Z"; + query + .mockResolvedValueOnce( + qr([ + { id: "t1", user_id: USER, created_at: ts, updated_at: ts }, + { id: "t2", user_id: USER, created_at: ts, updated_at: ts }, + ]), + ) + .mockResolvedValueOnce( + qr([ + { + thread_id: "t1", + id: "m1", + role: "user", + content: "a", + tool_call_id: null, + tool_calls: null, + created_at: ts, + }, + { + thread_id: "t1", + id: "m2", + role: "assistant", + content: "b", + tool_call_id: null, + tool_calls: null, + created_at: ts, + }, + ]), + ); + + const threads = await new LakebaseThreadStore({ pool }).list(USER); + + expect(sql(query.mock.calls[0])).toContain("WHERE user_id = $1"); + expect(sql(query.mock.calls[0])).toContain("ORDER BY updated_at DESC"); + expect(params(query.mock.calls[0])).toEqual([USER]); + expect(params(query.mock.calls[1])).toEqual([USER]); + expect(query).toHaveBeenCalledTimes(2); // no N+1 + expect(threads[0].messages.map((m) => m.id)).toEqual(["m1", "m2"]); + expect(threads[1].messages).toEqual([]); + }); +}); + +describe("LakebaseThreadStore.addMessage", () => { + const message: Message = { + id: "m1", + role: "user", + content: "hello", + createdAt: new Date(), + }; + + test("writes the message + bumps updated_at in one user-scoped statement", async () => { + const { pool, query } = makePool(); + query.mockResolvedValueOnce(qr([], 1)); + + await new LakebaseThreadStore({ pool }).addMessage("t1", USER, message); + + const call = query.mock.calls[0]; + expect(sql(call)).toContain("UPDATE agent_threads SET updated_at = now()"); + expect(sql(call)).toContain("WHERE id = $1 AND user_id = $2"); + expect(sql(call)).toContain("INSERT INTO agent_messages"); + const p = params(call); + expect(p[0]).toBe("t1"); + expect(p[1]).toBe(USER); + expect(p[5]).toBeNull(); // tool_call_id absent + expect(p[6]).toBeNull(); // tool_calls absent + }); + + test("serializes tool_calls to JSON for the jsonb column", async () => { + const { pool, query } = makePool(); + query.mockResolvedValueOnce(qr([], 1)); + const toolCalls: ToolCall[] = [ + { id: "c1", name: "run", args: {}, thoughtSignature: "Zm9v" }, + ]; + + await new LakebaseThreadStore({ pool }).addMessage("t1", USER, { + ...message, + toolCallId: "c1", + toolCalls, + }); + + const p = params(query.mock.calls[0]); + expect(p[5]).toBe("c1"); + expect(p[6]).toBe(JSON.stringify(toolCalls)); + }); + + test("throws when the thread does not exist for the user (0 rows written)", async () => { + const { pool, query } = makePool(); + query.mockResolvedValueOnce(qr([], 0)); + + await expect( + new LakebaseThreadStore({ pool }).addMessage("nope", OTHER, message), + ).rejects.toThrow("Thread nope not found"); + }); +}); + +describe("LakebaseThreadStore.delete", () => { + test("deletes scoped to the user and reports whether a row was removed", async () => { + const { pool, query } = makePool(); + query.mockResolvedValueOnce(qr([], 1)).mockResolvedValueOnce(qr([], 0)); + const store = new LakebaseThreadStore({ pool }); + + expect(await store.delete("t1", USER)).toBe(true); + expect(await store.delete("t1", OTHER)).toBe(false); + expect(sql(query.mock.calls[0])).toContain( + "DELETE FROM agent_threads WHERE id = $1 AND user_id = $2", + ); + expect(params(query.mock.calls[0])).toEqual(["t1", USER]); + }); +}); + +describe("LakebaseThreadStore.listSummaries", () => { + test("projects summaries user-scoped, derives title, counts, orders by recency", async () => { + const { pool, query } = makePool(); + const ts = "2026-04-04T00:00:00.000Z"; + query.mockResolvedValueOnce( + qr([ + { + id: "t1", + title: "Weather in Paris", + message_count: 3, + created_at: ts, + updated_at: ts, + }, + { + id: "t2", + title: "", + message_count: 0, + created_at: ts, + updated_at: ts, + }, + ]), + ); + + const summaries = await new LakebaseThreadStore({ pool }).listSummaries( + USER, + ); + + const s = sql(query.mock.calls[0]); + expect(s).toContain("COALESCE(t.title, left(fm.content, 80)"); + expect(s).toContain("count(*)::int"); + expect(s).toContain("LEFT JOIN LATERAL"); + expect(s).toContain("WHERE t.user_id = $1"); + expect(s).toContain("ORDER BY t.updated_at DESC"); + expect(params(query.mock.calls[0])).toEqual([USER]); + + expect(summaries[0]).toMatchObject({ + id: "t1", + title: "Weather in Paris", + messageCount: 3, + }); + expect(summaries[0].updatedAt).toBeInstanceOf(Date); + expect(summaries[0].createdAt).toBeInstanceOf(Date); + expect(summaries[1].title).toBe(""); + }); +}); + +describe("LakebaseThreadStore.rename", () => { + test("updates title user-scoped without bumping updated_at; reports found", async () => { + const { pool, query } = makePool(); + query.mockResolvedValueOnce(qr([], 1)).mockResolvedValueOnce(qr([], 0)); + const store = new LakebaseThreadStore({ pool }); + + expect(await store.rename("t1", USER, "New title")).toBe(true); + expect(await store.rename("t1", OTHER, "New title")).toBe(false); + + const s = sql(query.mock.calls[0]); + expect(s).toContain("UPDATE agent_threads SET title = $3"); + expect(s).toContain("WHERE id = $1 AND user_id = $2"); + expect(s).not.toContain("updated_at"); // rename must not reorder by recency + expect(params(query.mock.calls[0])).toEqual(["t1", USER, "New title"]); + }); +}); + +describe("LakebaseThreadStore user_id scoping (isolation boundary)", () => { + test("every CRUD query carries the user id in its params", async () => { + const { pool, query } = makePool(); + // Enough rows for each method to exercise both queries where applicable. + query.mockResolvedValue( + qr([ + { + id: "t1", + user_id: USER, + created_at: "2026-01-01T00:00:00.000Z", + updated_at: "2026-01-01T00:00:00.000Z", + }, + ]), + ); + const store = new LakebaseThreadStore({ pool }); + + await store.create(USER); + await store.get("t1", USER); + await store.list(USER); + await store.addMessage("t1", USER, { + id: "m1", + role: "user", + content: "x", + createdAt: new Date(), + }); + await store.delete("t1", USER); + + for (const call of query.mock.calls) { + expect(params(call)).toContain(USER); + } + }); +}); diff --git a/packages/appkit/src/plugins/agents/tests/thread-store.test.ts b/packages/appkit/src/plugins/agents/tests/thread-store.test.ts index 2dda38640..702d1b124 100644 --- a/packages/appkit/src/plugins/agents/tests/thread-store.test.ts +++ b/packages/appkit/src/plugins/agents/tests/thread-store.test.ts @@ -136,4 +136,73 @@ describe("InMemoryThreadStore", () => { expect(user1Threads).toHaveLength(2); expect(user2Threads).toHaveLength(1); }); + + test("listSummaries() derives title from first user message + counts, sorted", async () => { + const store = new InMemoryThreadStore(); + const empty = await store.create("user-1"); + const chatted = await store.create("user-1"); + await store.addMessage(chatted.id, "user-1", { + id: "m1", + role: "user", + content: "What is the weather in Paris today?", + createdAt: new Date(), + }); + await store.addMessage(chatted.id, "user-1", { + id: "m2", + role: "assistant", + content: "Sunny.", + createdAt: new Date(), + }); + + const summaries = await store.listSummaries("user-1"); + expect(summaries).toHaveLength(2); + // (ordering is covered by the list() test; two same-ms creates can tie) + const chattedSummary = summaries.find((s) => s.id === chatted.id); + expect(chattedSummary?.title).toBe("What is the weather in Paris today?"); + expect(chattedSummary?.messageCount).toBe(2); + expect(chattedSummary?.updatedAt).toBeInstanceOf(Date); + // Empty thread → empty derived title, zero count. + const emptySummary = summaries.find((s) => s.id === empty.id); + expect(emptySummary?.title).toBe(""); + expect(emptySummary?.messageCount).toBe(0); + }); + + test("listSummaries() truncates a long derived title to 80 chars", async () => { + const store = new InMemoryThreadStore(); + const t = await store.create("user-1"); + await store.addMessage(t.id, "user-1", { + id: "m1", + role: "user", + content: "x".repeat(200), + createdAt: new Date(), + }); + const [summary] = await store.listSummaries("user-1"); + expect(summary.title).toHaveLength(80); + }); + + test("rename() sets the title, wins over the derived default, and does not reorder", async () => { + const store = new InMemoryThreadStore(); + const t = await store.create("user-1"); + await store.addMessage(t.id, "user-1", { + id: "m1", + role: "user", + content: "original first message", + createdAt: new Date(), + }); + const before = (await store.get(t.id, "user-1"))?.updatedAt.getTime(); + + expect(await store.rename(t.id, "user-1", "My renamed thread")).toBe(true); + + const [summary] = await store.listSummaries("user-1"); + expect(summary.title).toBe("My renamed thread"); + // Rename must not bump updatedAt. + expect((await store.get(t.id, "user-1"))?.updatedAt.getTime()).toBe(before); + }); + + test("rename() returns false for wrong user or missing thread", async () => { + const store = new InMemoryThreadStore(); + const t = await store.create("user-1"); + expect(await store.rename(t.id, "user-2", "nope")).toBe(false); + expect(await store.rename("missing", "user-1", "nope")).toBe(false); + }); }); diff --git a/packages/appkit/src/plugins/agents/thread-store.ts b/packages/appkit/src/plugins/agents/thread-store.ts index aefcdc686..123575818 100644 --- a/packages/appkit/src/plugins/agents/thread-store.ts +++ b/packages/appkit/src/plugins/agents/thread-store.ts @@ -1,6 +1,31 @@ import { randomUUID } from "node:crypto"; -import type { Message, Thread, ThreadStore } from "shared"; +import type { Message, Thread, ThreadStore, ThreadSummary } from "shared"; + +/** Longest derived title (first user message is truncated to this). */ +const DERIVED_TITLE_MAX = 80; + +/** + * Project a full {@link Thread} to a {@link ThreadSummary}: explicit `title` + * wins, else the first user message truncated, else empty. Shared by + * {@link InMemoryThreadStore.listSummaries} and the agents plugin's fallback + * for stores that don't implement `listSummaries`. + */ +export function deriveThreadSummary(thread: Thread): ThreadSummary { + const explicit = thread.title?.trim(); + const firstUser = thread.messages.find((m) => m.role === "user")?.content; + const title = + explicit && explicit.length > 0 + ? explicit + : (firstUser?.slice(0, DERIVED_TITLE_MAX) ?? ""); + return { + id: thread.id, + title, + createdAt: thread.createdAt, + updatedAt: thread.updatedAt, + messageCount: thread.messages.length, + }; +} /** * In-memory thread store backed by a nested Map. @@ -41,6 +66,23 @@ export class InMemoryThreadStore implements ThreadStore { ); } + async listSummaries(userId: string): Promise { + return (await this.list(userId)).map(deriveThreadSummary); + } + + async rename( + threadId: string, + userId: string, + title: string, + ): Promise { + const thread = this.userMap(userId).get(threadId); + if (!thread) return false; + // Deliberately does NOT touch updatedAt — recency ordering reflects + // conversation activity, not renames. + thread.title = title; + return true; + } + async addMessage( threadId: string, userId: string, diff --git a/packages/shared/src/agent.ts b/packages/shared/src/agent.ts index 811092b84..e092e8f93 100644 --- a/packages/shared/src/agent.ts +++ b/packages/shared/src/agent.ts @@ -94,6 +94,26 @@ export interface Thread { messages: Message[]; createdAt: Date; updatedAt: Date; + /** + * Optional human title. Defaults to a value derived from the first user + * message (see {@link ThreadStore.listSummaries}); an explicit rename via + * {@link ThreadStore.rename} takes precedence. Undefined until renamed. + */ + title?: string; +} + +/** + * Lightweight thread projection for a history list — no message bodies, so a + * sidebar of many threads stays cheap. `title` is already resolved (explicit + * rename, else derived from the first user message; may be empty when neither + * exists). Returned by {@link ThreadStore.listSummaries}. + */ +export interface ThreadSummary { + id: string; + title: string; + createdAt: Date; + updatedAt: Date; + messageCount: number; } // --------------------------------------------------------------------------- @@ -106,6 +126,29 @@ export interface ThreadStore { list(userId: string): Promise; addMessage(threadId: string, userId: string, message: Message): Promise; delete(threadId: string, userId: string): Promise; + /** + * Optional cheap list projection for a history sidebar — summaries only, no + * message bodies. When a store omits it, the agents plugin falls back to + * deriving summaries from {@link list} (correct, just heavier). + */ + listSummaries?(userId: string): Promise; + /** + * Optional rename of a thread's title (user-scoped). Returns `false` when no + * matching thread exists for the user. When a store omits it, the rename + * route reports the operation as unsupported. + */ + rename?(threadId: string, userId: string, title: string): Promise; + /** + * Optional one-time initialization — e.g. verify connectivity and bootstrap + * a backing schema. Called once during agents-plugin setup, so a failure + * here fails boot fast. In-memory stores omit it. + */ + init?(): Promise; + /** + * Optional teardown — e.g. close an owned connection pool. Called during + * agents-plugin shutdown. In-memory stores omit it. + */ + close?(): Promise; } // --------------------------------------------------------------------------- diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json index 998dc8b06..cad324d59 100644 --- a/template/appkit.plugins.json +++ b/template/appkit.plugins.json @@ -56,6 +56,88 @@ "origin": "user" } } + }, + { + "type": "postgres", + "alias": "Agent thread storage", + "resourceKey": "postgres", + "description": "Optional Lakebase Postgres for persistent agent threads (LakebaseThreadStore). Unbound, agents fall back to the in-memory thread store.", + "permission": "CAN_CONNECT_AND_CREATE", + "fields": { + "project": { + "description": "Lakebase project resource name", + "examples": [ + "projects/{project-id}" + ], + "discovery": { + "type": "kind", + "resourceKind": "postgres_project", + "select": "name" + }, + "origin": "user" + }, + "branch": { + "description": "Lakebase branch resource name", + "examples": [ + "projects/{project-id}/branches/{branch-id}" + ], + "discovery": { + "type": "kind", + "resourceKind": "postgres_branch", + "select": "name", + "dependsOn": "project" + }, + "origin": "user" + }, + "database": { + "description": "Lakebase database resource name", + "examples": [ + "projects/{project-id}/branches/{branch-id}/databases/{database-id}" + ], + "discovery": { + "type": "kind", + "resourceKind": "postgres_database", + "select": "name", + "dependsOn": "branch" + }, + "origin": "user" + }, + "host": { + "env": "PGHOST", + "description": "Postgres host", + "localOnly": true, + "resolve": "postgres:host", + "origin": "platform" + }, + "databaseName": { + "env": "PGDATABASE", + "description": "Postgres database name", + "localOnly": true, + "resolve": "postgres:databaseName", + "origin": "platform" + }, + "endpointPath": { + "env": "LAKEBASE_ENDPOINT", + "description": "Lakebase endpoint resource name", + "bundleIgnore": true, + "resolve": "postgres:endpointPath", + "origin": "cli" + }, + "port": { + "env": "PGPORT", + "description": "Postgres port", + "localOnly": true, + "value": "5432", + "origin": "platform" + }, + "sslmode": { + "env": "PGSSLMODE", + "description": "Postgres SSL mode", + "localOnly": true, + "value": "require", + "origin": "platform" + } + } } ] },