feat: add unit tests, Docker polish, and workbench UX improvements

Unit tests: Consumer class (7), recursive-splitter (10), parseJsonResponse (11) — 28 total.
Docker: add 5 commented LLM provider services, dev compose override, .env.example.
Workbench: chat persistence, error boundary, disconnect banner, prompts error handling.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

ts/packages/base/src/__tests__/consumer.test.ts

@@ -0,0 +1,273 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+import { Consumer, type ConsumerOptions, type FlowContext } from "../messaging/consumer.js";
+import type {
+  PubSubBackend,
+  BackendConsumer,
+  Message,
+  BackendProducer,
+  CreateProducerOptions,
+  CreateConsumerOptions,
+} from "../backend/types.js";
+import { TooManyRequestsError } from "../errors.js";
+import type { Flow } from "../processor/flow.js";
+
+// ── Mock Message ──────────────────────────────────────────────────────
+function createMockMessage<T>(val: T, props: Record<string, string> = {}): Message<T> {
+  return {
+    value: () => val,
+    properties: () => props,
+  };
+}
+
+// ── Mock BackendConsumer ──────────────────────────────────────────────
+function createMockBackendConsumer<T>(): BackendConsumer<T> & {
+  receive: ReturnType<typeof vi.fn>;
+  acknowledge: ReturnType<typeof vi.fn>;
+  negativeAcknowledge: ReturnType<typeof vi.fn>;
+  unsubscribe: ReturnType<typeof vi.fn>;
+  close: ReturnType<typeof vi.fn>;
+} {
+  return {
+    receive: vi.fn().mockResolvedValue(null),
+    acknowledge: vi.fn().mockResolvedValue(undefined),
+    negativeAcknowledge: vi.fn().mockResolvedValue(undefined),
+    unsubscribe: vi.fn().mockResolvedValue(undefined),
+    close: vi.fn().mockResolvedValue(undefined),
+  };
+}
+
+// ── Mock PubSubBackend ───────────────────────────────────────────────
+function createMockPubSub<T>(
+  backendConsumer: BackendConsumer<T>,
+): PubSubBackend {
+  return {
+    createProducer: vi.fn().mockResolvedValue({} as BackendProducer<unknown>),
+    createConsumer: vi.fn().mockResolvedValue(backendConsumer),
+    close: vi.fn().mockResolvedValue(undefined),
+  };
+}
+
+// ── Minimal FlowContext stub ─────────────────────────────────────────
+function createFlowContext(): FlowContext {
+  return {
+    id: "test-flow-id",
+    name: "test-flow",
+    flow: {} as Flow,
+  };
+}
+
+describe("Consumer", () => {
+  let backendConsumer: ReturnType<typeof createMockBackendConsumer>;
+  let pubsub: PubSubBackend;
+  let flowCtx: FlowContext;
+
+  beforeEach(() => {
+    backendConsumer = createMockBackendConsumer();
+    pubsub = createMockPubSub(backendConsumer);
+    flowCtx = createFlowContext();
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  // ── Constructor ──────────────────────────────────────────────────
+  it("stores options and applies defaults", () => {
+    const handler = vi.fn();
+    const consumer = new Consumer({
+      pubsub,
+      topic: "my-topic",
+      subscription: "my-sub",
+      handler,
+    });
+
+    // Access private fields via any-cast to verify defaults
+    expect((consumer as any).concurrency).toBe(1);
+    expect((consumer as any).rateLimitRetryMs).toBe(10_000);
+  });
+
+  it("accepts custom concurrency and rateLimitRetryMs", () => {
+    const consumer = new Consumer({
+      pubsub,
+      topic: "t",
+      subscription: "s",
+      handler: vi.fn(),
+      concurrency: 4,
+      rateLimitRetryMs: 5_000,
+    });
+
+    expect((consumer as any).concurrency).toBe(4);
+    expect((consumer as any).rateLimitRetryMs).toBe(5_000);
+  });
+
+  // ── start() creates consumer and calls handler ─────────────────
+  it("creates a backend consumer and invokes handler for received messages", async () => {
+    const handler = vi.fn().mockResolvedValue(undefined);
+    const msg = createMockMessage({ data: "hello" }, { id: "1" });
+
+    // First call returns a message, second call triggers stop
+    let callCount = 0;
+    backendConsumer.receive.mockImplementation(async () => {
+      callCount++;
+      if (callCount === 1) return msg;
+      // Stop the consumer on second receive
+      await consumer.stop();
+      return null;
+    });
+
+    const consumer = new Consumer({
+      pubsub,
+      topic: "topic-a",
+      subscription: "sub-a",
+      handler,
+    });
+
+    // start() blocks until the consume loop ends, so we don't need to await separately
+    await consumer.start(flowCtx);
+
+    expect(pubsub.createConsumer).toHaveBeenCalledWith({
+      topic: "topic-a",
+      subscription: "sub-a",
+      initialPosition: "latest",
+    });
+    expect(handler).toHaveBeenCalledWith({ data: "hello" }, { id: "1" }, flowCtx);
+  });
+
+  // ── Messages are acknowledged after successful handling ────────
+  it("acknowledges messages after successful handling", async () => {
+    const handler = vi.fn().mockResolvedValue(undefined);
+    const msg = createMockMessage("payload");
+
+    let callCount = 0;
+    backendConsumer.receive.mockImplementation(async () => {
+      callCount++;
+      if (callCount === 1) return msg;
+      await consumer.stop();
+      return null;
+    });
+
+    const consumer = new Consumer({
+      pubsub,
+      topic: "t",
+      subscription: "s",
+      handler,
+    });
+
+    await consumer.start(flowCtx);
+
+    expect(backendConsumer.acknowledge).toHaveBeenCalledWith(msg);
+    expect(backendConsumer.negativeAcknowledge).not.toHaveBeenCalled();
+  });
+
+  // ── Messages are negatively acknowledged on handler error ──────
+  it("negatively acknowledges messages when the handler throws", async () => {
+    const handler = vi.fn().mockRejectedValue(new Error("handler boom"));
+    const msg = createMockMessage("bad-payload");
+
+    let callCount = 0;
+    backendConsumer.receive.mockImplementation(async () => {
+      callCount++;
+      if (callCount === 1) return msg;
+      // Stop on second call (after the 1s sleep from error handling)
+      await consumer.stop();
+      return null;
+    });
+
+    const consumer = new Consumer({
+      pubsub,
+      topic: "t",
+      subscription: "s",
+      handler,
+    });
+
+    // Suppress console.error noise
+    const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
+
+    // start() will block; the error path sleeps 1s, so we need to advance timers
+    const startPromise = consumer.start(flowCtx);
+    // Advance past the 1s sleep in the error handler
+    await vi.advanceTimersByTimeAsync(1500);
+    await startPromise;
+
+    expect(backendConsumer.negativeAcknowledge).toHaveBeenCalledWith(msg);
+    expect(backendConsumer.acknowledge).not.toHaveBeenCalled();
+
+    errorSpy.mockRestore();
+  });
+
+  // ── TooManyRequestsError triggers retry ────────────────────────
+  it("retries the handler on TooManyRequestsError", async () => {
+    let handlerCalls = 0;
+    const handler = vi.fn().mockImplementation(async () => {
+      handlerCalls++;
+      if (handlerCalls === 1) {
+        throw new TooManyRequestsError("rate limited");
+      }
+      // Second call succeeds
+    });
+
+    const msg = createMockMessage("rate-limited-payload");
+
+    let receiveCount = 0;
+    backendConsumer.receive.mockImplementation(async () => {
+      receiveCount++;
+      if (receiveCount === 1) return msg;
+      await consumer.stop();
+      return null;
+    });
+
+    const consumer = new Consumer({
+      pubsub,
+      topic: "t",
+      subscription: "s",
+      handler,
+      rateLimitRetryMs: 500,
+    });
+
+    const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
+
+    const startPromise = consumer.start(flowCtx);
+    // Advance past the rate-limit retry delay (500ms)
+    await vi.advanceTimersByTimeAsync(600);
+    await startPromise;
+
+    // Handler called twice: first throws TooManyRequestsError, second succeeds
+    expect(handler).toHaveBeenCalledTimes(2);
+    // Message should be acknowledged (retry succeeded)
+    expect(backendConsumer.acknowledge).toHaveBeenCalledWith(msg);
+
+    warnSpy.mockRestore();
+  });
+
+  // ── stop() closes the backend ──────────────────────────────────
+  it("stop() sets running=false and closes the backend", async () => {
+    // Make receive block forever (returns null) until stopped
+    backendConsumer.receive.mockImplementation(async () => {
+      // Yield control so stop() can run
+      await new Promise((r) => setTimeout(r, 100));
+      return null;
+    });
+
+    const consumer = new Consumer({
+      pubsub,
+      topic: "t",
+      subscription: "s",
+      handler: vi.fn(),
+    });
+
+    const startPromise = consumer.start(flowCtx);
+
+    // Advance timers to let the consume loop iterate once
+    await vi.advanceTimersByTimeAsync(200);
+
+    await consumer.stop();
+
+    // Advance timers further so the loop can exit
+    await vi.advanceTimersByTimeAsync(200);
+    await startPromise;
+
+    expect(backendConsumer.close).toHaveBeenCalled();
+    expect((consumer as any).running).toBe(false);
+  });
+});

ts/packages/base/tsconfig.json

@@ -5,5 +5,6 @@
     "rootDir": "src",
     "composite": true
   },
-  "include": ["src"]
+  "include": ["src"],
+  "exclude": ["src/**/*.test.ts", "src/**/*.spec.ts"]
 }

