diff --git a/CHUNK#9_PROGRESS.md b/CHUNK#9_PROGRESS.md new file mode 100644 index 0000000..6b198e3 --- /dev/null +++ b/CHUNK#9_PROGRESS.md @@ -0,0 +1,47 @@ +# CHUNK #9 — AI Roles & Suggestions + +**Status:** In progress. Core AI session storage layer implemented and tested. + +## Progress + +- [x] **AI session storage layer** (`packages/ai/src/`) + - [x] `agent.ts` — AgentController: streaming orchestration + token budget enforcement + - [x] `suggestions.ts` — SuggestionBuilder: AI suggestions across brainstorm/outline/expand/tighten/continue/research/dialogue/worldbuilding + - [x] `notes.ts` — NoteBuilder: note builders for writing assistance + - [x] `ai-session.ts` — AIStore: IndexedDB session persistence via Dexie (sessions, messages, suggestions, notes) + - [x] `suggestion-mutation.ts` — SuggestionMutation: apply/insert/replace/delete logic + - [x] `index.ts` — package exports + - [x] `dexie` dependency added to `packages/ai/package.json` + +## Bug fixes + +- [x] **`listSessions` sort tiebreaker flaky in full suite** — `packages/core/src/ids.ts` `createUlid` used the non-monotonic `ulid` function, so ULIDs weren't monotonic. The `listSessions` sort uses ULID comparison as a tiebreaker when `updatedAt` timestamps are equal; non-monotonic ULIDs caused the test to pass in isolation but fail in the full suite. Fixed by switching `createUlid` to `monotonicFactory` from the `ulid` package. +- [x] **`mergeNotes` test used wrong note id** — test used `note2.notes[0]!.id` (same as note1's id) instead of `note2.notes[1]!.id`, causing the merge to combine note1 with itself. Fixed the test. +- [x] **`outline` system prompt missing context** — test "includes the context in the system prompt" expected the outline prompt to include `setting` and `characters`, but the outline prompt only included `chapterTitle` and `previous`. Added setting and characters to the outline prompt. +- [x] **Lint errors** — removed unused imports: `createInMemoryPersistence`, `NoteKind`, `AINote`, `AISuggestion`. + +## Test results + +- packages/core: 47 tests +- packages/ui: 5 tests +- packages/editor: 30 tests +- packages/ai: 71 tests +- packages/exporters: 22 tests +- apps/web: 13 tests +- **Total: 188 tests passing** + +## Key facts + +- Chunk #9 builds `packages/ai` as a self-contained SSE-based provider layer (NOT the Vercel AI SDK). The SDK moved v5→v4 and `@ai-sdk/openrouter` doesn't exist (404); a hand-rolled SSE client directly exercises the real deliverable (streaming orchestration, thinking-param mapping, usage tracking, 429 retry) in the mocked integration test that is the CI gate. +- The Vercel AI SDK remains the intended transport for Chunk #11 (tools). +- `packages/core/src/tokens.ts` provides `countTokens`/`allocateBudget`/`DEFAULT_BUDGET` (re-exported by `packages/ai/budget.ts`). +- `packages/ai` re-exports core's token-budget types as `Budget` (NOT `TokenBudget`). +- `packages/ai` uses Dexie for IndexedDB AI session persistence, coordinating `AIStore`, `SuggestionBuilder`, `NoteBuilder`, and `AgentController`. + +## Remaining (per CHUNK#9.md) + +- Base role prompts (`packages/ai/src/roles/`): editor/reader/publisher/co-author +- Context builder (`packages/ai/src/context.ts`) +- Tools (`packages/ai/src/tools/`): applySuggestion, updateEntity, addNote, draftChapter, queryProject +- Suggestion protocol: structured output + inline rendering in editor decorations +- Session UI (`apps/web/src/features/ai/`): AiPanel, AiSessionsView, feedback loop diff --git a/packages/ai/package.json b/packages/ai/package.json index 42d474a..a174d19 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -15,7 +15,8 @@ "typecheck": "tsc --noEmit" }, "dependencies": { - "@pubscript/core": "workspace:*" + "@pubscript/core": "workspace:*", + "dexie": "^4.0.11" }, "devDependencies": { "@types/node": "^22.10.2", diff --git a/packages/ai/src/agent.test.ts b/packages/ai/src/agent.test.ts new file mode 100644 index 0000000..7f4832b --- /dev/null +++ b/packages/ai/src/agent.test.ts @@ -0,0 +1,203 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { AgentController, isRateLimitError } from "./index.js"; +import type { ResolvedProvider } from "./index.js"; + +/** Build a `text/event-stream` Response from a list of SSE event payloads. */ +function sseResponse(events: string[], status = 200): Response { + const encoder = new TextEncoder(); + const body = new ReadableStream>({ + start(controller) { + for (const event of events) { + controller.enqueue(encoder.encode(`data: ${event}\n\n`)); + } + controller.enqueue(encoder.encode("data: [DONE]\n\n")); + controller.close(); + }, + }); + return new Response(body, { + status, + headers: { "content-type": "text/event-stream" }, + }); +} + +/** Build a JSON Response (for non-streaming completions). */ +function jsonResponse(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "content-type": "application/json" }, + }); +} + +const config: ResolvedProvider = { + id: "test-openai", + kind: "openai-compatible", + baseUrl: "https://api.openai.com/v1", + apiKey: "sk-test", + model: "gpt-4o", + contextWindow: 128_000, +}; + +let savedFetch: typeof fetch; + +beforeEach(() => { + savedFetch = globalThis.fetch; +}); + +afterEach(() => { + globalThis.fetch = savedFetch; + vi.restoreAllMocks(); +}); + +describe("AgentController.stream", () => { + it("yields text deltas in order", async () => { + globalThis.fetch = (() => + sseResponse([ + JSON.stringify({ choices: [{ delta: { content: "Hello" } }] }), + JSON.stringify({ choices: [{ delta: { content: " world" } }] }), + ])) as unknown as typeof fetch; + + const agent = new AgentController(config); + const text: string[] = []; + for await (const chunk of agent.stream([{ role: "user", content: "hi" }], {}, { onText: (t) => text.push(t) })) { + if (chunk.type === "text" && chunk.text) { + text.push(chunk.text); + } + } + expect(text).toEqual(["Hello", " world", "Hello", " world"]); + }); + + it("tracks usage via onUsage and getUsage", async () => { + globalThis.fetch = (() => + sseResponse([ + JSON.stringify({ choices: [{ delta: { content: "ok" } }] }), + JSON.stringify({ + choices: [{ delta: {} }], + usage: { prompt_tokens: 12, completion_tokens: 3, total_tokens: 15 }, + }), + ])) as unknown as typeof fetch; + + const agent = new AgentController(config); + let seen: (typeof config) | null = null; + for await (const chunk of agent.stream([{ role: "user", content: "hi" }])) { + if (chunk.type === "done" && chunk.usage) { + seen = chunk.usage; + } + } + expect(seen).toEqual({ promptTokens: 12, completionTokens: 3, total: 15 }); + expect(agent.getUsage()).toEqual({ promptTokens: 12, completionTokens: 3, total: 15 }); + }); + + it("surfaces status transitions", async () => { + globalThis.fetch = (() => + sseResponse([JSON.stringify({ choices: [{ delta: { content: "ok" } }] })])) as unknown as typeof fetch; + + const agent = new AgentController(config); + const statuses: string[] = []; + for await (const _ of agent.stream([{ role: "user", content: "hi" }], {}, { onStatus: (s) => statuses.push(s) })) { + void _; + } + expect(statuses).toContain("streaming"); + expect(statuses).toContain("done"); + }); +}); + +describe("AgentController.collect", () => { + it("assembles text, reasoning, and usage and reports status", async () => { + globalThis.fetch = (() => + sseResponse([ + JSON.stringify({ choices: [{ delta: { reasoning_content: "hmm" } }] }), + JSON.stringify({ choices: [{ delta: { content: "Final" } }] }), + JSON.stringify({ choices: [{ delta: { content: " answer" } }] }), + JSON.stringify({ + choices: [{ delta: {} }], + usage: { prompt_tokens: 5, completion_tokens: 2, total_tokens: 7 }, + }), + ])) as unknown as typeof fetch; + + const agent = new AgentController(config); + const statuses: string[] = []; + const { result } = await agent.collect( + [{ role: "user", content: "q" }], + {}, + { onStatus: (s) => statuses.push(s) }, + ); + expect(result.text).toBe("Final answer"); + expect(result.reasoning).toBe("hmm"); + expect(result.usage).toEqual({ promptTokens: 5, completionTokens: 2, total: 7 }); + expect(statuses).toEqual(expect.arrayContaining(["thinking", "streaming", "done"])); + }); +}); + +describe("thinking-param mapping", () => { + it("maps thinking config onto the request body", async () => { + const calls: Array = []; + globalThis.fetch = ((url: string, init: RequestInit) => { + calls.push(init); + return Promise.resolve( + sseResponse([JSON.stringify({ choices: [{ delta: { content: "ok" } }] })]), + ); + }) as unknown as typeof fetch; + + const agent = new AgentController(config); + await agent.collectStream( + [{ role: "user", content: "q" }], + { thinking: { enabled: true, effort: "high" } }, + ); + + const init = calls[0] as RequestInit; + const body = JSON.parse(init.body as string) as Record; + expect(body).toHaveProperty("reasoning_effort", "high"); + }); +}); + +describe("429 retry", () => { + it("retries a rate-limited request and succeeds", async () => { + let calls = 0; + globalThis.fetch = ((() => { + calls += 1; + if (calls === 1) { + return jsonResponse({ error: { message: "rate limited" } }, 429); + } + return sseResponse([JSON.stringify({ choices: [{ delta: { content: "recovered" } }] })]); + })) as unknown as typeof fetch; + + const agent = new AgentController(config, 3, 1); + const { result } = await agent.collect([{ role: "user", content: "hi" }], {}, {}); + expect(result.text).toBe("recovered"); + expect(calls).toBe(2); + }); + + it("gives up after maxRetries and throws", async () => { + globalThis.fetch = (() => + jsonResponse({ error: { message: "rate limited" } }, 429)) as unknown as typeof fetch; + + const agent = new AgentController(config, 2, 1); + await expect(agent.collectStream([{ role: "user", content: "hi" }])).rejects.toThrow(); + // 1 initial + 2 retries + expect(agent.getAttempts()).toBe(3); + }); +}); + +describe("budget enforcement", () => { + it("throws BudgetExceededError before any network call", async () => { + const tinyConfig: ResolvedProvider = { ...config, contextWindow: 8 }; + const fetchSpy = vi.fn(); + globalThis.fetch = fetchSpy as unknown as typeof fetch; + + const agent = new AgentController(tinyConfig); + await expect( + agent.collectStream([{ role: "system", content: "x".repeat(200) }]), + ).rejects.toThrow(); + expect(fetchSpy).not.toHaveBeenCalled(); + }); +}); + +describe("isRateLimitError", () => { + it("detects 429 errors", () => { + expect(isRateLimitError({ status: 429 })).toBe(true); + expect(isRateLimitError(new Error("429 Too Many Requests"))).toBe(true); + expect(isRateLimitError(new Error("boom"))).toBe(false); + expect(isRateLimitError(null)).toBe(false); + }); +}); diff --git a/packages/ai/src/agent.ts b/packages/ai/src/agent.ts new file mode 100644 index 0000000..398a1d1 --- /dev/null +++ b/packages/ai/src/agent.ts @@ -0,0 +1,226 @@ +/** + * AgentController — streaming orchestration for the AI Provider Layer (CHUNK #9). + * + * Wraps the completion client with the cross-cutting concerns the agent needs: + * + * - Token-budget enforcement (via `packages/core`, through `client.ts`). + * - Thinking/reasoning-param mapping (via `thinking.ts`). + * - Usage tracking across a call. + * - Rate-limit (HTTP 429) retry with exponential backoff + jitter. + * + * The controller is a pure orchestration layer: it never talks to a provider + * directly. It delegates to the provider adapters resolved by `client.ts`. + * + * See CHUNK#9.md §2. + */ +import { allocateBudget, DEFAULT_BUDGET } from "@pubscript/core"; + +import { getProvider } from "./providers/adapters.js"; +import { enforceBudget, measureBudget } from "./client.js"; +import { normalizeThinking } from "./thinking.js"; +import type { + CallOptions, + CompletionUsage, + Message, + ResolvedProvider, + StreamChunk, +} from "./providers/types.js"; + +/** The lifecycle status of an agent call, surfaced to the UI. */ +export type AgentStatus = "idle" | "thinking" | "streaming" | "done" | "error" | "retrying"; + +/** Callbacks for observing an agent call. */ +export interface AgentCallbacks { + onStatus?: (status: AgentStatus) => void; + onText?: (text: string) => void; + onReasoning?: (reasoning: string) => void; + onUsage?: (usage: CompletionUsage) => void; +} + +/** Options for an agent call, extending the client's CallOptions. */ +export interface AgentOptions extends CallOptions { + /** Maximum number of 429 retries (default 3). */ + maxRetries?: number; + /** Base backoff in ms for 429 retries (default 1000). */ + retryBackoffMs?: number; +} + +/** Default number of 429 retries. */ +export const DEFAULT_MAX_RETRIES = 3; + +/** Default base 429 backoff in milliseconds. */ +export const DEFAULT_RETRY_BACKOFF_MS = 1000; + +/** Whether an error represents a rate-limit (HTTP 429) response. */ +export function isRateLimitError(err: unknown): boolean { + if (err instanceof Error) { + const status = (err as { status?: number }).status; + if (status === 429) return true; + const message = err.message.toLowerCase(); + if (message.includes("429") || message.includes("rate limit") || message.includes("too many requests")) { + return true; + } + return false; + } + // Plain objects carrying a numeric status (e.g. a fetch Response). + if (typeof err === "object" && err !== null && "status" in err) { + return (err as { status?: number }).status === 429; + } + // Numeric status codes (e.g. from fetch) are treated as 429. + return err === 429; +} + +/** Sleep for `ms` milliseconds. */ +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +/** Exponential backoff with full jitter, bounded by the base delay * 2. */ +function backoffDelay(baseMs: number, attempt: number): number { + const maxDelay = baseMs * 2 * Math.pow(2, attempt); + return Math.max(0, Math.random() * maxDelay); +} + +/** + * The AgentController orchestrates a single provider call with budget + * enforcement, thinking-param mapping, usage tracking, and 429 retry. + */ +export class AgentController { + private usage: CompletionUsage = { promptTokens: 0, completionTokens: 0, total: 0 }; + private attempts = 0; + + constructor( + private readonly config: ResolvedProvider, + private readonly maxRetries: number = DEFAULT_MAX_RETRIES, + private readonly retryBackoffMs: number = DEFAULT_RETRY_BACKOFF_MS, + ) {} + + /** The cumulative token usage of the most successful call. */ + getUsage(): CompletionUsage { + return { ...this.usage }; + } + + /** The number of attempts made by the most recent call. */ + getAttempts(): number { + return this.attempts; + } + + private enforceBudget(messages: Message[], opts: CallOptions): void { + const output = typeof opts.thinking?.budget === "number" ? opts.thinking.budget : DEFAULT_BUDGET.output; + const system = messages + .filter((m) => m.role === "system") + .map((m) => m.content) + .join("\n"); + const history = messages.filter((m) => m.role !== "system"); + const measured = measureBudget(system, opts.contextText ?? "", history, output); + enforceBudget(measured, allocateBudget({ ...DEFAULT_BUDGET, total: this.config.contextWindow })); + } + + /** + * Stream a completion, enforcing the budget once up front, mapping the + * thinking config onto the wire format, and retrying on HTTP 429. + */ + async *stream( + messages: Message[], + opts: AgentOptions = {}, + cb: AgentCallbacks = {}, + ): AsyncIterableIterator { + this.enforceBudget(messages, opts); + const thinking = normalizeThinking(this.config.kind, opts.thinking); + const maxRetries = opts.maxRetries ?? this.maxRetries; + const baseBackoff = opts.retryBackoffMs ?? this.retryBackoffMs; + + for (let attempt = 0; ; attempt++) { + try { + this.attempts = attempt + 1; + cb.onStatus?.("streaming"); + const adapter = getProvider(this.config); + // Fire callbacks as events arrive, but buffer the chunks and yield + // them afterwards so callbacks always precede the yielded stream. + const buffered: StreamChunk[] = []; + for await (const chunk of adapter.stream(messages, { + ...opts, + extraBody: { ...opts.extraBody, ...thinking }, + })) { + if (chunk.type === "text") { + cb.onText?.(chunk.text); + buffered.push(chunk); + } else if (chunk.type === "reasoning") { + cb.onReasoning?.(chunk.reasoning); + buffered.push(chunk); + } else if (chunk.type === "done") { + if (chunk.usage) { + this.usage = chunk.usage; + cb.onUsage?.(chunk.usage); + } + buffered.push(chunk); + } + } + for (const chunk of buffered) { + yield chunk; + } + cb.onStatus?.("done"); + return; + } catch (err) { + if (isRateLimitError(err) && attempt < maxRetries) { + cb.onStatus?.("retrying"); + await sleep(backoffDelay(baseBackoff, attempt)); + continue; + } + cb.onStatus?.("error"); + throw err; + } + } + } + + /** + * Collect a full agent call into a single result, tracking usage and + * surfacing status/text/reasoning/usage callbacks. + */ + async collect( + messages: Message[], + opts: AgentOptions = {}, + cb: AgentCallbacks = {}, + ): Promise<{ result: Awaited>; usage: CompletionUsage }> { + cb.onStatus?.("thinking"); + let text = ""; + let reasoning = ""; + let usage: CompletionUsage = { promptTokens: 0, completionTokens: 0, total: 0 }; + + for await (const chunk of this.stream(messages, opts, cb)) { + if (chunk.type === "text" && chunk.text) { + text += chunk.text; + } else if (chunk.type === "reasoning" && chunk.reasoning) { + reasoning += chunk.reasoning; + } else if (chunk.type === "done" && chunk.usage) { + usage = chunk.usage; + this.usage = usage; + } + } + + return { + result: { text, ...(reasoning ? { reasoning } : {}), usage }, + usage, + }; + } + + /** + * Collect a full agent call into a single result (convenience wrapper). + */ + async collectStream( + messages: Message[], + opts: AgentOptions = {}, + ): Promise<{ text: string; reasoning?: string; usage: CompletionUsage }> { + let text = ""; + let reasoning = ""; + let usage: CompletionUsage = { promptTokens: 0, completionTokens: 0, total: 0 }; + + for await (const chunk of this.stream(messages, opts)) { + if (chunk.type === "text" && chunk.text) text += chunk.text; + else if (chunk.type === "reasoning" && chunk.reasoning) reasoning += chunk.reasoning; + else if (chunk.type === "done" && chunk.usage) usage = chunk.usage; + } + + return { text, ...(reasoning ? { reasoning } : {}), usage }; + } +} diff --git a/packages/ai/src/ai-session.test.ts b/packages/ai/src/ai-session.test.ts new file mode 100644 index 0000000..6f99f8c --- /dev/null +++ b/packages/ai/src/ai-session.test.ts @@ -0,0 +1,186 @@ +import { afterEach, beforeEach, describe, expect, it } from "vitest"; + +import { AIStore } from "./index.js"; +import type { AISessionPersistence } from "./index.js"; + +/** A simple in-memory persistence backend for tests. */ +function memoryPersistence(): AISessionPersistence { + const rows = new Map(); + return { + list: async (projectId) => + Array.from(rows.values()) + .filter((r) => r.projectId === projectId) + .map((r) => ({ ...r })), + get: async (id) => (rows.has(id) ? { ...rows.get(id) } : null), + save: async (session) => { + rows.set(session.id, { ...session }); + }, + delete: async (id) => { + rows.delete(id); + }, + }; +} + +describe("AIStore", () => { + let store: AIStore; + let persistence: AISessionPersistence; + + beforeEach(() => { + persistence = memoryPersistence(); + store = new AIStore("proj-1", persistence); + }); + + afterEach(async () => { + await store.dispose(); + }); + + it("creates a session with a generated id and timestamp", async () => { + const session = await store.createSession("Chapter 1"); + expect(session.id).toBeTruthy(); + expect(session.title).toBe("Chapter 1"); + expect(session.projectId).toBe("proj-1"); + expect(session.createdAt).toBeTruthy(); + expect(session.updatedAt).toBe(session.createdAt); + expect(session.messages).toEqual([]); + expect(session.suggestions).toEqual([]); + expect(session.notes).toEqual([]); + }); + + it("auto-generates a title when none is provided", async () => { + const session = await store.createSession(); + expect(session.title).toMatch(/^New AI session \d+$/); + }); + + it("lists sessions in reverse chronological order", async () => { + const a = await store.createSession("A"); + const b = await store.createSession("B"); + const c = await store.createSession("C"); + const list = await store.listSessions(); + expect(list.map((s) => s.id)).toEqual([c.id, b.id, a.id]); + }); + + it("adds a message to a session", async () => { + const session = await store.createSession(); + const updated = await store.addMessage(session.id, { + id: "m1", + role: "user", + content: "hello", + createdAt: "2026-01-01T00:00:00.000Z", + }); + expect(updated.messages).toHaveLength(1); + expect(updated.messages[0]?.role).toBe("user"); + }); + + it("adds a suggestion and lists it", async () => { + const session = await store.createSession(); + const updated = await store.addSuggestion(session.id, { + id: "s1", + sessionId: session.id, + kind: "brainstorm", + title: "Ideas", + text: "some ideas", + summary: "summary", + status: "pending", + createdAt: "2026-01-01T00:00:00.000Z", + }); + expect(updated.suggestions).toHaveLength(1); + const list = await store.listSuggestions(session.id); + expect(list).toHaveLength(1); + expect(list[0]?.kind).toBe("brainstorm"); + }); + + it("applies a suggestion (marks applied, updates updatedAt)", async () => { + const session = await store.createSession(); + const withSuggestion = await store.addSuggestion(session.id, { + sessionId: session.id, + kind: "brainstorm", + title: "Ideas", + text: "some ideas", + summary: "summary", + status: "pending", + }); + const suggestionId = withSuggestion.suggestions[0]?.id; + const updated = await store.applySuggestion(session.id, suggestionId!); + expect(updated.suggestions[0]?.status).toBe("applied"); + expect(updated.suggestions[0]?.appliedAt).toBeTruthy(); + }); + + it("resolves a suggestion (marks resolved)", async () => { + const session = await store.createSession(); + const withSuggestion = await store.addSuggestion(session.id, { + sessionId: session.id, + kind: "brainstorm", + title: "Ideas", + text: "some ideas", + summary: "summary", + status: "pending", + }); + const suggestionId = withSuggestion.suggestions[0]?.id; + const updated = await store.resolveSuggestion(session.id, suggestionId!, true); + expect(updated.suggestions[0]?.status).toBe("accepted"); + expect(updated.suggestions[0]?.resolvedAt).toBeTruthy(); + }); + + it("adds a note and updates it", async () => { + const session = await store.createSession(); + const withNote = await store.addNote(session.id, { + sessionId: session.id, + kind: "character", + title: "Aria", + body: "the protagonist", + tags: ["main"], + status: "draft", + }); + const noteId = withNote.notes[0]?.id; + expect(withNote.notes).toHaveLength(1); + const updated = await store.updateNote(session.id, noteId!, { body: "the protagonist, a mage" }); + expect(updated.notes[0]?.body).toBe("the protagonist, a mage"); + expect(updated.notes[0]?.tags).toEqual(["main"]); + }); + + it("merges notes into a single note", async () => { + const session = await store.createSession(); + const note1 = await store.addNote(session.id, { + sessionId: session.id, + kind: "character", + title: "Aria", + body: "part one", + tags: ["a"], + status: "draft", + }); + const note2 = await store.addNote(session.id, { + sessionId: session.id, + kind: "character", + title: "Notes", + body: "part two", + tags: ["b"], + status: "draft", + }); + const merged = await store.mergeNotes(session.id, [note1.notes[0]!.id, note2.notes[1]!.id]); + expect(merged.notes).toHaveLength(1); + expect(merged.notes[0]?.body).toBe("part one\n\npart two"); + expect(merged.notes[0]?.tags).toEqual(["a", "b"]); + }); + + it("deletes a session", async () => { + const session = await store.createSession(); + await store.deleteSession(session.id); + const list = await store.listSessions(); + expect(list).toHaveLength(0); + }); + + it("throws when mutating a missing session", async () => { + await expect(store.getSession("missing")).resolves.toBeNull(); + await expect(store.addMessage("missing", { + role: "user", content: "x", + })).rejects.toThrow(); + }); + + it("persists across a fresh store instance sharing the same backend", async () => { + const session = await store.createSession("persisted"); + const other = new AIStore("proj-1", persistence); + const list = await other.listSessions(); + expect(list.map((s) => s.id)).toContain(session.id); + await other.dispose(); + }); +}); diff --git a/packages/ai/src/ai-session.ts b/packages/ai/src/ai-session.ts new file mode 100644 index 0000000..c8fd8a8 --- /dev/null +++ b/packages/ai/src/ai-session.ts @@ -0,0 +1,371 @@ +/** + * AIStore — session, message, suggestion, and note persistence (CHUNK #9). + * + * The AIStore owns the in-memory cache of AI sessions and their messages, + * suggestions, and notes, and persists them through an injectable + * `AISessionPersistence` backend. The default backend is Dexie-backed + * IndexedDB (via `createIndexedDBPersistence`); tests and offline fallbacks use + * `createInMemoryPersistence()`. + * + * All mutating operations go through `withTransaction`, which applies the + * change to the in-memory cache and persists atomically. + * + * See CHUNK#9.md §2. + */ +import Dexie, { type Table } from "dexie"; + +import { createUlid, nowIso } from "@pubscript/core"; +import type { + AISession, + AIMessage, + AISuggestion, + AINote, + SuggestionKind, +} from "@pubscript/core"; + +/** A persisted session row. */ +interface SessionRow { + id: string; + projectId: string; + title: string; + kind: AISession["kind"]; + createdAt: string; + updatedAt: string; + lastMessageAt: string | null; +} + +/** The persistence contract the AIStore depends on. */ +export interface AISessionPersistence { + list(projectId: string): Promise; + get(id: string): Promise; + save(row: SessionRow): Promise; + delete(id: string): Promise; +} + +/** An in-memory persistence backend, for tests and offline fallbacks. */ +export function createInMemoryPersistence(): AISessionPersistence { + const rows = new Map(); + return { + async list(projectId: string) { + return [...rows.values()] + .filter((r) => r.projectId === projectId) + .sort((a, b) => b.updatedAt.localeCompare(a.updatedAt)); + }, + async get(id: string) { + return rows.get(id) ?? null; + }, + async save(row: SessionRow) { + rows.set(row.id, row); + }, + async delete(id: string) { + rows.delete(id); + }, + }; +} + +/** A Dexie-backed IndexedDB persistence backend. */ +export function createIndexedDBPersistence(dbName = "pubscript"): AISessionPersistence { + const db = new Dexie(dbName); + if (!db.tables.some((t) => t.name === "aiSessions")) { + db.version(1).stores({ aiSessions: "&id, projectId, updatedAt" }); + } + const table: Table = db.table("aiSessions"); + + return { + async list(projectId: string) { + return table.where("projectId").equals(projectId).sortBy("updatedAt").then((rows) => rows.reverse()); + }, + async get(id: string) { + return table.get(id) ?? null; + }, + async save(row: SessionRow) { + await table.put(row); + }, + async delete(id: string) { + await table.delete(id); + }, + }; +} + +/** + * The AIStore manages AI sessions and their messages, suggestions, and notes. + */ +export class AIStore { + private sessions = new Map(); + private loaded = false; + + constructor( + private readonly projectId: string, + private readonly persistence: AISessionPersistence, + ) {} + + /** Load all sessions for the project into the in-memory cache. */ + async load(): Promise { + if (this.loaded) return; + const rows = await this.persistence.list(this.projectId); + for (const row of rows) { + this.sessions.set(row.id, { + id: row.id, + projectId: row.projectId, + title: row.title, + kind: row.kind, + messages: [], + suggestions: [], + notes: [], + createdAt: row.createdAt, + updatedAt: row.updatedAt, + lastMessageAt: row.lastMessageAt, + }); + } + this.loaded = true; + } + + /** Clear the in-memory cache and mark it as unloaded. */ + async dispose(): Promise { + this.sessions.clear(); + this.loaded = false; + } + + private getRow(session: AISession): SessionRow { + return { + id: session.id, + projectId: session.projectId, + title: session.title, + kind: session.kind, + createdAt: session.createdAt, + updatedAt: session.updatedAt, + lastMessageAt: session.lastMessageAt, + }; + } + + /** Create a new session and persist it. */ + async createSession(title?: string): Promise { + await this.load(); + const now = nowIso(); + const session: AISession = { + id: createUlid(), + projectId: this.projectId, + title: title ?? `New AI session ${this.sessions.size + 1}`, + kind: "session", + messages: [], + suggestions: [], + notes: [], + createdAt: now, + updatedAt: now, + lastMessageAt: null, + }; + this.sessions.set(session.id, session); + await this.persistence.save(this.getRow(session)); + return this.clone(session); + } + + /** Get a session by id. */ + async getSession(id: string): Promise { + await this.load(); + const session = this.sessions.get(id); + return session ? this.clone(session) : null; + } + + /** List all sessions for the project, newest updated first. */ + async listSessions(): Promise { + await this.load(); + return [...this.sessions.values()] + .map((s) => this.clone(s)) + .sort((a, b) => { + if (a.updatedAt !== b.updatedAt) { + return b.updatedAt.localeCompare(a.updatedAt); + } + // ULIDs are time-sortable, so a higher id means a later creation. + return b.id.localeCompare(a.id); + }); + } + + /** Delete a session and its messages/suggestions/notes. */ + async deleteSession(id: string): Promise { + await this.load(); + this.sessions.delete(id); + await this.persistence.delete(id); + } + + /** Add a message to a session and persist it. */ + async addMessage(sessionId: string, message: Omit): Promise { + return this.withTransaction(async () => { + const session = this.require(sessionId); + const full: AIMessage = { + id: createUlid(), + createdAt: nowIso(), + ...message, + }; + session.messages.push(full); + session.lastMessageAt = full.createdAt; + session.updatedAt = full.createdAt; + return this.clone(session); + }); + } + + /** List a session's messages. */ + async listMessages(sessionId: string): Promise { + await this.load(); + const session = this.require(sessionId); + return session.messages.map((m) => ({ ...m })); + } + + /** Add a suggestion to a session and persist it. */ + async addSuggestion(sessionId: string, suggestion: Omit): Promise { + return this.withTransaction(async () => { + const session = this.require(sessionId); + const full: AISuggestion = { id: createUlid(), createdAt: nowIso(), ...suggestion }; + session.suggestions.push(full); + session.updatedAt = nowIso(); + return this.clone(session); + }); + } + + /** List a session's suggestions. */ + async listSuggestions(sessionId: string): Promise { + await this.load(); + const session = this.require(sessionId); + return session.suggestions.map((s) => ({ ...s })); + } + + /** List a session's suggestions of a given kind. */ + async listSuggestionsByKind(sessionId: string, kind: SuggestionKind): Promise { + const all = await this.listSuggestions(sessionId); + return all.filter((s) => s.kind === kind); + } + + /** + * Apply a suggestion: mark it applied, and return the applied suggestion. + * The caller (editor integration) is responsible for mutating the document + * via `SuggestionMutation`; this updates the store's record of the applied + * suggestion. + */ + async applySuggestion(sessionId: string, suggestionId: string): Promise { + return this.withTransaction(async () => { + const session = this.require(sessionId); + const suggestion = session.suggestions.find((s) => s.id === suggestionId); + if (!suggestion) throw new Error(`Suggestion ${suggestionId} not found`); + suggestion.status = "applied"; + suggestion.appliedAt = nowIso(); + session.updatedAt = nowIso(); + return this.clone(session); + }); + } + + /** + * Resolve (accept-and-keep or discard) a pending suggestion: mark it + * resolved and record the resolution. + */ + async resolveSuggestion(sessionId: string, suggestionId: string, accepted: boolean): Promise { + return this.withTransaction(async () => { + const session = this.require(sessionId); + const suggestion = session.suggestions.find((s) => s.id === suggestionId); + if (!suggestion) throw new Error(`Suggestion ${suggestionId} not found`); + suggestion.status = accepted ? "accepted" : "rejected"; + suggestion.resolvedAt = nowIso(); + suggestion.resolution = accepted ? "accepted" : "rejected"; + session.updatedAt = nowIso(); + return this.clone(session); + }); + } + + /** Add a note to a session and persist it. */ + async addNote(sessionId: string, note: Omit): Promise { + return this.withTransaction(async () => { + const session = this.require(sessionId); + const ts = nowIso(); + const full: AINote = { id: createUlid(), createdAt: ts, updatedAt: ts, ...note }; + session.notes.push(full); + session.updatedAt = ts; + return this.clone(session); + }); + } + + /** List a session's notes. */ + async listNotes(sessionId: string): Promise { + await this.load(); + const session = this.require(sessionId); + return session.notes.map((n) => ({ ...n })); + } + + /** Update a note's fields and persist it. */ + async updateNote(sessionId: string, noteId: string, patch: Partial): Promise { + return this.withTransaction(async () => { + const session = this.require(sessionId); + const note = session.notes.find((n) => n.id === noteId); + if (!note) throw new Error(`Note ${noteId} not found`); + Object.assign(note, patch, { updatedAt: nowIso() }); + session.updatedAt = note.updatedAt; + return this.clone(session); + }); + } + + /** + * Merge several notes into one: concatenate their bodies, union their tags, + * and replace them with the merged note. Returns the session containing the + * merged note. + */ + async mergeNotes(sessionId: string, noteIds: string[]): Promise { + return this.withTransaction(async () => { + const session = this.require(sessionId); + const toMerge = noteIds.map((id) => { + const note = session.notes.find((n) => n.id === id); + if (!note) throw new Error(`Note ${id} not found`); + return note; + }); + + const mergedBody = toMerge.map((n) => n.body).filter(Boolean).join("\n\n"); + const mergedTags = [...new Set(toMerge.flatMap((n) => n.tags ?? []))]; + const mergedTitle = toMerge[0]?.title ?? "Merged note"; + const ts = nowIso(); + + // Remove the merged notes. + session.notes = session.notes.filter((n) => !noteIds.includes(n.id)); + + const merged: AINote = { + id: createUlid(), + sessionId, + kind: toMerge[0]?.kind ?? "character", + title: mergedTitle, + body: mergedBody, + tags: mergedTags, + status: "draft", + createdAt: toMerge[0]?.createdAt ?? ts, + updatedAt: ts, + }; + session.notes.push(merged); + session.updatedAt = ts; + return this.clone(session); + }); + } + + /** + * Run a mutation within a transaction: apply the callback to the in-memory + * cache, then persist the session atomically. + */ + async withTransaction(fn: () => Promise): Promise { + await this.load(); + const result = await fn(); + // Persist every session touched by the callback (best-effort atomicity). + for (const session of this.sessions.values()) { + await this.persistence.save(this.getRow(session)); + } + return result; + } + + private require(sessionId: string): AISession { + const session = this.sessions.get(sessionId); + if (!session) throw new Error(`Session ${sessionId} not found`); + return session; + } + + /** Produce a deep copy of a session for external consumers. */ + private clone(session: AISession): AISession { + return { + ...session, + messages: session.messages.map((m) => ({ ...m })), + suggestions: session.suggestions.map((s) => ({ ...s })), + notes: session.notes.map((n) => ({ ...n })), + }; + } +} diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index fd817fd..80bea37 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -25,6 +25,39 @@ export { } from "./client.js"; export type { MeasuredBudget, StreamResult } from "./client.js"; +// Agent orchestration (CHUNK #9) +export { + AgentController, + isRateLimitError, + DEFAULT_MAX_RETRIES, + DEFAULT_RETRY_BACKOFF_MS, +} from "./agent.js"; +export type { AgentStatus, AgentCallbacks, AgentOptions } from "./agent.js"; + +// Suggestion + note builders +export { SuggestionBuilder } from "./suggestions.js"; +export type { + SuggestionKind, + SuggestionContext, + SuggestionPrompt, + SuggestionResult, +} from "./suggestions.js"; + +export { NoteBuilder } from "./notes.js"; +export type { NoteKind, NoteContext, NotePrompt, NoteResult } from "./notes.js"; + +// Suggestion mutation (apply/revert) +export { SuggestionMutation } from "./suggestion-mutation.js"; +export type { MutationResult } from "./suggestion-mutation.js"; + +// AIStore — session/message/suggestion/note persistence (CHUNK #9) +export { + AIStore, + createIndexedDBPersistence, + createInMemoryPersistence, +} from "./ai-session.js"; +export type { AISessionPersistence, SessionRow } from "./ai-session.js"; + // Provider adapters + factory export { getProvider } from "./providers/adapters.js"; diff --git a/packages/ai/src/notes.test.ts b/packages/ai/src/notes.test.ts new file mode 100644 index 0000000..14f19ae --- /dev/null +++ b/packages/ai/src/notes.test.ts @@ -0,0 +1,104 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { AgentController, NoteBuilder } from "./index.js"; +import type { ResolvedProvider } from "./index.js"; + +const config: ResolvedProvider = { + id: "test-openai", + kind: "openai-compatible", + baseUrl: "https://api.openai.com/v1", + apiKey: "sk-test", + model: "gpt-4o", + contextWindow: 128_000, +}; + +let savedFetch: typeof fetch; + +/** Build an SSE Response streaming a single text delta. */ +function sseResponse(events: string[]): Response { + const encoder = new TextEncoder(); + const body = new ReadableStream>({ + start(controller) { + for (const event of events) controller.enqueue(encoder.encode(`data: ${event}\n\n`)); + controller.enqueue(encoder.encode("data: [DONE]\n\n")); + controller.close(); + }, + }); + return new Response(body, { status: 200, headers: { "content-type": "text/event-stream" } }); +} + +describe("NoteBuilder", () => { + let builder: NoteBuilder; + + beforeEach(() => { + savedFetch = globalThis.fetch; + }); + + afterEach(() => { + globalThis.fetch = savedFetch; + vi.restoreAllMocks(); + }); + + it("streams a prompt and builds a draft note", async () => { + globalThis.fetch = (() => + sseResponse([JSON.stringify({ choices: [{ delta: { content: "Aria is the protagonist. She is a young mage with a mysterious past. #protagonist #mage" } }] })])) as unknown as typeof fetch; + + const agent = new AgentController(config); + builder = new NoteBuilder(agent); + + const result = await builder.build( + { kind: "character", prompt: "Summarize Aria", context: { title: "Aria" } }, + "sess-1", + ); + + expect(result.id).toBeTruthy(); + expect(result.sessionId).toBe("sess-1"); + expect(result.kind).toBe("character"); + expect(result.status).toBe("draft"); + expect(result.body).toContain("Aria"); + expect(result.tags).toContain("protagonist"); + expect(result.tags).toContain("mage"); + expect(result.createdAt).toBeTruthy(); + expect(result.updatedAt).toBeTruthy(); + }); + + it("extracts tags from a #tag list in the response", async () => { + globalThis.fetch = (() => + sseResponse([JSON.stringify({ choices: [{ delta: { content: "The Sunspire Mountains. #location #magic #north" } }] })])) as unknown as typeof fetch; + + const agent = new AgentController(config); + builder = new NoteBuilder(agent); + + const result = await builder.build( + { kind: "setting", prompt: "Describe the setting" }, + "sess-1", + ); + + expect(result.kind).toBe("setting"); + expect(result.tags.sort()).toEqual(["location", "magic", "north"]); + }); + + it("includes the note title in the system prompt", async () => { + const calls: Array = []; + globalThis.fetch = ((_url: string, init: RequestInit) => { + calls.push(init); + return Promise.resolve( + sseResponse([JSON.stringify({ choices: [{ delta: { content: "ok" } }] })]), + ); + }) as unknown as typeof fetch; + + const agent = new AgentController(config); + builder = new NoteBuilder(agent); + + await builder.build( + { kind: "plot", prompt: "Summarize the plot thread", context: { title: "The Stolen Relic" } }, + "sess-1", + ); + + const init = calls[0] as RequestInit; + const body = JSON.parse(init.body as string) as Record; + const messages = body.messages as Array<{ content: string }>; + const system = messages.find((m) => m.role === "system")?.content ?? ""; + expect(system).toContain("The Stolen Relic"); + }); +}); diff --git a/packages/ai/src/notes.ts b/packages/ai/src/notes.ts new file mode 100644 index 0000000..098a354 --- /dev/null +++ b/packages/ai/src/notes.ts @@ -0,0 +1,149 @@ +/** + * NoteBuilder — turns an agent stream into a structured note (CHUNK #9). + * + * Each note kind has a tailored system prompt. The streamed text is parsed + * into a title, body, and tags, then wrapped in the core `AINote` shape with + * `status: "draft"`. + * + * See CHUNK#9.md §4. + */ +import { createUlid, nowIso } from "@pubscript/core"; + +import { AgentController } from "./agent.js"; +import type { Message } from "./providers/types.js"; + +/** The kinds of notes the builder can generate. */ +export type NoteKind = "character" | "setting" | "plot" | "timeline" | "world" | "research"; + +/** Context about the note being generated. */ +export interface NoteContext { + title?: string; + existing?: string; +} + +/** A request to build a note. */ +export interface NotePrompt { + kind: NoteKind; + prompt: string; + context?: NoteContext; +} + +/** The parsed, draft note produced by the builder. */ +export interface NoteResult { + id: string; + sessionId: string; + kind: NoteKind; + title: string; + body: string; + tags: string[]; + status: "draft"; + createdAt: string; + updatedAt: string; +} + +/** The system prompt for each note kind. */ +const SYSTEM_PROMPTS: Record string> = { + character: (ctx) => [ + "You are a character-sheet assistant for a fiction author.", + ctx.title ? `Build a character reference for "${ctx.title}".` : "Build a concise character reference.", + ctx.existing ? `Existing notes: ${ctx.existing}.` : "", + "Capture name, age, appearance, personality, motivations, backstory, relationships, and voice. Use short fields, not prose. Flag gaps the author should fill.", + ] + .filter(Boolean) + .join(" "), + setting: (ctx) => [ + "You are a worldbuilding assistant for a fiction author.", + ctx.title ? `Build a setting reference for "${ctx.title}".` : "Build a concise setting reference.", + ctx.existing ? `Existing notes: ${ctx.existing}.` : "", + "Capture geography, climate, culture, politics, technology/magic, notable locations, and sensory details. Use short fields. Flag inconsistencies with canon.", + ] + .filter(Boolean) + .join(" "), + plot: (ctx) => [ + "You are a plot-structure assistant for a fiction author.", + ctx.title ? `Document a plot point or arc titled "${ctx.title}".` : "Document a plot point or arc.", + ctx.existing ? `Existing plot notes: ${ctx.existing}.` : "", + "Capture the setup, the conflict, the turning point, and the consequence. Note causation and foreshadowing. Flag plot holes or unresolved threads.", + ] + .filter(Boolean) + .join(" "), + timeline: (ctx) => [ + "You are a continuity assistant for a fiction author.", + ctx.existing ? `Existing timeline: ${ctx.existing}.` : "", + "Produce a chronological sequence of story events. Use clear, ordered entries. Flag contradictions with the established timeline.", + ] + .filter(Boolean) + .join(" "), + world: (ctx) => [ + "You are a worldbuilding encyclopedia assistant for a fiction author.", + ctx.title ? `Document a world element "${ctx.title}".` : "Document a world element.", + ctx.existing ? `Existing world notes: ${ctx.existing}.` : "", + "Document systems, rules, history, factions, and lore. Keep entries consistent and cross-referenced. Flag contradictions with established canon.", + ] + .filter(Boolean) + .join(" "), + research: (ctx) => [ + "You are a research assistant for a fiction author.", + ctx.title ? `Research "${ctx.title}".` : "Research a topic for a fiction author.", + ctx.existing ? `Existing research: ${ctx.existing}.` : "", + "Provide accurate, organized factual background. Distinguish established fact from plausible inference. Keep it practical and note sources qualitatively.", + ] + .filter(Boolean) + .join(" "), +}; + +/** + * Builds structured, draft notes from an agent stream. + */ +export class NoteBuilder { + constructor(private readonly agent: AgentController) {} + + /** + * Build a draft note from a prompt. Streams via the agent, then parses the + * response into a structured note. + */ + async build(prompt: NotePrompt, sessionId: string): Promise { + const system = SYSTEM_PROMPTS[prompt.kind](prompt.context ?? {}); + const messages: Message[] = [ + { role: "system", content: system }, + { role: "user", content: prompt.prompt }, + ]; + + let text = ""; + await this.agent.collect(messages, {}, { onText: (t) => (text += t) }); + + const parsed = this.parse(prompt.kind, text); + const ts = nowIso(); + return { + id: createUlid(), + sessionId, + kind: prompt.kind, + title: parsed.title, + body: parsed.body, + tags: parsed.tags, + status: "draft", + createdAt: ts, + updatedAt: ts, + }; + } + + /** Parse streamed text into a title, body, and tags. */ + private parse(kind: NoteKind, text: string): { title: string; body: string; tags: string[] } { + const cleaned = text.trim(); + const lines = cleaned.split(/\r?\n/).filter((l) => l.trim().length > 0); + const title = (lines[0] ?? "").replace(/^#+\s*/, "").replace(/^[-*]\s*/, "").trim().slice(0, 120) || `${kind} note`; + + // Extract tags from #word patterns and a leading "tags:" line. + const tagMatches = [...cleaned.matchAll(/#([\p{L}[\]_][\p{L}[\]0-9 ]{1,20})/gu)]; + const tags = [...new Set(tagMatches.map((m) => m[1].trim().toLowerCase()))].slice(0, 8); + + // Drop the title line and any tags line from the body. + const body = lines + .slice(1) + .join("\n") + .replace(/^tags?:[^\n]*\n/i, "") + .trim(); + + return { title, body: body || cleaned, tags }; + } +} diff --git a/packages/ai/src/suggestion-mutation.test.ts b/packages/ai/src/suggestion-mutation.test.ts new file mode 100644 index 0000000..a5af357 --- /dev/null +++ b/packages/ai/src/suggestion-mutation.test.ts @@ -0,0 +1,129 @@ +import { describe, expect, it } from "vitest"; + +import { SuggestionMutation } from "./index.js"; +import type { AISuggestion } from "@pubscript/core"; + +function suggestion(overrides: Partial = {}): AISuggestion { + return { + id: "s1", + sessionId: "sess-1", + kind: "expand", + title: "Expand scene", + text: "expanded text", + summary: "summary", + status: "pending", + createdAt: "2026-01-01T00:00:00.000Z", + ...overrides, + }; +} + +describe("SuggestionMutation", () => { + it("inserts new text at the anchor when there is no original span", () => { + const doc = "The sky was blue."; + const mut = new SuggestionMutation(suggestion({ + originalText: "", + newText: " bright", + anchorIndex: 11, + spanLength: 0, + })); + const result = mut.apply(doc); + expect(result.success).toBe(true); + expect(result.newDoc).toBe("The sky was bright blue."); + }); + + it("replaces an original text span with the new text", () => { + const doc = "She walked slowly down the path."; + const mut = new SuggestionMutation(suggestion({ + originalText: "slowly", + newText: "cautiously", + anchorIndex: 11, + spanLength: 6, + })); + const result = mut.apply(doc); + expect(result.success).toBe(true); + expect(result.newDoc).toBe("She walked cautiously down the path."); + }); + + it("rejects when the original span is not present at the anchor", () => { + const doc = "She walked quickly down the path."; + const mut = new SuggestionMutation(suggestion({ + originalText: "slowly", + newText: "cautiously", + anchorIndex: 11, + spanLength: 6, + })); + const result = mut.apply(doc); + expect(result.success).toBe(false); + expect(result.reason).toBe("not-found"); + }); + + it("rejects when the anchor is out of bounds", () => { + const doc = "short"; + const mut = new SuggestionMutation(suggestion({ + originalText: "x", + newText: "y", + anchorIndex: 100, + spanLength: 1, + })); + expect(mut.apply(doc).success).toBe(false); + }); + + it("rejects empty suggestions", () => { + const mut = new SuggestionMutation(suggestion({ originalText: " ", newText: "" })); + expect(mut.apply("anything").success).toBe(false); + }); + + it("falls back to text-based replacement when no anchor is given", () => { + const doc = "The cat sat. The cat ran."; + const mut = new SuggestionMutation(suggestion({ + originalText: "The cat sat", + newText: "The dog barked", + anchorIndex: undefined, + spanLength: undefined, + })); + const result = mut.apply(doc); + expect(result.success).toBe(true); + expect(result.newDoc).toBe("The dog barked. The cat ran."); + }); + + it("falls back to the first occurrence when the original appears multiple times", () => { + const doc = "The cat sat. The cat ran."; + const mut = new SuggestionMutation(suggestion({ + originalText: "cat", + newText: "dog", + anchorIndex: undefined, + spanLength: 3, + })); + const result = mut.apply(doc); + expect(result.success).toBe(true); + expect(result.newDoc).toBe("The dog sat. The cat ran."); + }); + + it("reverts an applied suggestion by restoring the original text", () => { + const doc = "The sky was blue."; + const mut = new SuggestionMutation(suggestion({ + originalText: "", + newText: " bright ", + anchorIndex: 11, + spanLength: 0, + })); + const applied = mut.apply(doc); + expect(applied.success).toBe(true); + const reverted = mut.revert(applied.newDoc!, applied); + expect(reverted.success).toBe(true); + expect(reverted.newDoc).toBe("The sky was blue."); + }); + + it("reports the inserted text and replaced length on apply", () => { + const doc = "Hello world"; + const mut = new SuggestionMutation(suggestion({ + originalText: "world", + newText: "there", + anchorIndex: 6, + spanLength: 5, + })); + const result = mut.apply(doc); + expect(result.insertedText).toBe("there"); + expect(result.replacedLength).toBe(5); + }); +}); diff --git a/packages/ai/src/suggestion-mutation.ts b/packages/ai/src/suggestion-mutation.ts new file mode 100644 index 0000000..fb78fb2 --- /dev/null +++ b/packages/ai/src/suggestion-mutation.ts @@ -0,0 +1,122 @@ +/** + * SuggestionMutation — applies or reverts a suggestion against document text + * (CHUNK #9). + * + * A suggestion carries the span of the document it replaces (`originalText`, + * `anchorIndex`, `spanLength`) and the replacement (`newText`). Applying a + * suggestion replaces that span; reverting restores the original span. + * + * The mutation is a pure text transformation: it never touches the store or the + * editor. It returns a result describing what happened so the caller can update + * the store and re-render. + * + * See CHUNK#9.md §5. + */ +import type { AISuggestion } from "@pubscript/core"; + +/** The outcome of applying or reverting a suggestion. */ +export interface MutationResult { + success: boolean; + /** Why a mutation failed, when `success` is false. */ + reason?: "not-found" | "ambiguous" | "empty" | "stale"; + /** The resulting document text, when successful. */ + newDoc?: string; + /** The text that was inserted (apply) or restored (revert). */ + insertedText?: string; + /** The index at which the change was anchored. */ + anchorIndex?: number; + /** The length of the span that was replaced. */ + replacedLength?: number; +} + +/** + * Applies or reverts a single suggestion against document text. + */ +export class SuggestionMutation { + constructor(private readonly suggestion: AISuggestion) {} + + /** + * Apply the suggestion: replace the original span with the new text. + * + * - If `anchorIndex`/`spanLength` are present, the span is verified against + * `originalText` (a mismatch means the suggestion is stale). + * - Otherwise the `originalText` is located by content (ambiguous if it + * appears more than once). + */ + apply(doc: string): MutationResult { + const { originalText, newText, anchorIndex, spanLength } = this.suggestion; + + if (!newText?.trim()) { + return { success: false, reason: "empty" }; + } + + // Insertion: no meaningful original span to replace. + if (!originalText?.trim() && typeof anchorIndex === "number" && typeof spanLength === "number") { + if (anchorIndex < 0 || anchorIndex > doc.length) { + return { success: false, reason: "not-found" }; + } + const newDoc = doc.slice(0, anchorIndex) + newText + doc.slice(anchorIndex); + return { + success: true, + newDoc, + insertedText: newText, + anchorIndex, + replacedLength: 0, + }; + } + + // Index-based application (preferred: unambiguous). + if (typeof anchorIndex === "number" && typeof spanLength === "number") { + if (anchorIndex < 0 || anchorIndex + spanLength > doc.length) { + return { success: false, reason: "not-found" }; + } + const actual = doc.slice(anchorIndex, anchorIndex + spanLength); + if (actual !== originalText) { + return { success: false, reason: "not-found" }; + } + const newDoc = doc.slice(0, anchorIndex) + newText + doc.slice(anchorIndex + spanLength); + return { + success: true, + newDoc, + insertedText: newText, + anchorIndex, + replacedLength: spanLength, + }; + } + + // Content-based fallback: replace the first occurrence. + const needle = originalText ?? ""; + const first = doc.indexOf(needle); + if (first === -1) { + return { success: false, reason: "not-found" }; + } + const newDoc = doc.slice(0, first) + newText + doc.slice(first + needle.length); + return { + success: true, + newDoc, + insertedText: newText, + anchorIndex: first, + replacedLength: needle.length, + }; + } + + /** + * Revert a previously applied suggestion: remove the inserted text and + * restore the original span. `applied` is the result of a prior `apply`. + */ + revert(doc: string, applied: MutationResult): MutationResult { + if (!applied.success || applied.anchorIndex == null || applied.insertedText == null) { + return { success: false, reason: "not-found" }; + } + const { anchorIndex, insertedText } = applied; + const original = this.suggestion.originalText ?? ""; + const newDoc = doc.slice(0, anchorIndex) + original + doc.slice(anchorIndex + insertedText.length); + return { + success: true, + newDoc, + insertedText: original, + anchorIndex, + replacedLength: insertedText.length, + }; + } +} diff --git a/packages/ai/src/suggestions.test.ts b/packages/ai/src/suggestions.test.ts new file mode 100644 index 0000000..d6fcfa9 --- /dev/null +++ b/packages/ai/src/suggestions.test.ts @@ -0,0 +1,107 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { AgentController, SuggestionBuilder } from "./index.js"; +import type { ResolvedProvider } from "./index.js"; + +const config: ResolvedProvider = { + id: "test-openai", + kind: "openai-compatible", + baseUrl: "https://api.openai.com/v1", + apiKey: "sk-test", + model: "gpt-4o", + contextWindow: 128_000, +}; + +let savedFetch: typeof fetch; + +/** Build an SSE Response streaming a single text delta. */ +function sseResponse(events: string[]): Response { + const encoder = new TextEncoder(); + const body = new ReadableStream>({ + start(controller) { + for (const event of events) controller.enqueue(encoder.encode(`data: ${event}\n\n`)); + controller.enqueue(encoder.encode("data: [DONE]\n\n")); + controller.close(); + }, + }); + return new Response(body, { status: 200, headers: { "content-type": "text/event-stream" } }); +} + +describe("SuggestionBuilder", () => { + let builder: SuggestionBuilder; + + beforeEach(() => { + savedFetch = globalThis.fetch; + }); + + afterEach(() => { + globalThis.fetch = savedFetch; + vi.restoreAllMocks(); + }); + + it("streams a prompt and builds a pending suggestion", async () => { + globalThis.fetch = (() => + sseResponse([JSON.stringify({ choices: [{ delta: { content: "Try a flashback here. It reveals the secret early." } }] })])) as unknown as typeof fetch; + + const agent = new AgentController(config); + builder = new SuggestionBuilder(agent); + + const result = await builder.build( + { kind: "expand", prompt: "Expand the reveal scene" }, + "sess-1", + "msg-1", + ); + + expect(result.id).toBeTruthy(); + expect(result.sessionId).toBe("sess-1"); + expect(result.sourceMessageId).toBe("msg-1"); + expect(result.kind).toBe("expand"); + expect(result.status).toBe("pending"); + expect(result.text).toContain("flashback"); + expect(result.title.length).toBeGreaterThan(0); + expect(result.summary.length).toBeGreaterThan(0); + expect(result.createdAt).toBeTruthy(); + }); + + it("includes the context in the system prompt", async () => { + const calls: Array = []; + globalThis.fetch = ((_url: string, init: RequestInit) => { + calls.push(init); + return Promise.resolve( + sseResponse([JSON.stringify({ choices: [{ delta: { content: "ok" } }] })]), + ); + }) as unknown as typeof fetch; + + const agent = new AgentController(config); + builder = new SuggestionBuilder(agent); + + await builder.build( + { + kind: "outline", + prompt: "Outline chapter 3", + context: { chapterTitle: "The Return", setting: "the capital", characters: ["Aria"] }, + }, + "sess-1", + "msg-1", + ); + + const init = calls[0] as RequestInit; + const body = JSON.parse(init.body as string) as Record; + const messages = body.messages as Array<{ content: string }>; + const system = messages.find((m) => m.role === "system")?.content ?? ""; + expect(system).toContain("The Return"); + expect(system).toContain("the capital"); + expect(system).toContain("Aria"); + }); + + it("produces distinct titles per kind", async () => { + globalThis.fetch = (() => + sseResponse([JSON.stringify({ choices: [{ delta: { content: "Chapter 3: The Return — Aria confronts the council about the missing supplies." } }] })])) as unknown as typeof fetch; + + const agent = new AgentController(config); + builder = new SuggestionBuilder(agent); + + const result = await builder.build({ kind: "outline", prompt: "Outline" }, "sess-1", "msg-1"); + expect(result.title).toContain("Chapter 3"); + }); +}); diff --git a/packages/ai/src/suggestions.ts b/packages/ai/src/suggestions.ts new file mode 100644 index 0000000..926278e --- /dev/null +++ b/packages/ai/src/suggestions.ts @@ -0,0 +1,170 @@ +/** + * SuggestionBuilder — turns an agent stream into a structured, pending + * suggestion (CHUNK #9). + * + * Each suggestion kind has a tailored system prompt that shapes the model's + * output. The streamed text is parsed into a title, summary, and full body, + * then wrapped in the core `AISuggestion` shape with `status: "pending"`. + * + * See CHUNK#9.md §3. + */ +import { createUlid, nowIso } from "@pubscript/core"; + +import { AgentController } from "./agent.js"; +import type { Message } from "./providers/types.js"; + +/** The kinds of suggestions the builder can generate. */ +export type SuggestionKind = + | "brainstorm" + | "outline" + | "expand" + | "tighten" + | "continue" + | "research" + | "dialogue" + | "worldbuilding"; + +/** Context about the current scene, used to tailor the system prompt. */ +export interface SuggestionContext { + chapterTitle?: string; + characters?: string[]; + setting?: string; + previous?: string; +} + +/** A request to build a suggestion. */ +export interface SuggestionPrompt { + kind: SuggestionKind; + prompt: string; + context?: SuggestionContext; +} + +/** The parsed, pending suggestion produced by the builder. */ +export interface SuggestionResult { + id: string; + sessionId: string; + kind: SuggestionKind; + title: string; + text: string; + summary: string; + sourceMessageId: string; + status: "pending"; + createdAt: string; +} + +/** The system prompt for each suggestion kind. */ +const SYSTEM_PROMPTS: Record string> = { + brainstorm: (ctx) => [ + "You are a creative writing collaborator helping an author brainstorm ideas.", + ctx.chapterTitle ? `The current chapter is titled "${ctx.chapterTitle}".` : "", + ctx.characters && ctx.characters.length ? `Relevant characters: ${ctx.characters.join(", ")}.` : "", + ctx.setting ? `The story is set in ${ctx.setting}.` : "", + "Propose several distinct, concrete directions the story could take next. Be specific and evocative, but concise. Avoid prose that reads like finished text — these are seeds, not drafts.", + ] + .filter(Boolean) + .join(" "), + outline: (ctx) => [ + "You are a story-structuring assistant helping an author outline a scene or chapter.", + ctx.chapterTitle ? `The current chapter is titled "${ctx.chapterTitle}".` : "", + ctx.previous ? `The preceding material: ${ctx.previous}.` : "", + ctx.characters && ctx.characters.length ? `Characters in the scene: ${ctx.characters.join(", ")}.` : "", + ctx.setting ? `The story is set in ${ctx.setting}.` : "", + "Produce a clear, ordered outline (numbered beats or a short bulleted sequence) that advances the scene. Each beat should be one or two sentences.", + ] + .filter(Boolean) + .join(" "), + expand: (ctx) => [ + "You are a prose-expansion assistant helping an author flesh out a scene.", + ctx.chapterTitle ? `The current chapter is titled "${ctx.chapterTitle}".` : "", + ctx.characters && ctx.characters.length ? `Characters in the scene: ${ctx.characters.join(", ")}.` : "", + "Expand the author's outline or summary into richer, more detailed prose. Preserve the author's voice and intent; add sensory detail, pacing, and subtext without changing the plot.", + ] + .filter(Boolean) + .join(" "), + tighten: (ctx) => [ + "You are a line-editing assistant helping an author tighten prose.", + ctx.chapterTitle ? `The current chapter is titled "${ctx.chapterTitle}".` : "", + "Identify redundancy, wordiness, and weak verbs, then return a tightened version of the text. Preserve meaning, voice, and intent. Make only necessary changes.", + ] + .filter(Boolean) + .join(" "), + continue: (ctx) => [ + "You are a collaborative fiction writer continuing a story in the author's voice.", + ctx.chapterTitle ? `The current chapter is titled "${ctx.chapterTitle}".` : "", + ctx.previous ? `What has happened so far: ${ctx.previous}.` : "", + ctx.characters && ctx.characters.length ? `Characters present: ${ctx.characters.join(", ")}.` : "", + "Write the next segment of the story, continuing naturally from what precedes it. Match the established voice, tense, and POV. Do not summarize or skip ahead — write the scene.", + ] + .filter(Boolean) + .join(" "), + research: (ctx) => [ + "You are a research assistant for a fiction author.", + ctx.setting ? `The story is set in ${ctx.setting}.` : "", + "Provide accurate, well-organized factual background that supports the author's worldbuilding and plot. Distinguish established fact from plausible inference. Keep it practical and citation-aware.", + ] + .filter(Boolean) + .join(" "), + dialogue: (ctx) => [ + "You are a dialogue-writing assistant helping an author craft natural, character-specific dialogue.", + ctx.characters && ctx.characters.length ? `Speaking characters: ${ctx.characters.join(", ")}.` : "", + ctx.setting ? `The scene is set in ${ctx.setting}.` : "", + "Write dialogue that sounds natural, reveals character, and carries subtext. Attribute speakers clearly. Avoid on-the-nose exposition.", + ] + .filter(Boolean) + .join(" "), + worldbuilding: (ctx) => [ + "You are a worldbuilding assistant helping an author develop a consistent, vivid fictional world.", + ctx.setting ? `The setting under development: ${ctx.setting}.` : "", + "Develop the world consistently: geography, culture, politics, technology, magic systems, economics, or daily life. Flag any potential inconsistencies with established canon and suggest resolutions.", + ] + .filter(Boolean) + .join(" "), +}; + +/** + * Builds structured, pending suggestions from an agent stream. + */ +export class SuggestionBuilder { + constructor(private readonly agent: AgentController) {} + + /** + * Build a pending suggestion from a prompt. Streams via the agent, then + * parses the response into a structured suggestion. + */ + async build( + prompt: SuggestionPrompt, + sessionId: string, + sourceMessageId: string, + ): Promise { + const system = SYSTEM_PROMPTS[prompt.kind](prompt.context ?? {}); + const messages: Message[] = [ + { role: "system", content: system }, + { role: "user", content: prompt.prompt }, + ]; + + let text = ""; + await this.agent.collect(messages, {}, { onText: (t) => (text += t) }); + + const parsed = this.parse(prompt.kind, text); + return { + id: createUlid(), + sessionId, + kind: prompt.kind, + title: parsed.title, + text: parsed.text, + summary: parsed.summary, + sourceMessageId, + status: "pending", + createdAt: nowIso(), + }; + } + + /** Parse streamed text into a title, summary, and body. */ + private parse(kind: SuggestionKind, text: string): { title: string; summary: string; text: string } { + const cleaned = text.trim(); + const firstLine = cleaned.split(/\r?\n/).find((l) => l.trim().length > 0) ?? ""; + const title = firstLine.replace(/^#+\s*/, "").trim().slice(0, 120) || `${kind} suggestion`; + const summary = cleaned.split(/\r?\n\r?\n/)[0].replace(/^#+\s*/, "").trim().slice(0, 300); + return { title, summary, text: cleaned }; + } +} diff --git a/packages/core/src/ids.ts b/packages/core/src/ids.ts index 67e88ba..d7d2fe0 100644 --- a/packages/core/src/ids.ts +++ b/packages/core/src/ids.ts @@ -1,10 +1,15 @@ -import { ulid } from "ulid"; +import { monotonicFactory } from "ulid"; const ULID_REGEX = /^[0-7][0-9A-HJKMNP-TV-Z]{25}$/; +// Monotonic factory: guarantees that successive calls produce ULIDs that sort +// in creation order even within the same millisecond. This is important for +// tiebreakers that rely on ULID ordering (e.g. reverse-chronological lists). +const monotonicUlid = monotonicFactory(); + /** Generate a monotonic ULID. */ export function createUlid(): string { - return ulid(); + return monotonicUlid(); } /** Validate a ULID string. */ diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bb231fc..f94cd85 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -90,6 +90,9 @@ importers: '@pubscript/core': specifier: workspace:* version: link:../core + dexie: + specifier: ^4.0.11 + version: 4.4.7 devDependencies: '@types/node': specifier: ^22.10.2