chat-messages: render and batch-submit multiple HITL approval cards

This commit is contained in:
CREDO23 2026-05-13 21:00:01 +02:00
parent 0fd87ccb7f
commit 1bb9f435e5
5 changed files with 160 additions and 76 deletions

View file

@ -49,7 +49,11 @@ import {
type TokenUsageData,
TokenUsageProvider,
} from "@/components/assistant-ui/token-usage-context";
import { type HitlDecision, PendingInterruptProvider } from "@/features/chat-messages/hitl";
import {
type HitlDecision,
PendingInterruptProvider,
type PendingInterruptState,
} from "@/features/chat-messages/hitl";
import { TimelineDataUI } from "@/features/chat-messages/timeline";
import {
applyActionLogSse,
@ -208,7 +212,10 @@ const MentionedDocumentInfoSchema = z.object({
id: z.number(),
title: z.string(),
document_type: z.string(),
kind: z.union([z.literal("doc"), z.literal("folder")]).optional().default("doc"),
kind: z
.union([z.literal("doc"), z.literal("folder")])
.optional()
.default("doc"),
});
const MentionedDocumentsPartSchema = z.object({
@ -269,12 +276,16 @@ export default function NewChatPage() {
const [tokenUsageStore] = useState(() => createTokenUsageStore());
const abortControllerRef = useRef<AbortController | null>(null);
const recentCancelRequestedAtRef = useRef(0);
const [pendingInterrupt, setPendingInterrupt] = useState<{
threadId: number;
assistantMsgId: string;
interruptData: Record<string, unknown>;
bundleToolCallIds: string[];
} | null>(null);
// One entry per paused subagent, in receipt order (which matches the
// backend's ``state.interrupts`` traversal — and therefore the order
// ``slice_decisions_by_tool_call`` consumes on resume). Cleared on submit
// or on a fresh user turn.
const [pendingInterrupts, setPendingInterrupts] = useState<PendingInterruptState[]>([]);
// Per-card staged decisions held until every pending card has submitted,
// at which point we batch them into one ``hitl-decision`` event in the
// same order as ``pendingInterrupts``. Using a ref because partial
// progress should not re-render the page.
const stagedDecisionsByInterruptIdRef = useRef<Map<string, HitlDecision[]>>(new Map());
const toolsWithUI = TOOLS_WITH_UI_ALL;
const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom);
@ -1029,9 +1040,7 @@ export default function NewChatPage() {
mentioned_surfsense_doc_ids: hasSurfsenseDocIds
? mentionedDocumentIds.surfsense_doc_ids
: undefined,
mentioned_folder_ids: hasFolderIds
? mentionedDocumentIds.folder_ids
: undefined,
mentioned_folder_ids: hasFolderIds ? mentionedDocumentIds.folder_ids : undefined,
// Full mention metadata (docs + folders, with
// ``kind`` discriminator) so the BE can embed a
// ``mentioned-documents`` ContentPart on the
@ -1193,12 +1202,24 @@ export default function NewChatPage() {
)
);
if (currentThreadId) {
setPendingInterrupt({
threadId: currentThreadId,
assistantMsgId,
interruptData,
bundleToolCallIds,
});
// ``tool_call_id`` is stamped on the backend by
// ``checkpointed_subagent_middleware``. Without it we
// can't address the paused subagent on resume — skip
// rather than fabricate a synthetic key.
const interruptId = String(interruptData.tool_call_id ?? "");
if (interruptId) {
const incoming: PendingInterruptState = {
interruptId,
threadId: currentThreadId,
assistantMsgId,
interruptData,
bundleToolCallIds,
};
setPendingInterrupts((prev) => {
const without = prev.filter((p) => p.interruptId !== interruptId);
return [...without, incoming];
});
}
}
break;
}
@ -1274,7 +1295,7 @@ export default function NewChatPage() {
// by ``persist_assistant_shell``. Rename the optimistic
// id, migrate ``tokenUsageStore`` so any pending
// ``data-token-usage`` payload binds to the new id,
// remap any in-flight ``pendingInterrupt`` reference,
// remap any in-flight ``pendingInterrupts`` entries,
// and reassign the closure variable so the in-stream
// flush callback (line ~1074) keeps writing to the
// renamed message.
@ -1290,10 +1311,12 @@ export default function NewChatPage() {
: m
)
);
setPendingInterrupt((prev) =>
prev && prev.assistantMsgId === oldAssistantMsgId
? { ...prev, assistantMsgId: newAssistantMsgId }
: prev
setPendingInterrupts((prev) =>
prev.map((p) =>
p.assistantMsgId === oldAssistantMsgId
? { ...p, assistantMsgId: newAssistantMsgId }
: p
)
);
assistantMsgId = newAssistantMsgId;
break;
@ -1380,14 +1403,23 @@ export default function NewChatPage() {
edited_action?: { name: string; args: Record<string, unknown> };
}>
) => {
if (!pendingInterrupt) return;
const { threadId: resumeThreadId } = pendingInterrupt;
if (pendingInterrupts.length === 0) return;
// All cards in this turn share the same threadId/assistantMsgId
// (they're siblings of one parent agent step), so reading from
// the first entry is safe.
const resumeThreadId = pendingInterrupts[0].threadId;
// Destructured separately as ``let`` so the SSE
// ``data-assistant-message-id`` handler (resume always
// allocates a fresh server-side row) can rename it to
// the canonical ``msg-{db_id}`` mid-stream.
let assistantMsgId = pendingInterrupt.assistantMsgId;
setPendingInterrupt(null);
let assistantMsgId = pendingInterrupts[0].assistantMsgId;
// Concatenate every card's tool-call ids in pendingInterrupts order;
// this matches the ``decisions`` ordering produced by
// ``handleApprovalSubmit`` and the backend slicer's traversal of
// ``state.interrupts``.
const allBundleToolCallIds = pendingInterrupts.flatMap((p) => p.bundleToolCallIds);
setPendingInterrupts([]);
stagedDecisionsByInterruptIdRef.current.clear();
setIsRunning(true);
const token = getBearerToken();
@ -1464,7 +1496,7 @@ export default function NewChatPage() {
// collapse onto ``decisions[0]``. Cards outside the bundle are
// untouched. Mirrors the host ``hitl-decision`` handler.
const decisionByTcId = new Map<string, (typeof decisions)[number]>();
const tcIds = pendingInterrupt.bundleToolCallIds;
const tcIds = allBundleToolCallIds;
if (decisions.length === tcIds.length) {
for (let i = 0; i < tcIds.length; i++) decisionByTcId.set(tcIds[i], decisions[i]);
}
@ -1596,12 +1628,22 @@ export default function NewChatPage() {
: m
)
);
setPendingInterrupt({
threadId: resumeThreadId,
assistantMsgId,
interruptData,
bundleToolCallIds,
});
{
const interruptId = String(interruptData.tool_call_id ?? "");
if (interruptId) {
const incoming: PendingInterruptState = {
interruptId,
threadId: resumeThreadId,
assistantMsgId,
interruptData,
bundleToolCallIds,
};
setPendingInterrupts((prev) => {
const without = prev.filter((p) => p.interruptId !== interruptId);
return [...without, incoming];
});
}
}
break;
}
@ -1679,7 +1721,7 @@ export default function NewChatPage() {
}
},
[
pendingInterrupt,
pendingInterrupts,
messages,
searchSpaceId,
localFilesystemEnabled,
@ -1700,17 +1742,19 @@ export default function NewChatPage() {
edited_action?: { name: string; args: Record<string, unknown> };
}>;
};
if (!detail?.decisions || !pendingInterrupt) return;
if (!detail?.decisions || pendingInterrupts.length === 0) return;
const incoming = detail.decisions;
if (incoming.length === 0) return;
const tcIds = pendingInterrupt.bundleToolCallIds;
// Concatenated tool-call ids across every pending card, in the
// order ``handleApprovalSubmit`` produced ``incoming``.
const tcIds = pendingInterrupts.flatMap((p) => p.bundleToolCallIds);
const N = tcIds.length;
// Bundles must submit exactly one decision per action_request.
// Refuse rather than silently broadcast a single decision across
// the bundle (would mis-apply rejects/edits and diverge from
// what handleResume sends to /resume).
if (N > 1 && incoming.length !== N) {
// Refuse rather than silently broadcast or drop. The orchestrator
// only fires ``hitl-decision`` once every pending card has
// submitted, so a count mismatch indicates a contract drift
// (and would later make the backend slicer raise).
if (incoming.length !== N) {
toast.error(
`Cannot resume: ${incoming.length} decision(s) submitted for ${N} pending actions.`
);
@ -1721,9 +1765,12 @@ export default function NewChatPage() {
for (let i = 0; i < tcIds.length; i++) byTcId.set(tcIds[i], incoming[i]);
const submittedDecisions = tcIds.map((id) => byTcId.get(id)!);
// All pending cards belong to the same assistant message, so a
// single content-update pass suffices.
const targetAssistantMsgId = pendingInterrupts[0].assistantMsgId;
setMessages((prev) =>
prev.map((m) => {
if (m.id !== pendingInterrupt.assistantMsgId) return m;
if (m.id !== targetAssistantMsgId) return m;
const parts = m.content as unknown as Array<Record<string, unknown>>;
const newContent = parts.map((part) => {
const tcId = part.toolCallId as string | undefined;
@ -1760,7 +1807,7 @@ export default function NewChatPage() {
};
window.addEventListener("hitl-decision", handler);
return () => window.removeEventListener("hitl-decision", handler);
}, [handleResume, pendingInterrupt]);
}, [handleResume, pendingInterrupts]);
// Convert message (pass through since already in correct format)
const convertMessage = useCallback(
@ -1900,12 +1947,10 @@ export default function NewChatPage() {
filesystem_mode: selection.filesystem_mode,
client_platform: selection.client_platform,
local_filesystem_mounts: selection.local_filesystem_mounts,
mentioned_document_ids:
regenerateDocIds.length > 0 ? regenerateDocIds : undefined,
mentioned_document_ids: regenerateDocIds.length > 0 ? regenerateDocIds : undefined,
mentioned_surfsense_doc_ids:
regenerateSurfsenseDocIds.length > 0 ? regenerateSurfsenseDocIds : undefined,
mentioned_folder_ids:
regenerateFolderIds.length > 0 ? regenerateFolderIds : undefined,
mentioned_folder_ids: regenerateFolderIds.length > 0 ? regenerateFolderIds : undefined,
// Full mention metadata for the regenerate-specific
// source list. Only meaningful for edit (the BE only
// re-persists a user row when ``user_query`` is set);
@ -2284,11 +2329,32 @@ export default function NewChatPage() {
[handleRegenerate, messages, agentActionItems]
);
const handleApprovalSubmit = useCallback((orderedDecisions: HitlDecision[]) => {
window.dispatchEvent(
new CustomEvent("hitl-decision", { detail: { decisions: orderedDecisions } })
);
}, []);
const handleApprovalSubmit = useCallback(
(interruptId: string, decisions: HitlDecision[]) => {
// Stage this card's decisions; only fire the resume once every
// pending card in the current turn has submitted, so the
// backend slicer sees a single concatenated decisions list
// whose total matches the parent state's pending action count.
stagedDecisionsByInterruptIdRef.current.set(interruptId, decisions);
if (stagedDecisionsByInterruptIdRef.current.size < pendingInterrupts.length) {
return;
}
const ordered: HitlDecision[] = [];
for (const pi of pendingInterrupts) {
const staged = stagedDecisionsByInterruptIdRef.current.get(pi.interruptId);
if (!staged) {
// Defensive: a missing entry means the staging map and
// the pending list disagreed for one cycle. Bail rather
// than dispatch a count-mismatched batch.
return;
}
ordered.push(...staged);
}
stagedDecisionsByInterruptIdRef.current.clear();
window.dispatchEvent(new CustomEvent("hitl-decision", { detail: { decisions: ordered } }));
},
[pendingInterrupts]
);
const handleEditDialogChoice = useCallback(
async (choice: EditMessageDialogChoice) => {
@ -2361,7 +2427,7 @@ export default function NewChatPage() {
<TimelineDataUI />
<StepSeparatorDataUI />
<PendingInterruptProvider
pendingInterrupt={pendingInterrupt}
pendingInterrupts={pendingInterrupts}
onSubmit={handleApprovalSubmit}
>
<div key={searchSpaceId} className="flex h-full overflow-hidden">