SurfSense/surfsense_web/features/chat-messages/ARCHITECTURE.md

23 KiB

features/chat-messages/ — Architecture

Scope. This module owns everything between an assistant message arriving and its rendering inside the chat UI: the timeline (the agent's process — reasoning + every tool call), and the HITL primitives that per-tool components compose to render approval views.

It does NOT own: the thread shell, the composer, the streaming pipeline, the message frame (assistant-message.tsx, user-message.tsx, markdown renderer, citations), the comments sidebar, or any of the 63 individual tool-ui integration files under components/tool-ui/.


1. Mental model

Every assistant message has two regions:

Region What it shows
Timeline The agent's process. Reasoning, every tool call, grouped by delegation spanId into a tree. Each tool call is rendered by its registered component, which selects its own view (running, awaiting approval, success, error, etc.) by discriminating its result data.
Body The agent's product. Markdown text, citations, native reasoning blocks, and value-add deliverables (image viewer, chart, canvas). Connector tool cards do NOT render here.

Principle: timeline = process, body = product. No overlap.

A tool's UI lives in the body if and only if it produces a deliverable the user wants to interact with — view, scrub, copy, share. If the UI just shows that the tool ran and what it did, it lives in the timeline.

┌─ Assistant Message ─────────────────────────────────────────┐
│                                                             │
│  ╔═════════════════════════════════════════════════════╗   │
│  ║  TIMELINE  (process)                                ║   │
│  ║                                                     ║   │
│  ║  ▸ task: NotionAgent              [running]         ║   │
│  ║      ▸ search_workspace           [completed]       ║   │
│  ║      ▸ update_page                ← rendered by     ║   │
│  ║          (Notion-styled approval     UpdateNotion-  ║   │
│  ║           card OR Notion-styled      PageToolUI;    ║   │
│  ║           success/error card,        the component  ║   │
│  ║           per its own data           picks the view ║   │
│  ║           discrimination)            from result)   ║   │
│  ║  ▸ summarize                      [completed]       ║   │
│  ╚═════════════════════════════════════════════════════╝   │
│                                                             │
│  ╔═════════════════════════════════════════════════════╗   │
│  ║  BODY  (product)                                    ║   │
│  ║                                                     ║   │
│  ║  Markdown text, citations, value-add deliverables   ║   │
│  ║  only. Connector tool cards do NOT render here.     ║   │
│  ╚═════════════════════════════════════════════════════╝   │
│                                                             │
└─────────────────────────────────────────────────────────────┘

2. The single data model

The timeline reads ONE data structure: a TimelineItem[]. There are no parallel structures for "thinking steps", "tool calls", "HITL bundles", etc. Every visible piece of agent activity is a TimelineItem.

2.1 The discriminated union (outer discrimination)

Two kinds. The timeline does outer discrimination — it chooses reasoning view vs tool-call mounting based on kind.

type ItemStatus =
    | "pending" | "running" | "completed" | "cancelled" | "error";

interface BaseItem {
    id: string;
    spanId?: string;             // groups items into delegation tree (parent task + children)
    status: ItemStatus;
}

interface ReasoningItem extends BaseItem {
    kind: "reasoning";
    text: string;
}

interface ToolCallItem extends BaseItem {
    kind: "tool-call";
    toolName: string;
    args: Record<string, unknown>;
    argsText?: string;
    result?: unknown;            // per-tool component discriminates this internally
    langchainToolCallId?: string;
}

type TimelineItem = ReasoningItem | ToolCallItem;

interface TimelineGroup {
    parent: TimelineItem;
    children: TimelineItem[];
}

ToolCallItem has no approval field, no phase, no view. All of that is derived inside the per-tool component from the result data.

2.2 Inner discrimination (per-tool component)

Each tool registers a component that receives the tool-call data and decides what to render based on its own result-shape discriminators:

const UpdateNotionPageToolUI: TimelineToolComponent = (props) => {
    if (isInterruptResult(props.result))    return <NotionApprovalCard {...props} />;
    if (isAuthErrorResult(props.result))    return <NotionAuthErrorCard {...props} />;
    if (isErrorResult(props.result))        return <NotionErrorCard {...props} />;
    if (isInfoResult(props.result))         return <NotionNotFoundCard {...props} />;
    if (isSuccessResult(props.result))      return <NotionSuccessCard {...props} />;
    return <NotionPendingCard {...props} />;
};

The discriminators (isInterruptResult, isAuthErrorResult, etc.) are types, not centralized infrastructure. The component owns the dispatch. The timeline knows none of this.

2.3 The pure builder

function buildTimeline(
    content: MessageContent[],
    thinkingSteps: ThinkingStep[],
): TimelineGroup[]

Builds the timeline from existing message content + thinking-step data parts. Pure function. Sets kind and status on each item; preserves result verbatim for per-tool discrimination.

2.4 The dispatch (timeline-level)

Two cases. Exhaustive switch. No runtime guards in the timeline renderer.

function TimelineItemView({ item }: { item: TimelineItem }) {
    switch (item.kind) {
        case "reasoning": return <ReasoningItemView item={item} />;
        case "tool-call": return <ToolCallItemView item={item} />;
    }
}

function ToolCallItemView({ item }: { item: ToolCallItem }) {
    const ToolBody = getToolComponent(item.toolName) ?? FallbackToolBody;
    return <ToolBody {...adaptItemToProps(item)} />;
}

No card frame, header, body slot, approval area, or result panel at the timeline level. Each tool component owns its own visual presentation. This matches how every existing tool-ui component already works — they each render their own rounded card with their own header.


3. The timeline's tool-component contract

Tool components mounted by the timeline implement a subset of assistant-ui's ToolCallMessagePartProps — only the fields the timeline can supply:

interface TimelineToolProps {
    toolCallId: string;
    toolName: string;
    args: Record<string, unknown>;
    argsText?: string;
    result?: unknown;
    langchainToolCallId?: string;
    status: ItemStatus;          // simple enum, not assistant-ui's complex status object
}

type TimelineToolComponent = (props: TimelineToolProps) => ReactNode;

Notably absent (compared to ToolCallMessagePartProps):

  • addResult, resume — runtime-only, not needed; HITL decisions flow through useHitlDecision (a hook) which talks to the runtime directly.
  • The complex status: ToolCallMessagePartState["status"] object — replaced by our simple ItemStatus enum.

The 15 existing HITL-aware tool-ui components only use the subset above. They are retyped to TimelineToolComponent in the cutover commit (mechanical: ToolCallMessagePartComponentTimelineToolComponent).


4. Rendering topology — how the body opts out

The body uses assistant-ui's MessagePrimitive.Parts and registers a no-op fallback for tool calls so they don't render here:

<MessagePrimitive.Parts
    components={{
        Text: MarkdownText,
        Reasoning,
        Source,
    }}
    tools={{
        by_name: BODY_TOOLS,         // value-add deliverables only (image viewer, etc.)
        fallback: () => null,         // every other tool-call: render nothing in the body
    }}
/>

BODY_TOOLS starts empty (no value-add deliverables exist yet) and grows as we identify them. Every tool not in BODY_TOOLS renders nothing in the body.

The timeline reads message content via useAuiState(({ message }) => message?.content) and runs buildTimeline to produce the items it renders. Tool-call data IS in the message; the body just chooses not to render it.

Result: zero dual placement. Zero suppression HOC. Zero render-target context. Zero coordination.


5. Slice layout