ts/packages/base/vitest.config.ts

@@ -0,0 +1,6 @@
+import { defineConfig } from "vitest/config";
+export default defineConfig({
+  test: {
+    globals: true,
+  },
+});

ts/packages/flow/src/__tests__/parse-json.test.ts

@@ -0,0 +1,95 @@
+import { describe, it, expect, vi } from "vitest";
+import { parseJsonResponse } from "../extract/knowledge-extract.js";
+
+describe("parseJsonResponse", () => {
+  // Suppress console.warn from the function under test
+  beforeEach(() => {
+    vi.spyOn(console, "warn").mockImplementation(() => {});
+  });
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  // ── Valid JSON array ────────────────────────────────────────────
+  it("parses a valid JSON array", () => {
+    const result = parseJsonResponse<{ a: number }[]>('[{"a":1}]');
+    expect(result).toEqual([{ a: 1 }]);
+  });
+
+  // ── JSON with markdown fences ──────────────────────────────────
+  it("strips markdown fences and parses JSON", () => {
+    const input = '```json\n[{"a":1}]\n```';
+    const result = parseJsonResponse<{ a: number }[]>(input);
+    expect(result).toEqual([{ a: 1 }]);
+  });
+
+  // ── JSON embedded in surrounding text ──────────────────────────
+  it("extracts JSON array embedded in surrounding text", () => {
+    const input = 'Here is the result: [{"a":1}] hope that helps';
+    const result = parseJsonResponse<{ a: number }[]>(input);
+    expect(result).toEqual([{ a: 1 }]);
+  });
+
+  // ── Truncated array ────────────────────────────────────────────
+  it("repairs truncated array by closing at last complete object", () => {
+    const input = '[{"a":1},{"b":2';
+    const result = parseJsonResponse<Record<string, number>[]>(input);
+    expect(result).toEqual([{ a: 1 }]);
+  });
+
+  // ── Single object (not array) ──────────────────────────────────
+  it("parses a single object directly (valid JSON passes Attempt 1)", () => {
+    const input = '{"a":1}';
+    const result = parseJsonResponse<{ a: number }>(input);
+    // A bare object is valid JSON, so Attempt 1 (JSON.parse) succeeds directly
+    expect(result).toEqual({ a: 1 });
+  });
+
+  it("wraps a single object in an array when embedded in non-JSON text", () => {
+    // When the object is surrounded by garbage, Attempt 1 and 2 fail,
+    // so Attempt 4 extracts the object and wraps it in an array
+    const input = 'some text {"a":1} more text';
+    const result = parseJsonResponse<{ a: number }[]>(input);
+    expect(result).toEqual([{ a: 1 }]);
+  });
+
+  // ── Complete garbage ───────────────────────────────────────────
+  it("returns null for complete garbage", () => {
+    const result = parseJsonResponse("not json at all");
+    expect(result).toBeNull();
+  });
+
+  // ── Empty string ───────────────────────────────────────────────
+  it("returns null for empty string", () => {
+    const result = parseJsonResponse("");
+    expect(result).toBeNull();
+  });
+
+  // ── Nested fences with language tag ────────────────────────────
+  it("parses JSON inside fences with language tag (single object)", () => {
+    const input = '```json\n{"key":"value"}\n```';
+    const result = parseJsonResponse<{ key: string }[]>(input);
+    // The function first strips fences, then tries JSON.parse which yields an object,
+    // then if that fails as array extraction, falls back to wrapping in array
+    // Actually: JSON.parse of '{"key":"value"}' succeeds directly, returning the object
+    expect(result).toEqual({ key: "value" });
+  });
+
+  // ── Multiple objects in valid array ────────────────────────────
+  it("parses a multi-element array correctly", () => {
+    const input = '[{"name":"Alice"},{"name":"Bob"},{"name":"Carol"}]';
+    const result = parseJsonResponse<{ name: string }[]>(input);
+    expect(result).toEqual([
+      { name: "Alice" },
+      { name: "Bob" },
+      { name: "Carol" },
+    ]);
+  });
+
+  // ── Fences without language tag ────────────────────────────────
+  it("strips fences without a language tag", () => {
+    const input = '```\n[{"x":42}]\n```';
+    const result = parseJsonResponse<{ x: number }[]>(input);
+    expect(result).toEqual([{ x: 42 }]);
+  });
+});

