Skip to content

feat(agents): persistent Lakebase thread store - #576

Draft
MarioCadenas wants to merge 4 commits into
mainfrom
feat/agent-persistent-threads
Draft

feat(agents): persistent Lakebase thread store#576
MarioCadenas wants to merge 4 commits into
mainfrom
feat/agent-persistent-threads

Conversation

@MarioCadenas

Copy link
Copy Markdown
Collaborator

What

Adds LakebaseThreadStore — a persistent ThreadStore for AppKit agents backed by Databricks Lakebase (Postgres). Today agents use InMemoryThreadStore, which loses every thread on restart and grows without bound; this closes that parity gap. The injection seam already existed (agents({ threadStore })), so this is one new store implementation plus lifecycle wiring — no plumbing surgery.

Implements the approved spec in design-docs/agent-persistent-threads-plan.md exactly.

How it works

  • Raw parameterized pg SQL over a pool from createLakebasePool() (OAuth handled inside the pool). Not the DatabasePlugin data layer.
  • Two user_id-scoped tables — agent_threads + agent_messages (FK ON DELETE CASCADE) — self-bootstrapped with CREATE TABLE IF NOT EXISTS on init() (once-guarded, select 1 connectivity probe first).
  • Every query filters WHERE user_id = $ — the isolation boundary; a user can never read or mutate another user's threads.
  • tool_calls stored as jsonb, preserving ToolCall.thoughtSignature verbatim (node-pg parses jsonb on read; JSON.stringify + ::jsonb cast on write). Date fields revived on read.
  • Owns its pool by default (closed on shutdown); an injected pool is left alone.
  • Optional init?()/close?() added to the ThreadStore interface (in-memory store untouched); wired into the agents plugin setup() (fail-fast) and shutdown().
  • Optional postgres resource on the agents manifest (same field/env shape as the database plugin) so a deploy can bind Lakebase; unbound apps fall back to in-memory.

Commits

  1. add persistent Lakebase thread store (schema + 5 methods) + unit tests
  2. wire ThreadStore init/close lifecycle into agents plugin
  3. optional postgres resource on agents manifest + export store
  4. persistent threads example + docs

Verification (agent-run)

  • New unit tests green (17 cases): per-method SQL + params, user_id scoping on every query, thoughtSignature round-trip, addMessage throws on missing thread, pool ownership, once-guarded bootstrap.
  • pnpm -r typecheck clean (no new errors vs baseline).
  • pnpm check:fix (oxlint + oxfmt) clean.
  • pnpm build && pnpm docs:build succeed.
  • Opt-in dev-playground example (agents({ threadStore: new LakebaseThreadStore() }) gated on LAKEBASE_ENDPOINT).

NOT verified

This has not been run against a live Lakebase. The SQL correctness (casts, CTE atomicity, jsonb round-trip, cascade) is only exercised against a mocked pool in unit tests. Live dogfood — deploy against a real Lakebase, create a thread, restart the app, confirm persistence and per-user scoping — is the remaining human step. Draft until then.

Deviations from the plan (flagged, not expanded)

  • Commit scope appkit, not agents: commitlint's scope enum is [appkit, appkit-ui, shared, playground, docs, deps, release]agents is rejected. PR title keeps (agents) per the task.
  • Export via @databricks/appkit/beta, not the GA root index.ts: the agents plugin and the ThreadStore type are beta exports, so the store lives beside them ("beside the existing agents exports").
  • Full postgres field set copied from the database manifest (host/databaseName/endpointPath/port/sslmode), not just the 5 named in the plan, so a bound deploy injects every env var createLakebasePool() reads.

Persistent ThreadStore backed by Lakebase (Postgres) over raw parameterized
pg SQL from createLakebasePool(). Two user_id-scoped tables (agent_threads,
agent_messages, FK ON DELETE CASCADE), self-bootstrapping CREATE TABLE IF NOT
EXISTS via init(). Every query filters WHERE user_id = the caller's id (the
isolation boundary); tool_calls stored as jsonb preserving thoughtSignature
verbatim; Dates revived on read. Owns its pool by default, leaves an injected
pool alone on close().

Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>
Add optional init?()/close?() to the ThreadStore interface (in-memory store
untouched). The agents plugin awaits threadStore.init?.() first in setup()
for fail-fast connectivity and threadStore.close?.() in shutdown() to release
an owned pool. Field typed as ThreadStore so the optional hooks resolve.

Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>
…store

Declare an optional postgres resource on the agents manifest (same field/env
shape as the database plugin, CAN_CONNECT_AND_CREATE) so a deploy can bind
Lakebase for persistent threads; unbound apps are unaffected and fall back to
in-memory. Export LakebaseThreadStore (and its options type) via the agents
barrel and the beta surface, beside the other agents exports.

Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>
Add a Thread persistence section to the agents plugin docs (in-memory default
vs LakebaseThreadStore, self-bootstrap, per-user isolation, deploy binding,
custom-store contract). Wire an opt-in dev-playground example that uses
LakebaseThreadStore when LAKEBASE_ENDPOINT is set (same signal the lakebase
plugin uses) and falls back to in-memory otherwise. Regenerated API docs and
template appkit.plugins.json.

Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant