diff --git a/bun.lock b/bun.lock index ebda34d44df..fe6e3436d9d 100644 --- a/bun.lock +++ b/bun.lock @@ -585,6 +585,31 @@ "typescript": "^5.9.3", }, }, + "packages/durable-session": { + "name": "@superset/durable-session", + "version": "0.0.1", + "dependencies": { + "@durable-streams/state": "^0.2.0", + "@standard-schema/spec": "^1.0.0", + "@tanstack/ai": "^0.3.0", + "@tanstack/db": "^0.5.22", + "@tanstack/db-ivm": "^0.1.17", + "zod": "^4.1.12", + }, + "devDependencies": { + "@superset/typescript": "workspace:*", + "@superset/ui": "workspace:*", + "@types/react": "~19.1.0", + "lucide-react": "^0.563.0", + "typescript": "^5.9.3", + }, + "peerDependencies": { + "@superset/ui": "workspace:*", + "@tanstack/react-db": "^0.1.66", + "lucide-react": ">=0.300.0", + "react": "^18.0.0 || ^19.0.0", + }, + }, "packages/email": { "name": "@superset/email", "version": "0.1.0", @@ -1897,6 +1922,8 @@ "@superset/docs": ["@superset/docs@workspace:apps/docs"], + "@superset/durable-session": ["@superset/durable-session@workspace:packages/durable-session"], + "@superset/email": ["@superset/email@workspace:packages/email"], "@superset/local-db": ["@superset/local-db@workspace:packages/local-db"], @@ -1963,10 +1990,14 @@ "@tailwindcss/vite": ["@tailwindcss/vite@4.1.18", "", { "dependencies": { "@tailwindcss/node": "4.1.18", "@tailwindcss/oxide": "4.1.18", "tailwindcss": "4.1.18" }, "peerDependencies": { "vite": "^5.2.0 || ^6 || ^7" } }, "sha512-jVA+/UpKL1vRLg6Hkao5jldawNmRo7mQYrZtNHMIVpLfLhDml5nMRUo/8MwoX2vNXvnaXNNMedrMfMugAVX1nA=="], + "@tanstack/ai": ["@tanstack/ai@0.3.0", "", { "dependencies": { "@tanstack/devtools-event-client": "^0.4.0", "partial-json": "^0.1.7" } }, "sha512-+MmqvLFmM1XvzGzVm2Pwgj8TmN8l3AV5NikDgam+B1HZY0MGdj5Pz7QFFYDEjTmzcFrDSDOAIobxR7Qp2iuKog=="], + "@tanstack/db": ["@tanstack/db@0.5.22", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "@tanstack/db-ivm": "0.1.17", "@tanstack/pacer-lite": "^0.2.0" }, "peerDependencies": { "typescript": ">=4.7" } }, "sha512-Qh33d9Idw7lEa/sg7veJfG8EUmSMtBzRkk/ghAVIDA3MJWAjyGOzU29TVZaM7K36BegTL8T/yVVlAFqjn/G2pw=="], "@tanstack/db-ivm": ["@tanstack/db-ivm@0.1.17", "", { "dependencies": { "fractional-indexing": "^3.2.0", "sorted-btree": "^1.8.1" }, "peerDependencies": { "typescript": ">=4.7" } }, "sha512-DK7vm56CDxNuRAdsbiPs+gITJ+16tUtYgZg3BRTLYKGIDsy8sdIO7sQFq5zl7Y+aIKAPmMAbVp9UjJ75FTtwgQ=="], + "@tanstack/devtools-event-client": ["@tanstack/devtools-event-client@0.4.0", "", {}, "sha512-RPfGuk2bDZgcu9bAJodvO2lnZeHuz4/71HjZ0bGb/SPg8+lyTA+RLSKQvo7fSmPSi8/vcH3aKQ8EM9ywf1olaw=="], + "@tanstack/electric-db-collection": ["@tanstack/electric-db-collection@0.2.27", "", { "dependencies": { "@electric-sql/client": "^1.3.1", "@standard-schema/spec": "^1.1.0", "@tanstack/db": "0.5.22", "@tanstack/store": "^0.8.0", "debug": "^4.4.3" } }, "sha512-nCJosh3GV8iX0w5b1LNys8a6l8k0I28bU86Zw3lIN/xN/blZ/dOO7nVDWE+jFBfu7UWYK2qFSzHATPKezWrHfw=="], "@tanstack/history": ["@tanstack/history@1.154.14", "", {}, "sha512-xyIfof8eHBuub1CkBnbKNKQXeRZC4dClhmzePHVOEel4G7lk/dW+TQ16da7CFdeNLv6u6Owf5VoBQxoo6DFTSA=="], @@ -4147,6 +4178,8 @@ "parseurl": ["parseurl@1.3.3", "", {}, "sha512-CiyeOxFT/JZyN5m0z9PfXw4SCBJ6Sygz1Dpl0wqjlhDEGGBP1GnsUVEL0p63hoG1fcj3fHynXi9NYO4nWOL+qQ=="], + "partial-json": ["partial-json@0.1.7", "", {}, "sha512-Njv/59hHaokb/hRUjce3Hdv12wd60MtM9Z5Olmn+nehe0QDAsRtRbJPvJ0Z91TusF0SuZRIvnM+S4l6EIP8leA=="], + "patch-console": ["patch-console@2.0.0", "", {}, "sha512-0YNdUceMdaQwoKce1gatDScmMo5pu/tfABfnzEqeG0gtTmd7mh/WcwgUjtAeOU7N8nFFlbQBnFK2gXW5fGvmMA=="], "path-data-parser": ["path-data-parser@0.1.0", "", {}, "sha512-NOnmBpt5Y2RWbuv0LMzsayp3lVylAHLPUTut412ZA3l+C4uw4ZVkQbjShYCQ8TCpUMdPapr4YjUqLYD6v68j+w=="], diff --git a/docs/ai-chat-plan.md b/docs/ai-chat-plan.md index d17a6d599a9..40bd2225081 100644 --- a/docs/ai-chat-plan.md +++ b/docs/ai-chat-plan.md @@ -78,7 +78,7 @@ Any Client (Web/Desktop/Mobile) ## Key Design Decisions -1. **Vendor `@electric-sql/durable-session`** — Not published to npm. Vendored from [electric-sql/transport](https://github.com/electric-sql/transport) (~35 files, ~4500 LOC). Gives us reactive collections, optimistic mutations, TanStack AI compatibility. +1. **Vendor `@electric-sql/durable-session`** — Not published to npm. Vendored from [electric-sql/transport](https://github.com/electric-sql/transport) into `packages/durable-session/` (~20 files). Required compatibility fixes for unreleased `@tanstack/db` aggregates (`collect`, `minStr`) and `@tanstack/ai` types. Gives us reactive collections, optimistic mutations, TanStack AI compatibility. 2. **Proxy pattern** — Proxy handles message writing, agent invocation, stream fan-out. Clients never write to durable stream directly. 3. **Agent endpoint** — Claude SDK runs as an "agent" the proxy calls via HTTP. Agent handles entire tool loop server-side. Returns standard TanStack AI SSE chunks. 4. **TanStack AI message format** — Messages use `parts: MessagePart[]` (TextPart, ToolCallPart, ToolResultPart, ThinkingPart) not Anthropic-specific `BetaContentBlock[]`. SDK output converted at the agent boundary. @@ -112,21 +112,22 @@ The agent endpoint converts these to TanStack AI `StreamChunk` format before wri | Session manager (v1) | DONE — `apps/desktop/src/lib/trpc/routers/ai-chat/utils/session-manager/session-manager.ts` | | Desktop tRPC router | DONE — `apps/desktop/src/lib/trpc/routers/ai-chat/index.ts` | | Durable stream server (v1) | DONE — `apps/streams/` (custom HTTP proxy + session registry) | -| Stream client (v1) | DONE — `packages/ai-chat/src/stream/client.ts` (custom DurableChatClient) | -| Stream hook (v1) | DONE — `packages/ai-chat/src/stream/useChatSession.ts` (custom hook) | -| Custom materialization (v1) | DONE — `packages/ai-chat/src/stream/materialize.ts` | -| ChatInput component | DONE — `packages/ai-chat/src/components/ChatInput/` | -| PresenceBar component | DONE — `packages/ai-chat/src/components/PresenceBar/` | +| Vendored durable-session client | DONE — `packages/durable-session/` (vendored from electric-sql/transport) | +| React hook (useDurableChat) | DONE — `packages/durable-session/src/react/use-durable-chat.ts` | +| ChatInput component | DONE — `packages/durable-session/src/react/components/ChatInput/` | +| PresenceBar component | DONE — `packages/durable-session/src/react/components/PresenceBar/` | +| Old ai-chat package | REMOVED — replaced by `@superset/durable-session` | +| Vendored proxy (A2) | NOT BUILT — **Next phase** | +| Claude agent endpoint (B) | NOT BUILT | | Database schema | NOT BUILT | | API chat router | NOT BUILT | | Desktop chat UI (renderer) | NOT BUILT | | Web chat UI | NOT BUILT | | Message rendering component | NOT BUILT | -| Vendored durable-session | NOT BUILT — **Next phase** | --- -## Phase A: Vendor `@electric-sql/durable-session` ← NEXT +## Phase A: Vendor `@electric-sql/durable-session` Source: [electric-sql/transport](https://github.com/electric-sql/transport) (unpublished, Apache-2.0) @@ -135,7 +136,7 @@ Reference source cloned to `/tmp/electric-sql-transport/` via: git clone https://github.com/electric-sql/transport.git /tmp/electric-sql-transport ``` -### A1. Create `packages/durable-session/` +### A1. Create `packages/durable-session/` — DONE Vendor from `packages/durable-session` + `packages/react-durable-session` in the transport repo. @@ -258,9 +259,9 @@ export const sessionStateSchema = createStateSchema({ **Collection pipeline** (`collections/messages.ts`): ``` -chunks → groupBy(messageId) + collect(chunk) + minStr(createdAt) + count(chunk) +chunks → groupBy(messageId) + count(chunk) + min(createdAt) → orderBy(startedAt, 'asc') - → fn.select(materializeMessage(collected.rows)) + → fn.select(imperatively gather chunks → materializeMessage(rows)) → getKey: row.id ``` @@ -288,7 +289,28 @@ const rawDb = createStreamDB({ - Auto-connects on mount if `autoConnect: true` (default) - Returns TanStack AI-compatible API: messages, sendMessage, isLoading, etc. -### A2. Vendor proxy into `apps/streams/` +#### Compatibility Fixes Applied + +The transport repo uses `workspace:*` (unreleased local versions) of `@tanstack/db`, `@tanstack/ai`, and `@durable-streams/state`. The published npm versions differ, requiring these fixes: + +| Issue | Fix | +|-------|-----| +| `collect` aggregate not in `@tanstack/db` v0.5.22 | Rewrote `messages.ts`, `session-stats.ts`, `presence.ts` to use `groupBy + count` as change discriminator + `fn.select` with imperative collection filtering | +| `minStr` aggregate not in `@tanstack/db` v0.5.22 | Replaced with `min()` which handles strings at runtime | +| `DoneStreamChunk` not in `@tanstack/ai` v0.3.0 | Replaced with `chunk.type === 'RUN_FINISHED'` type guard | +| `LiveMode` not in `@durable-streams/state` v0.2.1 | Removed import and re-export (was already unused in practice) | + +#### UI Components Migrated + +`ChatInput` and `PresenceBar` from the old `packages/ai-chat` were moved into `packages/durable-session/src/react/components/`. They are exported from `@superset/durable-session/react`: + +```typescript +import { ChatInput, PresenceBar } from '@superset/durable-session/react' +``` + +The old `packages/ai-chat` package has been fully removed. + +### A2. Vendor proxy into `apps/streams/` — NEXT Vendor from `packages/durable-session-proxy` in the transport repo. @@ -592,40 +614,17 @@ type StreamChunk = ## Phase C: Update Client Packages -### C1. Update `packages/ai-chat` +### C1. ~~Update `packages/ai-chat`~~ — DONE -**Remove** (replaced by vendored `@superset/durable-session`): -- `src/stream/client.ts` -- `src/stream/schema.ts` -- `src/stream/materialize.ts` -- `src/stream/materialize.test.ts` -- `src/stream/useChatSession.ts` -- `src/stream/useCollectionData.ts` -- `src/stream/actions.ts` +`packages/ai-chat` has been fully removed. All stream client code, hooks, materialization, and UI components are now in `packages/durable-session`. Consumers import directly: -**Rewrite** `src/stream/index.ts`: ```typescript -export { - DurableChatClient, createDurableChatClient, - type MessageRow, type ConnectionStatus, type DurableChatCollections, - type DurableChatClientOptions, type AgentSpec, - sessionStateSchema, extractTextContent, - isUserMessage, isAssistantMessage, messageRowToUIMessage, -} from "@superset/durable-session" - -export { - useDurableChat, - type UseDurableChatOptions, type UseDurableChatReturn, -} from "@superset/durable-session/react" -``` - -**Update** `package.json`: -- Remove: `@durable-streams/client`, `@durable-streams/state`, `@anthropic-ai/claude-agent-sdk`, `@anthropic-ai/sdk` -- Add: `@superset/durable-session: workspace:*` +// Data layer +import { DurableChatClient, createDurableChatClient } from '@superset/durable-session' -**Update** `src/types.ts`: -- Remove `StreamEvent`, `StreamEntry`, `Draft` types -- Keep `PresenceUser`, `ChatMessage`, `ChatSession` +// React hooks + components +import { useDurableChat, ChatInput, PresenceBar } from '@superset/durable-session/react' +``` ### C2. Simplify desktop session manager @@ -774,8 +773,8 @@ apps/desktop/src/renderer/screens/chat/ │ ├── ChatSidebar.tsx │ ├── ChatMessageList.tsx │ ├── ChatMessage.tsx -- Renders MessageRow with parts: TextPart, ToolCallPart, etc. -│ ├── ChatInput.tsx -- Reuse from @superset/ai-chat -│ ├── PresenceBar.tsx -- Reuse from @superset/ai-chat +│ ├── ChatInput.tsx -- Reuse from @superset/durable-session/react +│ ├── PresenceBar.tsx -- Reuse from @superset/durable-session/react │ └── TypingIndicator.tsx └── stores/ └── chat-store.ts @@ -783,7 +782,7 @@ apps/desktop/src/renderer/screens/chat/ Usage in component: ```tsx -import { useDurableChat } from "@superset/ai-chat/stream" +import { useDurableChat, ChatInput, PresenceBar } from "@superset/durable-session/react" function ChatRoom({ sessionId }: { sessionId: string }) { const { @@ -868,28 +867,32 @@ DURABLE_STREAM_URL=https://stream.superset.sh ## Complete File Operations Summary -### Files to CREATE (vendored) +### Files CREATED (vendored client — Phase A1) ✅ -| Destination | Source | Lines | +All files below are created and typechecking. Compatibility fixes applied for unreleased `@tanstack/db` aggregates (`collect`, `minStr`) and `@tanstack/ai` types (`DoneStreamChunk`). + +| Destination | Source | Status | |---|---|---| -| `packages/durable-session/package.json` | NEW | ~30 | -| `packages/durable-session/tsconfig.json` | NEW | ~8 | -| `packages/durable-session/src/index.ts` | `durable-session/src/index.ts` | 197 | -| `packages/durable-session/src/client.ts` | `durable-session/src/client.ts` | ~830 | -| `packages/durable-session/src/collection.ts` | `durable-session/src/collection.ts` | 155 | -| `packages/durable-session/src/materialize.ts` | `durable-session/src/materialize.ts` | 251 | -| `packages/durable-session/src/schema.ts` | `durable-session/src/schema.ts` | 253 | -| `packages/durable-session/src/types.ts` | `durable-session/src/types.ts` | 422 | -| `packages/durable-session/src/collections/index.ts` | `durable-session/src/collections/index.ts` | 57 | -| `packages/durable-session/src/collections/messages.ts` | `durable-session/src/collections/messages.ts` | 225 | -| `packages/durable-session/src/collections/active-generations.ts` | `durable-session/src/collections/active-generations.ts` | 82 | -| `packages/durable-session/src/collections/session-meta.ts` | `durable-session/src/collections/session-meta.ts` | 111 | -| `packages/durable-session/src/collections/session-stats.ts` | `durable-session/src/collections/session-stats.ts` | 265 | -| `packages/durable-session/src/collections/model-messages.ts` | `durable-session/src/collections/model-messages.ts` | 83 | -| `packages/durable-session/src/collections/presence.ts` | `durable-session/src/collections/presence.ts` | 77 | -| `packages/durable-session/src/react/index.ts` | `react-durable-session/src/index.ts` | 88 | -| `packages/durable-session/src/react/types.ts` | `react-durable-session/src/types.ts` | 119 | -| `packages/durable-session/src/react/use-durable-chat.ts` | `react-durable-session/src/use-durable-chat.ts` | 341 | +| `packages/durable-session/package.json` | NEW | ✅ | +| `packages/durable-session/tsconfig.json` | NEW | ✅ | +| `packages/durable-session/src/index.ts` | `durable-session/src/index.ts` | ✅ | +| `packages/durable-session/src/client.ts` | `durable-session/src/client.ts` | ✅ (fixed) | +| `packages/durable-session/src/collection.ts` | `durable-session/src/collection.ts` | ✅ | +| `packages/durable-session/src/materialize.ts` | `durable-session/src/materialize.ts` | ✅ (fixed) | +| `packages/durable-session/src/schema.ts` | `durable-session/src/schema.ts` | ✅ | +| `packages/durable-session/src/types.ts` | `durable-session/src/types.ts` | ✅ (fixed) | +| `packages/durable-session/src/collections/index.ts` | `durable-session/src/collections/index.ts` | ✅ | +| `packages/durable-session/src/collections/messages.ts` | `durable-session/src/collections/messages.ts` | ✅ (rewritten) | +| `packages/durable-session/src/collections/active-generations.ts` | `durable-session/src/collections/active-generations.ts` | ✅ | +| `packages/durable-session/src/collections/session-meta.ts` | `durable-session/src/collections/session-meta.ts` | ✅ | +| `packages/durable-session/src/collections/session-stats.ts` | `durable-session/src/collections/session-stats.ts` | ✅ (rewritten) | +| `packages/durable-session/src/collections/model-messages.ts` | `durable-session/src/collections/model-messages.ts` | ✅ | +| `packages/durable-session/src/collections/presence.ts` | `durable-session/src/collections/presence.ts` | ✅ (rewritten) | +| `packages/durable-session/src/react/index.ts` | `react-durable-session/src/index.ts` | ✅ | +| `packages/durable-session/src/react/types.ts` | `react-durable-session/src/types.ts` | ✅ | +| `packages/durable-session/src/react/use-durable-chat.ts` | `react-durable-session/src/use-durable-chat.ts` | ✅ | +| `packages/durable-session/src/react/components/ChatInput/` | Migrated from `packages/ai-chat` | ✅ | +| `packages/durable-session/src/react/components/PresenceBar/` | Migrated from `packages/ai-chat` | ✅ | ### Files to CREATE (new code) @@ -920,61 +923,59 @@ DURABLE_STREAM_URL=https://stream.superset.sh | `apps/streams/src/routes/auth.ts` | `durable-session-proxy/src/routes/auth.ts` | 146 | | `apps/streams/src/routes/fork.ts` | `durable-session-proxy/src/routes/fork.ts` | 50 | -### Files to DELETE +### Files DELETED ✅ + +| File | Reason | Status | +|---|---|---| +| `packages/ai-chat/` (entire package) | Replaced by `@superset/durable-session` | ✅ Removed | + +### Files to DELETE (remaining) | File | Reason | |---|---| -| `apps/streams/src/session-registry.ts` | Replaced by proxy's built-in session management | -| `packages/ai-chat/src/stream/client.ts` | Replaced by vendored DurableChatClient | -| `packages/ai-chat/src/stream/schema.ts` | Replaced by vendored sessionStateSchema | -| `packages/ai-chat/src/stream/materialize.ts` | Replaced by vendored materializeMessage | -| `packages/ai-chat/src/stream/materialize.test.ts` | Tests for deleted file | -| `packages/ai-chat/src/stream/useChatSession.ts` | Replaced by vendored useDurableChat | -| `packages/ai-chat/src/stream/useCollectionData.ts` | Replaced (built into vendored hook) | -| `packages/ai-chat/src/stream/actions.ts` | Replaced by vendored optimistic actions | +| `apps/streams/src/session-registry.ts` | Replaced by proxy's built-in session management (Phase A2) | -### Files to REWRITE +### Files to REWRITE (remaining) | File | Description | |---|---| -| `apps/streams/src/index.ts` | New entrypoint with Hono proxy + DurableStreamTestServer | -| `packages/ai-chat/src/stream/index.ts` | Re-export from `@superset/durable-session` | -| `apps/desktop/.../session-manager.ts` | Thin HTTP orchestrator (no StreamWatcher/Producer) | +| `apps/streams/src/index.ts` | New entrypoint with Hono proxy + DurableStreamTestServer (Phase A2) | +| `apps/desktop/.../session-manager.ts` | Thin HTTP orchestrator (no StreamWatcher/Producer) (Phase C2) | -### Files to MODIFY +### Files to MODIFY (remaining) | File | Changes | |---|---| -| `packages/ai-chat/package.json` | Remove: @durable-streams/*, @anthropic-ai/*. Add: @superset/durable-session | -| `apps/streams/package.json` | Add: hono, @hono/node-server, @durable-streams/client, @superset/durable-session | -| `packages/ai-chat/src/types.ts` | Remove StreamEvent, StreamEntry, Draft types | -| `packages/ai-chat/src/index.ts` | Update exports to match new stream/index.ts | +| `apps/streams/package.json` | Add: hono, @hono/node-server, @durable-streams/client, @superset/durable-session (Phase A2) | --- ## Implementation Order -1. **Phase A1** — Vendor `@superset/durable-session` package (copy 18 files, adjust 3 import paths) -2. **Phase A2** — Vendor proxy into `apps/streams` (copy 17 files, adjust 3 import paths) -3. **Phase B** — Claude agent endpoint + SDK-to-AI chunk converter (2 new files) -4. **Phase C** — Update `packages/ai-chat` + simplify session manager (delete 7 files, rewrite 3) -5. **Phase D** — Database schema + migration -6. **Phase E** — API tRPC router -7. **Phase F** — Desktop chat UI -8. **Phase G** — Web chat UI +1. ~~**Phase A1** — Vendor `@superset/durable-session` package~~ ✅ DONE +2. ~~**Phase C1** — Remove old `packages/ai-chat`, migrate UI components~~ ✅ DONE +3. **Phase A2** — Vendor proxy into `apps/streams` (copy 17 files, adjust 3 import paths) ← NEXT +4. **Phase B** — Claude agent endpoint + SDK-to-AI chunk converter (2 new files) +5. **Phase C2** — Simplify desktop session manager +6. **Phase C3** — Handle drafts +7. **Phase D** — Database schema + migration +8. **Phase E** — API tRPC router +9. **Phase F** — Desktop chat UI +10. **Phase G** — Web chat UI --- ## Risks -| Risk | Impact | Mitigation | -|------|--------|------------| -| `@tanstack/ai` API mismatch with vendored code | Build breaks | Vendored code uses `workspace:*` — pin to compatible published versions, fix API differences | -| SDKMessage → AI chunk conversion errors | Broken rendering | Comprehensive unit tests with real Claude output fixtures | -| Claude binary path outside Electron | Agent can't start | `CLAUDE_BINARY_PATH` env var set by desktop at streams startup | -| Multi-turn resume state lost on restart | Context lost | In-memory map + optional file-based persistence in data dir | -| Interrupt via HTTP abort | Claude subprocess continues | Agent detects fetch abort → calls `query.interrupt()` + `abortController.abort()` | -| Proxy `workspace:*` TanStack DB deps | Import errors | Pin all `@tanstack/*` to compatible published versions across monorepo | +| Risk | Impact | Mitigation | Status | +|------|--------|------------|--------| +| `@tanstack/ai` API mismatch with vendored code | Build breaks | Vendored code uses `workspace:*` — pin to compatible published versions, fix API differences | ✅ Resolved — `DoneStreamChunk` → `RUN_FINISHED`, `LiveMode` removed | +| `@tanstack/db` unreleased aggregates | Build breaks | Rewrite collection pipelines with `groupBy + count + fn.select` workaround | ✅ Resolved — `collect`/`minStr` replaced | +| SDKMessage → AI chunk conversion errors | Broken rendering | Comprehensive unit tests with real Claude output fixtures | Pending (Phase B) | +| Claude binary path outside Electron | Agent can't start | `CLAUDE_BINARY_PATH` env var set by desktop at streams startup | Pending | +| Multi-turn resume state lost on restart | Context lost | In-memory map + optional file-based persistence in data dir | Pending | +| Interrupt via HTTP abort | Claude subprocess continues | Agent detects fetch abort → calls `query.interrupt()` + `abortController.abort()` | Pending | +| Proxy `workspace:*` TanStack DB deps | Import errors | Pin all `@tanstack/*` to compatible published versions across monorepo | Pending (Phase A2) | --- @@ -1138,14 +1139,14 @@ it('converts result to done chunk', () => { ## Verification -### Phase A Verification (Vendored Package) +### Phase A1 Verification (Vendored Package) ✅ PASSED ```bash # 1. Install deps cd packages/durable-session && bun install -# 2. Type check vendored package -bun run typecheck -# 3. Verify exports resolve -node -e "require('@superset/durable-session')" # or bun +# 2. Type check vendored package — 0 errors, 0 warnings +bunx tsc --noEmit +# 3. Lint — 0 errors, 0 warnings +bun run lint:fix ``` ### Phase A2 + B Verification (Proxy + Agent) diff --git a/packages/ai-chat/package.json b/packages/ai-chat/package.json deleted file mode 100644 index 4d67b51cf71..00000000000 --- a/packages/ai-chat/package.json +++ /dev/null @@ -1,38 +0,0 @@ -{ - "name": "@superset/ai-chat", - "version": "0.0.1", - "description": "Shared AI chat hooks and utilities", - "type": "module", - "main": "./src/index.ts", - "exports": { - ".": "./src/index.ts", - "./components": "./src/components/index.ts", - "./stream": "./src/stream/index.ts" - }, - "scripts": { - "test": "bun test", - "typecheck": "tsc --noEmit --emitDeclarationOnly false" - }, - "dependencies": { - "@anthropic-ai/claude-agent-sdk": "^0.2.19", - "@anthropic-ai/sdk": "^0.72.1", - "@durable-streams/client": "^0.2.0", - "@durable-streams/state": "^0.2.0", - "@superset/ui": "workspace:*", - "@tanstack/db": "0.5.22", - "@tanstack/react-db": "0.1.66", - "lucide-react": "^0.563.0", - "zod": "^4.3.5" - }, - "devDependencies": { - "@superset/typescript": "workspace:*", - "@types/node": "^24.9.1", - "@types/react": "~19.1.0", - "bun-types": "^1.3.1", - "react": "19.1.0", - "typescript": "^5.9.3" - }, - "peerDependencies": { - "react": "^18.0.0 || ^19.0.0" - } -} diff --git a/packages/ai-chat/src/components/PresenceBar/index.ts b/packages/ai-chat/src/components/PresenceBar/index.ts deleted file mode 100644 index accc01d8d92..00000000000 --- a/packages/ai-chat/src/components/PresenceBar/index.ts +++ /dev/null @@ -1 +0,0 @@ -export { PresenceBar, type PresenceBarProps } from "./PresenceBar"; diff --git a/packages/ai-chat/src/components/index.ts b/packages/ai-chat/src/components/index.ts deleted file mode 100644 index 134ece8df83..00000000000 --- a/packages/ai-chat/src/components/index.ts +++ /dev/null @@ -1,2 +0,0 @@ -export { ChatInput, type ChatInputProps } from "./ChatInput"; -export { PresenceBar, type PresenceBarProps } from "./PresenceBar"; diff --git a/packages/ai-chat/src/index.ts b/packages/ai-chat/src/index.ts deleted file mode 100644 index a44ec03885c..00000000000 --- a/packages/ai-chat/src/index.ts +++ /dev/null @@ -1,3 +0,0 @@ -export * from "./components"; -export * from "./stream"; -export * from "./types"; diff --git a/packages/ai-chat/src/stream/actions.ts b/packages/ai-chat/src/stream/actions.ts deleted file mode 100644 index 9ed5719aedb..00000000000 --- a/packages/ai-chat/src/stream/actions.ts +++ /dev/null @@ -1,129 +0,0 @@ -/** - * Stream Actions - * - * Simple async functions for writing to the durable stream. - * Follows the Electric SQL pattern of plain functions over complex hooks. - */ - -export interface SessionUser { - userId: string; - name: string; -} - -/** - * Create a new stream (PUT request) - * Returns true if created, false if already exists - */ -export async function createStream( - baseUrl: string, - sessionId: string, -): Promise { - const response = await fetch(`${baseUrl}/streams/${sessionId}`, { - method: "PUT", - headers: { "Content-Type": "application/json" }, - }); - - // 201 = created, 200 = already exists - return response.status === 201; -} - -/** - * Helper to POST events to the stream with correct content-type - */ -async function appendToStream(url: string, events: unknown[]): Promise { - const response = await fetch(url, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(events), - }); - - if (!response.ok) { - throw new Error(`Failed to append to stream: ${response.status}`); - } -} - -export interface SessionActions { - /** Announce presence when joining */ - join: () => Promise; - /** Remove presence when leaving */ - leave: () => Promise; - /** Update draft content (empty content = delete draft) */ - updateDraft: (content: string) => Promise; -} - -/** - * Creates action functions for a chat session - * - * @example - * ```ts - * const actions = createSessionActions({ - * baseUrl: "http://localhost:8080", - * sessionId: "abc123", - * user: { userId: "user-1", name: "Alice" } - * }); - * - * await actions.join(); - * await actions.updateDraft("Hello..."); - * await actions.leave(); - * ``` - */ -export function createSessionActions({ - baseUrl, - sessionId, - user, -}: { - baseUrl: string; - sessionId: string; - user: SessionUser; -}): SessionActions { - const streamUrl = `${baseUrl}/streams/${sessionId}`; - - return { - join: async () => { - await appendToStream(streamUrl, [ - { - type: "presence", - key: user.userId, - value: { - userId: user.userId, - userName: user.name, - joinedAt: new Date().toISOString(), - }, - headers: { operation: "upsert" }, - }, - ]); - }, - - leave: async () => { - // Delete presence and draft on leave - await appendToStream(streamUrl, [ - { - type: "presence", - key: user.userId, - headers: { operation: "delete" }, - }, - { - type: "draft", - key: user.userId, - headers: { operation: "delete" }, - }, - ]); - }, - - updateDraft: async (content: string) => { - await appendToStream(streamUrl, [ - { - type: "draft", - key: user.userId, - value: { - userId: user.userId, - userName: user.name, - content, - updatedAt: new Date().toISOString(), - }, - headers: { operation: content ? "upsert" : "delete" }, - }, - ]); - }, - }; -} diff --git a/packages/ai-chat/src/stream/client.ts b/packages/ai-chat/src/stream/client.ts deleted file mode 100644 index e62ac7f715f..00000000000 --- a/packages/ai-chat/src/stream/client.ts +++ /dev/null @@ -1,416 +0,0 @@ -/** - * DurableChatClient - Chat session client adapted from Electric SQL. - * - * Follows the same patterns as @electric-sql/durable-session: - * - Synchronous construction (collections available immediately) - * - Async connection via connect() - * - Proper dispose lifecycle - * - * Simplified for our schema: chunks, presence, drafts (no agents, tool calls, etc.) - */ - -import { createStreamDB, type StreamDB } from "@durable-streams/state"; -import type { Collection } from "@tanstack/db"; -import type { SessionUser } from "./actions"; -import { - type SessionStateSchema, - type StreamChunk, - type StreamDraft, - type StreamPresence, - sessionStateSchema, -} from "./schema"; - -// ============================================================================ -// Helpers -// ============================================================================ - -/** UUID generator with fallback for environments without crypto.randomUUID (React Native). */ -function generateUUID(): string { - if ( - typeof crypto !== "undefined" && - typeof crypto.randomUUID === "function" - ) { - return crypto.randomUUID(); - } - return "xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx".replace(/[xy]/g, (c) => { - const r = (Math.random() * 16) | 0; - const v = c === "x" ? r : (r & 0x3) | 0x8; - return v.toString(16); - }); -} - -// ============================================================================ -// Types -// ============================================================================ - -export type ConnectionStatus = - | "disconnected" - | "connecting" - | "connected" - | "error"; - -export interface DurableChatClientOptions { - sessionId: string; - proxyUrl: string; - actorId?: string; - user?: SessionUser | null; - onError?: (error: Error) => void; - onStatusChange?: (status: ConnectionStatus) => void; -} - -/** - * Collections map with correct row types. - * - * stream-db injects the primary key field at runtime, so our types - * include the `id` field for chunks. - */ -export interface SessionCollections { - chunks: Collection; - presence: Collection; - drafts: Collection; -} - -// ============================================================================ -// DurableChatClient -// ============================================================================ - -/** - * DurableChatClient - Adapted from @electric-sql/durable-session. - * - * Provides a simple chat interface backed by Durable Streams. - * All collections are available immediately after construction. - * - * @example - * ```typescript - * const client = new DurableChatClient({ - * sessionId: 'my-session', - * proxyUrl: 'http://localhost:8080', - * }) - * - * // Collections available immediately - * const chunks = client.collections.chunks - * - * // Connect to start syncing - * await client.connect() - * - * // Send messages - * await client.sendMessage('Hello!') - * - * // Cleanup - * client.dispose() - * ``` - */ -export class DurableChatClient { - readonly sessionId: string; - readonly actorId: string; - - private readonly options: DurableChatClientOptions; - - // Stream-db instance (created synchronously in constructor) - private readonly _db: StreamDB; - - // Collections are always available after construction - private readonly _collections: SessionCollections; - - private _isConnected = false; - private _isDisposed = false; - private _connectionStatus: ConnectionStatus = "disconnected"; - private _error: Error | undefined; - - // AbortController for canceling stream sync - private readonly _abortController: AbortController; - - // ========================================================================= - // Constructor - // ========================================================================= - - constructor(options: DurableChatClientOptions) { - this.options = options; - this.sessionId = options.sessionId; - this.actorId = options.actorId ?? options.user?.userId ?? generateUUID(); - - // Create abort controller before anything else - this._abortController = new AbortController(); - - // Create stream-db synchronously (connection happens on preload) - this._db = createStreamDB({ - streamOptions: { - url: `${options.proxyUrl}/streams/${options.sessionId}`, - signal: this._abortController.signal, - }, - state: sessionStateSchema, - }); - - // Collections are available immediately - this._collections = this._db.collections as unknown as SessionCollections; - } - - // ========================================================================= - // Getters - // ========================================================================= - - /** - * Get all collections for direct access. - * Collections are available immediately after construction. - */ - get collections(): SessionCollections { - return this._collections; - } - - /** - * Get current connection status. - */ - get connectionStatus(): ConnectionStatus { - return this._connectionStatus; - } - - /** - * Check if the client has been disposed. - */ - get isDisposed(): boolean { - return this._isDisposed; - } - - /** - * Get the current error, if any. - */ - get error(): Error | undefined { - return this._error; - } - - // ========================================================================= - // Status Management - // ========================================================================= - - private _setConnectionStatus(status: ConnectionStatus): void { - if (this._connectionStatus === status) return; - this._connectionStatus = status; - this.options.onStatusChange?.(status); - } - - // ========================================================================= - // Lifecycle - // ========================================================================= - - /** - * Connect to the durable stream and start syncing. - * - * This method handles network operations only - collections are already - * created synchronously in the constructor and are immediately available. - */ - async connect(): Promise { - if (this._isConnected) return; - if (this._isDisposed) { - throw new Error("Cannot connect disposed client"); - } - - try { - this._setConnectionStatus("connecting"); - console.log( - `[ai-chat/client] connect() sessionId=${this.sessionId} url=${this.options.proxyUrl}`, - ); - - // Preload stream data - await this._db.preload(); - - this._isConnected = true; - this._setConnectionStatus("connected"); - console.log(`[ai-chat/client] connected, preload complete`); - - // Announce presence if we have a user - if (this.options.user) { - await this._announcePresence(); - } - } catch (error) { - this._error = error instanceof Error ? error : new Error(String(error)); - this._setConnectionStatus("error"); - this.options.onError?.(this._error); - throw error; - } - } - - /** - * Disconnect from the stream. - */ - disconnect(): void { - // Remove presence before disconnecting - if (this.options.user && this._isConnected) { - this._removePresence().catch(() => {}); - } - - this._db.close(); - this._abortController.abort(); - this._isConnected = false; - this._setConnectionStatus("disconnected"); - } - - /** - * Dispose the client and clean up resources. - * - * Note: We only disconnect here - we don't manually cleanup collections. - * TanStack DB will GC collections automatically when they have no subscribers. - */ - dispose(): void { - if (this._isDisposed) return; - this._isDisposed = true; - this.disconnect(); - } - - // ========================================================================= - // Actions - // ========================================================================= - - /** - * Send a user message. - */ - async sendMessage(content: string): Promise { - if (!this._isConnected) { - throw new Error("Client not connected. Call connect() first."); - } - - const user = this.options.user; - if (!user) { - throw new Error("Cannot send message without user"); - } - - const uuid = generateUUID(); - console.log( - `[ai-chat/client] sendMessage uuid=${uuid} content="${content.slice(0, 50)}"`, - ); - - await this._appendToStream([ - { - type: "chunk", - key: uuid, - value: { - type: "user_input", - content, - actorId: user.userId, - createdAt: new Date().toISOString(), - }, - headers: { operation: "insert" }, - }, - ]); - console.log(`[ai-chat/client] sendMessage POST complete`); - } - - /** - * Update the user's draft. - */ - async updateDraft(content: string): Promise { - const user = this.options.user; - if (!user) return; - - const now = new Date().toISOString(); - - await this._appendToStream([ - { - type: "draft", - key: user.userId, - value: { - userId: user.userId, - userName: user.name, - content, - updatedAt: now, - }, - headers: { operation: content ? "upsert" : "delete" }, - }, - ]); - } - - /** - * Update the user associated with this client. - * Useful for late binding when user auth completes after client creation. - */ - setUser(user: SessionUser | null): void { - const previousUser = this.options.user; - (this.options as { user: SessionUser | null }).user = user; - - // If connected, update presence - if (this._isConnected) { - if (previousUser && !user) { - // User logged out - remove presence - this._removePresence().catch(this.options.onError ?? console.error); - } else if (user && !previousUser) { - // User logged in - announce presence - this._announcePresence().catch(this.options.onError ?? console.error); - } else if (user && previousUser && user.userId !== previousUser.userId) { - // User changed - update presence - this._removePresence() - .then(() => this._announcePresence()) - .catch(this.options.onError ?? console.error); - } - } - } - - // ========================================================================= - // Private Helpers - // ========================================================================= - - private async _announcePresence(): Promise { - const user = this.options.user; - if (!user) return; - - await this._appendToStream([ - { - type: "presence", - key: user.userId, - value: { - userId: user.userId, - userName: user.name, - joinedAt: new Date().toISOString(), - }, - headers: { operation: "upsert" }, - }, - ]); - } - - private async _removePresence(): Promise { - const user = this.options.user; - if (!user) return; - - await this._appendToStream([ - { - type: "presence", - key: user.userId, - headers: { operation: "delete" }, - }, - { - type: "draft", - key: user.userId, - headers: { operation: "delete" }, - }, - ]); - } - - private async _appendToStream(events: unknown[]): Promise { - const response = await fetch( - `${this.options.proxyUrl}/streams/${this.sessionId}`, - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(events), - }, - ); - - if (!response.ok) { - throw new Error(`Failed to append to stream: ${response.status}`); - } - } -} - -// ============================================================================ -// Factory Function -// ============================================================================ - -/** - * Create a new DurableChatClient instance. - * - * @param options - Client options - * @returns New client instance - */ -export function createDurableChatClient( - options: DurableChatClientOptions, -): DurableChatClient { - return new DurableChatClient(options); -} diff --git a/packages/ai-chat/src/stream/index.ts b/packages/ai-chat/src/stream/index.ts deleted file mode 100644 index e7273a9682e..00000000000 --- a/packages/ai-chat/src/stream/index.ts +++ /dev/null @@ -1,42 +0,0 @@ -/** - * Stream module exports - */ - -export type { - BetaContentBlock, - BetaToolUseBlock, -} from "@anthropic-ai/sdk/resources/beta/messages/messages"; -export { - createSessionActions, - createStream, - type SessionActions, - type SessionUser, -} from "./actions"; -export { - type ConnectionStatus, - createDurableChatClient, - DurableChatClient, - type DurableChatClientOptions, - type SessionCollections, -} from "./client"; -export { - type ChunkRow, - type MessageRole, - type MessageRow, - materializeMessages, - type ToolResult, -} from "./materialize"; -export { - type SessionStateSchema, - type StreamChunk, - type StreamDraft, - type StreamPresence, - sessionStateSchema, -} from "./schema"; -export { - type ChatUser, - type UseChatSessionOptions, - type UseChatSessionReturn, - useChatSession, -} from "./useChatSession"; -export { useCollectionData } from "./useCollectionData"; diff --git a/packages/ai-chat/src/stream/materialize.test.ts b/packages/ai-chat/src/stream/materialize.test.ts deleted file mode 100644 index 9e51b593227..00000000000 --- a/packages/ai-chat/src/stream/materialize.test.ts +++ /dev/null @@ -1,75 +0,0 @@ -import { describe, expect, it } from "bun:test"; -import { materializeMessages } from "./materialize"; - -type ChunkInput = { - id: string; - type: "user_input"; - content: string; - actorId: string; - createdAt?: string; - _seq?: number; -}; - -const makeUserChunk = ({ - id, - createdAt, - _seq, -}: { - id: string; - createdAt?: string; - _seq?: number; -}): ChunkInput => ({ - id, - type: "user_input", - content: `msg-${id}`, - actorId: "user-1", - createdAt, - _seq, -}); - -describe("materializeMessages ordering", () => { - it("orders by _seq when available", () => { - const chunks = [ - makeUserChunk({ - id: "a", - createdAt: "2024-01-01T00:00:02.000Z", - _seq: 2, - }), - makeUserChunk({ - id: "b", - createdAt: "2024-01-01T00:00:01.000Z", - _seq: 1, - }), - ]; - - const messages = materializeMessages(chunks); - expect(messages.map((m) => m.id)).toEqual(["b", "a"]); - }); - - it("falls back to createdAt when _seq is missing", () => { - const chunks = [ - makeUserChunk({ - id: "a", - createdAt: "2024-01-01T00:00:02.000Z", - }), - makeUserChunk({ - id: "b", - createdAt: "2024-01-01T00:00:01.000Z", - }), - ]; - - const messages = materializeMessages(chunks); - expect(messages.map((m) => m.id)).toEqual(["b", "a"]); - }); - - it("preserves input order when both _seq and createdAt are missing", () => { - const chunks = [ - makeUserChunk({ id: "a" }), - makeUserChunk({ id: "b" }), - makeUserChunk({ id: "c" }), - ]; - - const messages = materializeMessages(chunks); - expect(messages.map((m) => m.id)).toEqual(["a", "b", "c"]); - }); -}); diff --git a/packages/ai-chat/src/stream/materialize.ts b/packages/ai-chat/src/stream/materialize.ts deleted file mode 100644 index a2b12dcbbae..00000000000 --- a/packages/ai-chat/src/stream/materialize.ts +++ /dev/null @@ -1,373 +0,0 @@ -/** - * Message Materialization (SDK-Native, Zero Envelope) - * - * Processes raw SDK messages and user input chunks from the durable stream - * in collection order. Turn boundaries are detected from SDK message type - * transitions (stream_event after user = new turn). - * - * No envelope fields (messageId, seq, role, actorId) — SDK messages are - * stored as-is, user input uses { type: "user_input", content, actorId }. - */ - -import type { - SDKAssistantMessage, - SDKMessage, - SDKPartialAssistantMessage, - SDKUserMessage, -} from "@anthropic-ai/claude-agent-sdk"; -import type { - BetaContentBlock, - BetaRawContentBlockDeltaEvent, - BetaRawContentBlockStartEvent, - BetaTextBlock, - BetaToolUseBlock, -} from "@anthropic-ai/sdk/resources/beta/messages/messages"; - -// ============================================================================ -// Types -// ============================================================================ - -export type MessageRole = "user" | "assistant" | "system"; - -/** Pre-merged tool result — SDK has no joined tool_use + tool_result concept */ -export interface ToolResult { - output: string; - isError: boolean; -} - -/** Materialized UI state from durable stream chunks */ -export interface MessageRow { - id: string; - role: MessageRole; - content: string; - contentBlocks: BetaContentBlock[]; - toolResults: Map; - actorId: string; - isComplete: boolean; - isStreaming: boolean; - createdAt: Date; -} - -/** Raw chunk row from the durable stream collection */ -export type ChunkRow = Record & { id: string }; - -interface ChunkSortKey { - seq: number | null; - time: number | null; - index: number; -} - -// ============================================================================ -// Materialize All Messages -// ============================================================================ - -/** - * Materialize all messages from raw chunk rows. - * - * Chunks are sorted by createdAt + _seq before processing to ensure - * correct ordering regardless of collection insertion order (which - * may not match stream append order after Electric SQL sync/reconnect). - * - * User input chunks become user messages. SDK message chunks are grouped - * into assistant turns with automatic boundary detection. - */ -export function materializeMessages(chunks: ChunkRow[]): MessageRow[] { - if (chunks.length === 0) return []; - - // Sort by stream sequence when available, then createdAt, then original index. - // createdAt is not guaranteed on SDK messages, so _seq should be authoritative. - const sorted = [...chunks] - .map((chunk, index) => ({ - chunk, - key: getChunkSortKey(chunk, index), - })) - .sort((a, b) => compareChunkSortKeys(a.key, b.key)) - .map(({ chunk }) => chunk); - - console.log( - `[ai-chat/materialize] processing ${sorted.length} sorted chunks, types: ${sorted.map((c) => c.type).join(", ")}`, - ); - - const messages: MessageRow[] = []; - let currentTurnChunks: ChunkRow[] = []; - let lastRenderingType: string | null = null; - - for (const chunk of sorted) { - const chunkType = chunk.type as string | undefined; - - // User input from client - if (chunkType === "user_input") { - // Flush current assistant turn - if (currentTurnChunks.length > 0) { - messages.push(materializeTurn(currentTurnChunks)); - currentTurnChunks = []; - lastRenderingType = null; - } - messages.push({ - id: chunk.id, - role: "user", - content: String(chunk.content ?? ""), - contentBlocks: [], - toolResults: new Map(), - actorId: String(chunk.actorId ?? ""), - isComplete: true, - isStreaming: false, - createdAt: new Date( - String(chunk.createdAt ?? new Date().toISOString()), - ), - }); - continue; - } - - // SDK message — only process rendering-relevant types for turns - // Note: "result" is excluded — materializeTurn doesn't handle it, - // and including it creates ghost streaming messages from lone result - // chunks. It also masks the user→stream_event turn boundary. - const isRenderingType = - chunkType === "stream_event" || - chunkType === "assistant" || - chunkType === "user"; - - if (!isRenderingType) continue; - - // Turn boundary: stream_event or assistant after user (tool result) = new turn - if ( - lastRenderingType === "user" && - (chunkType === "stream_event" || chunkType === "assistant") - ) { - if (currentTurnChunks.length > 0) { - messages.push(materializeTurn(currentTurnChunks)); - currentTurnChunks = []; - } - } - - currentTurnChunks.push(chunk); - lastRenderingType = chunkType; - } - - // Flush remaining turn - if (currentTurnChunks.length > 0) { - messages.push(materializeTurn(currentTurnChunks)); - } - - return messages; -} - -// ============================================================================ -// Turn Materialization -// ============================================================================ - -/** - * Materialize a single assistant turn from its SDK message chunks. - */ -function materializeTurn(chunks: ChunkRow[]): MessageRow { - const firstChunk = chunks[0] as ChunkRow; - - console.log( - `[ai-chat/materialize] materializeTurn: ${chunks.length} chunks, types: ${chunks.map((c) => c.type).join(", ")}`, - ); - - let assistantMsg: SDKAssistantMessage | null = null; - const streamEvents: SDKPartialAssistantMessage[] = []; - const userMsgs: SDKUserMessage[] = []; - for (const chunk of chunks) { - const msg = chunk as unknown as SDKMessage; - switch (msg.type) { - case "assistant": - assistantMsg = msg; - break; - case "stream_event": - streamEvents.push(msg); - break; - case "user": - userMsgs.push(msg); - break; - } - } - - // Build content blocks: prefer authoritative assistant message - let contentBlocks: BetaContentBlock[]; - let isStreaming: boolean; - - if (assistantMsg) { - contentBlocks = assistantMsg.message.content; - isStreaming = false; - } else { - contentBlocks = buildBlocksFromStreamEvents(streamEvents); - isStreaming = true; - } - - // Build tool results from user messages (tool_result blocks) - const toolResults = buildToolResults(userMsgs); - - // Join text blocks for backward-compat content field - const content = contentBlocks - .filter((b): b is BetaTextBlock => b.type === "text") - .map((b) => b.text) - .join(""); - - console.log( - `[ai-chat/materialize] turn result: id=${firstChunk.id.slice(0, 8)} hasAssistant=${!!assistantMsg} streamEvents=${streamEvents.length} userMsgs=${userMsgs.length} blocks=${contentBlocks.length} contentLen=${content.length} isComplete=${assistantMsg !== null} isStreaming=${isStreaming}`, - ); - - return { - id: firstChunk.id, - role: "assistant", - content, - contentBlocks, - toolResults, - actorId: "claude", - isComplete: assistantMsg !== null, - isStreaming, - createdAt: firstChunk.createdAt - ? new Date(String(firstChunk.createdAt)) - : new Date(), - }; -} - -// ============================================================================ -// Tool Result Extraction -// ============================================================================ - -function buildToolResults(userMsgs: SDKUserMessage[]): Map { - const toolResults = new Map(); - - for (const userMsg of userMsgs) { - const msgContent = userMsg.message.content; - if (!Array.isArray(msgContent)) continue; - - for (const block of msgContent) { - if (typeof block !== "object" || block === null) continue; - if (!("type" in block) || block.type !== "tool_result") continue; - - const tr = block as { - tool_use_id?: string; - content?: string | Array<{ type: string; text?: string }>; - is_error?: boolean; - }; - if (!tr.tool_use_id) continue; - - let output = ""; - if (typeof tr.content === "string") { - output = tr.content; - } else if (Array.isArray(tr.content)) { - output = tr.content - .filter((c) => c.type === "text" && c.text) - .map((c) => c.text as string) - .join(""); - } - - toolResults.set(tr.tool_use_id, { - output, - isError: tr.is_error ?? false, - }); - } - } - - return toolResults; -} - -// ============================================================================ -// Stream Event Reconstruction -// ============================================================================ - -function buildBlocksFromStreamEvents( - events: SDKPartialAssistantMessage[], -): BetaContentBlock[] { - const blocks: BetaContentBlock[] = []; - const jsonAccumulators = new Map(); - - for (const { event } of events) { - switch (event.type) { - case "content_block_start": { - const e = event as BetaRawContentBlockStartEvent; - blocks[e.index] = { ...e.content_block }; - if (e.content_block.type === "tool_use") { - jsonAccumulators.set(e.index, ""); - } - break; - } - - case "content_block_delta": { - const e = event as BetaRawContentBlockDeltaEvent; - const block = blocks[e.index]; - if (!block) break; - - const delta = e.delta; - if (delta.type === "text_delta" && block.type === "text") { - (block as BetaTextBlock).text += delta.text; - } else if (delta.type === "input_json_delta") { - const accumulated = - (jsonAccumulators.get(e.index) ?? "") + delta.partial_json; - jsonAccumulators.set(e.index, accumulated); - if (block.type === "tool_use") { - try { - (block as BetaToolUseBlock).input = JSON.parse(accumulated); - } catch { - // Partial JSON - } - } - } else if ( - delta.type === "thinking_delta" && - block.type === "thinking" - ) { - (block as { thinking: string }).thinking += delta.thinking; - } - break; - } - - case "content_block_stop": { - const index = (event as { index: number }).index; - const accumulated = jsonAccumulators.get(index); - const block = blocks[index]; - if (accumulated && block?.type === "tool_use") { - try { - (block as BetaToolUseBlock).input = JSON.parse(accumulated); - } catch { - // Best effort - } - jsonAccumulators.delete(index); - } - break; - } - } - } - - return blocks.filter(Boolean); -} - -// ============================================================================ -// Helpers -// ============================================================================ - -export function isUserMessage(row: MessageRow): boolean { - return row.role === "user"; -} - -export function isAssistantMessage(row: MessageRow): boolean { - return row.role === "assistant"; -} - -function getChunkSortKey(chunk: ChunkRow, index: number): ChunkSortKey { - const seq = typeof chunk._seq === "number" ? chunk._seq : null; - const time = chunk.createdAt - ? new Date(String(chunk.createdAt)).getTime() - : null; - return { seq, time, index }; -} - -function compareChunkSortKeys(a: ChunkSortKey, b: ChunkSortKey): number { - if (a.seq !== null || b.seq !== null) { - if (a.seq === null) return 1; - if (b.seq === null) return -1; - if (a.seq !== b.seq) return a.seq - b.seq; - } - - if (a.time !== null || b.time !== null) { - if (a.time === null) return 1; - if (b.time === null) return -1; - if (a.time !== b.time) return a.time - b.time; - } - - return a.index - b.index; -} diff --git a/packages/ai-chat/src/stream/schema.ts b/packages/ai-chat/src/stream/schema.ts deleted file mode 100644 index c72444a10ef..00000000000 --- a/packages/ai-chat/src/stream/schema.ts +++ /dev/null @@ -1,79 +0,0 @@ -/** - * Durable Streams State Schema - * - * Defines the state protocol schemas for real-time chat streaming with TanStack DB. - * Uses @durable-streams/state for protocol-compliant state management. - * - * Simplified design: - * - Presence = who's in the room (no status field - typing is derived from drafts) - * - Drafts = content being typed (non-empty content = user is typing) - * - Chunks = streaming message tokens - */ - -import { createStateSchema } from "@durable-streams/state"; -import { z } from "zod"; - -/** - * Chunk schema - raw SDK messages and user input stored directly. - * - * No envelope wrapping — SDK messages are passthrough, user input uses - * { type: "user_input", content, actorId, createdAt }. - * The `id` field (primary key) is injected at runtime from the event's `key`. - */ -export const chunkSchema = z.record(z.string(), z.unknown()); - -export type StreamChunk = z.infer; - -/** - * Presence schema - user presence tracking - * - * Simplified: No status field. Typing is derived from drafts with non-empty content. - */ -export const presenceSchema = z.object({ - userId: z.string(), - userName: z.string(), - joinedAt: z.string(), -}); - -export type StreamPresence = z.infer; - -/** - * Draft schema - user draft messages - * - * Non-empty content = user is typing. - */ -export const draftSchema = z.object({ - userId: z.string(), - userName: z.string(), - content: z.string(), - cursorPosition: z.number().optional(), - updatedAt: z.string(), -}); - -export type StreamDraft = z.infer; - -/** - * Combined session state schema - * - * This creates a typed schema that routes different event types - * to their respective TanStack DB collections. - */ -export const sessionStateSchema = createStateSchema({ - chunks: { - schema: chunkSchema, - type: "chunk", - primaryKey: "id", - }, - presence: { - schema: presenceSchema, - type: "presence", - primaryKey: "userId", - }, - drafts: { - schema: draftSchema, - type: "draft", - primaryKey: "userId", - }, -}); - -export type SessionStateSchema = typeof sessionStateSchema; diff --git a/packages/ai-chat/src/stream/useChatSession.ts b/packages/ai-chat/src/stream/useChatSession.ts deleted file mode 100644 index 8284508191f..00000000000 --- a/packages/ai-chat/src/stream/useChatSession.ts +++ /dev/null @@ -1,314 +0,0 @@ -/** - * useChatSession - React hook for durable chat. - * - * Copied 1:1 from @electric-sql/react-durable-session/use-durable-chat.ts - * Adapted for our schema: chunks, presence, drafts (no agents, tool calls, etc.) - */ -import { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { - type ConnectionStatus, - DurableChatClient, - type DurableChatClientOptions, - type SessionCollections, -} from "./client"; -import { type MessageRow, materializeMessages } from "./materialize"; -import type { StreamDraft, StreamPresence } from "./schema"; -import { useCollectionData } from "./useCollectionData"; - -// ============================================================================ -// Types -// ============================================================================ - -export interface UseChatSessionOptions - extends Omit { - /** Auto-connect when hook mounts. Defaults to true. */ - autoConnect?: boolean; - /** Provide an existing client instead of creating one */ - client?: DurableChatClient; - /** Error handler */ - onError?: (error: Error) => void; -} - -export interface ChatUser { - userId: string; - name: string; -} - -export interface UseChatSessionReturn { - // Data - messages: MessageRow[]; - streamingMessage: MessageRow | null; - users: ChatUser[]; - draft: string; - drafts: StreamDraft[]; - - // Actions - sendMessage: (content: string) => Promise; - setDraft: (content: string) => void; // Sync for onChange compatibility - - // State - isLoading: boolean; - error: Error | undefined; - connectionStatus: ConnectionStatus; - - // Extensions - client: DurableChatClient; - collections: SessionCollections; - connect: () => Promise; - disconnect: () => void; -} - -// ============================================================================ -// Hook Implementation -// ============================================================================ - -/** - * React hook for durable chat with TanStack AI-compatible API. - * - * Provides reactive data binding with automatic updates when underlying - * collection data changes. Supports SSR through proper `useSyncExternalStore` - * integration. - * - * The client and collections are always available synchronously. - * Connection state is managed separately via `connectionStatus`. - * - * @example Basic usage - * ```typescript - * function Chat() { - * const { messages, sendMessage, isLoading, collections } = useChatSession({ - * sessionId: 'my-session', - * proxyUrl: 'http://localhost:8080', - * }) - * - * return ( - *
- * {messages.map(m => )} - * - *
- * ) - * } - * ``` - */ -export function useChatSession( - options: UseChatSessionOptions, -): UseChatSessionReturn { - const { - autoConnect = true, - client: providedClient, - onError: userOnError, - ...clientOptions - } = options; - - // Error handler ref - allows client's onError to call setError - const [error, setError] = useState(); - const onErrorRef = useRef<(err: Error) => void>(() => {}); - onErrorRef.current = (err) => { - setError(err); - userOnError?.(err); - }; - - // Connection status state - ensures React re-renders on status changes - const [connectionStatus, setConnectionStatus] = - useState("disconnected"); - const onStatusChangeRef = useRef<(status: ConnectionStatus) => void>( - () => {}, - ); - onStatusChangeRef.current = (status) => { - setConnectionStatus(status); - }; - - // Create client synchronously - always available immediately - const clientRef = useRef<{ - client: DurableChatClient; - key: string; - } | null>(null); - const key = `${clientOptions.sessionId}:${clientOptions.proxyUrl}`; - - // Create or recreate client when key changes or client was disposed - // The isDisposed check handles React Strict Mode: cleanup disposes the client, - // so the next render must create a fresh one with a new AbortController. - if (providedClient) { - if (!clientRef.current || clientRef.current.client !== providedClient) { - clientRef.current = { client: providedClient, key: "provided" }; - } - } else if ( - !clientRef.current || - clientRef.current.key !== key || - clientRef.current.client.isDisposed - ) { - // Dispose old client if exists (may already be disposed, which is fine) - clientRef.current?.client.dispose(); - clientRef.current = { - client: new DurableChatClient({ - ...clientOptions, - onError: (err) => onErrorRef.current(err), - onStatusChange: (status) => onStatusChangeRef.current(status), - }), - key, - }; - } - - const client = clientRef.current.client; - - useEffect(() => { - client.setUser(clientOptions.user ?? null); - }, [client]); - - // ========================================================================= - // Collection Subscriptions (1:1 from Electric SQL) - // ========================================================================= - - const chunkRows = useCollectionData(client.collections.chunks); - const presenceRows = useCollectionData(client.collections.presence); - const draftRows = useCollectionData(client.collections.drafts); - - // ========================================================================= - // Derived State - // ========================================================================= - - // Materialize messages from chunks (processed in collection order) - const { messages, streamingMessage } = useMemo(() => { - console.log(`[ai-chat/session] materialize: ${chunkRows.length} chunks`); - if (chunkRows.length === 0) { - return { messages: [], streamingMessage: null }; - } - - const all = materializeMessages( - chunkRows as Array & { id: string }>, - ); - - // Separate complete messages from streaming (incomplete) message - const complete = all.filter((m) => m.isComplete); - const streaming = all.find((m) => !m.isComplete) ?? null; - - console.log( - `[ai-chat/session] materialized: ${all.length} total, ${complete.length} complete, streaming=${streaming ? `id=${streaming.id} role=${streaming.role} content="${streaming.content.slice(0, 80)}" blocks=${streaming.contentBlocks.length} isComplete=${streaming.isComplete} isStreaming=${streaming.isStreaming}` : "null"}`, - ); - if (all.length > 0) { - console.log( - `[ai-chat/session] all messages:`, - all.map((m) => ({ - id: m.id.slice(0, 8), - role: m.role, - isComplete: m.isComplete, - isStreaming: m.isStreaming, - contentLen: m.content.length, - blocks: m.contentBlocks.length, - })), - ); - } - - return { messages: complete, streamingMessage: streaming }; - }, [chunkRows]); - - // Transform presence to ChatUser[] - const users = useMemo( - (): ChatUser[] => - presenceRows.map((p: StreamPresence) => ({ - userId: p.userId, - name: p.userName, - })), - [presenceRows], - ); - - // All drafts - const drafts = useMemo((): StreamDraft[] => draftRows, [draftRows]); - - // Current user's draft - const draft = useMemo((): string => { - const user = clientOptions.user; - if (!user) return ""; - const myDraft = draftRows.find( - (d: StreamDraft) => d.userId === user.userId, - ); - return myDraft?.content ?? ""; - }, [draftRows]); - - const isLoading = connectionStatus !== "connected"; - - // ========================================================================= - // Connection Lifecycle (1:1 from Electric SQL) - // ========================================================================= - - useEffect(() => { - if (autoConnect && client.connectionStatus === "disconnected") { - client.connect().catch((err) => { - setError(err instanceof Error ? err : new Error(String(err))); - }); - } - - // Cleanup: unsubscribe and dispose (disposal is idempotent) - return () => { - if (!providedClient) { - client.dispose(); - } - }; - }, [client, autoConnect, providedClient]); - - // ========================================================================= - // Action Callbacks (1:1 from Electric SQL) - // ========================================================================= - - const sendMessage = useCallback( - async (content: string) => { - try { - await client.sendMessage(content); - } catch (err) { - setError(err instanceof Error ? err : new Error(String(err))); - throw err; - } - }, - [client], - ); - - // setDraft is sync (fire-and-forget) for onChange compatibility - const setDraft = useCallback( - (content: string) => { - client.updateDraft(content).catch((err) => { - setError(err instanceof Error ? err : new Error(String(err))); - }); - }, - [client], - ); - - const connect = useCallback(async () => { - try { - await client.connect(); - } catch (err) { - setError(err instanceof Error ? err : new Error(String(err))); - throw err; - } - }, [client]); - - const disconnect = useCallback(() => { - client.disconnect(); - }, [client]); - - // ========================================================================= - // Return (1:1 structure from Electric SQL) - // ========================================================================= - - return { - // Data - messages, - streamingMessage, - users, - draft, - drafts, - - // Actions - sendMessage, - setDraft, - - // State - isLoading, - error, - connectionStatus, - - // Extensions - client, - collections: client.collections, - connect, - disconnect, - }; -} diff --git a/packages/ai-chat/src/stream/useCollectionData.ts b/packages/ai-chat/src/stream/useCollectionData.ts deleted file mode 100644 index 17d8d631470..00000000000 --- a/packages/ai-chat/src/stream/useCollectionData.ts +++ /dev/null @@ -1,98 +0,0 @@ -/** - * SSR-safe hook for subscribing to TanStack DB collection data. - * - * Copied verbatim from @electric-sql/react-durable-session. - * This is a workaround to useLiveQuery not yet supporting SSR - * as per https://github.com/TanStack/db/pull/709 - */ - -import type { Collection } from "@tanstack/db"; -import { useEffect, useRef, useSyncExternalStore } from "react"; - -/** - * Extract the item type from a Collection. - * - * TanStack DB's Collection has 5 type parameters: - * `Collection` - * - * This helper extracts `T` (the item type) from any Collection variant. - */ -type CollectionItem = - // biome-ignore lint/suspicious/noExplicitAny: Collection has constrained generic params that require any - C extends Collection ? T : never; - -/** - * SSR-safe hook for subscribing to TanStack DB collection data. - * This is a workaround to useLiveQuery not yet supporting SSR - * as per https://github.com/TanStack/db/pull/709 - */ -export function useCollectionData< - // biome-ignore lint/suspicious/noExplicitAny: Collection has constrained generic params that require any - C extends Collection, ->(collection: C): CollectionItem[] { - type T = CollectionItem; - - const collectionRef = useRef(collection); - - // Track version to know when to create a new snapshot. - // Incremented by subscription callback when collection changes. - const versionRef = useRef(0); - - // Cache the last snapshot to maintain stable reference. - // useSyncExternalStore requires getSnapshot to return the same reference - // when data hasn't changed, otherwise it triggers infinite re-renders. - const snapshotRef = useRef<{ version: number; data: T[] }>({ - version: -1, // Force initial snapshot creation - data: [], - }); - - useEffect(() => { - collectionRef.current = collection; - versionRef.current = 0; - snapshotRef.current = { version: -1, data: [] }; - }, [collection]); - - // Subscribe callback - increments version to signal data changed. - // Stored in ref to maintain stable reference for useSyncExternalStore. - const subscribeRef = useRef<(onStoreChange: () => void) => () => void>( - () => () => {}, - ); - subscribeRef.current = (onStoreChange: () => void): (() => void) => { - const currentCollection = collectionRef.current; - const subscription = currentCollection.subscribeChanges(() => { - versionRef.current++; - console.log( - `[ai-chat/collection] change detected, version=${versionRef.current}, size=${currentCollection.size}`, - ); - onStoreChange(); - }); - return () => subscription.unsubscribe(); - }; - - // Snapshot callback - returns cached data unless version changed. - // Stored in ref to maintain stable reference for useSyncExternalStore. - const getSnapshotRef = useRef<() => T[]>(() => []); - getSnapshotRef.current = (): T[] => { - const currentVersion = versionRef.current; - const cached = snapshotRef.current; - - if (cached.version === currentVersion) { - return cached.data; - } - - const data = [...collectionRef.current.values()] as T[]; - snapshotRef.current = { version: currentVersion, data }; - return data; - }; - - const getServerSnapshotRef = useRef<() => T[]>(() => []); - getServerSnapshotRef.current = (): T[] => snapshotRef.current.data; - - // Use a stable server snapshot to keep SSR output consistent - // and avoid hydration mismatches. - return useSyncExternalStore( - subscribeRef.current, - getSnapshotRef.current, - getServerSnapshotRef.current, - ); -} diff --git a/packages/ai-chat/src/types.ts b/packages/ai-chat/src/types.ts deleted file mode 100644 index ee2c9903b06..00000000000 --- a/packages/ai-chat/src/types.ts +++ /dev/null @@ -1,112 +0,0 @@ -/** - * Shared types for AI chat - */ - -export interface StreamTextDeltaEvent { - type: "text_delta"; - text: string; - timestamp: number; -} - -export interface StreamToolUseStartEvent { - type: "tool_use_start"; - toolName: string; - toolId: string; - timestamp: number; -} - -export interface StreamToolUseDeltaEvent { - type: "tool_use_delta"; - toolId: string; - partialJson: string; - timestamp: number; -} - -export interface StreamToolUseEndEvent { - type: "tool_use_end"; - toolId: string; - timestamp: number; -} - -export interface StreamMessageCompleteEvent { - type: "message_complete"; - inputTokens?: number; - outputTokens?: number; - timestamp: number; -} - -export interface StreamErrorEvent { - type: "error"; - error: string; - timestamp: number; -} - -export interface StreamSessionStartEvent { - type: "session_start"; - timestamp: number; -} - -export interface StreamSessionEndEvent { - type: "session_end"; - exitCode: number | null; - timestamp: number; -} - -export type StreamEvent = - | StreamTextDeltaEvent - | StreamToolUseStartEvent - | StreamToolUseDeltaEvent - | StreamToolUseEndEvent - | StreamMessageCompleteEvent - | StreamErrorEvent - | StreamSessionStartEvent - | StreamSessionEndEvent; - -export interface StreamEntry { - offset: number; - event: StreamEvent; -} - -export interface PresenceUser { - userId: string; - name: string; - image?: string; -} - -export interface PresenceState { - viewers: PresenceUser[]; - typingUsers: PresenceUser[]; -} - -export interface Draft { - userId: string; - userName: string; - content: string; - updatedAt: number; -} - -export interface ChatMessage { - id: string; - sessionId: string; - role: "user" | "assistant"; - content: string; - toolCalls?: unknown[]; - inputTokens?: number; - outputTokens?: number; - createdById: string; - createdAt: Date; -} - -export interface ChatSession { - id: string; - organizationId: string; - repositoryId?: string | null; - workspaceId?: string | null; - title: string; - claudeSessionId?: string | null; - cwd?: string | null; - createdById: string; - archivedAt?: Date | null; - createdAt: Date; - updatedAt: Date; -} diff --git a/packages/durable-session/package.json b/packages/durable-session/package.json new file mode 100644 index 00000000000..864e1d25126 --- /dev/null +++ b/packages/durable-session/package.json @@ -0,0 +1,37 @@ +{ + "name": "@superset/durable-session", + "version": "0.0.1", + "private": true, + "type": "module", + "exports": { + ".": { + "types": "./src/index.ts", + "default": "./src/index.ts" + }, + "./react": { + "types": "./src/react/index.ts", + "default": "./src/react/index.ts" + } + }, + "dependencies": { + "@durable-streams/state": "^0.2.0", + "@standard-schema/spec": "^1.0.0", + "@tanstack/ai": "^0.3.0", + "@tanstack/db": "^0.5.22", + "@tanstack/db-ivm": "^0.1.17", + "zod": "^4.1.12" + }, + "devDependencies": { + "@superset/typescript": "workspace:*", + "@superset/ui": "workspace:*", + "@types/react": "~19.1.0", + "lucide-react": "^0.563.0", + "typescript": "^5.9.3" + }, + "peerDependencies": { + "react": "^18.0.0 || ^19.0.0", + "@tanstack/react-db": "^0.1.66", + "@superset/ui": "workspace:*", + "lucide-react": ">=0.300.0" + } +} diff --git a/packages/durable-session/src/client.ts b/packages/durable-session/src/client.ts new file mode 100644 index 00000000000..c074c0e3db2 --- /dev/null +++ b/packages/durable-session/src/client.ts @@ -0,0 +1,851 @@ +/** + * DurableChatClient - Framework-agnostic durable chat client. + * + * Provides TanStack AI-compatible API backed by Durable Streams + * with real-time sync and multi-agent support. + * + * All derived collections contain fully materialized MessageRow objects. + * Consumers filter message.parts to access specific part types (ToolCallPart, etc.). + */ + +import type { AnyClientTool, ToolCallPart, UIMessage } from "@tanstack/ai"; +import type { Transaction } from "@tanstack/db"; +import { createCollection, createOptimisticAction } from "@tanstack/db"; +import { createSessionDB, type SessionDB } from "./collection"; +import { + createActiveGenerationsCollection, + createInitialSessionMeta, + createMessagesCollection, + createPendingApprovalsCollection, + createPresenceCollection, + createSessionMetaCollectionOptions, + createSessionStatsCollection, + createToolCallsCollection, + createToolResultsCollection, + updateConnectionStatus, +} from "./collections"; +import { extractTextContent, messageRowToUIMessage } from "./materialize"; +import type { + ActorType, + AgentSpec, + ApprovalResponseInput, + ClientToolResultInput, + ConnectionStatus, + DurableChatClientOptions, + ForkOptions, + ForkResult, + MessageRow, + SessionMetaRow, + ToolResultInput, +} from "./types"; + +/** + * Unified input for all message optimistic actions. + */ +interface MessageActionInput { + /** Message content */ + content: string; + /** Client-generated message ID */ + messageId: string; + /** Message role */ + role: "user" | "assistant" | "system"; + /** Optional agent to invoke (for user messages) */ + agent?: AgentSpec; +} + +/** + * DurableChatClient provides a TanStack AI-compatible chat interface + * backed by Durable Streams for persistence and real-time sync. + * + * All derived collections contain fully materialized objects. + * Access data directly from collections - no helper functions needed. + * + * @example + * ```typescript + * import { DurableChatClient } from '@superset/durable-session' + * + * const client = new DurableChatClient({ + * sessionId: 'my-session', + * proxyUrl: 'http://localhost:4000', + * }) + * + * await client.connect() + * + * // Use TanStack AI-compatible API + * await client.sendMessage('Hello!') + * console.log(client.messages) + * + * // Or use collections directly + * for (const message of client.collections.messages.values()) { + * console.log(message.id, message.role, message.parts) + * } + * + * // Filter tool calls + * const pending = [...client.collections.toolCalls.values()] + * .filter(tc => tc.state === 'pending') + * ``` + */ + +export class DurableChatClient< + TTools extends ReadonlyArray = AnyClientTool[], +> { + readonly sessionId: string; + readonly actorId: string; + readonly actorType: ActorType; + + private readonly options: DurableChatClientOptions; + + // Stream-db instance (created synchronously in constructor) + // Either from options.sessionDB (tests) or createSessionDB() (production) + private readonly _db: SessionDB; + + // Collections are typed via inference from createCollections() + // Created synchronously in constructor - always available + private readonly _collections: ReturnType< + DurableChatClient["createCollections"] + >; + + private _isConnected = false; + private _isDisposed = false; + private _error: Error | undefined; + + // AbortController created at construction time to pass signal to stream-db. + // Aborted on disconnect() to cancel the stream sync. + private readonly _abortController: AbortController; + + // Optimistic actions for mutations (created synchronously in constructor) + private readonly _messageAction: (input: MessageActionInput) => Transaction; + private readonly _addToolResultAction: ( + input: ClientToolResultInput, + ) => Transaction; + private readonly _addApprovalResponseAction: ( + input: ApprovalResponseInput, + ) => Transaction; + + // ═══════════════════════════════════════════════════════════════════════ + // Constructor + // ═══════════════════════════════════════════════════════════════════════ + + constructor(options: DurableChatClientOptions) { + this.options = options; + this.sessionId = options.sessionId; + this.actorId = options.actorId ?? crypto.randomUUID(); + this.actorType = options.actorType ?? "user"; + + // Create abort controller before anything else + this._abortController = new AbortController(); + + // Create stream-db synchronously (use injected sessionDB for tests) + this._db = + options.sessionDB ?? + createSessionDB({ + sessionId: this.sessionId, + baseUrl: options.proxyUrl, + headers: options.stream?.headers, + signal: this._abortController.signal, + }); + + // Create all collections synchronously (always from _db.collections) + this._collections = this.createCollections(); + + // Initialize session metadata + this._collections.sessionMeta.insert( + createInitialSessionMeta(this.sessionId), + ); + + // Create optimistic actions (they use collections) + this._messageAction = this.createMessageAction(); + this._addToolResultAction = this.createAddToolResultAction(); + this._addApprovalResponseAction = this.createApprovalResponseAction(); + } + + // ═══════════════════════════════════════════════════════════════════════ + // Collection Setup + // ═══════════════════════════════════════════════════════════════════════ + + /** + * Create all derived collections from the chunks collection. + * + * Pipeline architecture: + * - chunks → (subquery) → messages (root materialized collection) + * - Derived collections filter messages via .fn.where() on parts + * + * CRITICAL: Materialization happens inside fn.select(). No imperative code + * outside this pattern. + */ + private createCollections() { + // Get root collections from stream-db (always available - from real or mock SessionDB) + // Note: rawPresence contains per-device records; we expose aggregated presence + const { chunks, presence: rawPresence, agents } = this._db.collections; + + // Root materialized collection: chunks → messages + // Uses inline subquery for chunk aggregation + const messages = createMessagesCollection({ + chunksCollection: chunks, + }); + + // Derived collections filter on message parts (lazy evaluation) + const toolCalls = createToolCallsCollection({ + messagesCollection: messages, + }); + + const pendingApprovals = createPendingApprovalsCollection({ + messagesCollection: messages, + }); + + const toolResults = createToolResultsCollection({ + messagesCollection: messages, + }); + + const activeGenerations = createActiveGenerationsCollection({ + messagesCollection: messages, + }); + + // Session metadata collection (local state) + const sessionMeta = createCollection( + createSessionMetaCollectionOptions({ + sessionId: this.sessionId, + }), + ); + + // Session statistics collection (aggregated from chunks) + const sessionStats = createSessionStatsCollection({ + sessionId: this.sessionId, + chunksCollection: chunks, + }); + + // Create aggregated presence collection (groups by actorId, filters for online) + // This provides a "who's online" view rather than raw per-device records + const presence = createPresenceCollection({ + sessionId: this.sessionId, + rawPresenceCollection: rawPresence, + }); + + return { + chunks, + presence, + agents, + messages, + toolCalls, + pendingApprovals, + toolResults, + activeGenerations, + sessionMeta, + sessionStats, + }; + } + + // ═══════════════════════════════════════════════════════════════════════ + // Core API (TanStack AI ChatClient compatible) + // ═══════════════════════════════════════════════════════════════════════ + + /** + * Get all messages as UIMessage array. + * Messages are accessed directly from the materialized collection. + */ + get messages(): UIMessage[] { + return [...this._collections.messages.values()].map(messageRowToUIMessage); + } + + /** + * Check if any generation is currently active. + * Uses the activeGenerations collection size directly. + */ + get isLoading(): boolean { + return this._collections.activeGenerations.size > 0; + } + + /** + * Get the current error, if any. + */ + get error(): Error | undefined { + return this._error; + } + + /** + * Check if the client has been disposed. + */ + get isDisposed(): boolean { + return this._isDisposed; + } + + /** + * Send a user message and trigger agent response. + * + * Uses optimistic updates for instant UI feedback. The message appears + * immediately in the UI while the server request is in flight. + * + * @param content - Text content to send + */ + async sendMessage(content: string): Promise { + if (!this._isConnected) { + throw new Error("Client not connected. Call connect() first."); + } + + await this.executeAction(this._messageAction, { + content, + messageId: crypto.randomUUID(), + role: "user", + agent: this.options.agent, + }); + } + + /** + * Append a message to the conversation. + * + * Uses optimistic updates for instant UI feedback. + * For user messages, this triggers agent response if an agent is configured. + * + * @param message - UIMessage or ModelMessage to append + */ + async append( + message: UIMessage | { role: string; content: string }, + ): Promise { + if (!this._isConnected) { + throw new Error("Client not connected. Call connect() first."); + } + + const content = + "parts" in message + ? extractTextContent(message as MessageRow) + : (message as { content: string }).content; + + const role = message.role as "user" | "assistant" | "system"; + const messageId = "id" in message ? message.id : crypto.randomUUID(); + + await this.executeAction(this._messageAction, { + content, + messageId, + role, + agent: role === "user" ? this.options.agent : undefined, + }); + } + + /** + * Execute an optimistic action with unified error handling. + */ + private async executeAction( + action: (input: T) => Transaction, + input: T, + ): Promise { + try { + const transaction = action(input); + await transaction.isPersisted.promise; + } catch (error) { + this._error = error instanceof Error ? error : new Error(String(error)); + this.options.onError?.(this._error); + throw error; + } + } + + /** + * POST JSON to proxy endpoint with error handling. + */ + private async postToProxy( + path: string, + body: Record, + options?: { actorIdHeader?: boolean }, + ): Promise { + const headers: Record = { + "Content-Type": "application/json", + }; + if (options?.actorIdHeader) { + headers["X-Actor-Id"] = this.actorId; + } + + const response = await fetch(`${this.options.proxyUrl}${path}`, { + method: "POST", + headers, + body: JSON.stringify(body), + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error(`Request failed: ${response.status} ${errorText}`); + } + } + + /** + * Create the unified optimistic action for all message types. + * Handles user, assistant, and system messages with the same pattern. + * + * Optimistic updates insert into the messages collection directly. + * This ensures the optimistic state propagates to all derived collections + * (toolCalls, pendingApprovals, toolResults, activeGenerations). + */ + private createMessageAction() { + return createOptimisticAction({ + onMutate: ({ content, messageId, role }) => { + const createdAt = new Date(); + + // Insert into messages collection directly + // This propagates to all derived collections + this._collections.messages.insert({ + id: messageId, + role, + parts: [{ type: "text" as const, content }], + actorId: this.actorId, + isComplete: true, + createdAt, + }); + }, + mutationFn: async ({ content, messageId, role, agent }) => { + const txid = crypto.randomUUID(); + + await this.postToProxy(`/v1/sessions/${this.sessionId}/messages`, { + messageId, + content, + role, + actorId: this.actorId, + actorType: this.actorType, + txid, + ...(agent && { agent }), + }); + + // Wait for txid to appear in synced stream + await this._db.utils.awaitTxId(txid); + }, + }); + } + + /** + * Reload the last user message and regenerate response. + */ + async reload(): Promise { + const msgs = this.messages; + if (msgs.length === 0) return; + + // Find the last user message + let lastUserMessage: UIMessage | undefined; + for (let i = msgs.length - 1; i >= 0; i--) { + if (msgs[i]?.role === "user") { + lastUserMessage = msgs[i]; + break; + } + } + + if (!lastUserMessage) return; + + // Get content of last user message + const content = extractTextContent( + lastUserMessage as unknown as MessageRow, + ); + + // Call regenerate endpoint + const response = await fetch( + `${this.options.proxyUrl}/v1/sessions/${this.sessionId}/regenerate`, + { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + fromMessageId: lastUserMessage.id, + content, + actorId: this.actorId, + actorType: this.actorType, + }), + }, + ); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error(`Failed to reload: ${response.status} ${errorText}`); + } + } + + /** + * Stop all active generations. + */ + stop(): void { + // Call stop endpoint + fetch(`${this.options.proxyUrl}/v1/sessions/${this.sessionId}/stop`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ messageId: null }), // null = stop all + }).catch((err) => { + console.warn("Failed to stop generation:", err); + }); + } + + /** + * Clear all messages (local only - does not affect server). + */ + clear(): void { + // Note: This only clears local state, not the durable stream + // For full clear, use the proxy's clear endpoint + this.options.onMessagesChange?.([]); + } + + /** + * Add a tool result. + * + * Uses optimistic updates for instant UI feedback. + * + * @param result - Tool result to add + */ + async addToolResult(result: ToolResultInput): Promise { + if (!this._isConnected) { + throw new Error("Client not connected. Call connect() first."); + } + + // Ensure messageId is set for optimistic updates + const inputWithMessageId: ClientToolResultInput = { + ...result, + messageId: result.messageId ?? crypto.randomUUID(), + }; + await this.executeAction(this._addToolResultAction, inputWithMessageId); + } + + /** + * Create the optimistic action for adding tool results. + * + * Inserts a new message with a ToolResultPart into the messages collection. + * Uses client-generated messageId for predictable IDs. + */ + private createAddToolResultAction() { + return createOptimisticAction({ + onMutate: ({ messageId, toolCallId, output, error }) => { + const createdAt = new Date(); + + // Insert a new message with tool-result part + this._collections.messages.insert({ + id: messageId, + role: "assistant", + parts: [ + { + type: "tool-result" as const, + toolCallId, + content: + typeof output === "string" ? output : JSON.stringify(output), + state: error ? ("error" as const) : ("complete" as const), + ...(error && { error }), + }, + ], + actorId: this.actorId, + isComplete: true, + createdAt, + }); + }, + mutationFn: async ({ messageId, toolCallId, output, error }) => { + const txid = crypto.randomUUID(); + + await this.postToProxy( + `/v1/sessions/${this.sessionId}/tool-results`, + { messageId, toolCallId, output, error: error ?? null, txid }, + { actorIdHeader: true }, + ); + + // Wait for txid to appear in synced stream + await this._db.utils.awaitTxId(txid); + }, + }); + } + + /** + * Add an approval response. + * + * Uses optimistic updates for instant UI feedback. + * + * @param response - Approval response + */ + async addToolApprovalResponse( + response: ApprovalResponseInput, + ): Promise { + if (!this._isConnected) { + throw new Error("Client not connected. Call connect() first."); + } + + await this.executeAction(this._addApprovalResponseAction, response); + } + + /** + * Create the optimistic action for approval responses. + * + * Finds the message containing the tool call with the approval and updates + * the approval.approved field. This propagates to pendingApprovals collection. + */ + private createApprovalResponseAction() { + return createOptimisticAction({ + onMutate: ({ id, approved }) => { + // Find the message containing this approval + for (const message of this._collections.messages.values()) { + for (const part of message.parts) { + if (part.type === "tool-call" && part.approval?.id === id) { + // Update the message with the approval response + this._collections.messages.update(message.id, (draft) => { + for (const p of draft.parts) { + const toolCall = p as ToolCallPart; + if ( + p.type === "tool-call" && + toolCall.approval?.id === id && + toolCall.approval + ) { + toolCall.approval.approved = approved; + } + } + }); + return; + } + } + } + }, + mutationFn: async ({ id, approved }) => { + const txid = crypto.randomUUID(); + + await this.postToProxy( + `/v1/sessions/${this.sessionId}/approvals/${id}`, + { approved, txid }, + { actorIdHeader: true }, + ); + + // Wait for txid to appear in synced stream + await this._db.utils.awaitTxId(txid); + }, + }); + } + + // ═══════════════════════════════════════════════════════════════════════ + // Collections + // ═══════════════════════════════════════════════════════════════════════ + + /** + * Get all collections for custom queries. + * All collections contain fully materialized objects. + * Collections are available immediately after construction. + */ + get collections() { + return this._collections; + } + + // ═══════════════════════════════════════════════════════════════════════ + // Durable-specific features + // ═══════════════════════════════════════════════════════════════════════ + + /** + * Get current connection status. + */ + get connectionStatus(): ConnectionStatus { + const meta = this._collections.sessionMeta.get(this.sessionId); + return meta?.connectionStatus ?? "disconnected"; + } + + /** + * Fork session at a message boundary. + * + * @param options - Fork options + * @returns New session info + */ + async fork(options?: ForkOptions): Promise { + const response = await fetch( + `${this.options.proxyUrl}/v1/sessions/${this.sessionId}/fork`, + { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + atMessageId: options?.atMessageId ?? null, + newSessionId: options?.newSessionId ?? null, + }), + }, + ); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error( + `Failed to fork session: ${response.status} ${errorText}`, + ); + } + + return await response.json(); + } + + /** + * Register agents to respond to session messages. + * + * @param agents - Agent specifications + */ + async registerAgents(agents: AgentSpec[]): Promise { + const response = await fetch( + `${this.options.proxyUrl}/v1/sessions/${this.sessionId}/agents`, + { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ agents }), + }, + ); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error( + `Failed to register agents: ${response.status} ${errorText}`, + ); + } + } + + /** + * Unregister an agent. + * + * @param agentId - Agent identifier + */ + async unregisterAgent(agentId: string): Promise { + const response = await fetch( + `${this.options.proxyUrl}/v1/sessions/${this.sessionId}/agents/${agentId}`, + { + method: "DELETE", + }, + ); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error( + `Failed to unregister agent: ${response.status} ${errorText}`, + ); + } + } + + // ═══════════════════════════════════════════════════════════════════════ + // Lifecycle + // ═══════════════════════════════════════════════════════════════════════ + + /** + * Connect to the durable stream and start syncing. + * + * This method handles network operations only - collections are already + * created synchronously in the constructor and are immediately available. + */ + async connect(): Promise { + if (this._isConnected) return; + + try { + // Update connection status + this.updateSessionMeta((meta) => + updateConnectionStatus(meta, "connecting"), + ); + + // Skip server call when using injected sessionDB (test mode) + // This allows tests to use connect() without needing a real server + if (!this.options.sessionDB) { + // Create or get the session on the server + const response = await fetch( + `${this.options.proxyUrl}/v1/sessions/${this.sessionId}`, + { + method: "PUT", + headers: this.options.stream?.headers, + signal: this._abortController.signal, + }, + ); + + if ( + !response.ok && + response.status !== 200 && + response.status !== 201 + ) { + throw new Error(`Failed to create session: ${response.status}`); + } + } + + // Preload stream data (works for both real and mock sessionDB) + await this._db.preload(); + + this._isConnected = true; + + // Update connection status + this.updateSessionMeta((meta) => + updateConnectionStatus(meta, "connected"), + ); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + this._error = err; + this.updateSessionMeta((meta) => + updateConnectionStatus(meta, "error", { + message: err.message, + }), + ); + this.options.onError?.(this._error); + throw error; + } + } + + /** + * Pause stream sync. + */ + pause(): void { + // The stream-db handles pausing internally via the abort signal + } + + /** + * Resume stream sync. + */ + async resume(): Promise { + if (!this._isConnected) { + await this.connect(); + return; + } + + // The stream-db handles resuming internally + } + + /** + * Disconnect from the stream. + */ + disconnect(): void { + // Close stream-db (which aborts the stream) + this._db.close(); + + this._abortController.abort(); + this._isConnected = false; + + this.updateSessionMeta((meta) => + updateConnectionStatus(meta, "disconnected"), + ); + } + + /** + * Dispose the client and clean up resources. + * + * Note: We only disconnect here - we don't manually cleanup collections. + * All exposed collections could be used by application code via useLiveQuery, + * and manual cleanup would error: "Source collection was manually cleaned up + * while live query depends on it." + * + * TanStack DB will GC collections automatically when they have no subscribers. + */ + dispose(): void { + if (this._isDisposed) return; + this._isDisposed = true; + this.disconnect(); + } + + // ═══════════════════════════════════════════════════════════════════════ + // Private Helpers + // ═══════════════════════════════════════════════════════════════════════ + + /** + * Update session metadata. + */ + private updateSessionMeta( + updater: (meta: SessionMetaRow) => SessionMetaRow, + ): void { + const current = this._collections.sessionMeta.get(this.sessionId); + if (current) { + const updated = updater(current); + this._collections.sessionMeta.update(this.sessionId, (draft) => { + Object.assign(draft, updated); + }); + } + } +} + +// ═══════════════════════════════════════════════════════════════════════════ +// Factory Function +// ═══════════════════════════════════════════════════════════════════════════ + +/** + * Create a new DurableChatClient instance. + * + * @param options - Client options + * @returns New client instance + */ +export function createDurableChatClient< + TTools extends ReadonlyArray = AnyClientTool[], +>(options: DurableChatClientOptions): DurableChatClient { + return new DurableChatClient(options); +} diff --git a/packages/durable-session/src/collection.ts b/packages/durable-session/src/collection.ts new file mode 100644 index 00000000000..8d278c61b6c --- /dev/null +++ b/packages/durable-session/src/collection.ts @@ -0,0 +1,165 @@ +/** + * Session stream-db factory. + * + * Creates a stream-backed database using `@durable-streams/state` for syncing + * a session's data from Durable Streams. This replaces the previous + * `@tanstack/durable-stream-db-collection` approach. + * + * The resulting StreamDB provides typed collections (chunks, presence, agents) + * that are automatically populated from the STATE-PROTOCOL events on the stream. + */ + +import { + createStreamDB, + type StreamDB, + type StreamDBMethods, +} from "@durable-streams/state"; +import type { Collection } from "@tanstack/db"; +import { + type AgentRow, + type ChunkRow, + type RawPresenceRow, + sessionStateSchema, +} from "./schema"; +import type { SessionDBConfig } from "./types"; + +// ============================================================================ +// Session StreamDB Types +// ============================================================================ + +/** + * Collections map with correct row types. + * + * stream-db injects the primary key field at runtime, so ChunkRow and + * RawPresenceRow include the `id` field even though it's not in the schema. + * We define the correct types here. + * + * Note: The presence collection here is the raw per-device presence. + * The aggregated per-actor presence is created in client.ts. + */ +export interface SessionCollections { + chunks: Collection; + presence: Collection; + agents: Collection; +} + +/** + * Type alias for a session stream-db instance. + * + * Provides typed access to: + * - `db.collections.chunks` - All message chunks + * - `db.collections.presence` - User/agent presence + * - `db.collections.agents` - Registered agents + * + * Plus stream-db methods: + * - `db.preload()` - Wait for initial sync + * - `db.close()` - Cleanup resources + * - `db.utils.awaitTxId(txid)` - Wait for specific write to sync + */ +export type SessionDB = { + collections: SessionCollections; +} & StreamDBMethods; + +/** + * Internal type for the raw stream-db instance. + * @internal + */ +type RawSessionDB = StreamDB; + +// ============================================================================ +// Session StreamDB Factory +// ============================================================================ + +/** + * Create a stream-db instance for a session. + * + * This function is synchronous - it creates the stream handle and collections + * but does not start the stream connection. Call `db.preload()` to connect + * and wait for the initial sync to complete. + * + * The returned SessionDB instance provides: + * - `db.collections.chunks` - Root chunks collection (for messages) + * - `db.collections.presence` - Presence tracking + * - `db.collections.agents` - Registered agents + * + * @example + * ```typescript + * import { createSessionDB } from '@superset/durable-session' + * + * // Create stream-db for this session (synchronous) + * const db = createSessionDB({ + * sessionId: 'my-session', + * baseUrl: 'http://localhost:4000', + * }) + * + * // Wait for initial data sync + * await db.preload() + * + * // Access typed collections + * for (const chunk of db.collections.chunks.values()) { + * console.log(chunk.messageId, chunk.role, chunk.chunk) + * } + * + * // Cleanup when done + * db.close() + * ``` + */ +export function createSessionDB(config: SessionDBConfig): SessionDB { + const { sessionId, baseUrl, headers, signal /* liveMode */ } = config; + + // Build the stream URL for this session + const streamUrl = `${baseUrl}/v1/stream/sessions/${sessionId}`; + + // Create the stream-db instance with our session state schema (synchronous) + const rawDb: RawSessionDB = createStreamDB({ + streamOptions: { + url: streamUrl, + headers, + signal, + }, + state: sessionStateSchema, + // liveMode, + }); + + // Cast to our SessionDB type which has correctly typed collections + // (stream-db injects the primary key at runtime, so our types reflect that) + return rawDb as unknown as SessionDB; +} + +// ============================================================================ +// Utility Functions +// ============================================================================ + +/** + * Get the primary key for a chunk (used for collection lookups). + * + * Key format: `${messageId}:${seq}` + * + * @param messageId - Message identifier + * @param seq - Sequence number within message + * @returns Primary key string + */ +export function getChunkKey(messageId: string, seq: number): string { + return `${messageId}:${seq}`; +} + +/** + * Parse a chunk key into its components. + * + * @param key - Chunk key in format `${messageId}:${seq}` + * @returns Parsed components or null if invalid + */ +export function parseChunkKey( + key: string, +): { messageId: string; seq: number } | null { + const lastColonIndex = key.lastIndexOf(":"); + if (lastColonIndex === -1) return null; + + const messageId = key.slice(0, lastColonIndex); + const seqStr = key.slice(lastColonIndex + 1); + const seq = parseInt(seqStr, 10); + + if (Number.isNaN(seq)) return null; + + return { messageId, seq }; +} diff --git a/packages/durable-session/src/collections/active-generations.ts b/packages/durable-session/src/collections/active-generations.ts new file mode 100644 index 00000000000..27aaf5de0e4 --- /dev/null +++ b/packages/durable-session/src/collections/active-generations.ts @@ -0,0 +1,81 @@ +/** + * Active generations collection - derived from messages. + * + * Tracks messages that are currently being streamed (have chunks but no finish chunk). + * This is derived from the messages collection by filtering for incomplete messages. + * + * This follows the pattern: derive from materialized data with fn.select + */ + +import type { Collection } from "@tanstack/db"; +import { createLiveQueryCollection } from "@tanstack/db"; +import type { ActiveGenerationRow, MessageRow } from "../types"; + +// ============================================================================ +// Active Generations Collection +// ============================================================================ + +/** + * Options for creating an active generations collection. + */ +export interface ActiveGenerationsCollectionOptions { + /** Messages collection to derive from */ + messagesCollection: Collection; +} + +/** + * Convert an incomplete message to an active generation row. + */ +function messageToActiveGeneration(message: MessageRow): ActiveGenerationRow { + return { + messageId: message.id, + actorId: message.actorId, + startedAt: message.createdAt, + lastChunkSeq: 0, // We don't track seq in messages, so use 0 as placeholder + lastChunkAt: message.createdAt, + }; +} + +/** + * Creates the active generations collection from messages. + * + * Filters messages to only include incomplete ones (isComplete === false) + * and transforms them into ActiveGenerationRow format. + * + * Active generations are useful for: + * - Showing typing indicators + * - Tracking streaming progress + * - Resuming interrupted generations + * + * @example + * ```typescript + * const activeGenerations = createActiveGenerationsCollection({ + * sessionId: 'my-session', + * messagesCollection, + * }) + * + * // Check if anything is generating + * const isLoading = activeGenerations.size > 0 + * + * // Access active generations directly + * for (const gen of activeGenerations.values()) { + * console.log(gen.messageId, gen.actorId, gen.startedAt) + * } + * ``` + */ +export function createActiveGenerationsCollection( + options: ActiveGenerationsCollectionOptions, +): Collection { + const { messagesCollection } = options; + + // Filter messages for incomplete ones and transform to ActiveGenerationRow + // Order by createdAt to ensure chronological ordering + return createLiveQueryCollection({ + query: (q) => + q + .from({ message: messagesCollection }) + .orderBy(({ message }) => message.createdAt, "asc") + .fn.where(({ message }) => !message.isComplete) + .fn.select(({ message }) => messageToActiveGeneration(message)), + }); +} diff --git a/packages/durable-session/src/collections/index.ts b/packages/durable-session/src/collections/index.ts new file mode 100644 index 00000000000..3b4e169c362 --- /dev/null +++ b/packages/durable-session/src/collections/index.ts @@ -0,0 +1,51 @@ +/** + * Collection exports for @superset/durable-session + * + * Pipeline architecture: + * - chunks → (subquery) → messages (root materialized collection) + * - Derived collections filter messages via .fn.where() on parts + * + * All derived collections return MessageRow[], preserving full message context. + * Consumers filter message.parts to access specific part types (ToolCallPart, etc.). + */ + +// Active generations collection (derived from messages) +export { + type ActiveGenerationsCollectionOptions, + createActiveGenerationsCollection, +} from "./active-generations"; +// Messages collection (root) and derived collections +export { + createMessagesCollection, + createPendingApprovalsCollection, + createToolCallsCollection, + createToolResultsCollection, + type DerivedMessagesCollectionOptions, + type MessagesCollectionOptions, +} from "./messages"; +// Model messages collection (for LLM invocation) +export { + createModelMessagesCollection, + type ModelMessage, + type ModelMessagesCollectionOptions, +} from "./model-messages"; +// Aggregated presence collection (derived from raw per-device presence) +export { + createPresenceCollection, + type PresenceCollectionOptions, +} from "./presence"; +// Session metadata collection (local state) +export { + createInitialSessionMeta, + createSessionMetaCollectionOptions, + type SessionMetaCollectionOptions, + updateConnectionStatus, + updateSyncProgress, +} from "./session-meta"; +// Session statistics collection (aggregated from chunks) +export { + computeSessionStats, + createEmptyStats, + createSessionStatsCollection, + type SessionStatsCollectionOptions, +} from "./session-stats"; diff --git a/packages/durable-session/src/collections/messages.ts b/packages/durable-session/src/collections/messages.ts new file mode 100644 index 00000000000..8fdad1e450a --- /dev/null +++ b/packages/durable-session/src/collections/messages.ts @@ -0,0 +1,161 @@ +/** + * Messages collection - core live query pipeline. + * + * Architecture: + * - chunks → (groupBy messageId + count/min) → fn.select(materialize) + * - Derived collections use .fn.where() to filter by message parts + * + * Note: The upstream @tanstack/db `collect` aggregate is not yet published. + * Instead, we use groupBy + count as a change discriminator, then + * imperatively filter the chunks collection inside fn.select to gather + * all chunks for each message. + */ + +import type { ToolCallPart } from "@tanstack/ai"; +import type { Collection } from "@tanstack/db"; +import { count, createLiveQueryCollection, min } from "@tanstack/db"; +import { materializeMessage } from "../materialize"; +import type { ChunkRow } from "../schema"; +import type { MessageRow } from "../types"; + +// ============================================================================ +// Messages Collection (Root) +// ============================================================================ + +/** + * Options for creating a messages collection. + */ +export interface MessagesCollectionOptions { + /** Chunks collection from stream-db */ + chunksCollection: Collection; +} + +/** + * Creates the messages collection with inline subquery for chunk aggregation. + * + * This is the root materialized collection in the live query pipeline. + * All derived collections (toolCalls, pendingApprovals, etc.) derive from this. + * + * Uses groupBy + count/min as change discriminators, then fn.select + * imperatively gathers chunks from the collection per messageId. + */ +export function createMessagesCollection( + options: MessagesCollectionOptions, +): Collection { + const { chunksCollection } = options; + + return createLiveQueryCollection({ + query: (q) => { + // Subquery: group chunks by messageId with aggregates for change detection + const grouped = q + .from({ chunk: chunksCollection }) + .groupBy(({ chunk }) => chunk.messageId) + .select(({ chunk }) => ({ + messageId: chunk.messageId, + // min() handles strings (ISO 8601 sort lexicographically) + startedAt: min(chunk.createdAt), + // Count as discriminator to force re-evaluation when chunks change + rowCount: count(chunk), + })); + + // Main query: materialize messages from chunks + return q + .from({ grouped }) + .orderBy(({ grouped }) => grouped.startedAt, "asc") + .fn.select(({ grouped }) => { + // Imperatively gather all chunks for this messageId + const rows = [...chunksCollection.values()].filter( + (c) => (c as ChunkRow).messageId === grouped.messageId, + ) as ChunkRow[]; + return materializeMessage(rows); + }); + }, + getKey: (row) => row.id, + }); +} + +// ============================================================================ +// Derived Collections +// ============================================================================ + +/** + * Options for creating a derived collection from messages. + */ +export interface DerivedMessagesCollectionOptions { + /** Messages collection to derive from */ + messagesCollection: Collection; +} + +/** + * Creates a collection of messages that contain tool calls. + * + * Filters messages where at least one part has type 'tool-call'. + * The collection is lazy - filtering only runs when accessed. + */ +export function createToolCallsCollection( + options: DerivedMessagesCollectionOptions, +): Collection { + const { messagesCollection } = options; + + return createLiveQueryCollection({ + query: (q) => + q + .from({ message: messagesCollection }) + .fn.where(({ message }) => + message.parts.some((p): p is ToolCallPart => p.type === "tool-call"), + ) + .orderBy(({ message }) => message.createdAt, "asc"), + getKey: (row) => row.id, + }); +} + +/** + * Creates a collection of messages that have pending approval requests. + * + * Filters messages where at least one tool call part has: + * - approval.needsApproval === true + * - approval.approved === undefined (not yet responded) + */ +export function createPendingApprovalsCollection( + options: DerivedMessagesCollectionOptions, +): Collection { + const { messagesCollection } = options; + + return createLiveQueryCollection({ + query: (q) => + q + .from({ message: messagesCollection }) + .fn.where(({ message }) => + message.parts.some( + (p): p is ToolCallPart => + p.type === "tool-call" && + p.approval?.needsApproval === true && + p.approval.approved === undefined, + ), + ) + .orderBy(({ message }) => message.createdAt, "asc"), + getKey: (row) => row.id, + }); +} + +/** + * Creates a collection of messages that contain tool results. + * + * Filters messages where at least one part has type 'tool-result'. + */ +export function createToolResultsCollection( + options: DerivedMessagesCollectionOptions, +): Collection { + const { messagesCollection } = options; + + return createLiveQueryCollection({ + query: (q) => + q + .from({ message: messagesCollection }) + .fn.where(({ message }) => + message.parts.some((p) => p.type === "tool-result"), + ) + .orderBy(({ message }) => message.createdAt, "asc"), + getKey: (row) => row.id, + }); +} diff --git a/packages/durable-session/src/collections/model-messages.ts b/packages/durable-session/src/collections/model-messages.ts new file mode 100644 index 00000000000..d5dfb2fdea1 --- /dev/null +++ b/packages/durable-session/src/collections/model-messages.ts @@ -0,0 +1,82 @@ +/** + * Model messages collection - LLM-ready message history. + * + * Derives from the messages collection: + * 1. Filters to complete messages only (isComplete === true) + * 2. Converts to { role, content } format expected by LLMs + * 3. Orders chronologically + */ + +import type { Collection } from "@tanstack/db"; +import { createLiveQueryCollection, eq } from "@tanstack/db"; +import { extractTextContent } from "../materialize"; +import type { MessageRow } from "../types"; + +// ============================================================================ +// Types +// ============================================================================ + +/** + * Message format expected by LLMs (OpenAI/Anthropic compatible). + */ +export interface ModelMessage { + id: string; + role: "user" | "assistant" | "system"; + content: string; +} + +/** + * Options for creating a model messages collection. + */ +export interface ModelMessagesCollectionOptions { + /** Messages collection (from createMessagesPipeline) */ + messagesCollection: Collection; +} + +// ============================================================================ +// Factory +// ============================================================================ + +/** + * Creates a collection of LLM-ready messages. + * + * This derived collection: + * - Filters to complete messages only (streaming messages excluded) + * - Extracts text content from message parts + * - Provides chronologically ordered { role, content } objects + * + * @example + * ```typescript + * const { messages } = createMessagesPipeline({ ... }) + * const modelMessages = createModelMessagesCollection({ + * messagesCollection: messages, + * }) + * + * // Get LLM-ready history + * // Note: toArray is a getter (property), not a method + * const history = modelMessages.toArray.map(m => ({ + * role: m.role, + * content: m.content, + * })) + * ``` + */ +export function createModelMessagesCollection( + options: ModelMessagesCollectionOptions, +): Collection { + const { messagesCollection } = options; + + return createLiveQueryCollection({ + query: (q) => + q + .from({ message: messagesCollection }) + .where(({ message }) => eq(message.isComplete, true)) + .orderBy(({ message }) => message.createdAt, "asc") + .fn.select(({ message }) => ({ + id: message.id, + role: message.role, + content: extractTextContent(message), + })), + getKey: (row) => row.id, + startSync: true, + }); +} diff --git a/packages/durable-session/src/collections/presence.ts b/packages/durable-session/src/collections/presence.ts new file mode 100644 index 00000000000..99bcec52ca4 --- /dev/null +++ b/packages/durable-session/src/collections/presence.ts @@ -0,0 +1,79 @@ +/** + * Aggregated presence collection - derived from raw per-device presence. + * + * The raw presence from stream-db tracks each (actorId, deviceId) pair. + * This collection aggregates devices per actor, filtering for online status, + * to provide a simple "who's online" view. + * + * Note: The upstream @tanstack/db `collect` aggregate is not yet published. + * Instead, we use groupBy + count as a change discriminator, then + * imperatively gather device IDs inside fn.select. + */ + +import type { Collection } from "@tanstack/db"; +import { count, createLiveQueryCollection, eq } from "@tanstack/db"; +import type { PresenceRow, RawPresenceRow } from "../schema"; + +// ============================================================================ +// Aggregated Presence Collection +// ============================================================================ + +/** + * Options for creating an aggregated presence collection. + */ +export interface PresenceCollectionOptions { + /** Session identifier */ + sessionId: string; + /** Raw presence collection from stream-db (per-device records) */ + rawPresenceCollection: Collection; +} + +/** + * Creates the aggregated presence collection. + * + * Uses a live query pipeline to: + * 1. Filter raw presence for status='online' + * 2. Group by actorId + * 3. Use fn.select to imperatively collect device IDs + * + * The result is one row per online actor, with their device count. + */ +export function createPresenceCollection( + options: PresenceCollectionOptions, +): Collection { + const { rawPresenceCollection } = options; + + return createLiveQueryCollection({ + query: (q) => { + // Subquery: filter for online, group by actorId, count for change detection + const grouped = q + .from({ presence: rawPresenceCollection }) + .where(({ presence }) => eq(presence.status, "online")) + .groupBy(({ presence }) => presence.actorId) + .select(({ presence }) => ({ + actorId: presence.actorId, + deviceCount: count(presence.deviceId), + })); + + // Main query: imperatively gather device info per actor + return q.from({ grouped }).fn.select(({ grouped }) => { + // Get all online presence rows for this actor + const actorPresence = [...rawPresenceCollection.values()].filter( + (p) => + (p as RawPresenceRow).actorId === grouped.actorId && + (p as RawPresenceRow).status === "online", + ) as RawPresenceRow[]; + + const first = actorPresence[0]; + return { + actorId: grouped.actorId as string, + actorType: (first?.actorType ?? "user") as "user" | "agent", + name: first?.name, + deviceIds: actorPresence.map((p) => p.deviceId), + deviceCount: actorPresence.length, + }; + }); + }, + startSync: true, + }); +} diff --git a/packages/durable-session/src/collections/session-meta.ts b/packages/durable-session/src/collections/session-meta.ts new file mode 100644 index 00000000000..eb53fc114c5 --- /dev/null +++ b/packages/durable-session/src/collections/session-meta.ts @@ -0,0 +1,107 @@ +/** + * Session metadata collection - local state collection. + * + * Tracks connection state and sync progress. + * This is a local-only collection, not derived from stream. + */ + +import { localOnlyCollectionOptions } from "@tanstack/db"; +import type { ConnectionStatus, SessionMetaRow } from "../types"; + +/** + * Options for creating a session meta collection. + */ +export interface SessionMetaCollectionOptions { + /** Session identifier */ + sessionId: string; +} + +/** + * Creates collection config for the session metadata collection. + * + * This collection stores local state: + * - connectionStatus + * - lastSyncedOffset + * - lastSyncedAt + * - error + * + * The collection is a single-row collection keyed by sessionId. + * + * @example + * ```typescript + * import { createSessionMetaCollectionOptions } from '@superset/durable-session' + * import { createCollection } from '@tanstack/db' + * + * const sessionMetaCollection = createCollection( + * createSessionMetaCollectionOptions({ + * sessionId: 'my-session', + * }) + * ) + * ``` + */ +export function createSessionMetaCollectionOptions( + options: SessionMetaCollectionOptions, +) { + const { sessionId } = options; + + return localOnlyCollectionOptions({ + id: `session-meta:${sessionId}`, + getKey: (meta) => meta.sessionId, + }); +} + +/** + * Create initial session metadata. + * + * @param sessionId - Session identifier + * @returns Initial session metadata row + */ +export function createInitialSessionMeta(sessionId: string): SessionMetaRow { + return { + sessionId, + connectionStatus: "disconnected", + lastSyncedTxId: null, + lastSyncedAt: null, + error: null, + }; +} + +/** + * Update session metadata with new connection status. + * + * @param meta - Current metadata + * @param status - New connection status + * @param error - Optional error information + * @returns Updated metadata + */ +export function updateConnectionStatus( + meta: SessionMetaRow, + status: ConnectionStatus, + error?: { message: string; code?: string } | null, +): SessionMetaRow { + return { + ...meta, + connectionStatus: status, + error: error ?? (status === "connected" ? null : meta.error), + }; +} + +/** + * Update session metadata with sync progress. + * + * @param meta - Current metadata + * @param txId - Last synced transaction ID + * @returns Updated metadata + */ +export function updateSyncProgress( + meta: SessionMetaRow, + txId: string, +): SessionMetaRow { + return { + ...meta, + lastSyncedTxId: txId, + lastSyncedAt: new Date(), + connectionStatus: "connected", + error: null, + }; +} diff --git a/packages/durable-session/src/collections/session-stats.ts b/packages/durable-session/src/collections/session-stats.ts new file mode 100644 index 00000000000..d2244427a73 --- /dev/null +++ b/packages/durable-session/src/collections/session-stats.ts @@ -0,0 +1,244 @@ +/** + * Session statistics collection - aggregated from chunks. + * + * Computes aggregate statistics from the stream by: + * 1. Counting all chunks for the session (as change discriminator) + * 2. Imperatively grouping by messageId and computing stats + * + * Uses TanStack AI's MessagePart types for type-safe filtering. + */ + +import type { ToolCallPart } from "@tanstack/ai"; +import type { Collection } from "@tanstack/db"; +import { count, createLiveQueryCollection } from "@tanstack/db"; +import { materializeMessage, parseChunk } from "../materialize"; +import type { ChunkRow } from "../schema"; +import type { MessageRow, SessionStatsRow } from "../types"; + +// ============================================================================ +// Session Stats Collection +// ============================================================================ + +/** + * Options for creating a session stats collection. + */ +export interface SessionStatsCollectionOptions { + /** Session identifier */ + sessionId: string; + /** Chunks collection from stream-db */ + chunksCollection: Collection; +} + +/** + * Creates the session stats collection. + * + * Uses groupBy with count as a change discriminator, then fn.select + * imperatively gathers all chunks and computes stats. + */ +export function createSessionStatsCollection( + options: SessionStatsCollectionOptions, +): Collection { + const { sessionId, chunksCollection } = options; + + // Single-stage: group by sessionId (constant), count for change detection, compute stats in fn.select + const collectedRows = createLiveQueryCollection({ + query: (q) => + q + .from({ chunk: chunksCollection }) + .groupBy(() => sessionId) + .select(({ chunk }) => ({ + sessionId, + rowCount: count(chunk), + })), + }); + + return createLiveQueryCollection({ + query: (q) => + q.from({ collected: collectedRows }).fn.select(({ collected }) => { + // Imperatively gather all chunks + const rows = [...chunksCollection.values()] as ChunkRow[]; + return computeSessionStats(collected.sessionId as string, rows); + }), + }); +} + +/** + * Group chunk rows by messageId. + */ +function groupRowsByMessage(rows: ChunkRow[]): Map { + const grouped = new Map(); + + for (const row of rows) { + const existing = grouped.get(row.messageId); + if (existing) { + existing.push(row); + } else { + grouped.set(row.messageId, [row]); + } + } + + return grouped; +} + +/** + * Compute session statistics from chunk rows. + * + * Materializes messages and counts parts by type to derive statistics. + * + * @param sessionId - Session identifier + * @param rows - All chunk rows + * @returns Computed statistics + */ +export function computeSessionStats( + sessionId: string, + rows: ChunkRow[], +): SessionStatsRow { + if (rows.length === 0) { + return createEmptyStats(sessionId); + } + + // Group rows by message + const grouped = groupRowsByMessage(rows); + + // Materialize messages for counting + const messages: MessageRow[] = []; + for (const [, messageRows] of grouped) { + try { + messages.push(materializeMessage(messageRows)); + } catch { + // Skip invalid messages + } + } + + // Count message types and extract part counts + let userMessageCount = 0; + let assistantMessageCount = 0; + let toolCallCount = 0; + let pendingApprovalCount = 0; + let activeGenerationCount = 0; + let firstMessageAt: Date | null = null; + let lastMessageAt: Date | null = null; + + for (const msg of messages) { + // Count by role + if (msg.role === "user") { + userMessageCount++; + } else if (msg.role === "assistant") { + assistantMessageCount++; + } + + // Track timestamps + if (!firstMessageAt || msg.createdAt < firstMessageAt) { + firstMessageAt = msg.createdAt; + } + if (!lastMessageAt || msg.createdAt > lastMessageAt) { + lastMessageAt = msg.createdAt; + } + + // Count tool calls and pending approvals from parts + for (const part of msg.parts) { + if (part.type === "tool-call") { + toolCallCount++; + const toolCallPart = part as ToolCallPart; + if ( + toolCallPart.approval?.needsApproval === true && + toolCallPart.approval.approved === undefined + ) { + pendingApprovalCount++; + } + } + } + + // Count active generations (incomplete messages) + if (!msg.isComplete) { + activeGenerationCount++; + } + } + + // Extract token usage from chunks + const { totalTokens, promptTokens, completionTokens } = + extractTokenUsage(rows); + + return { + sessionId, + messageCount: messages.length, + userMessageCount, + assistantMessageCount, + toolCallCount, + approvalCount: pendingApprovalCount, + totalTokens, + promptTokens, + completionTokens, + activeGenerationCount, + firstMessageAt, + lastMessageAt, + }; +} + +/** + * Extract token usage from chunk rows. + * + * @param rows - Chunk rows to extract from + * @returns Token usage counts + */ +function extractTokenUsage(rows: ChunkRow[]): { + totalTokens: number; + promptTokens: number; + completionTokens: number; +} { + let totalTokens = 0; + let promptTokens = 0; + let completionTokens = 0; + + for (const row of rows) { + const chunk = parseChunk(row.chunk); + if (!chunk) continue; + + // Look for usage information in chunks + const usage = ( + chunk as { + usage?: { + totalTokens?: number; + promptTokens?: number; + completionTokens?: number; + total_tokens?: number; + prompt_tokens?: number; + completion_tokens?: number; + }; + } + ).usage; + + if (usage) { + // Handle both camelCase and snake_case formats + totalTokens += usage.totalTokens ?? usage.total_tokens ?? 0; + promptTokens += usage.promptTokens ?? usage.prompt_tokens ?? 0; + completionTokens += + usage.completionTokens ?? usage.completion_tokens ?? 0; + } + } + + return { totalTokens, promptTokens, completionTokens }; +} + +/** + * Create empty session statistics. + * + * @param sessionId - Session identifier + * @returns Empty statistics row + */ +export function createEmptyStats(sessionId: string): SessionStatsRow { + return { + sessionId, + messageCount: 0, + userMessageCount: 0, + assistantMessageCount: 0, + toolCallCount: 0, + approvalCount: 0, + totalTokens: 0, + promptTokens: 0, + completionTokens: 0, + activeGenerationCount: 0, + firstMessageAt: null, + lastMessageAt: null, + }; +} diff --git a/packages/durable-session/src/index.ts b/packages/durable-session/src/index.ts new file mode 100644 index 00000000000..de3243aa7da --- /dev/null +++ b/packages/durable-session/src/index.ts @@ -0,0 +1,181 @@ +/** + * @superset/durable-session + * + * Framework-agnostic durable chat client backed by TanStack DB and Durable Streams. + * + * This package provides: + * - TanStack AI-compatible API for chat applications + * - Durable persistence via Durable Streams + * - Real-time sync across tabs, devices, and users + * - Multi-agent support with webhook registration + * - Reactive collections for custom UI needs + * + * Architecture: + * - chunks → (subquery) → messages (root materialized collection) + * - Derived collections filter messages via .fn.where() on parts + * - All collections return MessageRow[], preserving full message context + * - Consumers filter message.parts to access specific part types + * + * @example + * ```typescript + * import { DurableChatClient } from '@superset/durable-session' + * + * const client = new DurableChatClient({ + * sessionId: 'my-session', + * proxyUrl: 'http://localhost:4000', + * }) + * + * await client.connect() + * + * // TanStack AI-compatible API + * await client.sendMessage('Hello!') + * console.log(client.messages) + * + * // Access collections directly + * for (const message of client.collections.messages.values()) { + * console.log(message.id, message.role, message.parts) + * } + * + * // Filter tool calls from message parts + * for (const message of client.collections.toolCalls.values()) { + * for (const part of message.parts) { + * if (part.type === 'tool-call') { + * console.log(part.name, part.state, part.arguments) + * } + * } + * } + * + * // Check for pending approvals + * for (const message of client.collections.pendingApprovals.values()) { + * for (const part of message.parts) { + * if (part.type === 'tool-call' && part.approval?.needsApproval) { + * console.log(`Approval needed: ${part.name}`) + * } + * } + * } + * ``` + * + * @packageDocumentation + */ + +// ============================================================================ +// Client +// ============================================================================ + +export { createDurableChatClient, DurableChatClient } from "./client"; + +// ============================================================================ +// Schema (STATE-PROTOCOL) +// ============================================================================ + +export { + type AgentRow, + type AgentValue, + agentValueSchema, + type ChunkRow, + type ChunkValue, + chunkValueSchema, + type PresenceRow, + type PresenceValue, + presenceValueSchema, + type RawPresenceRow, + type SessionStateSchema, + sessionStateSchema, +} from "./schema"; + +// ============================================================================ +// Types +// ============================================================================ + +export type { + // Active generation types + ActiveGenerationRow, + // Actor types + ActorType, + AgentSpec, + // Agent types + AgentTrigger, + ApprovalResponseInput, + // Session types + ConnectionStatus, + // Configuration types + DurableChatClientOptions, + // Collection types + DurableChatCollections, + // Fork types + ForkOptions, + ForkResult, + // Re-exported TanStack AI types for consumer convenience + MessagePart, + // Message types + MessageRole, + MessageRow, + SessionDBConfig, + SessionMetaRow, + SessionStatsRow, + TextPart, + ThinkingPart, + ToolCallPart, + // Input types + ToolResultInput, + ToolResultPart, +} from "./types"; + +// ============================================================================ +// Session DB Factory +// ============================================================================ + +export { + createSessionDB, + getChunkKey, + parseChunkKey, + type SessionDB, +} from "./collection"; + +// ============================================================================ +// Collection Factories +// ============================================================================ + +export { + type ActiveGenerationsCollectionOptions, + computeSessionStats, + // Active generations collection + createActiveGenerationsCollection, + createEmptyStats, + createInitialSessionMeta, + // Messages collection (root) and derived collections + createMessagesCollection, + // Model messages collection (for LLM invocation) + createModelMessagesCollection, + createPendingApprovalsCollection, + // Aggregated presence collection + createPresenceCollection, + // Session metadata collection (local state) + createSessionMetaCollectionOptions, + // Session statistics collection + createSessionStatsCollection, + createToolCallsCollection, + createToolResultsCollection, + type DerivedMessagesCollectionOptions, + type MessagesCollectionOptions, + type ModelMessage, + type ModelMessagesCollectionOptions, + type PresenceCollectionOptions, + type SessionMetaCollectionOptions, + type SessionStatsCollectionOptions, + updateConnectionStatus, + updateSyncProgress, +} from "./collections"; + +// ============================================================================ +// Materialization +// ============================================================================ + +export { + extractTextContent, + isAssistantMessage, + isUserMessage, + materializeMessage, + messageRowToUIMessage, + parseChunk, +} from "./materialize"; diff --git a/packages/durable-session/src/materialize.ts b/packages/durable-session/src/materialize.ts new file mode 100644 index 00000000000..d6472072797 --- /dev/null +++ b/packages/durable-session/src/materialize.ts @@ -0,0 +1,248 @@ +/** + * Message materialization from stream chunks. + * + * Handles two formats: + * 1. User messages: Single row with {type: 'whole-message', message: UIMessage} + * 2. Assistant messages: Multiple rows with TanStack AI StreamChunks + * + * Chunk processing is delegated to TanStack AI's StreamProcessor, which handles: + * - Text content accumulation + * - Tool call parsing (including arguments streaming) + * - Tool result handling + * - Approval request tracking + * + * The output is a MessageRow with parts using TanStack AI's MessagePart types. + * Derived collections filter on these parts rather than using separate extraction. + */ + +import type { StreamChunk, UIMessage } from "@tanstack/ai"; +import { StreamProcessor } from "@tanstack/ai"; +import type { ChunkRow } from "./schema"; +import type { + DurableStreamChunk, + MessageRole, + MessageRow, + WholeMessageChunk, +} from "./types"; + +// ============================================================================ +// Type Guards +// ============================================================================ + +function isDoneChunk(chunk: StreamChunk): boolean { + return chunk.type === "RUN_FINISHED"; +} + +/** + * Type guard for WholeMessageChunk. + */ +function isWholeMessageChunk( + chunk: DurableStreamChunk | null, +): chunk is WholeMessageChunk { + return chunk !== null && chunk.type === "whole-message"; +} + +// ============================================================================ +// Message Materialization +// ============================================================================ + +/** + * Parse a JSON-encoded chunk string. + * + * @param chunkJson - JSON string containing DurableStreamChunk + * @returns Parsed chunk or null if invalid + */ +export function parseChunk(chunkJson: string): DurableStreamChunk | null { + try { + return JSON.parse(chunkJson) as DurableStreamChunk; + } catch { + return null; + } +} + +/** + * Materialize a whole message from a single row. + * User messages are stored as complete UIMessage objects. + */ +function materializeWholeMessage( + row: ChunkRow, + chunk: WholeMessageChunk, +): MessageRow { + const { message } = chunk; + + return { + id: message.id, + role: message.role as MessageRole, + parts: message.parts, + actorId: row.actorId, + isComplete: true, + createdAt: message.createdAt + ? new Date(message.createdAt) + : new Date(row.createdAt), + }; +} + +/** + * Materialize an assistant message from streamed chunks. + * Uses TanStack AI's StreamProcessor to process chunks. + */ +function materializeAssistantMessage(rows: ChunkRow[]): MessageRow { + const sorted = [...rows].sort((a, b) => a.seq - b.seq); + const first = sorted[0] as ChunkRow; + + // Create processor and start assistant message + const processor = new StreamProcessor(); + processor.startAssistantMessage(); + + let isComplete = false; + + for (const row of sorted) { + const chunk = parseChunk(row.chunk); + if (!chunk) continue; + + // Extract type as string for legacy/proxy chunk type checks + const type = (chunk as { type: string }).type as string; + + // Skip legacy wrapper chunks (for backward compatibility) + if (type === "message-start" || type === "message-end") { + if (type === "message-end") { + isComplete = true; + } + continue; + } + + // Skip whole-message chunks (shouldn't be in assistant messages, but guard) + if (isWholeMessageChunk(chunk)) continue; + + // Process TanStack AI StreamChunk + try { + processor.processChunk(chunk as StreamChunk); + } catch { + // Skip chunks that can't be processed + } + + if (isDoneChunk(chunk as StreamChunk)) { + isComplete = true; + } + + // Also check for stop/error chunks (stop is from our proxy, not in TanStack AI types) + if (type === "stop" || type === "error" || type === "RUN_ERROR") { + isComplete = true; + } + } + + // Finalize if complete + if (isComplete) { + processor.finalizeStream(); + } + + // Get the materialized UIMessage + const messages = processor.getMessages(); + const message = messages[messages.length - 1]; + + return { + id: first.messageId, + role: first.role as MessageRole, + parts: message?.parts ?? [], + actorId: first.actorId, + isComplete, + createdAt: new Date(first.createdAt), + }; +} + +/** + * Materialize a MessageRow from collected chunk rows. + * + * Handles two formats: + * 1. User messages: Single row with {type: 'whole-message', message: UIMessage} + * 2. Assistant messages: Multiple rows with TanStack AI StreamChunks + * + * @param rows - Chunk rows for a single message + * @returns Materialized message row + */ +export function materializeMessage(rows: ChunkRow[]): MessageRow { + if (!rows || rows.length === 0) { + throw new Error("Cannot materialize message from empty rows"); + } + + // Sort by seq to ensure correct order + const sorted = [...rows].sort((a, b) => a.seq - b.seq); + const firstRow = sorted[0] as ChunkRow; + const firstChunk = parseChunk(firstRow.chunk); + + if (!firstChunk) { + throw new Error("Failed to parse first chunk"); + } + + // Check if this is a whole message + if (isWholeMessageChunk(firstChunk)) { + return materializeWholeMessage(firstRow, firstChunk); + } + + // Otherwise, process as streamed assistant message + return materializeAssistantMessage(sorted); +} + +// ============================================================================ +// Content Extraction Helpers +// ============================================================================ + +/** + * Extract text content from a UIMessage or MessageRow. + * + * @param message - Message to extract from + * @returns Combined text content + */ +export function extractTextContent(message: { + parts: Array<{ type: string; text?: string; content?: string }>; +}): string { + return message.parts + .filter((p) => p.type === "text") + .map((p) => p.text ?? p.content ?? "") + .join(""); +} + +/** + * Check if a message row is from a user. + * + * @param row - Message row to check + * @returns Whether the message is from a user + */ +export function isUserMessage(row: MessageRow): boolean { + return row.role === "user"; +} + +/** + * Check if a message row is from an assistant/agent. + * + * @param row - Message row to check + * @returns Whether the message is from an assistant + */ +export function isAssistantMessage(row: MessageRow): boolean { + return row.role === "assistant"; +} + +// ============================================================================ +// UIMessage Conversion +// ============================================================================ + +/** + * Convert a MessageRow to a TanStack AI UIMessage. + * + * This is a pure transformation function that maps the internal MessageRow + * representation to the public UIMessage interface expected by TanStack AI. + * + * @param row - MessageRow from the messages collection + * @returns UIMessage compatible with TanStack AI + */ +export function messageRowToUIMessage( + row: MessageRow, +): UIMessage & { actorId: string } { + return { + id: row.id, + role: row.role, + parts: row.parts, + createdAt: row.createdAt, + actorId: row.actorId, + }; +} diff --git a/packages/ai-chat/src/components/ChatInput/ChatInput.tsx b/packages/durable-session/src/react/components/ChatInput/ChatInput.tsx similarity index 100% rename from packages/ai-chat/src/components/ChatInput/ChatInput.tsx rename to packages/durable-session/src/react/components/ChatInput/ChatInput.tsx diff --git a/packages/ai-chat/src/components/ChatInput/index.ts b/packages/durable-session/src/react/components/ChatInput/index.ts similarity index 100% rename from packages/ai-chat/src/components/ChatInput/index.ts rename to packages/durable-session/src/react/components/ChatInput/index.ts diff --git a/packages/ai-chat/src/components/PresenceBar/PresenceBar.tsx b/packages/durable-session/src/react/components/PresenceBar/PresenceBar.tsx similarity index 96% rename from packages/ai-chat/src/components/PresenceBar/PresenceBar.tsx rename to packages/durable-session/src/react/components/PresenceBar/PresenceBar.tsx index b43af801c96..a29ceb8b010 100644 --- a/packages/ai-chat/src/components/PresenceBar/PresenceBar.tsx +++ b/packages/durable-session/src/react/components/PresenceBar/PresenceBar.tsx @@ -4,7 +4,12 @@ import { Avatar, AvatarFallback, AvatarImage } from "@superset/ui/avatar"; import { cn } from "@superset/ui/utils"; -import type { PresenceUser } from "../../types"; + +export interface PresenceUser { + userId: string; + name: string; + image?: string; +} export interface PresenceBarProps { viewers: PresenceUser[]; diff --git a/packages/durable-session/src/react/components/PresenceBar/index.ts b/packages/durable-session/src/react/components/PresenceBar/index.ts new file mode 100644 index 00000000000..fd8ff9b395c --- /dev/null +++ b/packages/durable-session/src/react/components/PresenceBar/index.ts @@ -0,0 +1,5 @@ +export { + PresenceBar, + type PresenceBarProps, + type PresenceUser, +} from "./PresenceBar"; diff --git a/packages/durable-session/src/react/components/index.ts b/packages/durable-session/src/react/components/index.ts new file mode 100644 index 00000000000..83c48c090a0 --- /dev/null +++ b/packages/durable-session/src/react/components/index.ts @@ -0,0 +1,6 @@ +export { ChatInput, type ChatInputProps } from "./ChatInput"; +export { + PresenceBar, + type PresenceBarProps, + type PresenceUser, +} from "./PresenceBar"; diff --git a/packages/durable-session/src/react/index.ts b/packages/durable-session/src/react/index.ts new file mode 100644 index 00000000000..8f6e6628049 --- /dev/null +++ b/packages/durable-session/src/react/index.ts @@ -0,0 +1,95 @@ +/** + * @superset/durable-session/react + * + * React bindings for durable chat client backed by TanStack DB and Durable Streams. + * + * This package provides React hooks for building durable chat applications with: + * - TanStack AI-compatible API (drop-in replacement for useChat) + * - Automatic React state management + * - Access to reactive collections for custom queries + * - Multi-agent support + * + * @example + * ```typescript + * import { useDurableChat } from '@superset/durable-session/react' + * + * function Chat() { + * const { messages, sendMessage, isLoading } = useDurableChat({ + * sessionId: 'my-session', + * proxyUrl: 'http://localhost:4000', + * }) + * + * return ( + *
+ * {messages.map(m => )} + * + *
+ * ) + * } + * ``` + * + * @packageDocumentation + */ + +// ============================================================================ +// Hooks +// ============================================================================ + +export { useDurableChat } from "./use-durable-chat"; + +// ============================================================================ +// Components +// ============================================================================ + +export { ChatInput, type ChatInputProps } from "./components/ChatInput"; +export { + PresenceBar, + type PresenceBarProps, + type PresenceUser, +} from "./components/PresenceBar"; + +// ============================================================================ +// Types +// ============================================================================ + +export type { UseDurableChatOptions, UseDurableChatReturn } from "./types"; + +// ============================================================================ +// Re-exports from durable-session +// ============================================================================ + +export { + type ActiveGenerationRow, + // Types + type ActorType, + type AgentRow, + type AgentSpec, + type AgentTrigger, + type ApprovalResponseInput, + type ChunkRow, + type ConnectionStatus, + createDurableChatClient, + // Client + DurableChatClient, + type DurableChatClientOptions, + type DurableChatCollections, + // Materialization helpers + extractTextContent, + type ForkOptions, + type ForkResult, + isAssistantMessage, + isUserMessage, + // Re-exported TanStack AI types for consumer convenience + type MessagePart, + type MessageRole, + type MessageRow, + type PresenceRow, + type RawPresenceRow, + type SessionMetaRow, + type SessionStatsRow, + type TextPart, + type ThinkingPart, + type ToolCallPart, + type ToolResultInput, + type ToolResultPart, +} from ".."; diff --git a/packages/durable-session/src/react/types.ts b/packages/durable-session/src/react/types.ts new file mode 100644 index 00000000000..cea457149ed --- /dev/null +++ b/packages/durable-session/src/react/types.ts @@ -0,0 +1,120 @@ +/** + * React-specific types for @superset/durable-session/react + */ + +import type { AnyClientTool, UIMessage } from "@tanstack/ai"; +import type { + AgentSpec, + ApprovalResponseInput, + ConnectionStatus, + DurableChatClient, + DurableChatClientOptions, + DurableChatCollections, + ForkOptions, + ForkResult, + ToolResultInput, +} from ".."; + +/** + * Options for the useDurableChat hook. + */ +export interface UseDurableChatOptions< + TTools extends ReadonlyArray = AnyClientTool[], +> extends Partial> { + /** + * Whether to automatically connect on mount. + * @default true + */ + autoConnect?: boolean; + + /** + * Pre-created client instance. + * If provided, the hook will use this client instead of creating a new one. + * Useful for testing or when you need to share a client between components. + */ + client?: DurableChatClient; +} + +/** + * Return value from useDurableChat hook. + */ +export interface UseDurableChatReturn< + TTools extends ReadonlyArray = AnyClientTool[], +> { + // ═══════════════════════════════════════════════════════════════════════ + // TanStack AI useChat compatible + // ═══════════════════════════════════════════════════════════════════════ + + /** All messages in the conversation */ + messages: UIMessage[]; + + /** Send a user message */ + sendMessage: (content: string) => Promise; + + /** Append a message to the conversation */ + append: ( + message: UIMessage | { role: string; content: string }, + ) => Promise; + + /** Reload and regenerate the last response */ + reload: () => Promise; + + /** Stop all active generations */ + stop: () => void; + + /** Clear all messages (local only) */ + clear: () => void; + + /** Whether any generation is currently active */ + isLoading: boolean; + + /** Current error, if any */ + error: Error | undefined; + + /** Add a tool result */ + addToolResult: (result: ToolResultInput) => Promise; + + /** Add an approval response */ + addToolApprovalResponse: (response: ApprovalResponseInput) => Promise; + + // ═══════════════════════════════════════════════════════════════════════ + // Durable extensions + // ═══════════════════════════════════════════════════════════════════════ + + /** + * The underlying DurableChatClient instance. + * Always available - created synchronously on hook initialization. + */ + client: DurableChatClient; + + /** + * All collections for custom queries. + * Always available - use directly with useLiveQuery. + * Data syncs when connectionStatus is 'connected'. + */ + collections: DurableChatCollections; + + /** Current connection status */ + connectionStatus: ConnectionStatus; + + /** Fork the session at a message boundary */ + fork: (options?: ForkOptions) => Promise; + + /** Register agents to respond to session messages */ + registerAgents: (agents: AgentSpec[]) => Promise; + + /** Unregister an agent */ + unregisterAgent: (agentId: string) => Promise; + + /** Connect to the stream (if not auto-connected) */ + connect: () => Promise; + + /** Disconnect from the stream */ + disconnect: () => void; + + /** Pause stream sync */ + pause: () => void; + + /** Resume stream sync */ + resume: () => Promise; +} diff --git a/packages/durable-session/src/react/use-durable-chat.ts b/packages/durable-session/src/react/use-durable-chat.ts new file mode 100644 index 00000000000..0c1ba07bece --- /dev/null +++ b/packages/durable-session/src/react/use-durable-chat.ts @@ -0,0 +1,370 @@ +/** + * useDurableChat - React hook for durable chat. + * + * Provides TanStack AI-compatible API backed by Durable Streams + */ + +import type { AnyClientTool, UIMessage } from "@tanstack/ai"; +import type { Collection } from "@tanstack/react-db"; +import { + useCallback, + useEffect, + useMemo, + useRef, + useState, + useSyncExternalStore, +} from "react"; +import type { DurableChatClientOptions } from ".."; +import { DurableChatClient, messageRowToUIMessage } from ".."; +import type { UseDurableChatOptions, UseDurableChatReturn } from "./types"; + +/** + * Extract the item type from a Collection. + * + * TanStack DB's Collection has 5 type parameters: + * `Collection` + * + * This helper extracts `T` (the item type) from any Collection variant. + */ +type CollectionItem = + // biome-ignore lint/suspicious/noExplicitAny: Collection generic params require any for conditional type inference + C extends Collection ? T : never; + +/** + * SSR-safe hook for subscribing to TanStack DB collection data. + * This is a workaround to useLiveQuery not yet supporting SSR + * as per https://github.com/TanStack/db/pull/709 + */ +// biome-ignore lint/suspicious/noExplicitAny: Collection generic params require any for type constraint +function useCollectionData>( + collection: C, +): CollectionItem[] { + type T = CollectionItem; + + // Track version to know when to create a new snapshot. + // Incremented by subscription callback when collection changes. + const versionRef = useRef(0); + + // Cache the last snapshot to maintain stable reference. + // useSyncExternalStore requires getSnapshot to return the same reference + // when data hasn't changed, otherwise it triggers infinite re-renders. + const snapshotRef = useRef<{ version: number; data: T[] }>({ + version: -1, // Force initial snapshot creation + data: [], + }); + + // Subscribe callback - increments version to signal data changed. + // Stored in ref to maintain stable reference for useSyncExternalStore. + const subscribeRef = useRef((onStoreChange: () => void): (() => void) => { + const subscription = collection.subscribeChanges(() => { + versionRef.current++; + onStoreChange(); + }); + return () => subscription.unsubscribe(); + }); + + // Update subscribe ref when collection changes + subscribeRef.current = (onStoreChange: () => void): (() => void) => { + const subscription = collection.subscribeChanges(() => { + versionRef.current++; + onStoreChange(); + }); + return () => subscription.unsubscribe(); + }; + + // Snapshot callback - returns cached data unless version changed. + // Stored in ref to maintain stable reference for useSyncExternalStore. + const getSnapshotRef = useRef((): T[] => { + const currentVersion = versionRef.current; + const cached = snapshotRef.current; + + // Return cached snapshot if version hasn't changed + if (cached.version === currentVersion) { + return cached.data; + } + + // Version changed - create new snapshot and cache it + const data = [...collection.values()] as T[]; + snapshotRef.current = { version: currentVersion, data }; + return data; + }); + + // Update getSnapshot ref when collection changes + getSnapshotRef.current = (): T[] => { + const currentVersion = versionRef.current; + const cached = snapshotRef.current; + + if (cached.version === currentVersion) { + return cached.data; + } + + const data = [...collection.values()] as T[]; + snapshotRef.current = { version: currentVersion, data }; + return data; + }; + + // Pass the same function for both getSnapshot and getServerSnapshot. + // This ensures server and client render the same initial state (empty array), + // preventing hydration mismatches while enabling proper SSR. + return useSyncExternalStore( + subscribeRef.current, + getSnapshotRef.current, + getSnapshotRef.current, + ); +} + +/** + * React hook for durable chat with TanStack AI-compatible API. + * + * Provides reactive data binding with automatic updates when underlying + * collection data changes. Supports SSR through proper `useSyncExternalStore` + * integration. + * + * The client and collections are always available synchronously. + * Connection state is managed separately via `connectionStatus`. + * + * @example Basic usage + * ```typescript + * function Chat() { + * const { messages, sendMessage, isLoading, collections } = useDurableChat({ + * sessionId: 'my-session', + * proxyUrl: 'http://localhost:4000', + * }) + * + * return ( + *
+ * {messages.map(m => )} + * + *
+ * ) + * } + * ``` + * + * @example Custom queries with useLiveQuery + * ```typescript + * import { useLiveQuery, eq } from '@tanstack/react-db' + * + * function Chat() { + * const { collections } = useDurableChat({ ... }) + * + * // Use collections with useLiveQuery for custom queries + * const pendingToolCalls = useLiveQuery(q => + * q.from({ tc: collections.toolCalls }) + * .where(({ tc }) => eq(tc.state, 'pending')) + * ) + * } + * ``` + */ +export function useDurableChat< + TTools extends ReadonlyArray = AnyClientTool[], +>(options: UseDurableChatOptions): UseDurableChatReturn { + const { + autoConnect = true, + client: providedClient, + ...clientOptions + } = options; + + // Error handler ref - allows client's onError to call setError + const [error, setError] = useState(); + const onErrorRef = useRef<(err: Error) => void>(() => {}); + onErrorRef.current = (err) => { + setError(err); + clientOptions.onError?.(err); + }; + + // Create client synchronously - always available immediately + const clientRef = useRef<{ + client: DurableChatClient; + key: string; + } | null>(null); + const key = `${clientOptions.sessionId}:${clientOptions.proxyUrl}`; + + // Create or recreate client when key changes or client was disposed + // The isDisposed check handles React Strict Mode: cleanup disposes the client, + // so the next render must create a fresh one with a new AbortController. + if (providedClient) { + if (!clientRef.current || clientRef.current.client !== providedClient) { + clientRef.current = { client: providedClient, key: "provided" }; + } + } else if ( + !clientRef.current || + clientRef.current.key !== key || + clientRef.current.client.isDisposed + ) { + // Dispose old client if exists (may already be disposed, which is fine) + clientRef.current?.client.dispose(); + clientRef.current = { + client: new DurableChatClient({ + ...clientOptions, + onError: (err) => onErrorRef.current(err), + } as DurableChatClientOptions), + key, + }; + } + + const client = clientRef.current.client; + + const messageRows = useCollectionData(client.collections.messages); + const activeGenerations = useCollectionData( + client.collections.activeGenerations, + ); + const sessionMetaRows = useCollectionData(client.collections.sessionMeta); + + const messages = useMemo( + // Transform MessageRow[] to UIMessage[] + () => messageRows.map(messageRowToUIMessage), + [messageRows], + ); + + const isLoading = activeGenerations.length > 0; + const connectionStatus = + sessionMetaRows[0]?.connectionStatus ?? "disconnected"; + + useEffect(() => { + if (autoConnect && client.connectionStatus === "disconnected") { + client.connect().catch((err) => { + setError(err instanceof Error ? err : new Error(String(err))); + }); + } + + // Cleanup: unsubscribe and dispose (disposal is idempotent) + return () => { + if (!providedClient) { + client.dispose(); + } + }; + }, [client, autoConnect, providedClient]); + + // Action Callbacks + + const sendMessage = useCallback( + async (content: string) => { + try { + await client.sendMessage(content); + } catch (err) { + setError(err instanceof Error ? err : new Error(String(err))); + throw err; + } + }, + [client], + ); + + const append = useCallback( + async (message: UIMessage | { role: string; content: string }) => { + try { + await client.append(message); + } catch (err) { + setError(err instanceof Error ? err : new Error(String(err))); + throw err; + } + }, + [client], + ); + + const reload = useCallback(async () => { + try { + await client.reload(); + } catch (err) { + setError(err instanceof Error ? err : new Error(String(err))); + throw err; + } + }, [client]); + + const stop = useCallback(() => { + client.stop(); + }, [client]); + + const clear = useCallback(() => { + client.clear(); + }, [client]); + + const addToolResult = useCallback( + async ( + result: Parameters["addToolResult"]>[0], + ) => { + await client.addToolResult(result); + }, + [client], + ); + + const addToolApprovalResponse = useCallback( + async ( + response: Parameters< + DurableChatClient["addToolApprovalResponse"] + >[0], + ) => { + await client.addToolApprovalResponse(response); + }, + [client], + ); + + const fork = useCallback( + async (opts?: Parameters["fork"]>[0]) => { + return client.fork(opts); + }, + [client], + ); + + const registerAgents = useCallback( + async ( + agents: Parameters["registerAgents"]>[0], + ) => { + await client.registerAgents(agents); + }, + [client], + ); + + const unregisterAgent = useCallback( + async (agentId: string) => { + await client.unregisterAgent(agentId); + }, + [client], + ); + + const connect = useCallback(async () => { + try { + await client.connect(); + } catch (err) { + setError(err instanceof Error ? err : new Error(String(err))); + throw err; + } + }, [client]); + + const disconnect = useCallback(() => { + client.disconnect(); + }, [client]); + + const pause = useCallback(() => { + client.pause(); + }, [client]); + + const resume = useCallback(async () => { + await client.resume(); + }, [client]); + + return { + // TanStack AI useChat compatible + messages, + sendMessage, + append, + reload, + stop, + clear, + isLoading, + error, + addToolResult, + addToolApprovalResponse, + + // Durable extensions + client, + collections: client.collections, + connectionStatus, + fork, + registerAgents, + unregisterAgent, + connect, + disconnect, + pause, + resume, + }; +} diff --git a/packages/durable-session/src/schema.ts b/packages/durable-session/src/schema.ts new file mode 100644 index 00000000000..10aff129a91 --- /dev/null +++ b/packages/durable-session/src/schema.ts @@ -0,0 +1,252 @@ +/** + * STATE-PROTOCOL schema definition for AI DB. + * + * Defines the collection schemas for a durable chat session using the + * `@durable-streams/state` package's STATE-PROTOCOL. + * + * Each collection maps to a `type` field in change events streamed from + * the Durable Streams server. The `primaryKey` specifies which field + * receives the event's `key` value (injected by stream-db). + * + * @example + * ```typescript + * import { sessionStateSchema } from '@superset/durable-session' + * + * // Create insert event for a chunk + * const event = sessionStateSchema.chunks.insert({ + * key: `${messageId}:${seq}`, + * value: { + * messageId: 'msg-1', + * actorId: 'user-123', + * role: 'user', + * chunk: JSON.stringify({ type: 'message', message: {...} }), + * seq: 0, + * createdAt: new Date().toISOString(), + * }, + * }) + * ``` + */ + +import { createStateSchema } from "@durable-streams/state"; +import { z } from "zod"; + +// ============================================================================ +// Chunk Schema +// ============================================================================ + +/** + * Chunk schema for all messages. + * + * Both complete messages (user input, cached messages) and streaming chunks + * (assistant responses) use this schema. The difference is: + * + * - Complete messages: Single chunk with `{type: 'message', message: UIMessage}` + * - Streaming chunks: Multiple chunks with TanStack AI StreamChunk types + * + * The `chunk` field is opaque JSON - parsing happens in materialize.ts. + * + * Key format: `${messageId}:${seq}` - e.g., "msg-1:0", "msg-2:5" + */ +export const chunkValueSchema = z.object({ + /** Message identifier - groups chunks belonging to the same message */ + messageId: z.string(), + /** Actor who wrote this chunk */ + actorId: z.string(), + /** Message role - aligns with TanStack AI UIMessage.role */ + role: z.enum(["user", "assistant", "system"]), + /** JSON-encoded chunk content - could be WholeMessageChunk or StreamChunk */ + chunk: z.string(), + /** Sequence number within message - monotonically increasing per messageId */ + seq: z.number(), + /** ISO 8601 timestamp when chunk was created */ + createdAt: z.string(), +}); + +/** Inferred type for chunk values (without the injected `id` field) */ +export type ChunkValue = z.infer; + +// ============================================================================ +// Presence Schema +// ============================================================================ + +/** + * Presence schema for tracking online users and agents. + * + * Uses upsert semantics with (actorId, deviceId) pairs. + * Each tab/page load gets a unique deviceId. + * + * Key format: `${actorId}:${deviceId}` - e.g., "user-123:device-abc" + */ +export const presenceValueSchema = z.object({ + /** Actor identifier */ + actorId: z.string(), + /** Device/tab identifier - unique per browser tab/page load */ + deviceId: z.string(), + /** Actor type - 'user' or 'agent' */ + actorType: z.enum(["user", "agent"]), + /** Optional display name */ + name: z.string().optional(), + /** Current status */ + status: z.enum(["online", "offline", "away"]), + /** ISO 8601 timestamp of last activity */ + lastSeenAt: z.string(), +}); + +/** Inferred type for presence values */ +export type PresenceValue = z.infer; + +// ============================================================================ +// Agent Schema +// ============================================================================ + +/** + * Agent registration schema. + * + * Registered agents are invoked by the proxy when messages are added to + * the session. The `triggers` field controls which messages trigger the agent. + * + * Key format: `${agentId}` - e.g., "claude-agent", "custom-tool" + */ +export const agentValueSchema = z.object({ + /** Agent identifier (same as key) */ + agentId: z.string(), + /** Optional display name */ + name: z.string().optional(), + /** Webhook endpoint URL */ + endpoint: z.string(), + /** Trigger mode - when to invoke this agent */ + triggers: z.enum(["all", "user-messages"]).optional(), +}); + +/** Inferred type for agent values */ +export type AgentValue = z.infer; + +// ============================================================================ +// Session State Schema +// ============================================================================ + +/** + * Session state schema - defines all collection types for a chat session. + * + * This schema is used by both: + * - **Client (ai-db)**: Passed to `createStreamDB()` to create typed collections + * - **Proxy (ai-db-proxy)**: Used to create properly formatted change events + * + * Each key maps to: + * - A STATE-PROTOCOL `type` field value + * - A TanStack DB collection in the resulting StreamDB + * + * @example + * ```typescript + * // Client-side: Create stream-db with these collections + * const db = createStreamDB({ + * streamOptions: { url: '/v1/stream/sessions/my-session' }, + * state: sessionStateSchema, + * }) + * + * // Access collections + * const chunks = db.collections.chunks + * const presence = db.collections.presence + * + * // Proxy-side: Create change events + * const event = sessionStateSchema.chunks.insert({ + * key: 'msg-1:0', + * value: { ... }, + * }) + * await stream.append(event) + * ``` + */ +export const sessionStateSchema = createStateSchema({ + /** + * Chunks collection - all message chunks (complete and streaming). + * + * Primary key `id` is injected from event.key = `${messageId}:${seq}` + */ + chunks: { + schema: chunkValueSchema, + type: "chunk", + primaryKey: "id", + allowSyncWhilePersisting: true, + }, + + /** + * Presence collection - online status of users and agents. + * + * Uses upsert semantics with (actorId, deviceId) pairs. + * Primary key `id` is injected from event.key = `${actorId}:${deviceId}` + * This follows the same pattern as chunks. + */ + presence: { + schema: presenceValueSchema, + type: "presence", + primaryKey: "id", + }, + + /** + * Agents collection - registered webhook agents. + * + * Uses upsert semantics for registration. Primary key `agentId` from event.key. + */ + agents: { + schema: agentValueSchema, + type: "agent", + primaryKey: "agentId", + }, +}); + +/** Type of the session state schema */ +export type SessionStateSchema = typeof sessionStateSchema; + +// ============================================================================ +// Row Types (with injected primary keys) +// ============================================================================ + +/** + * ChunkRow - a chunk value with the injected `id` primary key. + * + * This is the type of rows in `db.collections.chunks` after stream-db + * injects the primary key from the event key. + */ +export type ChunkRow = ChunkValue & { + /** Primary key - injected from event.key = `${messageId}:${seq}` */ + id: string; +}; + +/** + * RawPresenceRow - a presence value with the injected `id` primary key. + * + * This is the type of rows in raw presence collection after stream-db + * injects the primary key from the event key = `${actorId}:${deviceId}` + * + * This is the internal/raw type. For the public API, use PresenceRow + * which is an aggregated view per actor. + */ +export type RawPresenceRow = PresenceValue & { + /** Primary key - injected from event.key = `${actorId}:${deviceId}` */ + id: string; +}; + +/** + * PresenceRow - aggregated presence per actor. + * + * This is the public type exposed to components. It aggregates + * all devices for an actor into a single row showing who's online. + */ +export type PresenceRow = { + /** Actor identifier */ + actorId: string; + /** Actor type - 'user' or 'agent' */ + actorType: "user" | "agent"; + /** Optional display name */ + name?: string; + /** All online device IDs for this actor */ + deviceIds: string[]; + /** Number of online devices */ + deviceCount: number; +}; + +/** + * AgentRow - an agent value with the `agentId` key. + * (Note: agentId is already in the schema, so this is the same as AgentValue) + */ +export type AgentRow = AgentValue; diff --git a/packages/durable-session/src/types.ts b/packages/durable-session/src/types.ts new file mode 100644 index 00000000000..643fea0af56 --- /dev/null +++ b/packages/durable-session/src/types.ts @@ -0,0 +1,422 @@ +/** + * Core type definitions for @superset/durable-session + * + * Defines the stream protocol types, collection schemas, and API interfaces. + * + * Design principles: + * - Use TanStack AI types directly (MessagePart, ToolCallPart, etc.) + * - Single MessageRow type for all materialized messages + * - Derived collections filter on message parts, not custom row types + */ + +import type { + AnyClientTool, + MessagePart, + StreamChunk, + UIMessage, +} from "@tanstack/ai"; +import type { Collection } from "@tanstack/db"; +import type { SessionDB } from "./collection"; + +// Re-export TanStack AI message part types for consumer convenience +export type { + MessagePart, + TextPart, + ThinkingPart, + ToolCallPart, + ToolResultPart, +} from "@tanstack/ai"; +// Re-export schema types +export type { AgentRow, ChunkRow, ChunkValue, PresenceRow } from "./schema"; + +// ============================================================================ +// Stream Protocol Types +// ============================================================================ + +/** + * Whole message chunk - stored as single row in stream. + * Used for messages that are complete when written (user input, cached messages). + * + * This is different from TanStack AI's StreamChunk types, which are designed + * for streaming assistant responses. Whole messages are complete when sent, + * so we store them as complete UIMessage objects. + */ +export interface WholeMessageChunk { + type: "whole-message"; + message: UIMessage; +} + +/** + * Union of all chunk types we handle. + * - WholeMessageChunk: Complete messages (user input) + * - StreamChunk: TanStack AI streaming chunks (assistant responses) + */ +export type DurableStreamChunk = WholeMessageChunk | StreamChunk; + +/** + * Actor types in the chat session. + */ +export type ActorType = "user" | "agent"; + +/** + * Message role types (aligned with TanStack AI UIMessage.role). + */ +export type MessageRole = "user" | "assistant" | "system"; + +// ============================================================================ +// Message Collection Types +// ============================================================================ + +/** + * Materialized message row from stream. + * + * Extends TanStack AI's UIMessage with durable session metadata. + * Messages are materialized from ChunkRow arrays via the live query pipeline. + * + * Message parts use TanStack AI's discriminated union types directly: + * - TextPart: { type: 'text', content: string } + * - ToolCallPart: { type: 'tool-call', id, name, arguments, state, approval?, output? } + * - ToolResultPart: { type: 'tool-result', toolCallId, content, state, error? } + * - ThinkingPart: { type: 'thinking', content: string } + * + * @example + * ```typescript + * // Filter for tool calls in a message + * const toolCalls = message.parts.filter( + * (p): p is ToolCallPart => p.type === 'tool-call' + * ) + * + * // Check for pending approvals + * const pendingApprovals = toolCalls.filter( + * tc => tc.approval?.needsApproval && tc.approval.approved === undefined + * ) + * ``` + */ +export interface MessageRow { + /** Message identifier (same as messageId from chunks) */ + id: string; + /** Message role */ + role: MessageRole; + /** Message parts - uses TanStack AI's MessagePart type directly */ + parts: MessagePart[]; + /** Actor identifier who wrote this message */ + actorId: string; + /** Whether the message is complete (finish chunk received) */ + isComplete: boolean; + /** Message creation timestamp (from first chunk) */ + createdAt: Date; +} + +// ============================================================================ +// Active Generation Types +// ============================================================================ + +/** + * Messages currently being streamed (have chunks but no finish chunk). + */ +export interface ActiveGenerationRow { + /** The message being generated */ + messageId: string; + /** Actor identifier */ + actorId: string; + /** When generation started */ + startedAt: Date; + /** Last chunk sequence number */ + lastChunkSeq: number; + /** When last chunk was received */ + lastChunkAt: Date; +} + +// ============================================================================ +// Session Metadata Types +// ============================================================================ + +/** + * Connection status states. + */ +export type ConnectionStatus = + | "disconnected" + | "connecting" + | "connected" + | "error"; + +/** + * Session metadata row (local state only, not derived from stream). + */ +export interface SessionMetaRow { + /** Session identifier */ + sessionId: string; + /** Current connection status */ + connectionStatus: ConnectionStatus; + /** Last synced transaction ID (for txid tracking) */ + lastSyncedTxId: string | null; + /** Last sync timestamp */ + lastSyncedAt: Date | null; + /** Error information if any */ + error: { message: string; code?: string } | null; +} + +// ============================================================================ +// Session Statistics Types +// ============================================================================ + +/** + * Aggregate session statistics row. + */ +export interface SessionStatsRow { + /** Session identifier */ + sessionId: string; + /** Total message count */ + messageCount: number; + /** User message count */ + userMessageCount: number; + /** Assistant message count */ + assistantMessageCount: number; + /** Total tool call count */ + toolCallCount: number; + /** Total approval count */ + approvalCount: number; + /** Total tokens used */ + totalTokens: number; + /** Prompt tokens used */ + promptTokens: number; + /** Completion tokens used */ + completionTokens: number; + /** Currently active generation count */ + activeGenerationCount: number; + /** First message timestamp */ + firstMessageAt: Date | null; + /** Last message timestamp */ + lastMessageAt: Date | null; +} + +// ============================================================================ +// Agent Types +// ============================================================================ + +/** + * Agent trigger modes. + */ +export type AgentTrigger = "all" | "user-messages"; + +/** + * Unified structure for webhook registration and inline agent invocation. + */ +export interface AgentSpec { + /** Agent identifier */ + id: string; + /** Optional display name */ + name?: string; + /** Endpoint URL the proxy will call */ + endpoint: string; + /** HTTP method */ + method?: "POST"; + /** Additional headers for agent calls */ + headers?: Record; + /** Trigger mode (for registered agents) */ + triggers?: AgentTrigger; + /** Request body template */ + bodyTemplate?: Record; +} + +// ============================================================================ +// Collection Types +// ============================================================================ + +// Import types from schema +import type { AgentRow, ChunkRow, PresenceRow } from "./schema"; + +/** + * All collections exposed by DurableChatClient. + * + * Architecture: + * - `chunks`, `presence`, `agents`: Synced from Durable Stream via stream-db + * - `messages`: Root materialized collection (groupBy + collect → materialize) + * - `toolCalls`, `pendingApprovals`, `toolResults`: Derived from messages via .fn.where() + * - `activeGenerations`: Derived from messages (incomplete messages) + * - `sessionMeta`, `sessionStats`: Local/aggregated state + * + * The `chunks` and `agents` collections are synced directly from the Durable + * Stream via stream-db. The `presence` collection is aggregated from raw + * per-device presence records. Other collections are derived from chunks. + * + * Pipeline: + * ``` + * chunks → (subquery) → messages + * ↓ + * .fn.where(parts filtering) + * ↓ + * toolCalls (lazy) + * pendingApprovals (lazy) + * toolResults (lazy) + * activeGenerations (lazy) + * ``` + */ +export interface DurableChatCollections { + /** Root chunks collection synced from Durable Stream via stream-db */ + chunks: Collection; + /** Aggregated presence - one row per online actor with their device count */ + presence: Collection; + /** Agents collection - registered webhook agents (from stream-db) */ + agents: Collection; + /** All materialized messages (keyed by messageId) */ + messages: Collection; + /** Messages containing tool calls (keyed by messageId) */ + toolCalls: Collection; + /** Messages with pending approval requests (keyed by messageId) */ + pendingApprovals: Collection; + /** Messages containing tool results (keyed by messageId) */ + toolResults: Collection; + /** Active generations - incomplete messages (keyed by messageId) */ + activeGenerations: Collection; + /** Session metadata collection (local state) */ + sessionMeta: Collection; + /** Session statistics (keyed by sessionId) */ + sessionStats: Collection; +} + +// ============================================================================ +// Client Configuration Types +// ============================================================================ + +/** + * Configuration options for DurableChatClient. + */ +export interface DurableChatClientOptions< + TTools extends ReadonlyArray = AnyClientTool[], +> { + /** Session identifier */ + sessionId: string; + /** Proxy URL for API requests */ + proxyUrl: string; + /** Actor identifier (auto-generated if not provided) */ + actorId?: string; + /** Actor type */ + actorType?: ActorType; + /** Client tools */ + tools?: TTools; + /** Initial messages for SSR hydration */ + initialMessages?: UIMessage[]; + /** API endpoint */ + api?: string; + /** Additional request body fields */ + body?: Record; + /** + * Default agent to invoke for each user message. + * For single-agent scenarios, this provides a simpler alternative to registerAgents(). + * The agent spec is sent with each sendMessage request. + */ + agent?: AgentSpec; + + // Callbacks (TanStack AI compatible) + /** Called when response is received */ + onResponse?: (response?: Response) => void | Promise; + /** Called for each chunk */ + onChunk?: (chunk: StreamChunk) => void; + /** Called when message finishes */ + onFinish?: (message: UIMessage) => void; + /** Called on error */ + onError?: (error: Error) => void; + /** Called when messages change */ + onMessagesChange?: (messages: UIMessage[]) => void; + + /** Durable Streams configuration */ + stream?: { + /** Additional headers for stream requests */ + headers?: Record; + }; + + /** + * Pre-created SessionDB for testing. + * If provided, the client will use this instead of creating its own via createSessionDB(). + * This allows tests to inject mock collections with controlled data. + * @internal + */ + sessionDB?: SessionDB; +} + +// ============================================================================ +// Tool Result Input Types +// ============================================================================ + +/** + * Input for adding a tool result. + */ +export interface ToolResultInput { + /** Tool call identifier */ + toolCallId: string; + /** Tool output */ + output: unknown; + /** Error message if failed */ + error?: string; + /** Client-generated message ID for optimistic updates (auto-generated if not provided) */ + messageId?: string; +} + +/** + * Tool result input with required messageId (used internally for optimistic actions). + */ +export type ClientToolResultInput = Required< + Pick +> & + ToolResultInput; + +/** + * Input for adding an approval response. + */ +export interface ApprovalResponseInput { + /** Approval identifier */ + id: string; + /** Whether approved */ + approved: boolean; +} + +// ============================================================================ +// Fork Types +// ============================================================================ + +/** + * Options for forking a session. + */ +export interface ForkOptions { + /** Fork before this message (default: current end) */ + atMessageId?: string; + /** Custom session ID (default: auto-generated) */ + newSessionId?: string; +} + +/** + * Result of forking a session. + */ +export interface ForkResult { + /** New session identifier */ + sessionId: string; + /** Starting offset for new session */ + offset: string; +} + +// ============================================================================ +// Session DB Configuration Types +// ============================================================================ + +/** + * Configuration for creating a session stream-db. + */ +export interface SessionDBConfig { + /** Session identifier */ + sessionId: string; + /** Base URL for the proxy */ + baseUrl: string; + /** Additional headers for stream requests */ + headers?: Record; + /** + * AbortSignal to cancel the stream sync. + * When aborted, the sync will stop and cleanup will be called. + */ + signal?: AbortSignal; + // /** + // * Live mode for the stream connection. + // * Defaults to "sse" for efficient real-time updates. + // */ + // liveMode?: LiveMode +} diff --git a/packages/ai-chat/tsconfig.json b/packages/durable-session/tsconfig.json similarity index 81% rename from packages/ai-chat/tsconfig.json rename to packages/durable-session/tsconfig.json index 41102bd9f88..c23a37895a1 100644 --- a/packages/ai-chat/tsconfig.json +++ b/packages/durable-session/tsconfig.json @@ -4,8 +4,7 @@ "outDir": "./dist", "rootDir": "./src", "lib": ["ES2022", "DOM", "DOM.Iterable"], - "jsx": "react-jsx", - "types": ["bun-types", "node"] + "jsx": "react-jsx" }, "include": ["src/**/*.ts", "src/**/*.tsx"], "exclude": ["node_modules", "dist"]