ts/packages/flow/src/__tests__/recursive-splitter.test.ts

@@ -0,0 +1,97 @@
+import { describe, it, expect } from "vitest";
+import { recursiveSplit } from "../chunking/recursive-splitter.js";
+
+describe("recursiveSplit", () => {
+  // ── Short text returns single chunk ──────────────────────────────
+  it("returns single chunk when text is shorter than chunkSize", () => {
+    const result = recursiveSplit("Hello world", 100, 10);
+    expect(result).toEqual(["Hello world"]);
+  });
+
+  // ── Empty/whitespace text returns empty array ────────────────────
+  it("returns empty array for empty string", () => {
+    expect(recursiveSplit("", 100, 10)).toEqual([]);
+  });
+
+  it("returns empty array for whitespace-only text", () => {
+    expect(recursiveSplit("   \n\n  \n  ", 100, 10)).toEqual([]);
+  });
+
+  // ── Splits on paragraph boundary (\n\n) first ───────────────────
+  it("splits on paragraph boundary (\\n\\n) first", () => {
+    const text = "Paragraph one content here.\n\nParagraph two content here.";
+    const result = recursiveSplit(text, 30, 0);
+    expect(result.length).toBeGreaterThanOrEqual(2);
+    // Each chunk should contain content from its respective paragraph
+    expect(result[0]).toContain("Paragraph one");
+    expect(result[result.length - 1]).toContain("Paragraph two");
+  });
+
+  // ── Splits on \n when no \n\n present ────────────────────────────
+  it("splits on newline when no paragraph boundary present", () => {
+    const text = "Line one content.\nLine two content.\nLine three content.";
+    const result = recursiveSplit(text, 25, 0);
+    expect(result.length).toBeGreaterThanOrEqual(2);
+    expect(result[0]).toContain("Line one");
+  });
+
+  // ── Splits on spaces when no newlines present ────────────────────
+  it("splits on spaces when no newlines present", () => {
+    const text = "word1 word2 word3 word4 word5 word6 word7 word8 word9 word10";
+    const result = recursiveSplit(text, 20, 0);
+    expect(result.length).toBeGreaterThanOrEqual(2);
+    // Each chunk should be at most roughly chunkSize
+    for (const chunk of result) {
+      // Allow some tolerance for the splitting algorithm
+      expect(chunk.length).toBeLessThanOrEqual(30);
+    }
+  });
+
+  // ── Character-level split as last resort ─────────────────────────
+  it("splits at character level as last resort", () => {
+    // A single long word with no separators
+    const text = "abcdefghijklmnopqrstuvwxyz0123456789abcdefghijklmnopqrstuvwxyz";
+    const result = recursiveSplit(text, 10, 0);
+    expect(result.length).toBeGreaterThanOrEqual(2);
+    // Reassembled text should match original
+    expect(result.join("")).toBe(text);
+  });
+
+  // ── Overlap: second chunk starts with tail of first ──────────────
+  it("applies overlap so second chunk starts with tail of first", () => {
+    const text = "First paragraph here.\n\nSecond paragraph here.";
+    const result = recursiveSplit(text, 25, 5);
+    expect(result.length).toBeGreaterThanOrEqual(2);
+    if (result.length >= 2) {
+      // The second chunk should start with the last 5 chars of the first
+      const firstTail = result[0].slice(-5);
+      expect(result[1].startsWith(firstTail)).toBe(true);
+    }
+  });
+
+  // ── Large text produces multiple chunks ──────────────────────────
+  it("large text produces multiple chunks of approximately chunkSize", () => {
+    // Create a large block of text with paragraph separators
+    const paragraphs = Array.from(
+      { length: 20 },
+      (_, i) => `This is paragraph number ${i + 1} with some filler content to make it longer.`,
+    );
+    const text = paragraphs.join("\n\n");
+    const result = recursiveSplit(text, 100, 10);
+    expect(result.length).toBeGreaterThan(5);
+  });
+
+  // ── chunkOverlap=0 produces no overlap ───────────────────────────
+  it("chunkOverlap=0 produces no overlap between chunks", () => {
+    const text = "AAAA\n\nBBBB\n\nCCCC\n\nDDDD";
+    const result = recursiveSplit(text, 8, 0);
+    expect(result.length).toBeGreaterThanOrEqual(2);
+    // With zero overlap, no chunk (except possibly the first) should start with previous chunk's tail
+    for (let i = 1; i < result.length; i++) {
+      const prevTail = result[i - 1].slice(-3);
+      // The next chunk should NOT start with the previous chunk's tail
+      // (unless they happen to share content naturally, which won't happen with AAAA/BBBB/etc.)
+      expect(result[i].startsWith(prevTail)).toBe(false);
+    }
+  });
+});

