Skip to content

feat(workflows): extract ChatSessionHost and compose into stage chat - #1026

Merged
lavaman131 merged 9 commits into
mainfrom
refactor/workflow-chat-host-parity
May 24, 2026
Merged

feat(workflows): extract ChatSessionHost and compose into stage chat#1026
lavaman131 merged 9 commits into
mainfrom
refactor/workflow-chat-host-parity

Conversation

@flora131

@flora131 flora131 commented May 23, 2026

Copy link
Copy Markdown
Collaborator

Summary

Extracts a reusable ChatSessionHost component from the interactive-mode chat surface and composes it into StageChatView, eliminating ~800 lines of duplicated interactive-mode logic from stage chat. ctx.ui.custom requests are routed to the active stage node via a new StageUiBroker, and MCP session lifecycle and security are hardened throughout.

New abstractions

  • ChatSessionHost (packages/coding-agent/src/modes/interactive/components/chat-session-host.ts, ~1281 lines) — reusable component encapsulating the full interactive chat surface (slash commands, bash execution, external editor, image paste, queued messages, compaction, retries); interactive-mode.ts retains its legacy implementation — migration is follow-up work
  • StageUiBroker (packages/workflows/src/shared/stage-ui-broker.ts) — routes ctx.ui.custom factory requests to the correct active stage chat node with host registration, duplicate concurrent request rejection, AbortSignal support, and cleanup for throwing hosts
  • chat-input-actions.ts — standalone helpers for clipboard image paste, external editor launch, and queued-message coalescing; shared between interactive mode and stage chat
  • stage-chat-layout.ts — extracted frame-fitting and viewport helpers for the stage chat surface

Workflow stage chat

  • StageChatView now composes ChatSessionHost instead of duplicating interactive-mode logic
  • Stage awaiting-input state tracked in the store (recordStagePendingPrompt, resolveStagePendingPrompt, awaitStagePendingPrompt) with lifecycle cleanup on run end or removal
  • Executor defers live handle release until the SDK message queue drains (releaseLiveHandleWhenIdle), preventing premature session collapse while messages are still in flight; a subscription plus a documented 250 ms defensive poll covers the silent SDK drain path
  • Completed idle stage chats become read-only after drain
  • Exposes isDisposed getter on the live stage-control handle for accurate lifecycle queries
  • Run-detail, status-list, and workflow TUI panels refreshed with rounded, marker-free styling; render-inputs-schema migrated to the same renderer

MCP hardening

  • Status bar now counts only "connected" connections (previously counted all map entries regardless of status)
  • colorizeStatusText helper guards against themes that don't expose a .fg() function, preventing runtime errors on non-standard themes
  • OAuth callback error HTML rendered from a static renderCallbackErrorHtml function; raw OAuth query-string errors are no longer interpolated into HTML — they remain available only through rejection/logging paths
  • Completed UI-session handle lookup hardened with a null guard and a stable local binding to avoid TOCTOU races
  • Glimpse iframe src attribute HTML-escaped via escapeHtmlAttribute before injection

Tests

  • New: test/unit/chat-session-host.test.ts (543 lines) — full host lifecycle, slash commands, bash, compaction, retry, custom entry rendering
  • New: test/unit/mcp-init-statusbar.test.ts (82 lines) — connected-count filtering and theme-fallback safety
  • New: test/unit/stage-ui-broker.test.ts, store-pending-prompt.test.ts, workflow-attach-pane.test.ts, mcp-security.test.ts
  • Expanded coverage across stage-chat-view, wiring-adapters, overlay-graph, prompt-card, dispatch-confirm, inputs-picker, executor, store-widget-installer, and 5+ other unit test files

Validation

AGENT=1 bun test test/unit/stage-ui-broker.test.ts test/unit/store-pending-prompt.test.ts packages/coding-agent/test/chat-input-actions.test.ts
bun run typecheck
# tmux E2E smoke: two-stage dummy workflow via runWorkflow(..., { stubAgent: true })

🤖 Generated with Claude Code

flora131 added 2 commits May 23, 2026 13:27
Extract a reusable chat session host from interactive mode and wire workflow stage chat through stage-scoped UI brokers.\n\nUpdate workflow TUI rendering, input schema previews, docs, mockups, and tests for the embedded chat surface.

Assistant-model: GPT-5.5
Count only connected MCP servers in the status bar and fall back to plain text when theme colorization is unavailable.\n\nEscape callback authorization errors and retain a stable completed UI session handle when storing session messages.

Assistant-model: GPT-5.5
@claude claude Bot changed the title feat(workflows): embed custom UI in stage chat hosts feat(workflows): embed stage-scoped custom UI in chat session hosts May 23, 2026
@claude

claude Bot commented May 23, 2026

Copy link
Copy Markdown

PR Review: feat(workflows): embed custom UI in stage chat hosts

This is a sizable but well-structured refactor. The extraction of ChatSessionHost is a clear win for code reuse, and the stage UI broker is a clean abstraction. The MCP fixes (HTML escaping, connection counting, completed-session lookup) are nice incidental hardening. Tests added for the new behavior look thorough.

A few items below — most are non-blocking, but the executor lifetime change in particular is worth discussing before merging.

Bugs / correctness

1. Stage AgentSession is never disposed on normal settle (packages/workflows/src/runs/foreground/executor.ts:1559-1565)

The previous cleanup branch called releaseLiveHandle() either immediately or after the chat pane detached, which in turn called disposeInnerContext()innerCtx.__dispose() and unregisterStageHandle(). The PR removes both branches and replaces them with limiter.release() only:

// Old: schedule releaseLiveHandle() on detach (or call it now if not attached)
// New: keep the direct chat handle registered after the workflow-owned
//      stage operation settles.
limiter.release();

The new tests (executor.test.ts:2040-2089) intentionally assert disposeCalls === 0 after the stage completes — so this is by design. But the implications:

  • disposeInnerContext() is never called for any stage on normal completion.
  • unregisterStageHandle() is never called — the registry retains the handle for the lifetime of the process.
  • Nothing in recordRunEnd / finalizeKilled walks stage handles to release them either.

For long-running parent processes orchestrating workflows with many stages, this leaks AgentSession instances (file handles, model registries, persistence streams) indefinitely. Even if "kept alive for post-completion chat" is the right product call, there should be an explicit release path tied to either (a) the chat surface detaching, (b) the run being removed from the store, or (c) some idle/eviction policy. Right now I don't see one. Worth confirming the intended cleanup ownership and documenting it.

2. openExternalEditorForText mishandles editor commands with spaces (packages/coding-agent/src/modes/interactive/chat-input-actions.ts:78)

const [editor, ...editorArgs] = editorCommand.split(" ");

This breaks for EDITOR="/Applications/Visual Studio Code.app/Contents/Resources/app/bin/code -w" or quoted args. Pre-existing pattern, but as it's being newly factored into a shared module, this is a good moment to swap in shell-quote.parse or node:util.parseArgs. Otherwise users with editors at paths containing spaces silently get "No such file or directory."

3. Clipboard image temp file leaks (chat-input-actions.ts:48-50)

pasteClipboardImageToEditor writes pi-clipboard-<uuid>.<ext> to os.tmpdir() and inserts the path into the editor. There's no cleanup hook — every paste leaves the file there until the OS reaps tmp. Acceptable on most platforms, but worth at least scheduling a best-effort fs.unlink after the agent reads the file (e.g., on next prompt submit). Pre-existing behavior, just newly extracted.

4. StageUiBroker can permanently wedge a stage (packages/workflows/src/shared/stage-ui-broker.ts:60-62)

if (existing) {
  return Promise.reject(new Error(`pi-workflows: stage ${stageId} already has a pending custom UI request`));
}

The pending map only clears via resolve/reject. If a factory throws between requestCustomUi setting pending and the host mounting (e.g., mountStageCustomUi rejects after the entry is in pending), the entry is properly cleared via reject(...) in _showCustomUi's catch. Good. But there's still no defense against a host that loses the request reference without calling done() — that stage is then permanently blocked from showing further custom UI. Consider tying pending lifetime to the host registration (clear pending requests for a stage when its host unregisters), or expose an admin clearPending(runId, stageId).

5. Optimistic user-prompt dedup signature is text-only (packages/coding-agent/src/modes/interactive/components/chat-session-host.ts:215-227, 1091-1093)

function userMessageSignature(text: string): string {
  return text.trim();
}

If two consecutive prompts have the same trimmed text (e.g., "yes", "yes ", " yes"), the counter dedups them — but they're presumably distinct user actions. The current logic happens to be correct because each emits exactly one optimistic insertion + one message_start, so the counts balance. However, if a message_start is ever delivered without a preceding optimistic insertion (or vice-versa, e.g., on retry/abort), the counter desynchronizes and starts swallowing real user messages. A more robust signature would include a per-call nonce stored on the optimistic entry, or simply key on the messageId once one is assigned.

6. mountStageCustomUi focused-property bridging (stage-ui-broker.ts:139-148)

if ("focused" in rawComponent) {
  Object.defineProperty(component, "focused", {
    get: () => (rawComponent as Component & Partial<Focusable>).focused,
    set: (value: boolean) => { (rawComponent as Component & Partial<Focusable>).focused = value; },
    ...
  });
}

'focused' in rawComponent is true even if the factory set focused to undefined explicitly. The property accessor will then return/set undefined, which downstream setComponentFocused treats as a no-op. Minor — but consider typeof rawComponent.focused === 'boolean' instead.

Best practices / maintainability

7. chat-session-host.ts is 1248 lines

It absorbs editor wiring, bash execution, queue management, animation, render helpers, event normalization, and three legacy event adapters (legacyToolStartEvent, legacyToolResultEvent, legacyThinkingEvent). Worth splitting at least the legacy event normalization and the *EditorAccess helpers into adjacent modules; the host class is the consumer surface, the rest is implementation.

8. The mountStageCustomUi cast chain is awkward but unavoidable