features/chat-messages/
├── ARCHITECTURE.md
│
├── timeline/                                ← the process surface
│   ├── types.ts                             (TimelineItem union, ToolCallItem, ItemStatus, TimelineGroup)
│   ├── build-timeline.ts                    (pure: content + thinkingSteps → groups)
│   ├── grouping.ts                          (pure: group items by spanId)
│   ├── subagent-rename.ts                   (pure: parent task title from args.subagent_type)
│   ├── tool-registry/                       (PRIVATE to timeline; only timeline mounts tools)
│   │   ├── types.ts                         (TimelineToolComponent, TimelineToolProps)
│   │   ├── registry.ts                      (TOOLS_BY_NAME from components/tool-ui/*)
│   │   ├── adapt-props.ts                   (pure: ToolCallItem → TimelineToolProps)
│   │   ├── fallback/
│   │   │   ├── fallback-tool-body.tsx       (TimelineToolComponent for unregistered tools — discriminates internally)
│   │   │   ├── default-fallback-card.tsx    (the non-HITL fallback view: status icon + collapsible + JSON)
│   │   │   ├── revert-button.tsx            (revert affordance — used by default-fallback-card)
│   │   │   ├── use-tool-action.ts           (action lookup hook for revert)
│   │   │   └── index.ts
│   │   └── index.ts
│   ├── items/
│   │   ├── reasoning-item.tsx               (renders kind: "reasoning")
│   │   ├── tool-call-item.tsx               (lookup component + mount with adapted props — ~10 lines)
│   │   └── index.ts
│   ├── timeline.tsx                         (groups + iteration + 2-case dispatch)
│   ├── data-renderer.tsx                    (assistant-ui adapter; exports TimelineDataUI)
│   └── index.ts
│
├── hitl/                                    ← pure HITL primitives
│   ├── types.ts                             (InterruptResult, HitlPhase, HitlDecision, isInterruptResult)
│   ├── use-hitl-decision.ts                 (hook: dispatch approve/edit/reject — used by every approval card)
│   ├── use-hitl-phase.ts                    (hook: tracks pending → processing → approved/rejected/edited)
│   ├── approval-cards/                      (the FALLBACK-mounted approval views; per-tool components import from here OR build their own)
│   │   ├── generic-approval.tsx             (default approval UI — what FallbackToolBody mounts for interrupt results)
│   │   ├── doom-loop-approval.tsx           (special-case approval UI + isDoomLoopInterrupt)
│   │   └── index.ts
│   ├── edit-panel/
│   │   ├── edit-panel.atom.ts               (Jotai atoms for the panel state)
│   │   ├── edit-panel.tsx                   (root: atom wiring + desktop/mobile switch only)
│   │   ├── fields/
│   │   │   ├── email-tags-field.tsx         (EmailsTagField + parse/format helpers)
│   │   │   ├── calendar-field.tsx           (DateTimePickerField + parse/format helpers)
│   │   │   ├── extra-fields.tsx             (ExtraField switch renderer)
│   │   │   └── index.ts                     (private barrel)
│   │   └── index.ts
│   └── index.ts
│
└── (no body slice yet — body just registers `tools={{ fallback: () => null }}`)

5.1 Notable absences

Was Status Reason
tool-cards/ slice Folded into timeline/ Tool-call rendering happens in the timeline; the tool-registry is private to timeline.
bundleTool composer Deleted Body opts out via fallback: () => null. No HOCs to compose.
withDelegationSpanIndent HOC Deleted Tree indent is owned by the timeline's group renderer.
withBundleStep + HitlBundleProvider Deleted Multi-approval is just N inline renderings; no coordination needed.
withHitlInTimeline + HitlRenderTargetProvider Deleted Tool cards never render in body; no dual-placement to suppress.
pickApprovalCard central dispatcher Deleted Each tool component picks its own view via internal discrimination. The fallback has its OWN internal dispatcher (interrupt → generic-approval; doom-loop → doom-loop-approval).
getHitlToolComponent registry Deleted The tool-registry is just a Record<string, TimelineToolComponent>; lookup is TOOLS_BY_NAME[name].
Centralized approval-area.tsx in timeline Deleted The approval is a view the per-tool component renders, not an area the timeline composes.
ApprovalState on ToolCallItem Deleted Phase is local UI state inside per-tool approval cards (via useHitlPhase). The timeline doesn't track it.
ThinkingStepToolInfoMap Map join Deleted The unified TimelineItem union eliminates the join.

6. Public surfaces

timeline/index.ts

export { TimelineDataUI };               // the assistant-ui registration
export { Timeline };                     // exposed for tests
export type { TimelineItem, ReasoningItem, ToolCallItem, TimelineGroup, ItemStatus };
export type { TimelineToolComponent, TimelineToolProps };

hitl/index.ts

export type { InterruptResult, InterruptActionRequest, InterruptReviewConfig, HitlDecision, HitlPhase };
export { isInterruptResult };

export { useHitlDecision };
export { useHitlPhase };

export { GenericHitlApprovalToolUI };    // for tool-ui integrations that want to compose on top
export { DoomLoopApprovalToolUI, isDoomLoopInterrupt };

export { HitlEditPanel, MobileHitlEditPanel };
export { openHitlEditPanelAtom, closeHitlEditPanelAtom, hitlEditPanelAtom };
export type { ExtraField };

The 63 components/tool-ui/* integrations consume hitl/'s public surface (types, hooks, edit-panel atom, optionally the fallback approval cards). Nothing else.


7. Layering & SRP rules

7.1 The "what knows about what" rule

Component Knows about
timeline/ Itself + hitl/ (via the fallback) + components/tool-ui/* (via the registry)
timeline/tool-registry/ The TimelineToolComponent contract, components/tool-ui/*, and hitl/ (for the fallback's approval views)
hitl/ Itself only — no knowledge of timeline, tool-call types, registry
components/tool-ui/* hitl/ only (for HITL primitives + optional fallback approval cards); never reaches into timeline/
Body (assistant-message.tsx) The BODY_TOOLS registry and TimelineDataUI from timeline/index.ts

hitl/ does NOT import from timeline/. The dependency arrow is one-way.

7.2 Render policy belongs to the surface, not the primitive

  • hitl/ exposes hooks, types, and the fallback approval cards.
  • timeline/ decides WHEN and WHERE tool components mount (inside tool-call-item.tsx).
  • A hitl/ primitive must never assume it's being rendered in the timeline, the body, or anywhere else. It receives props, renders UI, returns. No environment sniffing, no context.
  • Per-tool components in components/tool-ui/* decide WHICH view to render based on result-shape discriminators. The timeline does not know these discriminators exist.

7.3 Single Responsibility

Rules in priority order:

  1. One responsibility per file. Need "and" to describe it? Split it.
  2. One responsibility per function. Same.
  3. Line count is a smell, not a budget. ~250 lines = pause and ask "still one responsibility?"; ~500 lines = strong presumption of split needed unless explicitly justified at the top of the file.

Notable splits driven by SRP during the port:

  • hitl-edit-panel.tsx (current 405 lines, 4 responsibilities) → 5 files: edit-panel.tsx (root + layout switch), email-tags-field.tsx, calendar-field.tsx, extra-fields.tsx, edit-panel.atom.ts.
  • tool-fallback.tsx (current 533 lines, 3 responsibilities) → split across fallback-tool-body.tsx, default-fallback-card.tsx, revert-button.tsx, use-tool-action.ts.
  • thinking-steps.tsx (current 434 lines, 5 responsibilities) → folded into the new timeline/ slice across types.ts, build-timeline.ts, grouping.ts, subagent-rename.ts, timeline.tsx, items/*, data-renderer.tsx.

8. Tested behaviors

Unit tests live next to the file they cover (*.test.ts(x)).

  • timeline/build-timeline.test.ts — content + thinkingSteps → correct items, correct kind, correct status, correct ordering. result preserved verbatim.
  • timeline/grouping.test.ts — items group correctly by spanId; first item with a spanId is the parent; orphaned children are promoted defensively.
  • timeline/subagent-rename.test.tstask step's display title resolves to args.subagent_type (title-cased); falls back to "Task" when subagent type is missing.
  • timeline/tool-registry/registry.test.tsTOOLS_BY_NAME includes every named tool; FallbackToolBody is returned for unknown names; the fallback dispatches correctly (interrupt → generic, doom-loop → doom-loop, otherwise → default fallback).
  • timeline/tool-registry/adapt-props.test.tsToolCallItemTimelineToolProps mapping is lossless; status mapping is correct.
  • hitl/use-hitl-phase.test.ts — phase transitions through pending → processing → approved/rejected/edited correctly.
  • hitl/approval-cards/doom-loop-approval.test.tsxisDoomLoopInterrupt matches doom-loop-shape interrupts only.

Smoke test after cutover:

  • Assistant message renders; markdown + citations work in body.
  • All connector tool calls render in timeline only (none in body).
  • Reasoning steps render in timeline.
  • Single HITL flow (Notion update): approve, edit, reject — each transitions through the phases correctly.
  • Multiple pending HITL cards: each renders inline at its position; deciding one doesn't affect the others.
  • Doom-loop approval renders the special card.
  • Revert button works on completed default-fallback cards and survives reload.
  • Subagent name renaming on task parent step.

9. Migration plan (strangler fig, single atomic cutover)

Phase A — Build the new slice in parallel

In dependency order: hitl/ first (leaf), then timeline/. The existing code (thinking-steps.tsx, tool-fallback.tsx, assistant-message.tsx's tool registry, etc.) remains fully functional throughout Phase A.

  1. Port hitl/ primitives. Apply SRP splits (edit panel into 5 files). hitl/approval-cards/{generic,doom-loop}-approval.tsx are ported as standalone components — they're what the fallback mounts and what per-tool integrations may compose on top of.
  2. Build timeline/ slice. Implement buildTimeline from scratch (do NOT copy thinking-steps logic verbatim — design the pure function around the new union). Build the tool-registry/ with TimelineToolComponent contract; the registry imports from components/tool-ui/* (no file moves yet).
  3. Add unit tests as listed in §8.
  4. Verify: tsc clean, biome clean, no consumer file touched, no linter regressions.

Phase B — Atomic cutover (single commit)

File Change
components/assistant-ui/assistant-message.tsx Replace TOOLS_BY_NAME/TOOLS_FALLBACK definitions with BODY_TOOLS (initially empty) + tools={{ fallback: () => null }}. Replace ThinkingStepsDataUI registration with TimelineDataUI.
components/public-chat/public-thread.tsx Same registry + data UI swap.
app/dashboard/.../new-chat/page.tsx Switch ThinkingStepsDataUITimelineDataUI. Drop HitlBundleProvider (no longer needed).
components/free-chat/free-chat-page.tsx Switch ThinkingStepsDataUITimelineDataUI.
components/public-chat/public-chat-view.tsx Same.
components/layout/ui/right-panel/RightPanel.tsx Switch HitlEditPanel import to @/features/chat-messages/hitl.
The 15 components/tool-ui/* HITL-aware integration files (a) Switch HITL imports from @/lib/hitl, @/hooks/use-hitl-phase, @/atoms/chat/hitl-edit-panel.atom@/features/chat-messages/hitl. (b) Retype from ToolCallMessagePartComponentTimelineToolComponent (mechanical type rename).

Phase C — Delete legacy

After cutover passes smoke tests:

  • components/assistant-ui/thinking-steps.tsx
  • components/assistant-ui/tool-fallback.tsx
  • lib/chat/delegation-span-indent.ts
  • lib/hitl/ (entire folder)
  • components/hitl-bundle-pager/ (entire folder)
  • components/tool-ui/generic-hitl-approval.tsx
  • components/tool-ui/doom-loop-approval.tsx
  • components/hitl-edit-panel/ (entire folder)
  • hooks/use-hitl-phase.ts
  • atoms/chat/hitl-edit-panel.atom.ts

Verify: no orphan files, no dead imports, no test regressions.


10. Out of scope (and one consumer relationship)

10.1 The 63 components/tool-ui/* integrations

These are first-class consumers of hitl/ and the TimelineToolComponent contract. They are imported by timeline/tool-registry/registry.ts to build TOOLS_BY_NAME. They never reach into timeline/ themselves.

They stay where they are. Future option to move them is a separate, mechanical follow-up refactor.

10.2 Not touched by this refactor

  • The composer (input bar, mention picker, prompt picker, tool toggles).
  • The streaming pipeline (lib/chat/streaming-state.ts, stream-pipeline.ts, thread-persistence.ts).
  • The chat-comments sidebar.
  • The message frame (assistant-message.tsx, user-message.tsx, markdown-text.tsx, inline-citation.tsx) beyond swapping the registry imports.

If any of these become a blocker for the refactor (e.g. the streaming pipeline needs a metadata field that doesn't exist), surface it explicitly and decide whether to expand scope before touching it.