ts/packages/flow/src/extract/knowledge-extract.ts

@@ -261,7 +261,7 @@ function literalTerm(value: string): Term {
  * Parse JSON from LLM output, handling markdown code fences and malformed output.
  * Uses progressive fallback: direct parse, array extraction, truncated array repair, single object wrap.
  */
-function parseJsonResponse<T>(raw: string): T | null {
+export function parseJsonResponse<T>(raw: string): T | null {
   // Attempt 1: direct parse after stripping fences
   let cleaned = raw.trim();
   const fenceMatch = cleaned.match(/^```(?:json)?\s*\n?([\s\S]*?)\n?```$/);

ts/packages/flow/tsconfig.json

@@ -6,6 +6,7 @@
     "composite": true
   },
   "include": ["src"],
+  "exclude": ["src/**/*.test.ts", "src/**/*.spec.ts"],
   "references": [
     { "path": "../base" }
   ]

ts/packages/flow/vitest.config.ts

@@ -0,0 +1,6 @@
+import { defineConfig } from "vitest/config";
+export default defineConfig({
+  test: {
+    globals: true,
+  },
+});

ts/packages/workbench/src/App.tsx

@@ -1,5 +1,6 @@
 import { BrowserRouter, Routes, Route, Navigate } from "react-router";
 import { RootLayout } from "@/components/layout/root-layout";
+import { ErrorBoundary } from "@/components/error-boundary";
 import ChatPage from "@/pages/chat";
 import LibraryPage from "@/pages/library";
 import GraphPage from "@/pages/graph";
@@ -16,14 +17,14 @@ export default function App() {
       <Routes>
         <Route element={<RootLayout />}>
           <Route index element={<Navigate to="/chat" replace />} />
-          <Route path="/chat" element={<ChatPage />} />
-          <Route path="/library" element={<LibraryPage />} />
-          <Route path="/graph" element={<GraphPage />} />
-          <Route path="/prompts" element={<PromptsPage />} />
-          <Route path="/token-cost" element={<TokenCostPage />} />
-          <Route path="/knowledge-cores" element={<KnowledgeCoresPage />} />
-          <Route path="/flows" element={<FlowsPage />} />
-          <Route path="/settings" element={<SettingsPage />} />
+          <Route path="/chat" element={<ErrorBoundary><ChatPage /></ErrorBoundary>} />
+          <Route path="/library" element={<ErrorBoundary><LibraryPage /></ErrorBoundary>} />
+          <Route path="/graph" element={<ErrorBoundary><GraphPage /></ErrorBoundary>} />
+          <Route path="/prompts" element={<ErrorBoundary><PromptsPage /></ErrorBoundary>} />
+          <Route path="/token-cost" element={<ErrorBoundary><TokenCostPage /></ErrorBoundary>} />
+          <Route path="/knowledge-cores" element={<ErrorBoundary><KnowledgeCoresPage /></ErrorBoundary>} />
+          <Route path="/flows" element={<ErrorBoundary><FlowsPage /></ErrorBoundary>} />
+          <Route path="/settings" element={<ErrorBoundary><SettingsPage /></ErrorBoundary>} />
         </Route>
       </Routes>
 

ts/packages/workbench/src/components/error-boundary.tsx

@@ -0,0 +1,61 @@
+import { Component, type ErrorInfo, type ReactNode } from "react";
+import { AlertTriangle, RefreshCw } from "lucide-react";
+
+interface Props {
+  children: ReactNode;
+  /** Optional fallback -- if omitted, a default card is shown */
+  fallback?: ReactNode;
+}
+
+interface State {
+  hasError: boolean;
+  error: Error | null;
+}
+
+export class ErrorBoundary extends Component<Props, State> {
+  constructor(props: Props) {
+    super(props);
+    this.state = { hasError: false, error: null };
+  }
+
+  static getDerivedStateFromError(error: Error): State {
+    return { hasError: true, error };
+  }
+
+  componentDidCatch(error: Error, info: ErrorInfo) {
+    console.error("[ErrorBoundary]", error, info.componentStack);
+  }
+
+  handleReset = () => {
+    this.setState({ hasError: false, error: null });
+  };
+
+  render() {
+    if (this.state.hasError) {
+      if (this.props.fallback) return this.props.fallback;
+
+      return (
+        <div className="flex h-full items-center justify-center p-8">
+          <div className="max-w-md rounded-lg border border-error/30 bg-error/5 p-6 text-center">
+            <AlertTriangle className="mx-auto mb-3 h-8 w-8 text-error" />
+            <h2 className="mb-2 text-lg font-semibold text-fg">
+              Something went wrong
+            </h2>
+            <p className="mb-4 text-sm text-fg-muted">
+              {this.state.error?.message || "An unexpected error occurred."}
+            </p>
+            <button
+              onClick={this.handleReset}
+              className="inline-flex items-center gap-2 rounded-lg bg-brand-600 px-4 py-2 text-sm font-medium text-white hover:bg-brand-500"
+            >
+              <RefreshCw className="h-3.5 w-3.5" />
+              Try Again
+            </button>
+          </div>
+        </div>
+      );
+    }
+
+    return this.props.children;
+  }
+}

ts/packages/workbench/src/components/layout/root-layout.tsx

@@ -1,7 +1,9 @@
 import { Outlet } from "react-router";
+import { WifiOff } from "lucide-react";
 import { Sidebar } from "./sidebar";
 import { FlowSelector } from "./flow-selector";
 import { useProgressStore } from "@/hooks/use-progress-store";
+import { useConnectionState } from "@/providers/socket-provider";
 
 /**
  * Top loading bar -- shown when any global activity is in progress.
@@ -22,6 +24,11 @@ function LoadingBar() {
  * Root layout: fixed sidebar + scrollable main content area with a top bar.
  */
 export function RootLayout() {
+  const connectionState = useConnectionState();
+  const isDisconnected =
+    connectionState.status === "failed" ||
+    connectionState.status === "reconnecting";
+
   return (
     <div className="relative flex h-screen w-full overflow-hidden bg-surface-0">
       {/* Global loading bar */}
@@ -35,6 +42,14 @@ export function RootLayout() {
           <FlowSelector />
         </header>
 
+        {/* Connection lost banner */}
+        {isDisconnected && (
+          <div className="flex items-center gap-2 border-b border-amber-500/30 bg-amber-500/10 px-4 py-2 text-xs text-amber-400">
+            <WifiOff className="h-3.5 w-3.5" />
+            <span>Connection lost. Attempting to reconnect...</span>
+          </div>
+        )}
+
         {/* Page content */}
         <main className="flex-1 overflow-y-auto p-6">
           <Outlet />

ts/packages/workbench/src/hooks/use-conversation.ts

@@ -1,4 +1,5 @@
 import { create } from "zustand";
+import { persist } from "zustand/middleware";
 
 // ---------------------------------------------------------------------------
 // Types
@@ -66,26 +67,38 @@ export function nextMessageId(): string {
   return `msg-${++_nextMsgId}-${Date.now()}`;
 }
 
-export const useConversation = create<ConversationState>()((set) => ({
-  messages: [],
-  input: "",
-  chatMode: "graph-rag",
+export const useConversation = create<ConversationState>()(
+  persist(
+    (set) => ({
+      messages: [],
+      input: "",
+      chatMode: "graph-rag",
 
-  setInput: (value) => set({ input: value }),
-  setChatMode: (mode) => set({ chatMode: mode }),
+      setInput: (value) => set({ input: value }),
+      setChatMode: (mode) => set({ chatMode: mode }),
 
-  addMessage: (message) =>
-    set((state) => ({ messages: [...state.messages, message] })),
+      addMessage: (message) =>
+        set((state) => ({ messages: [...state.messages, message] })),
 
-  updateLastMessage: (updater) =>
-    set((state) => {
-      if (state.messages.length === 0) return state;
-      const last = state.messages[state.messages.length - 1]!;
-      const updated = updater(last);
-      return {
-        messages: [...state.messages.slice(0, -1), updated],
-      };
+      updateLastMessage: (updater) =>
+        set((state) => {
+          if (state.messages.length === 0) return state;
+          const last = state.messages[state.messages.length - 1]!;
+          const updated = updater(last);
+          return {
+            messages: [...state.messages.slice(0, -1), updated],
+          };
+        }),
+
+      clearMessages: () => set({ messages: [] }),
     }),
-
-  clearMessages: () => set({ messages: [] }),
-}));
+    {
+      name: "tg-conversation",
+      // Only persist messages and chatMode, not input or transient state
+      partialize: (state) => ({
+        messages: state.messages.filter((m) => !m.isStreaming),
+        chatMode: state.chatMode,
+      }),
+    },
+  ),
+);

ts/packages/workbench/src/hooks/use-prompts.ts

@@ -8,13 +8,17 @@ export function usePrompts() {
   const [prompts, setPrompts] = useState<Array<{ id: string; name?: string; description?: string }>>([]);
   const [systemPrompt, setSystemPrompt] = useState<string>("");
   const [loading, setLoading] = useState(false);
+  const [error, setError] = useState<string | null>(null);
 
   const loadPrompts = useCallback(async () => {
     try {
       setLoading(true);
+      setError(null);
       const list = await socket.config().getPrompts();
       setPrompts(Array.isArray(list) ? list : []);
     } catch (err) {
+      const msg = err instanceof Error ? err.message : String(err);
+      setError(msg);
       console.error("Failed to load prompts:", err);
     } finally {
       setLoading(false);
@@ -46,5 +50,5 @@ export function usePrompts() {
     }
   }, [connectionState.status, loadPrompts, loadSystemPrompt]);
 
-  return { prompts, systemPrompt, loading, loadPrompts, loadSystemPrompt, getPrompt };
+  return { prompts, systemPrompt, loading, error, loadPrompts, loadSystemPrompt, getPrompt };
 }

ts/packages/workbench/src/pages/prompts.tsx

@@ -18,7 +18,7 @@ import { usePrompts } from "@/hooks/use-prompts";
 type Tab = "templates" | "system";
 
 export default function PromptsPage() {
-  const { prompts, systemPrompt, loading, loadPrompts, loadSystemPrompt, getPrompt } = usePrompts();
+  const { prompts, systemPrompt, loading, error, loadPrompts, loadSystemPrompt, getPrompt } = usePrompts();
 
   const [activeTab, setActiveTab] = useState<Tab>("templates");
   const [selectedPromptId, setSelectedPromptId] = useState<string | null>(null);
@@ -96,6 +96,13 @@ export default function PromptsPage() {
         </button>
       </div>
 
+      {/* Error display */}
+      {error && (
+        <p className="mb-4 rounded-lg bg-error/10 px-4 py-2 text-sm text-error">
+          {error}
+        </p>
+      )}
+
       {/* Templates tab */}
       {activeTab === "templates" && (
         <div className="flex flex-1 flex-col gap-4 overflow-hidden">