tui as unknown as Parameters<StageCustomUiRequest[\"factory\"]>[0]

The PiCustomOverlayFactoryTui is structurally compatible with TUI but the types don't reconcile statically. Consider defining the broker's factory parameter as PiCustomOverlayFactoryTui directly so the cast disappears (this would also tighten what's actually passed to factories).

9. setEditorBorderColor assigns borderColor only if it already exists (chat-session-host.ts:1009-1017)

if (candidate.borderColor !== undefined) candidate.borderColor = borderColor;

This is a no-op for editors that don't pre-initialize the field. Intentional (test-stub editors don't have it), but a try/catch around the assignment would be more idiomatic than the structural sniff. Not blocking.

Good things

  • The chat-message-renderer.ts:221-227 dedup that absorbs the message_start/toolResult echo after tool_execution_end is a clean fix for a real duplication bug.
  • HTML escaping in mcp-callback-server.ts correctly mitigates a reflected XSS via error_description — good defense-in-depth even though the localhost server is single-user.
  • updateStatusBar now counts only connected connections (vs. all entries in the map) — fixes a real visual bug where needs-auth / closed showed as connected.
  • The completedHandle aliasing in ui-session.ts:267-271 is a clean fix for the lookup race where handle could be reassigned to null between the === handle check and the field accesses.
  • planStageChatFrame reserving editor rows before body rows is the right call — guarantees the composer stays visible under aggressive shrink.
  • rejectAllStagePrompts on run end / removal closes a real leak.

Test coverage

Strong coverage of the new chat host (487 lines), stage chat view (363), and MCP status bar (82). Gaps:

  • No test for the executor cleanup change beyond the "handle remains chat-capable" assertion — would be useful to add one that runs N stages and verifies a documented retention/cleanup policy.
  • No test for the StageUiBroker wedge scenario in Flora131/feat/add skills #4.

Summary

The architectural direction (extracted chat host + stage UI broker) is solid. My one substantive concern is #1 — the executor change leaks AgentSessions for the run's lifetime with no documented release point. Everything else is small-to-medium polish.

@claude

claude Bot commented May 23, 2026

Copy link
Copy Markdown

Code Review

Nice extraction. The ChatSessionHost + StageUiBroker split eliminates a lot of duplication between interactive mode and stage chat, and the MCP fixes are good catches. A few things worth a closer look before merge:

🚨 Likely resource leak: stage SDK sessions no longer disposed

packages/workflows/src/runs/foreground/executor.ts (around L1559) replaces the previous detach-driven release with just dropStageControlHandle() + limiter.release(). The comment explains the intent (keep the chat handle live after the stage settles), but I can't find any path that disposes the SDK session afterwards:

  • releaseLiveHandle is now only called inside waitForStageRelease on barrier rejection (line 1449), i.e. only during cascade-pause cancellation.
  • recordRunEnd / removeRun in store.ts reject pending prompts but don't iterate stages to release SDK sessions.
  • The try { … } finally { opts.cancellation?.unregister(runId); } block at the end of run() doesn't call releaseLiveHandle either.

Net effect: every stage's AgentSession (with its MCP connections, file handles, etc.) lives until process exit even after the workflow completes. For long-lived daemons / runDetached workflows that's an unbounded leak. A run-end cleanup pass that calls releaseLiveHandle on every stage (or that defers release behind "no chat pane attached AND run terminal") would close the gap. If this is intentional and disposal happens further out (e.g. when the run is dropped from store history), it'd be worth a code comment pointing at the owner.

Predictable tmpfile path in openExternalEditorForText

packages/coding-agent/src/modes/interactive/chat-input-actions.ts:73:

const tmpFile = path.join(os.tmpdir(), `pi-editor-${Date.now()}.pi.md`);

This is a TOCTOU/symlink-attack vector on shared /tmp systems — an attacker who can guess the millisecond stamp can pre-create a symlink and have writeFileSync clobber an arbitrary file the user owns. protected_symlinks=1 on modern Linux mitigates most variants, but the sibling helper pasteClipboardImageToEditor already uses crypto.randomUUID() (L48); doing the same here costs nothing. This is pre-existing code being moved out of interactive-mode.ts, but the extraction is a clean opportunity to harden it.

Other observations

  • StageUiBroker abort wiring (stage-ui-broker.ts:55-89) — the onAbort handler calls this.reject(request, …), and reject correctly bails when pending.get(hostKey)?.id !== request.id, so stale aborts are safe. The early if (signal?.aborted) check at the top handles the already-aborted case before addEventListener is called. 👍
  • MCP XSS hardening (mcp-callback-server.ts:58-65)escapeHtml covers &<>\"'. Good. The success template has no user input.
  • updateStatusBar connected-count fix (init.ts:277-279) — clear improvement: previously map size lied during reconnect/needs-auth states. The new mcp-init-statusbar.test.ts cases lock that down.
  • chat-session-host.ts optimistic user dedup (L213-227) — signature is text.trim(), so two consecutive identical prompts use a counter. That's correct, but a comment near optimisticUserSignatureCounts explaining "counter dedupes the optimistic row against the real message_start echo" would save the next reader some archaeology.
  • chat-session-host.ts:626this.editor = undefined in dispose() but the field is EditorComponent | undefined; if the editor exposes its own dispose() it's not being called. Consider this.editor?.dispose?.(); this.editor = undefined; to match the symmetry with mountedCustomUi?.component.dispose?.() in StageChatView.dispose.
  • Test coveragechat-session-host.test.ts (487 lines) is thorough across slash commands, bash, compaction queueing, retry, and custom-entry rendering. stage-chat-view.test.ts adds 363 lines around the broker integration. Solid.
  • Date.now() in nextRequestId (stage-ui-broker.ts:26) — fine for an in-process ID; just noting it's not collision-proof under concurrent calls within the same millisecond. The 6-char random suffix makes that vanishingly unlikely.

Style nits

  • stage-chat-view.ts:163 defines ANSI constants (ITALIC, FG_RESET, etc.) as module locals; color-utils.js already exports BOLD/RESET. Worth consolidating to keep ANSI plumbing in one place.
  • stage-chat-view.ts:559shortenId returns the first 8 chars when length > 10 (so a 12-char ID gets cut to 8). The > 10 threshold reads off-by-one; intended? length > SHORT_ID_LEN would be more obvious.

Nothing here is blocking — the refactor is well-structured and the test coverage is strong. The stage-session disposal question is the only thing I'd want to confirm before merging.

@claude

claude Bot commented May 23, 2026

Copy link
Copy Markdown

PR Review

Substantial refactor (4.4k+ / -1.7k across 64 files) that extracts a reusable ChatSessionHost from interactive mode and wires it into workflow stage chat via a new StageUiBroker. The new abstractions look well thought out and the test coverage is solid. A few things worth addressing before merge.

Security

packages/mcp/mcp-callback-server.ts — HTML escape fix (positive)
Good catch turning HTML_ERROR into a function with HTML-escaped message. The escape order is correct (& first, then </>/\"/'), and covers the relevant reflected-XSS surface for the OAuth error path.

However: no test covers the new XSS hardening. For a defense-in-depth security fix, a small unit test asserting HTML_ERROR('<script>alert(1)</script>') produces escaped output would prevent regression. The error_description query param is attacker-influenced (a malicious authorization server could redirect with crafted error_description), so this is worth pinning down.

Bugs / correctness

packages/coding-agent/src/modes/interactive/chat-input-actions.ts:73

const tmpFile = path.join(os.tmpdir(), `pi-editor-${Date.now()}.pi.md`);

Two concurrent invocations within the same millisecond would collide and one would clobber the other's file. The neighbouring pasteClipboardImageToEditor already uses crypto.randomUUID() (line 48) — use the same here for consistency and collision safety.

packages/coding-agent/src/modes/interactive/chat-input-actions.ts:39-60 (pasteClipboardImageToEditor)
The catch {} block silently swallows all errors. It can't distinguish "no clipboard image" (expected) from "fs.writeFileSync failed" (should surface). Also: clipboard images written to os.tmpdir() are never cleaned up — over a long session this leaks temp files. Consider either a periodic cleanup or registering a cleanup callback when the editor consumes the path.

packages/workflows/src/shared/stage-ui-broker.ts:26-27 (nextRequestId)

return \`stage-ui-\${Date.now().toString(36)}-\${Math.random().toString(36).slice(2, 8)}\`;

Only ~6 chars of randomness. Collisions within the same ms are possible and would cause resolve/reject to mis-route via the pending.get(hostKey)?.id !== request.id guard (silent no-op — the original requester's promise would never settle). crypto.randomUUID() would be safer here.

packages/workflows/src/runs/foreground/executor.ts (releaseLiveHandleWhenIdle)
The function is declared async but the queued-work branch does not await the eventual release — it fires the subscription path and returns. So callers that await releaseLiveHandleWhenIdle() get no guarantee the handle is released. That's intentional given the immediately-following limiter.release(), but the contract is subtle. Worth a one-line comment so future maintainers don't accidentally rely on the await.

Also, this is a behaviour change: previously, an attached chat pane kept the live handle until the user detached. Now the handle drops as soon as the SDK queue drains — even if a pane is still attached. The PR description doesn't call this out; users attached to a completed stage may notice a switch to read-only behaviour they didn't trigger.

Style / code quality

ChatSessionHost (1248 lines) is large but the boundary feels right and the tests demonstrate the surface is exercisable. A few smells:

  • Many unknown casts (e.g. as Component & Partial<Focusable>, as { setPlaceholder?: ...} patterns at stage-ui-broker.ts:131-148, throughout chat-session-host.ts:1000-1022). CLAUDE.md asks to avoid unknown/any. The pattern is mostly forced by structural compatibility with pi-tui, but a few of the local helper signatures (e.g. editorAccess) could be typed.
  • applyAgentEvent at line 212 has a ~100-line switch with several as Extract<...> casts. Consider a tagged-event normaliser to keep it tighter.

packages/coding-agent/src/modes/interactive/chat-input-actions.ts:78
editorCommand.split(\" \") doesn't handle quoted paths or arguments with spaces (e.g. \$EDITOR=\"/Applications/My Editor.app/bin/edit --wait\"). Pre-existing limitation worth a TODO if not in scope.

Test coverage

  • ChatSessionHost (487 lines of test) covers slash commands, bash exclusion, optimistic dedup, compaction, retry, queue restoration. Solid.
  • StageUiBroker is exercised via wiring-adapters.test.ts and stage-chat-view.test.ts. The single-pending-request rejection path and signal.abort() flow could use direct unit tests.
  • Missing: no test for the HTML_ERROR escape (see Security above).
  • Missing: no test for releaseLiveHandleWhenIdle's async-subscription path — the previous attach/detach race was load-bearing and the new behaviour deserves a regression test (executor + a fake stage handle with pending messages).

Minor

  • packages/workflows/src/shared/stage-ui-broker.ts:131-148 mounts a custom-UI component by manually re-exposing each method (render, handleInput, invalidate, dispose, focused). If pi-tui ever adds a new Component method, this wrapper silently drops it. A spread + targeted overrides would be more forward-compatible.
  • packages/workflows/src/shared/store.ts:562 (recordStageAwaitingInput) — early-returns false for paused/blocked, but the broker calls recordStageAwaitingInput(runId, stageId, true, request.createdAt) unconditionally when a request lands. If a stage is paused when a custom-UI request arrives, the awaiting flag silently won't be set, and the resolve path won't reset anything. The pending request itself is still recorded in this.pending, so it'll fire when a host registers, but the store snapshot will misrepresent the state. Worth confirming this scenario can't occur (or guarding).

Overall: solid refactor — abstractions are well-named and tests give confidence in the core flows. Main asks: cover the XSS fix with a test, swap Date.now()/Math.random() ID generators for crypto.randomUUID(), and either document or guard the behaviour change in releaseLiveHandleWhenIdle.

@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

Code review — PR #1026 (refactor/workflow-chat-host-parity)

Reviewed against CLAUDE.md conventions and the Bun/raw-TS conventions for packages/workflows. Solid refactor — bringing stage chat to parity with interactive mode by extracting ChatSessionHost (~1248 LOC) and routing ctx.ui.custom via StageUiBroker is the right shape, and the net effect is ~800 LOC removed from stage-chat-view.ts. Test coverage is genuinely good. Most findings below are minor.

Strengths

  • Architecture: clean extraction. StageChatView now composes ChatSessionHost rather than duplicating logic. StageUiBroker is small, focused, and side-steps the awkward ctx.ui.custom routing problem with a per-stage queue, host registration, and signal-based abort cleanup.
  • MCP hardening:
    • updateStatusBar now counts only connection.status === "connected" (packages/mcp/init.ts:277-279) — the previous .size count was a real bug. Test covers the connected/needs-auth/closed mix.
    • colorizeStatusText defends against themes without .fg() and against .fg() throwing (packages/mcp/init.ts:284-292).
    • OAuth callback error template now interpolates ${escapeHtml(message)} — closes a small reflected-XSS surface in the local callback HTML.
    • state.uiServer === handle && handle !== null plus the stable completedHandle binding in ui-session.ts:267-300 removes a TOCTOU-style race where handle could be reassigned during the close cascade.
  • Executor: releaseLiveHandleWhenIdle (packages/workflows/src/runs/foreground/executor.ts:1354-1368) is a real correctness improvement — keeps the live handle alive until the SDK queue drains rather than collapsing the session mid-stream. Backed by a dedicated executor test (completed stage handle is kept only until queued messages drain).
  • Store: stage pendingPrompt lifecycle is consistent with the existing run-level prompt model; rejectAllStagePrompts is invoked from both recordRunEnd and removeRun, so prompt awaiters won't dangle.
  • Tests: 487 lines of new chat-session-host.test.ts covering the important branches (idle Enter, mid-stream steer, follow-up, bash !/!!, escape-while-bash, compaction queue + flush, slash-command short-circuit, optimistic-user dedup). Good.

Findings

  1. releaseLiveHandleWhenIdle return-promise semantics are misleadingexecutor.ts:1354-1368. When work is pending it subscribes and resolves immediately; the actual release happens later via the innerCtx.subscribe callback. The caller await releaseLiveHandleWhenIdle().catch(...) (line 1581) doesn't actually wait. Two concerns:

    • The Promise<void> return type suggests it resolves on release; it doesn't.
    • If the session never emits another event after draining (no queue-drain event), the handle leaks. The behaviour relies on the contract that subscribe fires on every relevant state change. Consider either documenting this contract explicitly in a comment, or defensively adding a microtask re-check loop with a bounded timeout.
  2. openExternalEditorForText tmp filename can collidechat-input-actions.ts:73. Uses path.join(os.tmpdir(), pi-editor-${Date.now()}.pi.md). Two simultaneous opens within the same millisecond (rare but possible across stage chat + interactive at once) would clobber each other and one would read the other's content. The sibling pasteClipboardImageToEditor uses crypto.randomUUID() — do the same here.

  3. Glimpse iframe HTML interpolates the URL unescapedpackages/mcp/ui-session.ts:327. <iframe src=\"${handle.url}\"> is built without escapeHtmlAttribute. The URL is internal/localhost so the blast radius is small, but packages/mcp/host-html-template.ts already has escapeHtmlAttribute — using it here would be consistent and defensive against any future change that lets tool names or params flow into the URL.

  4. Duplicate escapeHtml helperpackages/mcp/mcp-callback-server.ts:58-65 duplicates the helper already exported (privately) in host-html-template.ts. Either lift it to a shared module or export from one. Minor.

  5. StageUiBroker.registerHost silently overwritesstage-ui-broker.ts:38-46. Re-registering for the same (runId, stageId) replaces the prior host with no notification; the first host's returned unregister becomes a no-op via the === host guard. The replacement semantics are reasonable, but if two StageChatViews mount for the same stage (shouldn't happen but possible during transient detach/reattach races) the older one will silently lose UI events. A debug log on overwrite would help future debugging.

  6. Missing broker test casestest/unit/wiring-adapters.test.ts exercises the happy path well, but I didn't find tests for:

    • Re-entrancy: second requestCustomUi with one already pending returns the "already has a pending custom UI request" rejection.
    • Signal-driven abort: passing an already-aborted or later-aborted AbortSignal correctly rejects and clears the pending entry.
    • registerHost replacement while a pending request exists (does the new host receive showCustomUi?).
      These are all small but cover the broker's contract more fully.
  7. combineQueuedMessagesForEditor inconsistent trimmingchat-input-actions.ts:23-31. currentText is trimmed (text.trim() ? [...] : []) but each queued message is included verbatim. If a queued message is whitespace-only it still becomes a \\n\\n-separated empty paragraph. Minor; probably not user-visible.

  8. pasteClipboardImageToEditor does not clean up its temp PNGchat-input-actions.ts:48-50. Pre-existing behaviour copied from interactive-mode.ts, not introduced by this PR, so not a blocker — but worth tracking as a separate cleanup task since the path is sent to chat and the file may live in os.tmpdir() indefinitely.

  9. ChatSessionHost.dispose() does not detach editor handlerschat-session-host.ts:617-627. Only the timers and this.editor reference are cleared. The editor's onChange/onSubmit/onAction/onPasteImage/onEscape closures still hold this. If the host outlives its editor (or vice versa) and the editor is re-used by another consumer, those callbacks would fire into a disposed host. In the current flow the editor is created and owned by the host, so practically safe — but worth a sentence in dispose explaining why the closures don't need explicit teardown.

Conventions / housekeeping

  • bun usage, no dist//tsconfig.build.json introduced for workflows — ✅ compliant with the EXTREMELY_IMPORTANT block in CLAUDE.md.
  • .js import specifiers used consistently in new files — ✅.
  • No any introduced; well-typed throughout. ✅.
  • New code reads cleanly with minimal speculative abstraction. The size of chat-session-host.ts (1248 LOC) is justified by the surface it consolidates, though future cleanups could split rendering, input handling, and agent-event ingestion into distinct files.

Suggested follow-ups (non-blocking)

  • Add the three broker test cases listed in (6).
  • Switch openExternalEditorForText to crypto.randomUUID() for tmp filenames (2).
  • Wrap handle.url in escapeHtmlAttribute in the Glimpse template (3).
  • Clarify the releaseLiveHandleWhenIdle contract with a code comment or bounded timeout (1).

Net: this is a strong refactor with good test scaffolding and several real bug fixes mixed in. The findings above are mostly minor; (1) and (2) are the only ones that could plausibly cause observable issues in production, and both are easily addressed.

@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

PR #1026 review — embed stage-scoped custom UI in chat session hosts

Thanks for the big and otherwise well-factored refactor — extracting ChatSessionHost and pairing it with StageUiBroker is a clean win, and the MCP hardening edits travel well alongside it. A few things worth a closer look before merge.

Bugs / behavior

  • StageChatView doesn't react to an externally-rejected custom UI request (packages/workflows/src/tui/stage-chat-view.ts:344). When the broker rejects a request (signal abort from the executor, or another stage chat view registering as the host later), the view's mountedCustomUi stays non-null because the onDone in mountStageCustomUi is only fired from the factory's resolution path. The user is left with a live-looking form whose submit becomes a silent no-op (broker.resolve returns early because the pending entry is gone). A second listener — either via signal.addEventListener mirrored on the view, or a broker.onRequestRejected(request.id, …) seam — would let the view unmount the component and call requestRender to flip back to the chat composer.

  • registerHost can leave a stale host with a pending request (packages/workflows/src/shared/stage-ui-broker.ts:38-46). If a second host registers for the same (runId, stageId) while a pending request exists, both hosts will receive showCustomUi(request). The old host's unregister() becomes a no-op (because this.hosts.get(hostKey) === host is now false), but the old host already mounted the component. Consider either rejecting the second registerHost call, or actively evicting the previous host via a registered cleanup. As-is, this is mostly defensive — the realistic call sites are single-host — but it's an easy footgun if attach behavior ever changes.

  • openExternalEditorForText tmp filename collides under concurrent use (packages/coding-agent/src/modes/interactive/chat-input-actions.ts:73). pi-editor-${Date.now()}.pi.md is millisecond-resolution; two near-simultaneous editor invocations (interactive + stage chat opening editors back-to-back) can race. The sibling clipboard-paste helper already uses crypto.randomUUID(); suggest the same here for consistency and to avoid one user clobbering the other's draft.

  • editorCommand.split(" ") breaks on quoted paths (chat-input-actions.ts:78). EDITOR="/Applications/My Editor.app/.../bin/editor" will be split into /Applications/My + args. Probably acceptable as a pre-existing inherited behavior, but worth flagging because the new shared helper means more code paths hit it. A small shell-quote-aware parser, or at minimum a clearer warning, would help.

  • releaseLiveHandleWhenIdle relies on innerCtx.subscribe always firing on queue drain (packages/workflows/src/runs/foreground/executor.ts:1354-1368). The idempotent flag protects against double-release, and the queueMicrotask defer is sensible. The concern: if __pendingMessageCount decrements without emitting a subscribe event (e.g., a path that silently drops a queued message during teardown), the handle is never released. The new test at test/unit/executor.test.ts:2137 exercises the happy queue_update path; an extra test simulating a silent drop (or a small belt-and-suspenders timer/check) would protect against future SDK regressions.

Code quality

  • makeStageExtensionUiContext has a subtle behavioral change when meta is undefined and ui.custom is also undefined: previously the no-ui.custom branch threw "ask_user_question UI is unavailable" synchronously from inside the factory wrapper; the new flow throws inside an async function, so it now becomes a rejected promise instead. That's the correct change for awaiters, but worth a quick sweep for any call sites that did synchronous try/catch.

  • stageBuiltinPackagePaths filter uses basename(path) !== "workflows" (packages/workflows/src/extension/wiring.ts:166-174). The comment explains why clearly, which is great — but the basename match is fragile (a sibling fork named workflows-experimental would also be filtered if you ever rename, etc.). Not a blocker; consider keying off a sentinel field on the package manifest if you anticipate variants.

  • ChatSessionHost.handleInput's ASCII fallback (packages/coding-agent/src/modes/interactive/components/chat-session-host.ts:464) only accepts >= " " && <= "~". That's fine for the headless/no-editor path, but consider documenting it — anyone wiring a host without an editor (tests, CI) will silently lose Unicode keystrokes.

Security

  • OAuth callback HTML escaping is the right call. escapeHtml covers the standard five entities and is applied at the only string interpolation site. Note: the state parameter is not echoed back to the page anywhere; if you ever add it to the error page (or to telemetry) make sure that path is escaped too.

  • pasteClipboardImageToEditor writes a fresh file to os.tmpdir() per paste; nothing in the new helper deletes it after the editor reads it. Long sessions will accumulate these. Not new in this PR, but the freshly-extracted helper is a good place to add a cleanup pass.

Tests

New coverage for ChatSessionHost, MCP status bar filtering/theme fallback, and the stage-scoped pending-prompt lifecycle is substantial and reads cleanly. A few gaps worth filling:

  • No direct test for StageUiBroker itself (request lifecycle, abort, replacement). The store-pending-prompt tests cover the store side, but the broker's de-dup / abort / replacement semantics are unverified.
  • No test for the externally-rejected-while-mounted case described above. Once that bug is fixed, a regression test would be straightforward.
  • The new releaseLiveHandleWhenIdle test in executor.test.ts:2137 only exercises the happy drain — consider one for the immediate-release case (no queued work) and one for cancellation while draining.

Nit

  • packages/workflows/src/shared/stage-ui-broker.ts:65let request!: StageCustomUiRequest<T> with the definite-assignment assertion works, but assigning inside a new Promise constructor body is a slightly subtle pattern. Promise.withResolvers() (used elsewhere in executor.ts) would let you avoid the bang and keep the assignment linear.

Overall this is a strong refactor — the size mostly comes from the legitimate code-motion of extracting ChatSessionHost, and the deletions in stage-chat-view.ts (~800 lines) more than justify the new module. Worth resolving the externally-rejected custom-UI issue and the tmpfile collision before merging.


Reviewed by Claude (Opus 4.7).

Comment thread packages/mcp/mcp-callback-server.ts Fixed
Comment thread test/unit/mcp-security.test.ts Fixed
@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

ping test

@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

Review: feat(workflows) embed stage-scoped custom UI in chat session hosts

Solid refactor with meaningful security hardening alongside it. The extraction of ChatSessionHost is a real win — stage-chat-view.ts shedding ~500 net lines is the kind of duplication-collapse you want — and the StageUiBroker is well-shaped (single-owner pending request, host swap, abort wiring all covered by tests).

Likely bugs / behavior to revisit

  1. Double APP_NAME in editor temp filename — packages/coding-agent/src/modes/interactive/chat-input-actions.ts:172. The template literal interpolates APP_NAME twice into the temp file name (once as prefix, once before the .md extension), yielding paths like /tmp/atomic-editor-UUID.atomic.md (extension is .atomic.md, not .md). Editors keying syntax highlighting off the extension fall back to plain text. Looks unintentional — likely meant a single .md suffix.

  2. releaseLiveHandleWhenIdle poll timer is not cleared on abort — packages/workflows/src/runs/foreground/executor.ts:1354-1387. The 250 ms setInterval fallback only stops when releaseIfIdle observes either liveHandleReleased true or an empty queue. If the run is killed but releaseLiveHandle is never called via another path (e.g. because the inner SDK gets stuck and pendingMessageCount stays positive), the watcher polls indefinitely. The unref keeps the process from being pinned, but in long-lived servers or tests this is a slow leak. Suggest hooking ownController.signal so abort triggers cleanupWatcher and a final releaseLiveHandle.

  3. Broker ignores store rejection — packages/workflows/src/shared/stage-ui-broker.ts:88. The store refuses recordStageAwaitingInput when the stage is paused/blocked/completed/failed (store.ts:568). When refused, the broker still sets pending and shows the request to whichever host is registered. Net effect: a paused or terminal stage can mount a custom UI that nobody can resolve through the store-driven status display. Consider rejecting the request when recordStageAwaitingInput returns false, or at minimum log/warn.

  4. recordStagePendingPrompt clobbers awaitingInputSince set by ask_user_question — packages/workflows/src/shared/store.ts:436-450. If a stage already has awaitingInputSince set by the tool-event watcher (executor.ts:1325), recording a pending prompt overwrites it with prompt.createdAt. The two awaiting-input sources do not coordinate. Probably benign today (they do not overlap in practice), but worth a comment or guard before this grows a third caller.

  5. OAuth error path returns HTTP 200 — packages/mcp/mcp-callback-server.ts:117. All other error branches return 400; the error branch returns 200 with the error HTML. Maybe intentional for browser UX (avoids the this-site-cannot-be-reached overlays), but it is inconsistent with the rest of the file. A short comment would help future readers.

Security — what landed is good

  • HTML escaping in renderCallbackErrorHtml (mcp-callback-server.ts:53) closes a reflected-XSS path via error_description. The new mcp-security.test.ts exercises both the OAuth body and the iframe src escape — nice.
  • chat-input-actions.ts writes temp files with flag wx and mode 0o600, so existing-file overwrite attacks and other-user reads are blocked.
  • spawnSync with stdio inherit and shell only on Windows — args are pre-tokenized by parseEditorCommand; shell use is restricted to Windows where it is necessary for .cmd/.bat resolution. Reasonable trade-off.

One small follow-up: parseEditorCommand does not reject EDITOR values that contain shell metacharacters before the Windows shell path. Low-risk (user owns their EDITOR), but if you want defense-in-depth, sanitize before shell true.

Performance

  • The host extraction preserves the 80 ms streaming throttle and tail-capping — no regressions visible.
  • 250 ms poll cadence in releaseLiveHandleWhenIdle is reasonable, but combined with the subscriber-driven queueMicrotask(releaseIfIdle) it is belt-and-suspenders; once the abort cleanup above is in place, the poll could probably stretch to 500-1000 ms.

Tests

Coverage is strong for this size of change:

  • stage-ui-broker.test.ts covers id format, duplicate rejection, abort and host-replace
  • mcp-security.test.ts and mcp-init-statusbar.test.ts lock in the security and status-bar fixes
  • chat-session-host.test.ts (506 lines) and the additional 108 lines in executor.test.ts exercise the new lifecycle

Gaps worth filling:

  • A test that hits releaseLiveHandleWhenIdle while the run is aborted, asserting both the poll timer and subscription are cleaned up.
  • A test for the broker-received-a-request-the-store-refused case (item 3 above) — even just to pin down current behavior.

Style / CLAUDE.md alignment

  • All test files use bun:test and node:assert/strict — ok
  • No new dist/, no build step added to packages/workflows — ok
  • Uses Bun-native APIs throughout, no Node/npm leakage — ok
  • One minor: chat-session-host.ts is 1266 lines. It is a clean extraction so this is fine for landing, but consider splitting input-handling, render, and lifecycle into sibling files in a follow-up — the file is at the edge of one-screenful-in-the-editor territory.

Verdict

Approve with the editor-extension-typo fix (item 1) addressed; items 2 through 4 are good follow-up tickets but not blockers given the scope. Nice work on collapsing the duplication and tightening the MCP surfaces in the same PR.

Generated with Claude Code (https://claude.com/claude-code)

@flora131
flora131 force-pushed the refactor/workflow-chat-host-parity branch from 23064c5 to 495763b Compare May 24, 2026 01:18
@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

Code Review

Solid refactor — extracting ChatSessionHost and StageUiBroker removes a meaningful amount of duplicated logic and the new tests cover the important paths (host lifecycle, broker registration races, store-scoped prompt resolvers, MCP status filtering). A few findings, ordered by importance.

Bugs / correctness

1. PR description doesn't match the MCP callback fix (medium). The description says “OAuth callback error template made dynamic; error messages are HTML-escaped (escapeHtml) to prevent reflected XSS.” The actual change in packages/mcp/mcp-callback-server.ts does the opposite — it makes the template static (no user input is interpolated) and there's no escapeHtml call in this file. The XSS is correctly mitigated by removing the unsafe interpolation, but please update the description so reviewers/future readers aren't looking for an escapeHtml helper that doesn't exist.

2. Dead variables in mcp-callback-server.ts (low). Now that renderCallbackErrorHtml() takes no argument, errorMsg is computed but never used at lines 98 and 106:

const errorMsg = "Missing required state parameter - potential CSRF attack"
res.writeHead(400, { "Content-Type": "text/html" })
res.end(renderCallbackErrorHtml())   // errorMsg dropped on the floor

This compiles under noUnusedLocals only because errorMsg is a const initializer with a side-effect-free RHS in your tsconfig — please delete the bindings to avoid masking future drift, and consider sending the message to logger.debug for operator visibility (the user-facing HTML no longer carries that diagnostic).

3. releaseLiveHandleWhenIdle removes the “attached keeps handle alive” invariant (medium — behavioral). The previous executor logic kept the live SDK handle alive while a chat pane was attached; the new version drops it once the SDK queue drains once, then never re-arms. After that point, handle.prompt(text) from StageChatViewChatSessionHostcommands.prompt will fail because innerCtx is disposed. The PR description (“prevent premature session collapse while messages are still in flight”) and the inline comment (“the node reopens as a read-only archived session”) suggest this is intentional — but if so, StageChatView should also flip into a clearly read-only mode once handle.isDisposed === true, otherwise users can type into an editor whose Enter silently surfaces an exception. Worth either documenting the regression in CHANGELOG.md under “Changed”, or wiring isDisposed into ChatSessionHost.isDisabled so the input visibly grays out.

4. releaseLiveHandleWhenIdle polling has no upper bound (low). The 250 ms setInterval is unref()'d so it won't keep the process alive, but if innerCtx.__pendingMessageCount() is wedged (SDK bug, async leak), the watcher runs for the lifetime of the process and the SDK session is never disposed. A bounded retry budget (e.g. ~30 s) that falls through to a forced releaseLiveHandle() would be more defensive than the current “poll forever” behavior.

Code quality

5. ChatSessionHost shadow state for the input buffer. inputBuffer is set both from setEditorText and from the editor's onChange handler. They currently stay in sync because setEditorText writes inputBuffer first then calls editor.setText(...), but every new edit path needs to remember to go through setEditorTextsubmit() itself uses submittedText ?? this.inputBuffer, so a stale buffer leaks straight into the SDK. Worth either making editor.getText() the single source of truth, or adding a comment by inputBuffer warning new callers to use setEditorText only.

6. CLAUDE.md says avoid unknown (low). StageChatViewOpts.piTheme: unknown and piKeybindings: unknown in packages/workflows/src/tui/stage-chat-view.ts:101-102, plus keybindings: unknown on ChatSessionHostOpts in chat-session-host.ts:97. These are then handed straight into editorFactory(tui, theme, keybindings) where the factory expects concrete types — the host-side cast is doing work the types should be doing. The pi-tui types for Keybindings and the editor theme are already imported elsewhere in the same files; using them here would catch wiring regressions at compile time.

7. stage-ui-broker.ts uses an opaque cast for the factory's TUI argument. request.factory(tui as unknown as Parameters<StageCustomUiRequest["factory"]>[0], ...) in mountStageCustomUi is fighting the type system. Since the factory signature is PiCustomOverlayFactory<T>, taking the actual TUI parameter type from there would be clearer than the double-cast.

Security (positive)

  • escapeHtmlAttribute(handle.url) in ui-session.ts is a real fix — handle.url was being interpolated unescaped into iframe src=\"...\".
  • Connected-count fix in updateStatusBar (only connection.status === \"connected\" is counted) is correct; the previous .size counted disconnected/auth-pending connections.
  • colorizeStatusText guard against missing theme.fg looks fine.

Test coverage

Good. The new tests cover request-id collision resistance, duplicate-pending rejection, abort propagation, host replacement, and stage-scoped pending prompts. Two gaps worth adding:

  • StageUiBroker: no test for the case where mountStageCustomUi's factory throws — the broker's _showCustomUi callsite calls broker.reject(request, error) but that path is uncovered.
  • releaseLiveHandleWhenIdle in executor.ts: the test suite has executor.test.ts updates but I didn't see one that asserts the handle stays alive across a single queued steering message and is then released once the SDK reports the queue is drained. The expanded +108/-9 in that file is mostly about other paths.

Nits

  • chat-session-host.ts is 1266 lines — extracting bash execution and queued-message restoration into small helpers would make the host class easier to scan.
  • chat-input-actions.ts:30 defines CLIPBOARD_STALE_AGE_MS (24h) and the cleanup walks os.tmpdir() filtering by appTempPrefix(\"clipboard\") — that's fine, but the prefix uses APP_NAME which could shift across app builds; consider hardcoding a stable prefix for cleanup so an older atomic binary doesn't leave orphaned files an upgraded one ignores.

Overall the refactor looks worth merging once #1/#3 are addressed (or #3 is explicitly opted into in the changelog).

@flora131
flora131 force-pushed the refactor/workflow-chat-host-parity branch from 495763b to 3a3b0de Compare May 24, 2026 01:31
@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

Code Review — PR #1026

Substantial refactor — the ChatSessionHost extraction is a clean win and the test coverage growth is real. Findings below from a verification pass against the diff and the previous tree at b1a11b1.


Bugs (verified)

1. recordStageEnd doesn't reject pending stage-prompt waiters — packages/workflows/src/shared/store.ts:286-307
Only recordRunEnd (line 344) and removeRun (line 358) call rejectAllStagePrompts. If a single stage ends (failed/completed) while holding a pending stage prompt, stage.pendingPrompt stays set and the resolver remains in _resolvers until the whole run terminates. Mirror the pendingPrompt cleanup from recordRunEnd (lines 339-343) inside recordStageEnd after the delete existing.awaitingInputSince.

2. forceRelease unconditionally disposes the live handle after 30 s — packages/workflows/src/runs/foreground/executor.ts:1389-1396
There is no hasQueuedLiveWork() re-check before disposing. Any stage whose post-completion drain legitimately exceeds 30 s (long generation tail, slow tool, paused-on-user) will have the SDK session torn down mid-stream and queued messages dropped. Either re-check hasQueuedLiveWork() inside forceRelease and re-arm, or drop the timer (the subscribe + 250 ms poll already cover the silent-drop case the comment cites).

3. flushCompactionQueue drops queued messages on the first failure — packages/coding-agent/src/modes/interactive/components/chat-session-host.ts:908-915
The queue is cleared before await requiredCommand(\"prompt\")(first). If that prompt rejects, the for…of rest never runs and the remaining queued messages vanish. The invoker is void this.flushCompactionQueue() (line 306), so the rejection is also swallowed. Wrap in try/catch and re-queue (or push to follow-up) on error.

4. StageUiBroker leaks the pending request when host unregisters — packages/workflows/src/shared/stage-ui-broker.ts:52-54
The unsubscriber only removes the host map entry. If a stage pane closes with a request still pending and no signal was supplied, the request sits in this.pending with no host to display it and will never settle. Either reject on unregister-with-pending or make signal a required parameter and document it.

5. StageUiBroker abort TOCTOU — packages/workflows/src/shared/stage-ui-broker.ts:64-94
signal?.aborted is checked at line 64; the addEventListener is installed at line 94 after the request is registered and the host shown. If the signal aborts between those points, the listener will never fire (DOM addEventListener(\"abort\",…,{once}) on an already-aborted signal does NOT auto-invoke), and the request hangs. Re-check signal?.aborted after attaching the listener and abort manually if it flipped.

6. submit clears localPaused in finally on workflow resume — packages/workflows/src/tui/stage-chat-view.ts:248-257
On resume failure the stage may still be paused, but localPaused is set false. The store-subscribe path self-heals on the next tick, but until then the UI shows the wrong state. Old _resume (b1a11b1, lines 1073-1097) only cleared localPaused on success.


Concerns

7. PR description overstates the MCP "reflected XSS" fix.
The previous HTML_ERROR = () => … (b1a11b1 packages/mcp/mcp-callback-server.ts:37) took zero parameters — callers passed errorMsg into a no-arg function where it was silently discarded. There was no reflected XSS to fix. The rename to renderCallbackErrorHtml + moving the error string into logger.debug is a reasonable cleanup, but test/unit/mcp-security.test.ts:7-13 only asserts the static body contains no <script> — it would have passed against the old code too. Either reframe in the PR description, or add an integration test driving handleRequest with ?error=<script>alert(1)</script> to make the regression guard real.

8. No disposed guard on ChatSessionHost async paths — chat-session-host.ts:631-641
dispose() clears two timers and the editor reference but does not flip an internal disposed flag. In-flight submit, runBashCommand, flushCompactionQueue, pasteClipboardImageToEditor, openExternalEditorForText will continue to mutate this.transcript/this.statusMessage and call this.requestRender?.() after teardown. runBashCommand's onChunk (line 851-906) is the most exposed — it fires per-chunk and writes to bashMessage.output. Add a private disposed = false and bail at each await resume point.

9. Slash commands bypass the compaction guard — chat-session-host.ts:500-513
Slash handler runs before the compacting check, so /compact while compaction is in flight can double-invoke handle.agentSession.compact. Gate slashes during compaction or have the /compact handler itself check the flag.

10. recordStageAwaitingInput(false) from broker can stomp executor's state — stage-ui-broker.ts:105/114
activeAskUserQuestionCalls in the executor (lines ~1319) manages the same awaitingInputSince field. If both paths are concurrently active, the broker's blind flip to false on settle can clear an awaiting-input set by the other path. Consider a refcount or a discriminated key.

11. unknown proliferation violates the CLAUDE.md "avoid any and unknown" rule.
chat-session-host.ts declares keybindings: unknown (lines 97, 101, 645, 721) where the only consumer casts to ConstructorParameters<typeof CustomEditor>[2]. Same in stage-chat-view.ts (piTheme?: unknown, piKeybindings?: unknown, lines 101-102). Use the structural types from wiring.ts rather than unknown.


Test gaps

  • test/unit/stage-ui-broker.test.ts — does not cover (a) host registers then immediately unregisters with pending; (b) signal aborted between request creation and listener install; (c) host-replacement double-settle path.
  • test/unit/store-pending-prompt.test.ts — only exercises rejection on run end, not stage end (the bug above is uncovered).
  • test/unit/executor.test.ts — covers isStreaming/pendingMessageCount gating and queue_update drain, but not the 30 s force-release path (so the unconditional dispose is uncovered) nor the subscribe-after-drain race that the 250 ms poll exists to catch.
  • test/unit/mcp-security.test.ts — see updated readme #7.

Verified clean / good

  • packages/mcp/init.ts:277-279connection.status === \"connected\" filter is correct; old code counted needs-auth/closed too.
  • packages/mcp/init.ts:284-292colorizeStatusText theme guard with typeof theme?.fg !== \"function\" + try/catch is sound.
  • packages/mcp/ui-session.ts:268-269 — local completedHandle binding correctly fixes the closure-capture bug on the outer let handle.
  • packages/mcp/host-html-template.ts:421escapeHtmlAttribute is defense-in-depth (handle.url is internally generated in ui-server.ts:511, not attacker-controlled), but harmless and good hygiene.
  • chat-input-actions.ts — helpers are properly standalone; only deps are node:*, the clipboard-image util, and APP_NAME. Safe to share between interactive and stage chat.
  • stage-chat-view.ts 875 → 367 deduplication checks out — compaction, retry, steer/followUp, optimistic dedupe (now a counted map, an improvement), and animation tick are all preserved via ChatSessionHost.

Priority order for follow-up

  1. Bugs add agent instructions #1, updates to readme and instructions #2, update readme and mcp servers #3, Flora131/feat/add skills #4 — all small fixes with clear paths.
  2. Bug Flora131/feat/refine readme #5 (broker TOCTOU) — narrow but real.
  3. Concern updated readme #7 — reframe the PR description (no code change needed).
  4. Test coverage for add agent instructions #1, updates to readme and instructions #2, Flora131/feat/add skills #4, Flora131/feat/refine readme #5.

@flora131
flora131 force-pushed the refactor/workflow-chat-host-parity branch from 3a3b0de to 167e3cf Compare May 24, 2026 02:14
@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

Review — feat(workflows): embed stage-scoped custom UI in chat session hosts

Skimmed the extraction, the broker, the executor lifecycle change, and the MCP hardening. Overall this looks like a careful, well-tested refactor. A few things worth flagging.

Behavior change worth confirming intent

packages/workflows/src/runs/foreground/executor.ts:1352-1387 — the live stage-control handle release policy was changed from release on user detach (old) to release when SDK queue drains (new). The old test attached completed stage handle remains chat-capable until detach was removed and replaced with completed idle stage handle is released after settle.

Net effect for users: a chat pane that's still attached to a completed stage will transition from interactive → read-only as soon as the SDK queue empties, instead of waiting for the user to press Ctrl+D. If the user is in the middle of a conversation after the workflow finished, the composer can go cold under their cursor. The PR description doesn't call this out explicitly — is this the intended product behavior, or is the policy still meant to wait for detach when a pane is attached?

If intentional, please mention it in the PR description; if not, the predicate at executor.ts:1352-1353 probably needs to also consider stageSnapshot.attached === true.

StageUiBroker — description / implementation mismatch

packages/workflows/src/shared/stage-ui-broker.ts:76-79 — a second concurrent requestCustomUi for the same (runId, stageId) is rejected, not queued. The PR description says the broker supports "pending-request queuing." Either the description should say "rejects duplicate concurrent requests" or the implementation should actually queue. Current behavior is reasonable for ask_user_question (only one at a time per stage), but if a workflow stage ever fires two custom UIs back-to-back, the second one explodes immediately — and there isn't a clear way for the caller to know it should retry.

Brittle marker for "exclude workflows from child sessions"

packages/workflows/src/extension/wiring.ts:169-175:

function stageBuiltinPackagePaths(paths: readonly string[]): string[] {
  return paths.filter((path) => basename(path) !== "workflows");
}

The fix itself (avoid recursive workflows extension load in child stage sessions) is correct and the comment explains it nicely. The mechanism — pattern-matching on basename(path) === "workflows" — is brittle. If BUILTIN_PACKAGES ever exposes the workflows package under a different directory name (a dist layout variant, a renamed workspace dir), this filter silently fails open and the recursive-load bug returns.

Cheap improvement: have getBuiltinPackagePaths return { name, path } pairs, or export a WORKFLOWS_BUILTIN_DIR constant from builtin-packages.ts and import it here, so the marker is structural rather than a string match.

MCP hardening — minor

  • packages/mcp/mcp-callback-server.ts — static renderCallbackErrorHtml eliminates the reflected XSS cleanly. Note escapeHtmlAttribute only escapes for double-quoted attribute contexts (it intentionally skips '); the call site at ui-session.ts:329 uses double quotes, so it's fine, but worth a one-liner doc on escapeHtmlAttribute so future callers don't drop it into a single-quoted attribute.
  • packages/mcp/init.ts:284-292colorizeStatusText resolves the theme via (ui as { theme?: { fg?: ... } }). Since state.ui is typed in the file already, an explicit theme?: PiTheme field on the surface type would be cleaner than the ad-hoc cast.
  • The "counts only connected" status-bar fix (init.ts:277-279) is a real correctness improvement — previously needs-auth and closed entries inflated the count. Good catch.

ChatSessionHost extraction & tests

packages/coding-agent/src/modes/interactive/components/chat-session-host.ts — the host is well-encapsulated and the 543-line test file covers slash commands, bash, compaction queue flush, retries, and custom-entry rendering. Two small things:

  • dispose() (line 631) clears timers but doesn't notify the editor to release its own resources; if CustomEditor ever acquires async handles, this would leak.
  • The polling-interval fallback in releaseLiveHandleWhenIdle (executor.ts:1384) isn't covered by tests — the existing queued-drain test exercises only the subscription path. A defensive setInterval running every 250 ms is fine for one stage, but if a workflow with N stages all settle simultaneously, that's N intervals firing while queues drain. A test that verifies the interval path closes cleanly would be worth adding.

Worth a follow-up, not a blocker

packages/coding-agent/src/modes/interactive/chat-input-actions.ts:188-191spawnSync(editor, [...editorArgs, tmpFile], { shell: process.platform === "win32" }). On Windows, passing shell: true re-joins the parsed args through cmd.exe, which means parseEditorCommand's quoting work can be partially undone by shell interpretation. The trust boundary is $VISUAL/$EDITOR, so this is acceptable, but worth a note in the function so a future contributor doesn't widen the input source without remembering.

Things that look right

  • recordStagePendingPrompt / resolveStagePendingPrompt / awaitStagePendingPrompt plus the rejectAllStagePrompts cleanup in recordRunEnd and removeRun — the lifecycle coverage looks complete.
  • The TOCTOU hardening on state.uiServer === handle && handle !== null with const completedHandle = handle in ui-session.ts:268-307 is correct.
  • The isDisposed getter on the live handle is a clean addition for StageChatView._liveHandle() to gate on.

Nice work overall.

@flora131
flora131 force-pushed the refactor/workflow-chat-host-parity branch from 167e3cf to 2fd0a8b Compare May 24, 2026 02:23
@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

Review: PR #1026feat(workflows): embed stage-scoped custom UI in chat session hosts

A substantial refactor (84 files, +5239/-1758). High-quality test coverage in the new abstractions and clearly intentional MCP hardening. Below is a constructive review focused on items worth a second look.

MCP hardening — well-targeted

  • renderCallbackErrorHtml() no longer interpolates the OAuth error/description into the HTML — eliminates the reflected-XSS surface in mcp-callback-server.ts. The new mcp-security.test.ts correctly verifies that error=<script>... payloads don't reach the response body and that the rejection still carries the raw message for logs.
  • escapeHtmlAttribute(handle.url) before string interpolation into the Glimpse iframe src (ui-session.ts:328-329) — good. Note: escapeHtmlAttribute does not escape single quotes; since the attribute is double-quoted (src="${iframeUrl}"), this is safe today, but worth a comment so future edits don't switch to src='...'.
  • updateStatusBar now filters to status === "connected" — fixes a real bug where needs-auth/closed entries inflated the count. colorizeStatusText theme-fallback is a good defensive guard.

StageUiBroker (packages/workflows/src/shared/stage-ui-broker.ts) — a few subtle issues

  1. registerHost partial failure (lines 40–50): when replacing a host that has a pending request, if previousHost.hideCustomUi(...) throws, the new host is never installed and the pending request is orphaned. Consider wrapping the hide call in a try/catch (or scheduling it via queueMicrotask).
  2. Throwing showCustomUi: requestCustomUi calls this.hosts.get(hostKey)?.showCustomUi(request) synchronously (line 102). If the host throws, the pending entry remains and recordStageAwaitingInput(..., true) has already been recorded — the stage will be stuck in awaiting_input forever. The same applies in registerHost (line 51). A try/catch that calls this.reject(request, err) on failure would close the loop.
  3. Redundant abort wiring: addEventListener("abort", onAbort, { once: true }) + the post-set if (signal?.aborted) onAbort(); is belt-and-suspenders. Correct but the comments would help — at a glance it reads like a double-fire risk (it isn't, because reject is keyed on request.id).

Executor live-handle release (packages/workflows/src/runs/foreground/executor.ts:1354-1387)

The releaseLiveHandleWhenIdle adds a setInterval(releaseIfIdle, 250) "defensive fallback for silent queue drops during teardown." A few concerns:

  • Documenting which queue-drain events fail to fire would help — if the subscription path is reliable for current code, the timer is dead code; if it isn't, the underlying event emitter should be fixed rather than papered over with a polling fallback. Leaving a perpetual 250ms tick per ended stage is a smell, even with unref().
  • void releaseLiveHandle().catch(() => {}) (line 1381) silently swallows any release error. At minimum, a logger.debug would help debugging future hangs.

ChatSessionHost extraction (packages/coding-agent/src/modes/interactive/components/chat-session-host.ts, ~1280 lines)

The PR description states this was "extracted from interactive-mode.ts to eliminate duplication with stage chat." A quick grep shows the new class is consumed only by stage-chat-view.ts; interactive-mode.ts still contains its own (large) chat-driving implementation. Net result: the duplicated surface area has grown, not shrunk, until interactive mode is also migrated.

Two suggestions:

  • If migrating interactive-mode is intentionally out of scope for this PR, please update the description so reviewers don't assume the cleanup is done.
  • The ChatSessionHost API surface is quite wide (ChatSessionHostStyle, ChatSessionHostCommands, ChatSessionHostOpts, ChatSessionHostEntry, ChatSessionHostBashRequest — all exported from src/index.ts). Worth confirming this is intended as a public-stable API for external extension authors vs. an internal seam. If the latter, consider a internal/ namespace or non-exported re-export to avoid committing to the shape long-term.

chat-input-actions.ts

  • openExternalEditorForText writes the temp file as ${APP_NAME}-editor-${uuid}.${APP_NAME}.md (e.g., atomic-editor-<uuid>.atomic.md) — the doubled token reads like a typo; if it's intentional (matching legacy pi.md), worth a one-line comment.
  • cleanupStaleClipboardFiles uses fs.unlinkSync on every matching prefix entry; if a same-prefixed directory exists, the unlink throws and is swallowed by the empty catch. Minor, but fs.statSync(...).isFile() first would be more robust.
  • parseEditorCommand's end-of-input handling (line 109) treats a trailing backslash as a literal \. Most shells would treat this as a parse error or line continuation — non-critical given the source is $VISUAL/$EDITOR, but documenting the deviation would help.

Tests

Good coverage on the new code paths:

  • stage-ui-broker.test.ts — solid; exercises host replacement, abort-during-show, host unregister, duplicate request rejection. Consider adding: showCustomUi host that throws synchronously, and a recordStagePendingPrompt → recordRunEnd integration to confirm waiters reject.
  • mcp-security.test.ts — concise and on-point. The integration test using a live createServer(handleRequest) is the right level.
  • chat-session-host.test.ts (543 lines) — broad coverage of host lifecycle and slash/bash/compaction paths.

Style / repo conventions

  • cleanupStaleClipboardFiles(now = Date.now()) adds an optional parameter for testability — good pattern. The test file chat-input-actions.test.ts doesn't appear to exercise it; would be a cheap addition.
  • Several as unknown as ... casts in stage-ui-broker.ts:143-167 to bridge Component/Focusable — acceptable given the mount-time generic adapter, but worth a future cleanup pass to narrow the types.

Performance

  • The stage chat executor pre-allocates timers and broker entries per stage. With many concurrent stages the per-stage 250 ms poll noted above is the only one that scales with stage count; everything else is event-driven.

Summary

Net positive change: real security fixes, a meaningful re-routing of ctx.ui.custom into the correct stage's UI, and good test coverage. Main asks before merge:

  1. Resolve the "extracted but still duplicated" ChatSessionHost state (either migrate interactive-mode in the same PR or update the description).
  2. Add try/catch around previousHost.hideCustomUi and host.showCustomUi in StageUiBroker so a throwing host can't orphan the request.
  3. Reconsider or document the 250 ms poll in releaseLiveHandleWhenIdle.

Otherwise this looks ready to land.

— Review by Claude (claude-opus-4-7)

@flora131
flora131 force-pushed the refactor/workflow-chat-host-parity branch from 2fd0a8b to cd24204 Compare May 24, 2026 03:37
@claude claude Bot changed the title feat(workflows): embed stage-scoped custom UI in chat session hosts feat(workflows): extract ChatSessionHost and compose into stage chat May 24, 2026
@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

Code Review — feat(workflows): embed stage-scoped custom UI in chat session hosts

Reviewed the +5,362 / −1,758 diff across 84 files. The PR is a solid refactor: extracting ChatSessionHost, introducing StageUiBroker, and hardening MCP. Tests are extensive (especially stage-ui-broker.test.ts, chat-session-host.test.ts, mcp-security.test.ts, executor's drain/release tests). Below are findings worth a second look before merge.

🔒 Security — positive notes

  • OAuth callback HTML (packages/mcp/mcp-callback-server.ts:38-57): renderCallbackErrorHtml is now a static template; error / error_description query params are not interpolated into the response. Reflection-XSS path is closed. Good. mcp-security.test.ts covers the regression with a <script>alert(1)</script> payload.
  • Glimpse iframe src (packages/mcp/ui-session.ts:329): wrapping handle.url in escapeHtmlAttribute() before inline injection is correct. escapeHtmlAttribute (host-html-template.ts:421) escapes &, ", <, >. Since the attribute uses double quotes, single quotes don't need escaping. ✓
  • colorizeStatusText (packages/mcp/init.ts:284-292): the typeof theme?.fg !== \"function\" guard plus try/catch is a sensible belt-and-braces fix for non-standard themes.

🐛 Potential bugs / nits

  1. ChatSessionHost.animationTimer keeps polling when requestRender is undefined
    chat-session-host.ts:618-622setInterval is started even if this.requestRender is undefined, in which case the timer ticks every 80 ms and does nothing. Cheap to guard:

    if (shouldAnimate && !this.animationTimer && this.requestRender) { ... }
  2. Slash-command fall-through silently submits as a prompt
    chat-session-host.ts:500-507 and stage-chat-view.ts:461-479StageChatView._handleSlashCommand only handles /compact, /quit, /exit. Anything else (/model, /help, …) returns false and submit() then sends the literal text (e.g. /model claude-opus) as a prompt to the agent. That's unlikely to be what a user wanted in a stage chat. Consider:

    • logging an unknown command status, or
    • intercepting all ^/[a-z] inputs in stage-chat and warning the user that only a subset is available.
  3. releaseLiveHandleWhenIdle 250 ms defensive poll (packages/workflows/src/runs/foreground/executor.ts:1389-1390)
    The polling fallback is documented and pollTimer.unref?.() is called, so it can't keep the process alive. But: if hasQueuedLiveWork() is permanently stuck (e.g. a buggy SDK never decrements __pendingMessageCount), the timer fires every 250 ms for the life of the process. Consider a deadline (e.g. clear after N seconds, log a warning, and force-release) — currently a stuck stage silently leaks a handle. Minor.

  4. mountStageCustomUi snapshots component methods at wrap time (packages/workflows/src/shared/stage-ui-broker.ts:172-189)
    The wrapped component captures rawComponent.handleInput, dispose, etc. once. If a factory mutates these after construction (uncommon, but the Focusable.focused setter relies on this exact pattern), they won't be reflected. Worth a comment, or a getter-based proxy.

  5. StageUiBroker.requestCustomUi — abort listener race window
    packages/workflows/src/shared/stage-ui-broker.ts:113-128 — between signal.addEventListener(\"abort\", onAbort) (line 116) and this.pending.set(hostKey, request) (line 118), if onAbort fires, this.reject(request, ...) short-circuits because pending.get(hostKey)?.id !== request.id. In single-threaded JS this can't actually happen synchronously (the abort handler dispatches via dispatchEvent synchronously, but no other code runs between two sync statements). So this is not a real bug — just flagging that the post-set if (signal?.aborted) onAbort() recheck is the only thing covering the case of abort() being triggered during showHostOrReject. Worth a sentence in the comment.

  6. openExternalEditorForText Windows shell: true (packages/coding-agent/src/modes/interactive/chat-input-actions.ts:194)
    shell: process.platform === \"win32\" is needed for .cmd / .bat resolution, but editorCommand ultimately comes from $VISUAL / $EDITOR. The existing comment acknowledges this. Worth double-checking that nothing user-controlled (e.g. via slash-command argument or extension API) ever flows into editorCommand — currently it's safe, but if a future change wires this to user input, the Windows path will execute shell metacharacters.

  7. isCallbackServerRunning / getPendingAuthCount exported (packages/mcp/mcp-callback-server.ts:276-285)
    Not new in this PR, but the test file imports handleRequest and renderCallbackErrorHtml. Just confirming those are intentionally part of the public surface now (they are exported).

✅ Test coverage

  • stage-ui-broker.test.ts — 8 tests cover: collision-resistant IDs, duplicate rejection, abort, host unregister, sync-abort during show, host throw, register-after-pending, host replace. Very thorough.
  • mcp-security.test.ts — XSS regression on OAuth callback and escapeHtmlAttribute semantics. ✓
  • mcp-init-statusbar.test.ts — connected-vs-total counting and the theme-fallback path. ✓
  • chat-session-host.test.ts (543 lines) — covers idle prompt, steer, follow-up, bash, escape paths. Good.
  • chat-input-actions.test.ts — only runs on non-Windows (correct), validates atomic-branded tmp file format.
  • One small gap: I didn't see a test for chat-session-host's compaction-queued message replay path (flushCompactionQueue, line 908-930). If requiredCommand(\"prompt\") throws mid-queue, the rest is re-prepended; that try/catch is worth a unit test.

📦 Conventions

  • Bun-only, no node/npm/yarn. ✓
  • @bastani/workflows ships raw .ts — no dist/. ✓
  • Uses beforeAll/afterAll (Bun-compatible) in new tests. ✓
  • .js import specifiers in TS sources. ✓
  • Internal-only emoji removal in DESIGN.md aligns with the Unicode-icon policy.

Verdict

Solid PR. Items 1–3 above are worth addressing before merge; 4–7 are nits / documentation. The MCP hardening + StageUiBroker abstraction are net positives for both safety and code reuse.

@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

Code Review

Thorough refactor — the ChatSessionHost extraction and the broker-based stage UI plumbing are the right shapes. Below are findings worth addressing, ordered by impact. File:line references throughout.

Major

1. escapeHtmlAttribute is missing the apostrophe — latent XSSpackages/mcp/host-html-template.ts:421-427
The helper escapes &, ", <, > but not '. Sibling escapeHtml (line 414-418) escapes ' correctly — the asymmetry is the smoking gun. Today the only call site (ui-session.ts:328) uses a double-quoted attribute, so it's not exploitable, but the helper is exported, named generically, and mcp-security.test.ts:54 validates it only against double-quote/angle-bracket payloads, masking the gap. Any future single-quoted attribute caller is vulnerable. Add .replace(/'/g, \"&#39;\") and extend the test with a single-quoted payload.

2. releaseLiveHandleWhenIdle's 250 ms poll is a band-aidpackages/workflows/src/runs/foreground/executor.ts:1355-1390
The inline comment is candid: "SDK prompt/tool cleanup can also drain after the stage has stopped emitting workflow-visible events." Root cause is __pendingMessageCount() decrementing without an accompanying observable event. Correctness now silently depends on a heartbeat. Two issues: (a) executor.test.ts:2137 only exercises the event-driven path, so the poll fallback is untested; (b) the right fix is to emit a queue_update (or expose a drain promise) at every decrement site in coding-agent so the subscription completes deterministically. Recommend tracking the upstream fix and adding a regression test that asserts the silent-drain path resolves without firing any events.

3. error_description flows verbatim into the rejection Error.messagepackages/mcp/mcp-callback-server.ts:107, 116
The HTML fix is correct (no interpolation). But errorMsg = errorDescription || error becomes new Error(errorMsg), and mcp-security.test.ts:48-49 explicitly asserts that <SCRIPT>alert(2)</SCRIPT> survives. Audit every downstream consumer of waitForCallback's rejection — if it ever reaches an HTML surface (e.g. ui.notify), the fix is incomplete. At minimum, length-cap and strip control chars before constructing the Error.

4. Re-registering the same host while a request is pending re-invokes showCustomUipackages/workflows/src/shared/stage-ui-broker.ts:58-69
The previousHost && previousHost !== host && request guard skips the hide-notify when previousHost === host, but then unconditionally calls showHostOrReject(host, activeRequest). Non-idempotent hosts (allocating child components on each showCustomUi) leak. Add a previousHost === host early-return for the show path.

5. mountStageCustomUi race — async factory can leak componentspackages/workflows/src/shared/stage-ui-broker.ts:163, stage-chat-view.ts:344-374
If the request is rejected/aborted while await request.factory(...) is in flight, the resolved component is built but never displayed. done() calls broker.resolve, which no-ops via the id guard, and _showCustomUi never re-checks whether the request is still pending after the await. Add a post-await guard: if the broker no longer holds the request, dispose() immediately.

6. appendMessages bypasses optimistic-signature deduppackages/coding-agent/src/modes/interactive/components/chat-session-host.ts:188
Goes straight to liveChat.appendMessages without consulting optimisticUserSignatureCounts. After a re-attach, a snapshot replay followed by a message_start for the same user text will double-render the prompt.

7. optimisticUserSignatureCounts grows unbounded under failure pathschat-session-host.ts:157
Incremented on every optimistic send; decremented only via submit's catch and message_start. Any code path that swallows message_start (handle detached mid-send, agent failure that bypasses the event) leaks an entry permanently and silently suppresses a future identical prompt's echo. Add a TTL/cap or a deterministic decrement on session-end.

8. Timers can fire (and re-arm) after dispose()chat-session-host.ts:1018, 1270
requestEventRender's 80 ms timer and syncAnimationTick's interval are cleared on dispose, but neither callback (nor applyAgentEvent) sets/checks a disposed flag, so a late event from a still-subscribed handle resurrects them. Add a single disposed boolean checked at the top of applyAgentEvent, submit, and the timer callbacks.

Minor

  • Double-drop in idle-release pathexecutor.ts:1355 calls dropStageControlHandle() eagerly and releaseLiveHandle calls it again at :1348. Currently safe via the stageControlDropped guard, but the idempotency contract should be commented; a future guard removal would silently break the registry.
  • onAbort re-entry into reject()stage-ui-broker.ts:124 after the host's showCustomUi synchronously aborts the signal, the once-listener fires and the post-call re-check fires onAbort again. reject() no-ops via id guard, but gate with a local aborted flag for clarity.
  • recordStageAwaitingInput flip race with ask_user_question watcher — broker resolution and the executor watcher at :1320-1335 both flip the same boolean. Concurrent broker+tool flows can land the stage in running while a tool call is still pending. Coordinate via a counter.
  • StageChatView._showCustomUi doesn't reject the previous request when a stale mountedCustomUi exists (stage-chat-view.ts:345-346). Defense-in-depth; broker dedup prevents this today.
  • External-editor spawnSync is uninterruptible by host dispose()chat-input-actions.ts:190. If the host is torn down mid-edit, the finally block runs host.start()/requestRender(true) against a disposed host. Wire an abort hook.
  • colorizeStatusText cast bypasses typed uimcp/init.ts:285 uses (ui as { theme?: { fg?: ... } }). Fallback is correct, but the right fix is a typed helper interface on ExtensionAPI.
  • Comments restating code in mcp-callback-server.ts — lines 105, 108, 121, 128, 138 (// Handle OAuth errors, // Require authorization code). Per CLAUDE.md, no comments unless the WHY is non-obvious.

Nits

  • ChatSessionHost is 1280 lines mixing keybinding, editor wiring, bash lifecycle, transcript, optimistic dedup, animation timers, render, two-SDK event normalization, ANSI helpers. Split candidates: ChatTranscriptModel, ChatAnimationClock, ChatEditorBindings, ChatBashRunner, leaving the host as a coordinator.
  • ChatSessionHostStyle.editorRuleColor(disabled, agentSession, state) leaks isBashMode-awareness back to callers — host knows both pieces; pass only a colour callback.
  • isBashRunningOverride / isStreamingOverride — "override" is misleading; it's an OR with internal state.
  • renderCallbackErrorHtml() takes no arguments — CALLBACK_ERROR_HTML (constant) communicates the no-interpolation guarantee structurally rather than by convention.

Test gaps worth filling

  • escapeHtmlAttribute with single-quoted-attribute payload (will currently pass — exposes finding add agent instructions #1).
  • Silent SDK-drain path for releaseLiveHandleWhenIdle — no event emitted, only the poll fires.
  • Host re-registration while a broker request is pending (finding Flora131/feat/add skills #4).
  • mountStageCustomUi aborted during async factory (finding Flora131/feat/refine readme #5).
  • Stale resolveStagePendingPrompt after a previous prompt resolved.
  • dispose() preventing a late-arriving event from resurrecting timers.
  • Snapshot-replay + optimistic-dedup interaction (finding update readme instructions #6).
  • pasteClipboardImageToEditor integration through the host (only direct helper tested).

Conventions

escapeHtml (line 417) escapes ' as &#39;; copy the same pattern into escapeHtmlAttribute for symmetry. Otherwise the PR follows Bun/.ts conventions correctly — no build steps added to packages/workflows, no new any usage worth flagging, and the unknown casts in legacy event normalizers are appropriately scoped.

Overall: ship-worthy after #1, #3, #4, #5 are addressed; #2 and #6-#8 should land as a follow-up with regression tests.

Render workflow inputs as a single ask-style active field with tab status, footer rules, and list-style select/boolean choices. Normalize TUI key handling through typed key ids and shared action constants, and cover the updated navigation/rendering behavior in unit tests.

Assistant-model: GPT-5.5
@lavaman131
lavaman131 merged commit 22de7b9 into main May 24, 2026
10 checks passed
@lavaman131
lavaman131 deleted the refactor/workflow-chat-host-parity branch May 24, 2026 05:15
@claude

claude Bot commented May 24, 2026

Copy link
Copy Markdown

Review: feat(workflows): extract ChatSessionHost and compose into stage chat

Big, well-structured refactor. The duplication that previously lived in StageChatView is now consolidated, the new abstractions are well-tested, and the MCP hardening is solid. A few findings below.

Bugs / correctness

1. Race in StageChatView._showCustomUi — settled requests can leave a mounted UI orphaned. packages/workflows/src/tui/stage-chat-view.ts:344-374

If a factory invokes its done callback synchronously (or before mountStageCustomUi resolves), the sequence is:

  1. mountStageCustomUi calls request.factory(..., done).
  2. Factory immediately calls done(result)broker.resolve(...)hideHost fires hideCustomUi(request) on us.
  3. _hideMountedCustomUi sees this.mountedCustomUi === null and returns early.
  4. Factory returns the component; await resolves; this.mountedCustomUi = result is assigned.
  5. Result: the component is mounted (visible) but the request has already settled and no further hide path will fire — and onDone's guard this.mountedCustomUi?.request.id !== request.id was tripped at step 4 too, since mountedCustomUi was null.

Suggested fix: after await mountStageCustomUi(...), re-check whether the request is still the current one (e.g., track a currentRequestId set before await, and dispose immediately if it differs); or pass a cancellation token in and tear down inside mountStageCustomUi if the host hide fired before return.

2. chat-session-host.ts loadSessionFile early-returns when transcript.length > 0, but it's first called from StageChatView after _snapshotMessagesFromHandle may have already populated the transcript from handle.messages. That's intentional, but worth verifying that a stage which boots with no live handle but a sessionFile still gets its archived transcript — the empty-transcript precondition is essentially the only entry path for archive playback.

3. Hardcoded calloutRows = 6 in _renderReadOnlyArchiveBody. stage-chat-view.ts (~line 590). The actual callout (blank + banner + 2 text rows) is closer to 5; the trailing pad-or-truncate masks the off-by-one, but if the banner ever grows the transcript budget will silently steal the affordance. Computing the callout height from the rendered output would remove the magic number.

Performance

4. releaseLiveHandleWhenIdle 250 ms poll fallback (executor.ts:1387-1391) is necessary for the documented silent-SDK-drain path and the timer is .unref()-ed, so it can't keep the process alive. Reasonable, but worth noting: the subscribe-callback path is the primary signal and the poller is purely defensive. If you have telemetry, it would be useful to know how often the poll actually wins the race vs. the subscription — a metric/log line at debug level could surface dead code.

Security — well done

The MCP hardening is a meaningful upgrade and the mcp-security.test.ts coverage is on point:

  • renderCallbackErrorHtml no longer interpolates raw OAuth error / error_description into HTML — the test confirms <script> payloads are not reflected.
  • Glimpse iframe src now goes through escapeHtmlAttribute (packages/mcp/ui-session.ts:328), closing the obvious attribute-injection vector if a tool URL ever contained a ".
  • The TOCTOU race in the completed UI session handle is closed with the stable local binding.
  • The status-bar connectedCount fix is correct — previously it counted needs-auth / failed entries.

One minor note: colorizeStatusText swallows arbitrary errors from theme.fg. Defensible (you don't want a status-bar crash to bubble up), but a debug log on the catch would help diagnose broken themes.

Code quality / style

  • Both new files (chat-session-host.ts, stage-ui-broker.ts, chat-input-actions.ts) adhere to repo conventions: strict types, .js import specifiers, no any/unknown slipping in, no build-step assumptions for workflows.
  • chat-input-actions.openExternalEditorForText correctly uses flag: "wx" + mode: 0o600 and crypto.randomUUID() for the temp filename — symlink/race-safe. pasteClipboardImageToEditor does the same. Good.
  • parseEditorCommand is a small, sane shell-token parser; the shell: process.platform === "win32" is appropriately commented.
  • stageBuiltinPackagePaths (wiring.ts) filters workflows out of recursively-loaded builtin packages with a clear explanation — nice defensive fix.

Tests

  • stage-ui-broker.test.ts covers the important edge cases (abort during show, duplicate pending, host replacement with throwing hideCustomUi, listener-replay quirk).
  • chat-session-host.test.ts is thorough — covers slash routing, bash parsing, optimistic-user dedup, compaction queueing/flushing/failure-preservation, and the queue-update tolerance for malformed payloads.
  • mcp-security.test.ts exercises the actual HTTP path end-to-end against a live server — the right level for this kind of injection-regression test.
  • store-pending-prompt.test.ts confirms the new per-stage prompt rejection paths on stage-end and run-end.

The async race in finding #1 is the only item I'd consider blocking. Everything else is a suggestion or polish.

🤖 Generated with Claude Code

lavaman131 pushed a commit that referenced this pull request Jun 29, 2026
…1026)

* feat(workflows): embed stage chat custom UI hosts

Extract a reusable chat session host from interactive mode and wire workflow stage chat through stage-scoped UI brokers.\n\nUpdate workflow TUI rendering, input schema previews, docs, mockups, and tests for the embedded chat surface.

Assistant-model: GPT-5.5

* fix(mcp): harden status and callback handling

Count only connected MCP servers in the status bar and fall back to plain text when theme colorization is unavailable.\n\nEscape callback authorization errors and retain a stable completed UI session handle when storing session messages.

Assistant-model: GPT-5.5

* fix(workflows): preserve workflow name in empty inputs output

Assistant-model: GPT-5.5

* fix(workflows): release idle completed stage sessions

Assistant-model: GPT-5.5

* fix(workflows): freeze background widget timers

Assistant-model: GPT-5.5

* fix(workflows): pause stage chat on escape

Assistant-model: GPT-5.5

* fix: address PR review security and lifecycle bugs

* feat(workflows): align inline forms with question UI

Assistant-model: GPT-5.5

* feat(workflows): align inputs picker with ask UI

Render workflow inputs as a single ask-style active field with tab status, footer rules, and list-style select/boolean choices. Normalize TUI key handling through typed key ids and shared action constants, and cover the updated navigation/rendering behavior in unit tests.

Assistant-model: GPT-5.5
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants