mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-21 18:55:16 +02:00
Merge remote-tracking branch 'upstream/dev' into feat/split-auto-free-premium
This commit is contained in:
Anish Sarkar
commit
872065f90d
15 changed files with 1857 additions and 545 deletions
|
|
@ -14,13 +14,6 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
|||
import { toast } from "sonner";
|
||||
import { z } from "zod";
|
||||
import { disabledToolsAtom } from "@/atoms/agent-tools/agent-tools.atoms";
|
||||
import {
|
||||
agentActionsByChatTurnIdAtom,
|
||||
markAgentActionRevertedAtom,
|
||||
resetAgentActionMapAtom,
|
||||
updateAgentActionReversibleAtom,
|
||||
upsertAgentActionAtom,
|
||||
} from "@/atoms/chat/agent-actions.atom";
|
||||
import {
|
||||
clearTargetCommentIdAtom,
|
||||
currentThreadAtom,
|
||||
|
|
@ -56,6 +49,12 @@ import {
|
|||
type TokenUsageData,
|
||||
TokenUsageProvider,
|
||||
} from "@/components/assistant-ui/token-usage-context";
|
||||
import {
|
||||
applyActionLogSse,
|
||||
applyActionLogUpdatedSse,
|
||||
markActionRevertedInCache,
|
||||
useAgentActionsQuery,
|
||||
} from "@/hooks/use-agent-actions-query";
|
||||
import { useChatSessionStateSync } from "@/hooks/use-chat-session-state";
|
||||
import { useMessagesSync } from "@/hooks/use-messages-sync";
|
||||
import { getAgentFilesystemSelection } from "@/lib/agent-filesystem";
|
||||
|
|
@ -76,12 +75,12 @@ import {
|
|||
addToolCall,
|
||||
appendReasoning,
|
||||
appendText,
|
||||
appendToolInputDelta,
|
||||
buildContentForPersistence,
|
||||
buildContentForUI,
|
||||
type ContentPartsState,
|
||||
endReasoning,
|
||||
FrameBatchedUpdater,
|
||||
findToolCallIdByLcId,
|
||||
readSSEStream,
|
||||
type SSEEvent,
|
||||
type ThinkingStepData,
|
||||
|
|
@ -511,14 +510,6 @@ export default function NewChatPage() {
|
|||
const setAgentCreatedDocuments = useSetAtom(agentCreatedDocumentsAtom);
|
||||
const pendingUserImageUrls = useAtomValue(pendingUserImageDataUrlsAtom);
|
||||
const setPendingUserImageUrls = useSetAtom(pendingUserImageDataUrlsAtom);
|
||||
// Agent action log SSE side-channel.
|
||||
const upsertAgentAction = useSetAtom(upsertAgentActionAtom);
|
||||
const updateAgentActionReversible = useSetAtom(updateAgentActionReversibleAtom);
|
||||
const markAgentActionReverted = useSetAtom(markAgentActionRevertedAtom);
|
||||
const resetAgentActionMap = useSetAtom(resetAgentActionMapAtom);
|
||||
// Chat-turn-keyed action map for the edit-from-position pre-flight
|
||||
// that decides whether to show the confirmation dialog.
|
||||
const agentActionsByChatTurnId = useAtomValue(agentActionsByChatTurnIdAtom);
|
||||
// Edit dialog state. Holds the message id being edited and
|
||||
// the (already extracted) regenerate args so we can resume the edit
|
||||
// after the user picks "revert all" / "continue" / "cancel".
|
||||
|
|
@ -547,6 +538,11 @@ export default function NewChatPage() {
|
|||
content: unknown;
|
||||
author_id: string | null;
|
||||
created_at: string;
|
||||
// Forwarded so ``convertToThreadMessage`` can rebuild the
|
||||
// ``metadata.custom.chatTurnId`` on the
|
||||
// ``ThreadMessageLike``. Required by the inline Revert
|
||||
// button's per-turn fallback.
|
||||
turn_id?: string | null;
|
||||
}[]
|
||||
) => {
|
||||
if (isRunning) {
|
||||
|
|
@ -579,6 +575,11 @@ export default function NewChatPage() {
|
|||
created_at: msg.created_at,
|
||||
author_display_name: member?.user_display_name ?? existingAuthor?.displayName ?? null,
|
||||
author_avatar_url: member?.user_avatar_url ?? existingAuthor?.avatarUrl ?? null,
|
||||
// Forward the per-turn correlation id so the
|
||||
// inline Revert button's ``(chat_turn_id,
|
||||
// tool_name, position)`` fallback survives the
|
||||
// post-stream Zero re-sync.
|
||||
turn_id: msg.turn_id ?? null,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -595,6 +596,13 @@ export default function NewChatPage() {
|
|||
return Number.isNaN(parsed) ? 0 : parsed;
|
||||
}, [params.search_space_id]);
|
||||
|
||||
// Unified store for agent-action rows (the same react-query cache
|
||||
// the agent-actions sheet, the inline Revert button, and the
|
||||
// per-turn Revert button all read). Hydrates from
|
||||
// ``GET /threads/{id}/actions`` and is updated incrementally by the
|
||||
// SSE handlers + revert-batch results below — no atom side-channel.
|
||||
const { items: agentActionItems } = useAgentActionsQuery(threadId);
|
||||
|
||||
// Extract chat_id from URL params
|
||||
const urlChatId = useMemo(() => {
|
||||
const id = params.chat_id;
|
||||
|
|
@ -738,7 +746,8 @@ export default function NewChatPage() {
|
|||
clearPlanOwnerRegistry();
|
||||
closeReportPanel();
|
||||
closeEditorPanel();
|
||||
resetAgentActionMap();
|
||||
// Note: agent-action data is keyed by threadId in react-query so
|
||||
// switching threads naturally swaps caches; no explicit reset.
|
||||
|
||||
try {
|
||||
if (urlChatId > 0) {
|
||||
|
|
@ -807,7 +816,6 @@ export default function NewChatPage() {
|
|||
removeChatTab,
|
||||
searchSpaceId,
|
||||
tokenUsageStore,
|
||||
resetAgentActionMap,
|
||||
]);
|
||||
|
||||
// Initialize on mount, and re-init when switching search spaces (even if urlChatId is the same)
|
||||
|
|
@ -1145,6 +1153,15 @@ export default function NewChatPage() {
|
|||
);
|
||||
};
|
||||
const scheduleFlush = () => batcher.schedule(flushMessages);
|
||||
// Force-flush helper: ``batcher.flush()`` is a no-op when
|
||||
// ``dirty=false`` (e.g. a tool starts before any text
|
||||
// streamed). ``scheduleFlush(); batcher.flush()`` sets
|
||||
// the dirty bit FIRST so terminal events render
|
||||
// promptly without the 50ms throttle delay.
|
||||
const forceFlush = () => {
|
||||
scheduleFlush();
|
||||
batcher.flush();
|
||||
};
|
||||
|
||||
for await (const parsed of readSSEStream(response)) {
|
||||
switch (parsed.type) {
|
||||
|
|
@ -1181,13 +1198,23 @@ export default function NewChatPage() {
|
|||
false,
|
||||
parsed.langchainToolCallId
|
||||
);
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
|
||||
case "tool-input-delta":
|
||||
// High-frequency event: deltas can fire dozens
|
||||
// of times per call, so use throttled
|
||||
// scheduleFlush (NOT forceFlush) to coalesce.
|
||||
appendToolInputDelta(contentPartsState, parsed.toolCallId, parsed.inputTextDelta);
|
||||
scheduleFlush();
|
||||
break;
|
||||
|
||||
case "tool-input-available": {
|
||||
const finalArgsText = JSON.stringify(parsed.input ?? {}, null, 2);
|
||||
if (toolCallIndices.has(parsed.toolCallId)) {
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
args: parsed.input || {},
|
||||
argsText: finalArgsText,
|
||||
langchainToolCallId: parsed.langchainToolCallId,
|
||||
});
|
||||
} else {
|
||||
|
|
@ -1200,8 +1227,14 @@ export default function NewChatPage() {
|
|||
false,
|
||||
parsed.langchainToolCallId
|
||||
);
|
||||
// addToolCall doesn't accept argsText today;
|
||||
// backfill via updateToolCall so the new card
|
||||
// renders pretty-printed JSON.
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
argsText: finalArgsText,
|
||||
});
|
||||
}
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
}
|
||||
|
||||
|
|
@ -1220,7 +1253,7 @@ export default function NewChatPage() {
|
|||
}
|
||||
}
|
||||
}
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
}
|
||||
|
||||
|
|
@ -1316,34 +1349,17 @@ export default function NewChatPage() {
|
|||
}
|
||||
|
||||
case "data-action-log": {
|
||||
const al = parsed.data;
|
||||
const matchedToolCallId = al.lc_tool_call_id
|
||||
? findToolCallIdByLcId(contentPartsState, al.lc_tool_call_id)
|
||||
: null;
|
||||
upsertAgentAction({
|
||||
action: {
|
||||
id: al.id,
|
||||
threadId: currentThreadId,
|
||||
lcToolCallId: al.lc_tool_call_id,
|
||||
chatTurnId: al.chat_turn_id,
|
||||
toolName: al.tool_name,
|
||||
reversible: al.reversible,
|
||||
reverseDescriptorPresent: al.reverse_descriptor_present,
|
||||
error: al.error,
|
||||
revertedByActionId: null,
|
||||
isRevertAction: false,
|
||||
createdAt: al.created_at,
|
||||
},
|
||||
toolCallId: matchedToolCallId,
|
||||
});
|
||||
applyActionLogSse(queryClient, currentThreadId, searchSpaceId, parsed.data);
|
||||
break;
|
||||
}
|
||||
|
||||
case "data-action-log-updated": {
|
||||
updateAgentActionReversible({
|
||||
id: parsed.data.id,
|
||||
reversible: parsed.data.reversible,
|
||||
});
|
||||
applyActionLogUpdatedSse(
|
||||
queryClient,
|
||||
currentThreadId,
|
||||
parsed.data.id,
|
||||
parsed.data.reversible
|
||||
);
|
||||
break;
|
||||
}
|
||||
|
||||
|
|
@ -1559,6 +1575,15 @@ export default function NewChatPage() {
|
|||
toolName: String(p.toolName),
|
||||
args: (p.args as Record<string, unknown>) ?? {},
|
||||
result: p.result as unknown,
|
||||
// Restore argsText so persisted pretty-printed
|
||||
// JSON survives reloads (assistant-ui prefers
|
||||
// supplied argsText over JSON.stringify(args)).
|
||||
// langchainToolCallId restoration also fixes a
|
||||
// pre-existing dropped-id bug on resume.
|
||||
...(typeof p.argsText === "string" ? { argsText: p.argsText } : {}),
|
||||
...(typeof p.langchainToolCallId === "string"
|
||||
? { langchainToolCallId: p.langchainToolCallId }
|
||||
: {}),
|
||||
});
|
||||
contentPartsState.currentTextPartIndex = -1;
|
||||
} else if (p.type === "data-thinking-steps") {
|
||||
|
|
@ -1580,7 +1605,12 @@ export default function NewChatPage() {
|
|||
const editedAction = decisions[0].edited_action;
|
||||
for (const part of contentParts) {
|
||||
if (part.type === "tool-call" && part.toolName === editedAction.name) {
|
||||
part.args = { ...part.args, ...editedAction.args };
|
||||
const mergedArgs = { ...part.args, ...editedAction.args };
|
||||
part.args = mergedArgs;
|
||||
// Sync argsText so the rendered card shows the
|
||||
// edited inputs — assistant-ui prefers caller-
|
||||
// supplied argsText over JSON.stringify(args).
|
||||
part.argsText = JSON.stringify(mergedArgs, null, 2);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
|
@ -1637,6 +1667,10 @@ export default function NewChatPage() {
|
|||
);
|
||||
};
|
||||
const scheduleFlush = () => batcher.schedule(flushMessages);
|
||||
const forceFlush = () => {
|
||||
scheduleFlush();
|
||||
batcher.flush();
|
||||
};
|
||||
|
||||
for await (const parsed of readSSEStream(response)) {
|
||||
switch (parsed.type) {
|
||||
|
|
@ -1673,13 +1707,20 @@ export default function NewChatPage() {
|
|||
false,
|
||||
parsed.langchainToolCallId
|
||||
);
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
|
||||
case "tool-input-available":
|
||||
case "tool-input-delta":
|
||||
appendToolInputDelta(contentPartsState, parsed.toolCallId, parsed.inputTextDelta);
|
||||
scheduleFlush();
|
||||
break;
|
||||
|
||||
case "tool-input-available": {
|
||||
const finalArgsText = JSON.stringify(parsed.input ?? {}, null, 2);
|
||||
if (toolCallIndices.has(parsed.toolCallId)) {
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
args: parsed.input || {},
|
||||
argsText: finalArgsText,
|
||||
langchainToolCallId: parsed.langchainToolCallId,
|
||||
});
|
||||
} else {
|
||||
|
|
@ -1692,9 +1733,13 @@ export default function NewChatPage() {
|
|||
false,
|
||||
parsed.langchainToolCallId
|
||||
);
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
argsText: finalArgsText,
|
||||
});
|
||||
}
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
}
|
||||
|
||||
case "tool-output-available":
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
|
|
@ -1702,7 +1747,7 @@ export default function NewChatPage() {
|
|||
langchainToolCallId: parsed.langchainToolCallId,
|
||||
});
|
||||
markInterruptsCompleted(contentParts);
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
|
||||
case "data-thinking-step": {
|
||||
|
|
@ -1762,34 +1807,17 @@ export default function NewChatPage() {
|
|||
}
|
||||
|
||||
case "data-action-log": {
|
||||
const al = parsed.data;
|
||||
const matchedToolCallId = al.lc_tool_call_id
|
||||
? findToolCallIdByLcId(contentPartsState, al.lc_tool_call_id)
|
||||
: null;
|
||||
upsertAgentAction({
|
||||
action: {
|
||||
id: al.id,
|
||||
threadId: resumeThreadId,
|
||||
lcToolCallId: al.lc_tool_call_id,
|
||||
chatTurnId: al.chat_turn_id,
|
||||
toolName: al.tool_name,
|
||||
reversible: al.reversible,
|
||||
reverseDescriptorPresent: al.reverse_descriptor_present,
|
||||
error: al.error,
|
||||
revertedByActionId: null,
|
||||
isRevertAction: false,
|
||||
createdAt: al.created_at,
|
||||
},
|
||||
toolCallId: matchedToolCallId,
|
||||
});
|
||||
applyActionLogSse(queryClient, resumeThreadId, searchSpaceId, parsed.data);
|
||||
break;
|
||||
}
|
||||
|
||||
case "data-action-log-updated": {
|
||||
updateAgentActionReversible({
|
||||
id: parsed.data.id,
|
||||
reversible: parsed.data.reversible,
|
||||
});
|
||||
applyActionLogUpdatedSse(
|
||||
queryClient,
|
||||
resumeThreadId,
|
||||
parsed.data.id,
|
||||
parsed.data.reversible
|
||||
);
|
||||
break;
|
||||
}
|
||||
|
||||
|
|
@ -1902,6 +1930,11 @@ export default function NewChatPage() {
|
|||
return {
|
||||
...part,
|
||||
args: decision.edited_action.args, // Update displayed args
|
||||
// Sync argsText so the rendered card shows
|
||||
// the edited inputs — assistant-ui prefers
|
||||
// caller-supplied argsText over
|
||||
// JSON.stringify(args).
|
||||
argsText: JSON.stringify(decision.edited_action.args, null, 2),
|
||||
result: {
|
||||
...(part.result as Record<string, unknown>),
|
||||
__decided__: decisionType,
|
||||
|
|
@ -2127,6 +2160,10 @@ export default function NewChatPage() {
|
|||
);
|
||||
};
|
||||
const scheduleFlush = () => batcher.schedule(flushMessages);
|
||||
const forceFlush = () => {
|
||||
scheduleFlush();
|
||||
batcher.flush();
|
||||
};
|
||||
|
||||
for await (const parsed of readSSEStream(response)) {
|
||||
switch (parsed.type) {
|
||||
|
|
@ -2163,13 +2200,20 @@ export default function NewChatPage() {
|
|||
false,
|
||||
parsed.langchainToolCallId
|
||||
);
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
|
||||
case "tool-input-available":
|
||||
case "tool-input-delta":
|
||||
appendToolInputDelta(contentPartsState, parsed.toolCallId, parsed.inputTextDelta);
|
||||
scheduleFlush();
|
||||
break;
|
||||
|
||||
case "tool-input-available": {
|
||||
const finalArgsText = JSON.stringify(parsed.input ?? {}, null, 2);
|
||||
if (toolCallIndices.has(parsed.toolCallId)) {
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
args: parsed.input || {},
|
||||
argsText: finalArgsText,
|
||||
langchainToolCallId: parsed.langchainToolCallId,
|
||||
});
|
||||
} else {
|
||||
|
|
@ -2182,9 +2226,13 @@ export default function NewChatPage() {
|
|||
false,
|
||||
parsed.langchainToolCallId
|
||||
);
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
argsText: finalArgsText,
|
||||
});
|
||||
}
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
}
|
||||
|
||||
case "tool-output-available":
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
|
|
@ -2201,7 +2249,7 @@ export default function NewChatPage() {
|
|||
}
|
||||
}
|
||||
}
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
|
||||
case "data-thinking-step": {
|
||||
|
|
@ -2217,34 +2265,21 @@ export default function NewChatPage() {
|
|||
}
|
||||
|
||||
case "data-action-log": {
|
||||
const al = parsed.data;
|
||||
const matchedToolCallId = al.lc_tool_call_id
|
||||
? findToolCallIdByLcId(contentPartsState, al.lc_tool_call_id)
|
||||
: null;
|
||||
upsertAgentAction({
|
||||
action: {
|
||||
id: al.id,
|
||||
threadId,
|
||||
lcToolCallId: al.lc_tool_call_id,
|
||||
chatTurnId: al.chat_turn_id,
|
||||
toolName: al.tool_name,
|
||||
reversible: al.reversible,
|
||||
reverseDescriptorPresent: al.reverse_descriptor_present,
|
||||
error: al.error,
|
||||
revertedByActionId: null,
|
||||
isRevertAction: false,
|
||||
createdAt: al.created_at,
|
||||
},
|
||||
toolCallId: matchedToolCallId,
|
||||
});
|
||||
if (threadId !== null) {
|
||||
applyActionLogSse(queryClient, threadId, searchSpaceId, parsed.data);
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
case "data-action-log-updated": {
|
||||
updateAgentActionReversible({
|
||||
id: parsed.data.id,
|
||||
reversible: parsed.data.reversible,
|
||||
});
|
||||
if (threadId !== null) {
|
||||
applyActionLogUpdatedSse(
|
||||
queryClient,
|
||||
threadId,
|
||||
parsed.data.id,
|
||||
parsed.data.reversible
|
||||
);
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
|
|
@ -2281,12 +2316,16 @@ export default function NewChatPage() {
|
|||
: `Reverted ${summary.reverted} downstream actions before regenerating.`
|
||||
);
|
||||
}
|
||||
for (const r of summary.results) {
|
||||
if (r.status === "reverted" || r.status === "already_reverted") {
|
||||
markAgentActionReverted({
|
||||
id: r.action_id,
|
||||
newActionId: r.new_action_id ?? null,
|
||||
});
|
||||
if (threadId !== null) {
|
||||
for (const r of summary.results) {
|
||||
if (r.status === "reverted" || r.status === "already_reverted") {
|
||||
markActionRevertedInCache(
|
||||
queryClient,
|
||||
threadId,
|
||||
r.action_id,
|
||||
r.new_action_id ?? null
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
break;
|
||||
|
|
@ -2459,16 +2498,26 @@ export default function NewChatPage() {
|
|||
const downstream = messages.slice(editedIndex + 1);
|
||||
downstreamTotalCount = downstream.length;
|
||||
const seenTurns = new Set<string>();
|
||||
const downstreamTurnIds = new Set<string>();
|
||||
for (const m of downstream) {
|
||||
const meta = (m.metadata ?? {}) as { custom?: { chatTurnId?: string } };
|
||||
const tid = meta.custom?.chatTurnId;
|
||||
if (!tid || seenTurns.has(tid)) continue;
|
||||
seenTurns.add(tid);
|
||||
const turnActions = agentActionsByChatTurnId.get(tid) ?? [];
|
||||
for (const a of turnActions) {
|
||||
if (a.reversible && a.revertedByActionId === null && !a.isRevertAction && !a.error) {
|
||||
downstreamReversibleCount += 1;
|
||||
}
|
||||
downstreamTurnIds.add(tid);
|
||||
}
|
||||
// Source of truth: the unified react-query cache. Every
|
||||
// action whose ``chat_turn_id`` belongs to the slice we're
|
||||
// about to drop counts toward the prompt.
|
||||
for (const a of agentActionItems) {
|
||||
if (!a.chat_turn_id || !downstreamTurnIds.has(a.chat_turn_id)) continue;
|
||||
if (
|
||||
a.reversible &&
|
||||
(a.reverted_by_action_id === null || a.reverted_by_action_id === undefined) &&
|
||||
!a.is_revert_action &&
|
||||
(a.error === null || a.error === undefined)
|
||||
) {
|
||||
downstreamReversibleCount += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -2492,7 +2541,7 @@ export default function NewChatPage() {
|
|||
downstreamTotalCount,
|
||||
});
|
||||
},
|
||||
[handleRegenerate, messages, agentActionsByChatTurnId]
|
||||
[handleRegenerate, messages, agentActionItems]
|
||||
);
|
||||
|
||||
const handleEditDialogChoice = useCallback(
|
||||
|
|
|
|||
|
|
@ -1,194 +0,0 @@
|
|||
"use client";
|
||||
|
||||
import { atom } from "jotai";
|
||||
|
||||
/**
|
||||
* Minimal per-row projection of ``AgentActionLog`` that the tool card
|
||||
* needs to decide whether to render a Revert button.
|
||||
*
|
||||
* Fields are deliberately a subset of the full ``AgentAction`` so the
|
||||
* SSE side-channel (``data-action-log`` / ``data-action-log-updated``)
|
||||
* can populate them without depending on the REST endpoint
|
||||
* ``GET /threads/.../actions`` (which 503s when
|
||||
* ``SURFSENSE_ENABLE_ACTION_LOG`` is off).
|
||||
*/
|
||||
export interface AgentActionLite {
|
||||
id: number;
|
||||
threadId: number | null;
|
||||
lcToolCallId: string | null;
|
||||
chatTurnId: string | null;
|
||||
toolName: string;
|
||||
reversible: boolean;
|
||||
reverseDescriptorPresent: boolean;
|
||||
error: boolean;
|
||||
revertedByActionId: number | null;
|
||||
isRevertAction: boolean;
|
||||
createdAt: string | null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Map keyed off the LangChain ``tool_call.id`` (mirrors ``ContentPart
|
||||
* tool-call.langchainToolCallId``).
|
||||
*/
|
||||
export const agentActionByLcIdAtom = atom<Map<string, AgentActionLite>>(new Map());
|
||||
|
||||
/**
|
||||
* Parallel map keyed off the synthetic chat-card ``toolCallId``
|
||||
* (``call_<run-id>``) so ``ToolFallback`` (which only receives the
|
||||
* synthetic id from assistant-ui) can join its card to the action log.
|
||||
*
|
||||
* Both maps are kept in sync by ``upsertAgentActionAtom``.
|
||||
*/
|
||||
export const agentActionByToolCallIdAtom = atom<Map<string, AgentActionLite>>(new Map());
|
||||
|
||||
/**
|
||||
* Index keyed by ``chat_turn_id`` so the per-turn revert UI can answer
|
||||
* "how many reversible actions does this assistant turn contain?" in
|
||||
* O(1). Each entry's array is ordered by insertion (which
|
||||
* for a single turn matches ``created_at`` because action-log writes
|
||||
* happen synchronously).
|
||||
*/
|
||||
export const agentActionsByChatTurnIdAtom = atom<Map<string, AgentActionLite[]>>(new Map());
|
||||
|
||||
/**
|
||||
* Action to upsert one ``AgentActionLite`` row.
|
||||
*
|
||||
* ``toolCallId`` is the synthetic card id (``call_<run-id>`` from
|
||||
* ``stream_new_chat.py``). When provided alongside ``lcToolCallId``, the
|
||||
* action is indexed under BOTH ids so the tool card can perform the
|
||||
* lookup without going via the streaming state.
|
||||
*/
|
||||
export const upsertAgentActionAtom = atom(
|
||||
null,
|
||||
(_get, set, payload: { action: AgentActionLite; toolCallId?: string | null }) => {
|
||||
const { action, toolCallId } = payload;
|
||||
const upsertInto = (
|
||||
prev: Map<string, AgentActionLite>,
|
||||
key: string
|
||||
): Map<string, AgentActionLite> => {
|
||||
const next = new Map(prev);
|
||||
const existing = next.get(key);
|
||||
next.set(key, {
|
||||
...action,
|
||||
// Preserve the local "reverted" bookkeeping if a reversibility
|
||||
// flip arrives AFTER the user already reverted via the REST
|
||||
// route. We never want a stale ``reversible=true`` event to
|
||||
// resurrect a Reverted card.
|
||||
revertedByActionId: existing?.revertedByActionId ?? action.revertedByActionId,
|
||||
isRevertAction: existing?.isRevertAction ?? action.isRevertAction,
|
||||
});
|
||||
return next;
|
||||
};
|
||||
if (action.lcToolCallId) {
|
||||
set(agentActionByLcIdAtom, (prev) => upsertInto(prev, action.lcToolCallId as string));
|
||||
}
|
||||
if (toolCallId) {
|
||||
set(agentActionByToolCallIdAtom, (prev) => upsertInto(prev, toolCallId));
|
||||
}
|
||||
if (action.chatTurnId) {
|
||||
set(agentActionsByChatTurnIdAtom, (prev) => {
|
||||
const next = new Map(prev);
|
||||
const turnId = action.chatTurnId as string;
|
||||
const existing = next.get(turnId) ?? [];
|
||||
const priorEntry = existing.find((row) => row.id === action.id);
|
||||
const merged: AgentActionLite = {
|
||||
...action,
|
||||
revertedByActionId: priorEntry?.revertedByActionId ?? action.revertedByActionId,
|
||||
isRevertAction: priorEntry?.isRevertAction ?? action.isRevertAction,
|
||||
};
|
||||
const others = existing.filter((row) => row.id !== action.id);
|
||||
next.set(turnId, [...others, merged]);
|
||||
return next;
|
||||
});
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
function mutateById(
|
||||
prev: Map<string, AgentActionLite>,
|
||||
id: number,
|
||||
mutator: (entry: AgentActionLite) => AgentActionLite
|
||||
): Map<string, AgentActionLite> {
|
||||
let mutated = false;
|
||||
const next = new Map(prev);
|
||||
for (const [key, value] of next) {
|
||||
if (value.id === id) {
|
||||
next.set(key, mutator(value));
|
||||
mutated = true;
|
||||
}
|
||||
}
|
||||
return mutated ? next : prev;
|
||||
}
|
||||
|
||||
function mutateByIdInTurnIndex(
|
||||
prev: Map<string, AgentActionLite[]>,
|
||||
id: number,
|
||||
mutator: (entry: AgentActionLite) => AgentActionLite
|
||||
): Map<string, AgentActionLite[]> {
|
||||
let mutated = false;
|
||||
const next = new Map(prev);
|
||||
for (const [key, list] of next) {
|
||||
let listMutated = false;
|
||||
const updated = list.map((row) => {
|
||||
if (row.id === id) {
|
||||
listMutated = true;
|
||||
return mutator(row);
|
||||
}
|
||||
return row;
|
||||
});
|
||||
if (listMutated) {
|
||||
next.set(key, updated);
|
||||
mutated = true;
|
||||
}
|
||||
}
|
||||
return mutated ? next : prev;
|
||||
}
|
||||
|
||||
/**
|
||||
* Action to flip an existing entry's ``reversible`` flag, keyed by the
|
||||
* AgentActionLog row id (the SSE ``data-action-log-updated`` payload
|
||||
* does NOT carry ``lcToolCallId``).
|
||||
*/
|
||||
export const updateAgentActionReversibleAtom = atom(
|
||||
null,
|
||||
(_get, set, payload: { id: number; reversible: boolean }) => {
|
||||
const apply = (entry: AgentActionLite): AgentActionLite => ({
|
||||
...entry,
|
||||
reversible: payload.reversible,
|
||||
});
|
||||
set(agentActionByLcIdAtom, (prev) => mutateById(prev, payload.id, apply));
|
||||
set(agentActionByToolCallIdAtom, (prev) => mutateById(prev, payload.id, apply));
|
||||
set(agentActionsByChatTurnIdAtom, (prev) => mutateByIdInTurnIndex(prev, payload.id, apply));
|
||||
}
|
||||
);
|
||||
|
||||
/** Action to mark an existing entry as reverted (post-revert call). */
|
||||
export const markAgentActionRevertedAtom = atom(
|
||||
null,
|
||||
(_get, set, payload: { id: number; newActionId: number | null }) => {
|
||||
const apply = (entry: AgentActionLite): AgentActionLite => ({
|
||||
...entry,
|
||||
revertedByActionId: payload.newActionId ?? -1,
|
||||
});
|
||||
set(agentActionByLcIdAtom, (prev) => mutateById(prev, payload.id, apply));
|
||||
set(agentActionByToolCallIdAtom, (prev) => mutateById(prev, payload.id, apply));
|
||||
set(agentActionsByChatTurnIdAtom, (prev) => mutateByIdInTurnIndex(prev, payload.id, apply));
|
||||
}
|
||||
);
|
||||
|
||||
/** Mark every action in a turn as reverted, given a list of (id, newActionId) pairs. */
|
||||
export const markAgentActionsRevertedBatchAtom = atom(
|
||||
null,
|
||||
(_get, set, payload: { entries: Array<{ id: number; newActionId: number | null }> }) => {
|
||||
for (const entry of payload.entries) {
|
||||
set(markAgentActionRevertedAtom, entry);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
/** Reset all maps (e.g. when the active thread changes). */
|
||||
export const resetAgentActionMapAtom = atom(null, (_get, set) => {
|
||||
set(agentActionByLcIdAtom, new Map());
|
||||
set(agentActionByToolCallIdAtom, new Map());
|
||||
set(agentActionsByChatTurnIdAtom, new Map());
|
||||
});
|
||||
|
|
@ -1,9 +1,9 @@
|
|||
"use client";
|
||||
|
||||
import { useQuery, useQueryClient } from "@tanstack/react-query";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
import { useAtom, useAtomValue } from "jotai";
|
||||
import { Activity, RefreshCcw } from "lucide-react";
|
||||
import { useCallback, useMemo } from "react";
|
||||
import { useCallback } from "react";
|
||||
import { actionLogSheetAtom } from "@/atoms/agent/action-log-sheet.atom";
|
||||
import { agentFlagsAtom } from "@/atoms/agent/agent-flags-query.atom";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
|
|
@ -17,15 +17,12 @@ import {
|
|||
SheetTitle,
|
||||
} from "@/components/ui/sheet";
|
||||
import { Skeleton } from "@/components/ui/skeleton";
|
||||
import { agentActionsApiService } from "@/lib/apis/agent-actions-api.service";
|
||||
import {
|
||||
agentActionsQueryKey,
|
||||
useAgentActionsQuery,
|
||||
} from "@/hooks/use-agent-actions-query";
|
||||
import { ActionLogItem } from "./action-log-item";
|
||||
|
||||
const ACTION_LOG_PAGE_SIZE = 50;
|
||||
|
||||
function actionLogQueryKey(threadId: number) {
|
||||
return ["agent-actions", threadId] as const;
|
||||
}
|
||||
|
||||
function EmptyState() {
|
||||
return (
|
||||
<div className="flex flex-1 flex-col items-center justify-center gap-3 px-6 text-center">
|
||||
|
|
@ -85,25 +82,17 @@ export function ActionLogSheet() {
|
|||
|
||||
const threadId = state.threadId;
|
||||
|
||||
const { data, isLoading, isFetching, isError, error, refetch } = useQuery({
|
||||
queryKey: threadId !== null ? actionLogQueryKey(threadId) : ["agent-actions", "none"],
|
||||
queryFn: () =>
|
||||
agentActionsApiService.listForThread(threadId as number, {
|
||||
page: 0,
|
||||
pageSize: ACTION_LOG_PAGE_SIZE,
|
||||
}),
|
||||
enabled: state.open && threadId !== null && actionLogEnabled,
|
||||
staleTime: 15 * 1000,
|
||||
});
|
||||
const { data, items, isLoading, isFetching, isError, error, refetch } = useAgentActionsQuery(
|
||||
threadId,
|
||||
{ enabled: state.open && actionLogEnabled }
|
||||
);
|
||||
|
||||
const handleRevertSuccess = useCallback(() => {
|
||||
if (threadId !== null) {
|
||||
queryClient.invalidateQueries({ queryKey: actionLogQueryKey(threadId) });
|
||||
queryClient.invalidateQueries({ queryKey: agentActionsQueryKey(threadId) });
|
||||
}
|
||||
}, [queryClient, threadId]);
|
||||
|
||||
const items = useMemo(() => data?.items ?? [], [data]);
|
||||
|
||||
return (
|
||||
<Sheet open={state.open} onOpenChange={(open) => setState((s) => ({ ...s, open }))}>
|
||||
<SheetContent
|
||||
|
|
|
|||
|
|
@ -4,26 +4,22 @@
|
|||
* "Revert turn" button rendered at the bottom of every completed
|
||||
* assistant turn that has at least one reversible action.
|
||||
*
|
||||
* The button reads the action map keyed by ``chat_turn_id`` from the
|
||||
* SSE side-channel (``data-action-log`` events). It shows a confirmation
|
||||
* dialog summarising "N reversible / M total" and, on confirm, calls
|
||||
* ``POST /threads/{id}/revert-turn/{chat_turn_id}``.
|
||||
* The button reads from the unified ``useAgentActionsQuery`` cache
|
||||
* (the SAME react-query cache the agent-actions sheet and the inline
|
||||
* Revert button consume) filtered by ``chat_turn_id``. It shows a
|
||||
* confirmation dialog summarising "N reversible / M total" and, on
|
||||
* confirm, calls ``POST /threads/{id}/revert-turn/{chat_turn_id}``.
|
||||
*
|
||||
* The route returns a per-action result list and never collapses the
|
||||
* batch into a 4xx — so we render any failed/not_reversible rows inline
|
||||
* with their messages.
|
||||
*/
|
||||
|
||||
import { useAtomValue, useSetAtom } from "jotai";
|
||||
import { selectAtom } from "jotai/utils";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
import { useAtomValue } from "jotai";
|
||||
import { CheckIcon, RotateCcw, XCircleIcon } from "lucide-react";
|
||||
import { useMemo, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
import {
|
||||
type AgentActionLite,
|
||||
agentActionsByChatTurnIdAtom,
|
||||
markAgentActionsRevertedBatchAtom,
|
||||
} from "@/atoms/chat/agent-actions.atom";
|
||||
import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom";
|
||||
import {
|
||||
AlertDialog,
|
||||
|
|
@ -38,6 +34,10 @@ import {
|
|||
} from "@/components/ui/alert-dialog";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { getToolDisplayName } from "@/contracts/enums/toolIcons";
|
||||
import {
|
||||
applyRevertTurnResultsToCache,
|
||||
useAgentActionsQuery,
|
||||
} from "@/hooks/use-agent-actions-query";
|
||||
import {
|
||||
agentActionsApiService,
|
||||
type RevertTurnActionResult,
|
||||
|
|
@ -49,49 +49,33 @@ interface RevertTurnButtonProps {
|
|||
chatTurnId: string | null | undefined;
|
||||
}
|
||||
|
||||
// Empty-array sentinel so the per-turn ``selectAtom`` slice returns a
|
||||
// stable reference when the turn has no recorded actions yet. Without
|
||||
// this every render allocates a fresh ``[]`` and Jotai's
|
||||
// equality check would re-render the button on unrelated turn updates.
|
||||
const EMPTY_ACTIONS: readonly AgentActionLite[] = Object.freeze([]);
|
||||
|
||||
export function RevertTurnButton({ chatTurnId }: RevertTurnButtonProps) {
|
||||
const session = useAtomValue(chatSessionStateAtom);
|
||||
const markRevertedBatch = useSetAtom(markAgentActionsRevertedBatchAtom);
|
||||
const threadId = session?.threadId ?? null;
|
||||
const queryClient = useQueryClient();
|
||||
const { findByChatTurnId } = useAgentActionsQuery(threadId);
|
||||
const [isReverting, setIsReverting] = useState(false);
|
||||
const [confirmOpen, setConfirmOpen] = useState(false);
|
||||
const [resultsOpen, setResultsOpen] = useState(false);
|
||||
const [results, setResults] = useState<RevertTurnActionResult[]>([]);
|
||||
|
||||
// Subscribe ONLY to the slice of the global action map that belongs
|
||||
// to ``chatTurnId``. Previously the button read the whole
|
||||
// ``agentActionsByChatTurnIdAtom``, which meant every action
|
||||
// upsert (one per tool call) re-rendered every Revert button on
|
||||
// the page. With ``selectAtom`` we re-render only when our turn's
|
||||
// list reference changes — and the upsert/mark atoms produce a
|
||||
// fresh list reference for the affected turn only.
|
||||
const sliceAtom = useMemo(
|
||||
() =>
|
||||
selectAtom(
|
||||
agentActionsByChatTurnIdAtom,
|
||||
(turnIndex) => (chatTurnId ? turnIndex.get(chatTurnId) : undefined) ?? EMPTY_ACTIONS
|
||||
),
|
||||
[chatTurnId]
|
||||
);
|
||||
const actions = useAtomValue(sliceAtom);
|
||||
const actions = useMemo(() => findByChatTurnId(chatTurnId), [findByChatTurnId, chatTurnId]);
|
||||
|
||||
const reversibleCount = useMemo(
|
||||
() =>
|
||||
actions.filter(
|
||||
(a) => a.reversible && a.revertedByActionId === null && !a.isRevertAction && !a.error
|
||||
(a) =>
|
||||
a.reversible &&
|
||||
(a.reverted_by_action_id === null || a.reverted_by_action_id === undefined) &&
|
||||
!a.is_revert_action &&
|
||||
(a.error === null || a.error === undefined)
|
||||
).length,
|
||||
[actions]
|
||||
);
|
||||
const totalCount = useMemo(() => actions.filter((a) => !a.isRevertAction).length, [actions]);
|
||||
const totalCount = useMemo(() => actions.filter((a) => !a.is_revert_action).length, [actions]);
|
||||
|
||||
if (!chatTurnId) return null;
|
||||
if (reversibleCount === 0) return null;
|
||||
const threadId = session?.threadId;
|
||||
if (!threadId) return null;
|
||||
|
||||
const handleRevertTurn = async () => {
|
||||
|
|
@ -103,7 +87,7 @@ export function RevertTurnButton({ chatTurnId }: RevertTurnButtonProps) {
|
|||
.filter((r) => r.status === "reverted" || r.status === "already_reverted")
|
||||
.map((r) => ({ id: r.action_id, newActionId: r.new_action_id ?? null }));
|
||||
if (revertedEntries.length > 0) {
|
||||
markRevertedBatch({ entries: revertedEntries });
|
||||
applyRevertTurnResultsToCache(queryClient, threadId, revertedEntries);
|
||||
}
|
||||
if (response.status === "ok") {
|
||||
toast.success(
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
|
||||
import { useAtomValue, useSetAtom } from "jotai";
|
||||
import { CheckIcon, ChevronDownIcon, ChevronUpIcon, RotateCcw, XCircleIcon } from "lucide-react";
|
||||
import { useMemo, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
import {
|
||||
agentActionByToolCallIdAtom,
|
||||
markAgentActionRevertedAtom,
|
||||
} from "@/atoms/chat/agent-actions.atom";
|
||||
type ToolCallMessagePartComponent,
|
||||
useAuiState,
|
||||
} from "@assistant-ui/react";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
import { useAtomValue } from "jotai";
|
||||
import { CheckIcon, ChevronDownIcon, RotateCcw, XCircleIcon } from "lucide-react";
|
||||
import { useEffect, useMemo, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom";
|
||||
import {
|
||||
DoomLoopApprovalToolUI,
|
||||
|
|
@ -24,8 +24,17 @@ import {
|
|||
AlertDialogTitle,
|
||||
AlertDialogTrigger,
|
||||
} from "@/components/ui/alert-dialog";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { getToolDisplayName, getToolIcon } from "@/contracts/enums/toolIcons";
|
||||
import { Card } from "@/components/ui/card";
|
||||
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
|
||||
import { Separator } from "@/components/ui/separator";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import { getToolDisplayName } from "@/contracts/enums/toolIcons";
|
||||
import {
|
||||
markActionRevertedInCache,
|
||||
useAgentActionsQuery,
|
||||
} from "@/hooks/use-agent-actions-query";
|
||||
import { agentActionsApiService } from "@/lib/apis/agent-actions-api.service";
|
||||
import { AppError } from "@/lib/error";
|
||||
import { isInterruptResult } from "@/lib/hitl";
|
||||
|
|
@ -34,31 +43,128 @@ import { cn } from "@/lib/utils";
|
|||
/**
|
||||
* Inline Revert button rendered on a tool card when the matching
|
||||
* ``AgentActionLog`` row is reversible and hasn't been reverted yet.
|
||||
* Reads from the SSE side-channel atom keyed by the synthetic
|
||||
* ``toolCallId`` so it lights up even when ``GET /threads/.../actions``
|
||||
* is gated behind ``SURFSENSE_ENABLE_ACTION_LOG=False`` (503).
|
||||
*
|
||||
* Reads from the unified ``useAgentActionsQuery`` cache — the SAME
|
||||
* react-query cache the agent-actions sheet consumes. SSE events
|
||||
* (``data-action-log`` / ``data-action-log-updated``) and
|
||||
* ``POST /threads/{id}/revert/{id}`` responses both flow through the
|
||||
* cache via ``setQueryData`` helpers, so the card and the sheet stay
|
||||
* in lockstep on every code path: page reload, navigation, live
|
||||
* stream, post-stream reversibility flip, and explicit revert clicks.
|
||||
*
|
||||
* Match key (in priority order):
|
||||
* 1. ``a.tool_call_id === toolCallId`` — direct hit in parity_v2 when
|
||||
* the model streamed ``tool_call_chunks`` so the card's synthetic
|
||||
* id IS the LangChain id.
|
||||
* 2. ``a.tool_call_id === langchainToolCallId`` — legacy mode (or
|
||||
* parity_v2 with provider-side chunk emission) where the card's
|
||||
* synthetic id is ``call_<run_id>`` and the LangChain id is
|
||||
* backfilled onto the part by ``tool-output-available``.
|
||||
* 3. ``(chat_turn_id, tool_name, position-within-turn)`` — fallback
|
||||
* for cards whose synthetic id is ``call_<run_id>`` AND whose
|
||||
* ``langchainToolCallId`` never got backfilled (provider emitted
|
||||
* the tool_call as a single payload with no chunks AND streaming
|
||||
* pre-dated the ``tool-output-available langchainToolCallId``
|
||||
* backfill, e.g. older threads). Reads the parent message's
|
||||
* ``chatTurnId`` and ``content`` via ``useAuiState`` so we can
|
||||
* match position-by-tool-name within the turn against the
|
||||
* action_log rows the server returned in ``created_at`` order.
|
||||
*/
|
||||
function ToolCardRevertButton({ toolCallId }: { toolCallId: string }) {
|
||||
function ToolCardRevertButton({
|
||||
toolCallId,
|
||||
toolName,
|
||||
langchainToolCallId,
|
||||
}: {
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
langchainToolCallId?: string;
|
||||
}) {
|
||||
const session = useAtomValue(chatSessionStateAtom);
|
||||
const actionMap = useAtomValue(agentActionByToolCallIdAtom);
|
||||
const markReverted = useSetAtom(markAgentActionRevertedAtom);
|
||||
const action = actionMap.get(toolCallId);
|
||||
const threadId = session?.threadId ?? null;
|
||||
const queryClient = useQueryClient();
|
||||
const { findByToolCallId, findByChatTurnAndTool } = useAgentActionsQuery(threadId);
|
||||
|
||||
// Parent message metadata, read via the narrowest possible
|
||||
// selectors so this card doesn't re-render on every text-delta of
|
||||
// every other part in the same message during streaming.
|
||||
//
|
||||
// IMPORTANT — ``useAuiState`` re-renders the component whenever the
|
||||
// returned slice's identity changes. Returning ``message?.content``
|
||||
// (an array) would re-render on every token because the runtime
|
||||
// rebuilds the parts array. Returning a PRIMITIVE (the position
|
||||
// number) lets ``useAuiState``'s ``Object.is`` check short-circuit
|
||||
// when the position hasn't actually moved — which is the common
|
||||
// case during text streaming, when only ``text``/``reasoning``
|
||||
// parts are mutating and the same-toolName tool-call ordering is
|
||||
// stable. (See Vercel React rule ``rerender-defer-reads``.)
|
||||
const chatTurnId = useAuiState(({ message }) => {
|
||||
const meta = message?.metadata as { custom?: { chatTurnId?: string } } | undefined;
|
||||
return meta?.custom?.chatTurnId ?? null;
|
||||
});
|
||||
const positionInTurn = useAuiState(({ message }) => {
|
||||
const content = message?.content;
|
||||
if (!Array.isArray(content)) return -1;
|
||||
let n = -1;
|
||||
for (const part of content) {
|
||||
if (
|
||||
part &&
|
||||
typeof part === "object" &&
|
||||
(part as { type?: string }).type === "tool-call" &&
|
||||
(part as { toolName?: string }).toolName === toolName
|
||||
) {
|
||||
n += 1;
|
||||
if ((part as { toolCallId?: string }).toolCallId === toolCallId) return n;
|
||||
}
|
||||
}
|
||||
return -1;
|
||||
});
|
||||
|
||||
const action = useMemo(() => {
|
||||
// Tier 1 + 2: O(1) Map-backed direct id match. Covers
|
||||
// ~all parity_v2 streams and any legacy stream that backfilled
|
||||
// ``langchainToolCallId`` via ``tool-output-available``.
|
||||
const direct =
|
||||
findByToolCallId(toolCallId) ?? findByToolCallId(langchainToolCallId);
|
||||
if (direct) return direct;
|
||||
// Tier 3: position-within-turn fallback. Only kicks in when the
|
||||
// card has a synthetic ``call_<run_id>`` id AND no
|
||||
// ``langchainToolCallId`` was ever backfilled — i.e. the tool
|
||||
// was emitted as a single non-chunked payload AND streaming
|
||||
// pre-dated the on_tool_end backfill.
|
||||
if (!chatTurnId || positionInTurn < 0) return null;
|
||||
const turnSameTool = findByChatTurnAndTool(chatTurnId, toolName);
|
||||
return turnSameTool[positionInTurn] ?? null;
|
||||
}, [
|
||||
findByToolCallId,
|
||||
findByChatTurnAndTool,
|
||||
toolCallId,
|
||||
langchainToolCallId,
|
||||
chatTurnId,
|
||||
toolName,
|
||||
positionInTurn,
|
||||
]);
|
||||
|
||||
const [isReverting, setIsReverting] = useState(false);
|
||||
const [confirmOpen, setConfirmOpen] = useState(false);
|
||||
|
||||
if (!action) return null;
|
||||
if (!action.reversible) return null;
|
||||
if (action.revertedByActionId !== null) return null;
|
||||
if (action.isRevertAction) return null;
|
||||
if (action.error) return null;
|
||||
const threadId = session?.threadId;
|
||||
if (action.reverted_by_action_id !== null && action.reverted_by_action_id !== undefined)
|
||||
return null;
|
||||
if (action.is_revert_action) return null;
|
||||
if (action.error !== null && action.error !== undefined) return null;
|
||||
if (!threadId) return null;
|
||||
|
||||
const handleRevert = async () => {
|
||||
setIsReverting(true);
|
||||
try {
|
||||
const response = await agentActionsApiService.revert(threadId, action.id);
|
||||
markReverted({ id: action.id, newActionId: response.new_action_id ?? null });
|
||||
markActionRevertedInCache(
|
||||
queryClient,
|
||||
threadId,
|
||||
action.id,
|
||||
response.new_action_id ?? null
|
||||
);
|
||||
toast.success(response.message || "Action reverted.");
|
||||
} catch (err) {
|
||||
// 503 means revert is gated off on this deployment — hide the
|
||||
|
|
@ -91,8 +197,17 @@ function ToolCardRevertButton({ toolCallId }: { toolCallId: string }) {
|
|||
e.stopPropagation();
|
||||
setConfirmOpen(true);
|
||||
}}
|
||||
disabled={isReverting}
|
||||
>
|
||||
<RotateCcw className="size-3.5" />
|
||||
{isReverting ? (
|
||||
// Spinner's typed props don't accept ``data-icon`` and
|
||||
// it renders an <output>, not an <svg>, so Button's
|
||||
// auto-sizing rule doesn't apply. Bare spinner +
|
||||
// Button's gap handle layout.
|
||||
<Spinner size="xs" />
|
||||
) : (
|
||||
<RotateCcw data-icon="inline-start" />
|
||||
)}
|
||||
Revert
|
||||
</Button>
|
||||
</AlertDialogTrigger>
|
||||
|
|
@ -101,7 +216,7 @@ function ToolCardRevertButton({ toolCallId }: { toolCallId: string }) {
|
|||
<AlertDialogTitle>Revert this action?</AlertDialogTitle>
|
||||
<AlertDialogDescription>
|
||||
This will undo{" "}
|
||||
<span className="font-medium">{getToolDisplayName(action.toolName)}</span> and add a
|
||||
<span className="font-medium">{getToolDisplayName(action.tool_name)}</span> and add a
|
||||
new entry to the history. Your chat is preserved — only the changes the agent made to
|
||||
your knowledge base or connected apps will be rolled back where possible.
|
||||
</AlertDialogDescription>
|
||||
|
|
@ -114,8 +229,10 @@ function ToolCardRevertButton({ toolCallId }: { toolCallId: string }) {
|
|||
handleRevert();
|
||||
}}
|
||||
disabled={isReverting}
|
||||
className="gap-1.5"
|
||||
>
|
||||
{isReverting ? "Reverting…" : "Revert"}
|
||||
{isReverting && <Spinner size="xs" />}
|
||||
Revert
|
||||
</AlertDialogAction>
|
||||
</AlertDialogFooter>
|
||||
</AlertDialogContent>
|
||||
|
|
@ -123,18 +240,49 @@ function ToolCardRevertButton({ toolCallId }: { toolCallId: string }) {
|
|||
);
|
||||
}
|
||||
|
||||
const DefaultToolFallbackInner: ToolCallMessagePartComponent = ({
|
||||
toolCallId,
|
||||
toolName,
|
||||
argsText,
|
||||
result,
|
||||
status,
|
||||
}) => {
|
||||
const [isExpanded, setIsExpanded] = useState(false);
|
||||
/**
|
||||
* Compact tool-call card.
|
||||
*
|
||||
* shadcn composition note: we intentionally use ``Card`` as a visual
|
||||
* frame WITHOUT ``CardHeader / CardContent``. The full composition's
|
||||
* ``p-6`` padding doesn't fit a compact collapsible header that IS the
|
||||
* trigger; using ``Card`` alone preserves the rounded border, shadow,
|
||||
* and ``bg-card`` token (semantic colors) without forcing a layout
|
||||
* that doesn't fit. All status colors use semantic tokens — no manual
|
||||
* dark-mode overrides, no raw hex.
|
||||
*/
|
||||
const DefaultToolFallbackInner: ToolCallMessagePartComponent = (props) => {
|
||||
const { toolCallId, toolName, argsText, result, status } = props;
|
||||
// ``langchainToolCallId`` is a SurfSense-specific extension the
|
||||
// streaming pipeline attaches to the tool-call content part so
|
||||
// the Revert button can resolve its ``AgentActionLog`` row even
|
||||
// when only the LC id is known. assistant-ui's
|
||||
// ``ToolCallMessagePartProps`` doesn't list it, but the runtime
|
||||
// spreads ``{...part}`` so the prop reaches us at runtime.
|
||||
const langchainToolCallId = (props as { langchainToolCallId?: string }).langchainToolCallId;
|
||||
|
||||
const isCancelled = status?.type === "incomplete" && status.reason === "cancelled";
|
||||
const isError = status?.type === "incomplete" && status.reason === "error";
|
||||
const isRunning = status?.type === "running" || status?.type === "requires-action";
|
||||
|
||||
/*
|
||||
Per-card expansion state. Initial value is ``isRunning`` so a
|
||||
card streaming in mounts already-expanded (no flash of
|
||||
collapsed → expanded on first paint), while a card loaded from
|
||||
history (status="complete") mounts collapsed. The useEffect
|
||||
below keeps this in lockstep with this card's own ``isRunning``
|
||||
when it transitions: false → true auto-expands (e.g. a tool
|
||||
that re-runs after edit), true → false auto-collapses once the
|
||||
tool finishes. Because the dep is per-card ``isRunning`` and
|
||||
not the chat-level streaming flag, sibling cards on the same
|
||||
assistant turn each manage their own expansion independently.
|
||||
Once ``isRunning`` is false the user controls expansion via
|
||||
``onOpenChange``.
|
||||
*/
|
||||
const [isExpanded, setIsExpanded] = useState(isRunning);
|
||||
useEffect(() => {
|
||||
setIsExpanded(isRunning);
|
||||
}, [isRunning]);
|
||||
const errorData = status?.type === "incomplete" ? status.error : undefined;
|
||||
const serializedError = useMemo(
|
||||
() => (errorData && typeof errorData !== "string" ? JSON.stringify(errorData) : null),
|
||||
|
|
@ -160,108 +308,207 @@ const DefaultToolFallbackInner: ToolCallMessagePartComponent = ({
|
|||
: serializedError
|
||||
: null;
|
||||
|
||||
const Icon = getToolIcon(toolName);
|
||||
const displayName = getToolDisplayName(toolName);
|
||||
const subtitle = errorReason ?? cancelledReason;
|
||||
|
||||
return (
|
||||
<div
|
||||
<Card
|
||||
className={cn(
|
||||
"my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none",
|
||||
"my-4 max-w-lg overflow-hidden",
|
||||
isCancelled && "opacity-60",
|
||||
isError && "border-destructive/20 bg-destructive/5"
|
||||
isError && "border-destructive/30"
|
||||
)}
|
||||
>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setIsExpanded((prev) => !prev)}
|
||||
className="flex w-full items-center gap-3 px-5 py-4 text-left transition-colors hover:bg-muted/50 focus:outline-none focus-visible:outline-none"
|
||||
{/*
|
||||
``group`` lets the chevron (rendered as a sibling of the
|
||||
main trigger button) read the Collapsible Root's
|
||||
``data-[state=open]`` for rotation. The Collapsible is
|
||||
fully controlled via ``isExpanded`` — the useEffect
|
||||
above syncs it to ``isRunning`` so the card auto-opens
|
||||
while a tool streams in and auto-collapses once it
|
||||
finishes. We deliberately DON'T pass ``disabled`` so
|
||||
both triggers stay clickable; ``onOpenChange`` is wired
|
||||
to a setter that no-ops while ``isRunning`` (see
|
||||
``handleOpenChange`` below) which keeps the card pinned
|
||||
open mid-stream without losing keyboard / pointer
|
||||
affordance the moment streaming ends.
|
||||
*/}
|
||||
<Collapsible
|
||||
className="group"
|
||||
open={isExpanded}
|
||||
onOpenChange={(next) => {
|
||||
// Block manual collapse while the tool is still
|
||||
// streaming — otherwise a stray click on either
|
||||
// trigger would close the card and hide the live
|
||||
// ``argsText`` panel mid-run. After streaming the
|
||||
// user has full control again.
|
||||
if (isRunning) return;
|
||||
setIsExpanded(next);
|
||||
}}
|
||||
>
|
||||
<div
|
||||
className={cn(
|
||||
"flex size-8 shrink-0 items-center justify-center rounded-lg",
|
||||
isError ? "bg-destructive/10" : isCancelled ? "bg-muted" : "bg-primary/10"
|
||||
)}
|
||||
>
|
||||
{isError ? (
|
||||
<XCircleIcon className="size-4 text-destructive" />
|
||||
) : isCancelled ? (
|
||||
<XCircleIcon className="size-4 text-muted-foreground" />
|
||||
) : isRunning ? (
|
||||
<Icon className="size-4 text-primary animate-pulse" />
|
||||
) : (
|
||||
<CheckIcon className="size-4 text-primary" />
|
||||
)}
|
||||
</div>
|
||||
{/*
|
||||
Header row: main trigger on the left (icon + title
|
||||
col), Revert + chevron-trigger on the right as
|
||||
siblings of the main trigger. The chevron is wrapped
|
||||
in its OWN ``CollapsibleTrigger`` (Radix supports
|
||||
multiple triggers per Root) so clicking the chevron
|
||||
toggles the same state as clicking the title row.
|
||||
The Revert button stays a separate AlertDialog
|
||||
trigger and stops propagation in its onClick so it
|
||||
doesn't toggle the collapsible while opening the
|
||||
confirm dialog. Keeping these as flat siblings —
|
||||
rather than nesting Revert / chevron inside the
|
||||
title trigger — avoids invalid HTML
|
||||
(button-in-button) and lets the Revert button
|
||||
render in BOTH the collapsed and expanded states.
|
||||
*/}
|
||||
<div className="flex items-stretch transition-colors hover:bg-muted/50">
|
||||
<CollapsibleTrigger asChild>
|
||||
<button
|
||||
type="button"
|
||||
className={cn(
|
||||
"flex flex-1 min-w-0 items-center gap-3 py-4 pl-5 pr-2 text-left",
|
||||
// Inset ring — Card's ``overflow-hidden`` would
|
||||
// clip an ``offset-2`` ring; ``ring-inset``
|
||||
// paints inside the button box.
|
||||
"focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset",
|
||||
"disabled:cursor-default"
|
||||
)}
|
||||
>
|
||||
<div
|
||||
className={cn(
|
||||
"flex size-8 shrink-0 items-center justify-center rounded-lg",
|
||||
isError ? "bg-destructive/10" : isCancelled ? "bg-muted" : "bg-primary/10"
|
||||
)}
|
||||
>
|
||||
{isError ? (
|
||||
<XCircleIcon className="size-4 text-destructive" />
|
||||
) : isCancelled ? (
|
||||
<XCircleIcon className="size-4 text-muted-foreground" />
|
||||
) : isRunning ? (
|
||||
<Spinner size="sm" className="text-primary" />
|
||||
) : (
|
||||
<CheckIcon className="size-4 text-primary" />
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="flex-1 min-w-0">
|
||||
<p
|
||||
className={cn(
|
||||
"text-sm font-semibold",
|
||||
isError
|
||||
? "text-destructive"
|
||||
: isCancelled
|
||||
? "text-muted-foreground line-through"
|
||||
: "text-foreground"
|
||||
)}
|
||||
>
|
||||
{isRunning
|
||||
? displayName
|
||||
: isCancelled
|
||||
? `Cancelled: ${displayName}`
|
||||
: isError
|
||||
? `Failed: ${displayName}`
|
||||
: displayName}
|
||||
</p>
|
||||
{isRunning && <p className="text-xs text-muted-foreground mt-0.5">Working…</p>}
|
||||
{cancelledReason && (
|
||||
<p className="text-xs text-muted-foreground mt-0.5 truncate">{cancelledReason}</p>
|
||||
)}
|
||||
{errorReason && (
|
||||
<p className="text-xs text-destructive/80 mt-0.5 truncate">{errorReason}</p>
|
||||
)}
|
||||
</div>
|
||||
<div className="flex flex-1 min-w-0 flex-col gap-0.5">
|
||||
<div className="flex items-center gap-2">
|
||||
<p
|
||||
className={cn(
|
||||
"text-sm font-semibold truncate",
|
||||
isCancelled && "text-muted-foreground line-through",
|
||||
isError && "text-destructive"
|
||||
)}
|
||||
>
|
||||
{displayName}
|
||||
</p>
|
||||
{isRunning && <Badge variant="secondary">Running</Badge>}
|
||||
{isError && <Badge variant="destructive">Failed</Badge>}
|
||||
{isCancelled && <Badge variant="outline">Cancelled</Badge>}
|
||||
</div>
|
||||
{subtitle && (
|
||||
<p
|
||||
className={cn(
|
||||
"text-xs truncate",
|
||||
isError ? "text-destructive/80" : "text-muted-foreground"
|
||||
)}
|
||||
>
|
||||
{subtitle}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
</button>
|
||||
</CollapsibleTrigger>
|
||||
|
||||
{!isRunning && (
|
||||
<div className="shrink-0 text-muted-foreground">
|
||||
{isExpanded ? (
|
||||
<ChevronDownIcon className="size-4" />
|
||||
) : (
|
||||
<ChevronUpIcon className="size-4" />
|
||||
)}
|
||||
{/*
|
||||
Right-side controls. The Revert button is
|
||||
visible whenever the matching action is
|
||||
reversible — including the collapsed state —
|
||||
but ``ToolCardRevertButton`` itself returns
|
||||
``null`` while a tool is still running because
|
||||
no action-log row exists yet, so it doesn't
|
||||
need an explicit ``isRunning`` gate here.
|
||||
*/}
|
||||
<div className="flex shrink-0 items-center gap-2 pl-2 pr-5">
|
||||
<ToolCardRevertButton
|
||||
toolCallId={toolCallId}
|
||||
toolName={toolName}
|
||||
langchainToolCallId={langchainToolCallId}
|
||||
/>
|
||||
<CollapsibleTrigger asChild>
|
||||
<button
|
||||
type="button"
|
||||
aria-label={isExpanded ? "Collapse details" : "Expand details"}
|
||||
className={cn(
|
||||
"flex size-7 shrink-0 items-center justify-center rounded-md",
|
||||
"text-muted-foreground hover:bg-muted hover:text-foreground",
|
||||
"focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset",
|
||||
"disabled:cursor-default"
|
||||
)}
|
||||
>
|
||||
<ChevronDownIcon
|
||||
className={cn(
|
||||
"size-4 transition-transform duration-200",
|
||||
"group-data-[state=open]:rotate-180"
|
||||
)}
|
||||
/>
|
||||
</button>
|
||||
</CollapsibleTrigger>
|
||||
</div>
|
||||
)}
|
||||
</button>
|
||||
</div>
|
||||
|
||||
{isExpanded && !isRunning && (
|
||||
<>
|
||||
<div className="mx-5 h-px bg-border/50" />
|
||||
<div className="px-5 py-3 space-y-3">
|
||||
{argsText && (
|
||||
<div>
|
||||
<p className="text-xs font-medium text-muted-foreground mb-1">Inputs</p>
|
||||
<pre className="text-xs text-foreground/80 whitespace-pre-wrap break-all">
|
||||
{argsText}
|
||||
</pre>
|
||||
{/*
|
||||
CollapsibleContent body — auto-open while streaming
|
||||
(see ``open`` prop above) so the live ``argsText``
|
||||
streams into the Inputs panel directly, no need for
|
||||
a separate "Live input" panel. Native
|
||||
``overflow-auto`` instead of ``ScrollArea`` because
|
||||
Radix's Viewport can let content bleed past
|
||||
``max-h-*`` in dynamic flex layouts. ``min-w-0`` on
|
||||
the column wrappers guarantees ``break-all`` wraps
|
||||
correctly within the bounded ``max-w-lg`` Card.
|
||||
*/}
|
||||
<CollapsibleContent>
|
||||
<Separator />
|
||||
<div className="flex flex-col gap-3 px-5 py-3">
|
||||
{(argsText || isRunning) && (
|
||||
<div className="flex flex-col gap-1 min-w-0">
|
||||
<p className="text-xs font-medium text-muted-foreground">Inputs</p>
|
||||
<div className="max-h-48 overflow-auto rounded-md bg-muted/40">
|
||||
{argsText ? (
|
||||
<pre className="px-3 py-2 text-xs text-foreground/80 whitespace-pre-wrap break-all font-mono">
|
||||
{argsText}
|
||||
</pre>
|
||||
) : (
|
||||
// Bridges the brief gap between
|
||||
// ``tool-input-start`` (creates the
|
||||
// card, ``argsText`` undefined) and
|
||||
// the first ``tool-input-delta``.
|
||||
<p className="px-3 py-2 text-xs italic text-muted-foreground">
|
||||
Waiting for input…
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
{!isCancelled && result !== undefined && (
|
||||
<>
|
||||
<div className="h-px bg-border/30" />
|
||||
<div>
|
||||
<p className="text-xs font-medium text-muted-foreground mb-1">Result</p>
|
||||
<pre className="text-xs text-foreground/80 whitespace-pre-wrap break-all">
|
||||
{typeof result === "string" ? result : serializedResult}
|
||||
</pre>
|
||||
<Separator />
|
||||
<div className="flex flex-col gap-1 min-w-0">
|
||||
<p className="text-xs font-medium text-muted-foreground">Result</p>
|
||||
<div className="max-h-64 overflow-auto rounded-md bg-muted/40">
|
||||
<pre className="px-3 py-2 text-xs text-foreground/80 whitespace-pre-wrap break-all font-mono">
|
||||
{typeof result === "string" ? result : serializedResult}
|
||||
</pre>
|
||||
</div>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
<div className="flex justify-end">
|
||||
<ToolCardRevertButton toolCallId={toolCallId} />
|
||||
</div>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
</CollapsibleContent>
|
||||
</Collapsible>
|
||||
</Card>
|
||||
);
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ import {
|
|||
addToolCall,
|
||||
appendReasoning,
|
||||
appendText,
|
||||
appendToolInputDelta,
|
||||
buildContentForUI,
|
||||
type ContentPartsState,
|
||||
endReasoning,
|
||||
|
|
@ -188,6 +189,10 @@ export function FreeChatPage() {
|
|||
);
|
||||
};
|
||||
const scheduleFlush = () => batcher.schedule(flushMessages);
|
||||
const forceFlush = () => {
|
||||
scheduleFlush();
|
||||
batcher.flush();
|
||||
};
|
||||
|
||||
try {
|
||||
for await (const parsed of readSSEStream(response)) {
|
||||
|
|
@ -225,13 +230,20 @@ export function FreeChatPage() {
|
|||
false,
|
||||
parsed.langchainToolCallId
|
||||
);
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
|
||||
case "tool-input-available":
|
||||
case "tool-input-delta":
|
||||
appendToolInputDelta(contentPartsState, parsed.toolCallId, parsed.inputTextDelta);
|
||||
scheduleFlush();
|
||||
break;
|
||||
|
||||
case "tool-input-available": {
|
||||
const finalArgsText = JSON.stringify(parsed.input ?? {}, null, 2);
|
||||
if (toolCallIndices.has(parsed.toolCallId)) {
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
args: parsed.input || {},
|
||||
argsText: finalArgsText,
|
||||
langchainToolCallId: parsed.langchainToolCallId,
|
||||
});
|
||||
} else {
|
||||
|
|
@ -244,16 +256,20 @@ export function FreeChatPage() {
|
|||
false,
|
||||
parsed.langchainToolCallId
|
||||
);
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
argsText: finalArgsText,
|
||||
});
|
||||
}
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
}
|
||||
|
||||
case "tool-output-available":
|
||||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||||
result: parsed.output,
|
||||
langchainToolCallId: parsed.langchainToolCallId,
|
||||
});
|
||||
batcher.flush();
|
||||
forceFlush();
|
||||
break;
|
||||
|
||||
case "data-thinking-step": {
|
||||
|
|
|
|||
|
|
@ -1,7 +1,13 @@
|
|||
import { z } from "zod";
|
||||
|
||||
/**
|
||||
* Raw message from database (real-time sync)
|
||||
* Raw message from database (real-time sync).
|
||||
*
|
||||
* ``turn_id`` is included so consumers (e.g. ``convertToThreadMessage``)
|
||||
* can populate ``metadata.custom.chatTurnId`` on the
|
||||
* ``ThreadMessageLike`` even after the live-collab Zero re-sync. The
|
||||
* inline Revert button's ``(chat_turn_id, tool_name, position)``
|
||||
* fallback in tool-fallback.tsx depends on it.
|
||||
*/
|
||||
export const rawMessage = z.object({
|
||||
id: z.number(),
|
||||
|
|
@ -10,6 +16,7 @@ export const rawMessage = z.object({
|
|||
content: z.unknown(),
|
||||
author_id: z.string().nullable(),
|
||||
created_at: z.string(),
|
||||
turn_id: z.string().nullable().optional(),
|
||||
});
|
||||
|
||||
export type RawMessage = z.infer<typeof rawMessage>;
|
||||
|
|
|
|||
416
surfsense_web/hooks/use-agent-actions-query.ts
Normal file
416
surfsense_web/hooks/use-agent-actions-query.ts
Normal file
|
|
@ -0,0 +1,416 @@
|
|||
"use client";
|
||||
|
||||
import { type QueryClient, useQuery } from "@tanstack/react-query";
|
||||
import { useCallback, useEffect, useMemo, useRef } from "react";
|
||||
import {
|
||||
type AgentAction,
|
||||
type AgentActionListResponse,
|
||||
agentActionsApiService,
|
||||
} from "@/lib/apis/agent-actions-api.service";
|
||||
|
||||
// =============================================================================
|
||||
// DIAGNOSTIC LOGGING — gated behind a single switch. Flip ``RevertDebug``
|
||||
// to ``true`` to trace the full SSE → cache → card → button pipeline in
|
||||
// the browser console. Off by default so we don't spam production. The
|
||||
// infrastructure stays in place because the underlying id-mismatch
|
||||
// failure mode is rare-but-real and surfaces only at runtime.
|
||||
// =============================================================================
|
||||
const RevertDebug = false;
|
||||
const dbg = (...args: unknown[]) => {
|
||||
if (RevertDebug && typeof window !== "undefined") {
|
||||
// eslint-disable-next-line no-console
|
||||
console.log("[RevertDebug]", ...args);
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Unified store for ``AgentActionLog`` rows scoped to one thread.
|
||||
*
|
||||
* Replaces the previous SSE side-channel atom mess
|
||||
* (``agentActionByLcIdAtom`` / ``agentActionByToolCallIdAtom`` /
|
||||
* ``agentActionsByChatTurnIdAtom``) and the standalone hydration hook.
|
||||
* One react-query cache entry is now the single source of truth for:
|
||||
*
|
||||
* * the inline Revert button on every tool-call card
|
||||
* * the per-turn "Revert turn" button under each assistant message
|
||||
* * the edit-from-position pre-flight that decides whether to show
|
||||
* the confirmation dialog
|
||||
* * the agent-actions sheet
|
||||
*
|
||||
* The cache is hydrated by ``GET /threads/{id}/actions`` (sized to
|
||||
* 200, the server max) and updated incrementally by helpers that turn
|
||||
* SSE events / revert RPC responses into ``setQueryData`` mutations.
|
||||
* That keeps the card and the sheet in lockstep on every code path —
|
||||
* page reload, navigation, live stream, post-stream reversibility flip,
|
||||
* and explicit revert clicks.
|
||||
*/
|
||||
|
||||
export const ACTION_LOG_PAGE_SIZE = 200;
|
||||
|
||||
/** Stable react-query key for the per-thread action list. */
|
||||
export function agentActionsQueryKey(threadId: number | null) {
|
||||
return threadId !== null
|
||||
? (["agent-actions", threadId] as const)
|
||||
: (["agent-actions", "none"] as const);
|
||||
}
|
||||
|
||||
/** Subset of the SSE ``data-action-log`` payload we care about. */
|
||||
export interface ActionLogSseEvent {
|
||||
id: number;
|
||||
lc_tool_call_id: string | null;
|
||||
chat_turn_id: string | null;
|
||||
tool_name: string;
|
||||
reversible: boolean;
|
||||
reverse_descriptor_present: boolean;
|
||||
error: boolean;
|
||||
created_at: string | null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Append or upsert a freshly-emitted ``AgentActionLog`` row into the
|
||||
* thread-scoped query cache.
|
||||
*
|
||||
* The SSE payload is a strict subset of ``AgentAction``; missing
|
||||
* fields (``args``, ``reverse_descriptor``, ``user_id``) are filled
|
||||
* with ``null`` placeholders. The next refetch (sheet open, user
|
||||
* focus, route stale) backfills them — but the inline Revert button
|
||||
* only reads the fields the SSE payload carries, so it lights up
|
||||
* immediately.
|
||||
*/
|
||||
export function applyActionLogSse(
|
||||
queryClient: QueryClient,
|
||||
threadId: number,
|
||||
searchSpaceId: number,
|
||||
event: ActionLogSseEvent
|
||||
): void {
|
||||
dbg("applyActionLogSse: incoming SSE event", {
|
||||
threadId,
|
||||
searchSpaceId,
|
||||
event,
|
||||
});
|
||||
queryClient.setQueryData<AgentActionListResponse>(
|
||||
agentActionsQueryKey(threadId),
|
||||
(prev) => {
|
||||
const placeholder: AgentAction = {
|
||||
id: event.id,
|
||||
thread_id: threadId,
|
||||
user_id: null,
|
||||
search_space_id: searchSpaceId,
|
||||
tool_name: event.tool_name,
|
||||
args: null,
|
||||
result_id: null,
|
||||
reversible: event.reversible,
|
||||
reverse_descriptor: event.reverse_descriptor_present ? {} : null,
|
||||
error: event.error ? {} : null,
|
||||
reverse_of: null,
|
||||
reverted_by_action_id: null,
|
||||
is_revert_action: false,
|
||||
tool_call_id: event.lc_tool_call_id,
|
||||
chat_turn_id: event.chat_turn_id,
|
||||
created_at: event.created_at ?? new Date().toISOString(),
|
||||
};
|
||||
if (!prev) {
|
||||
return {
|
||||
items: [placeholder],
|
||||
total: 1,
|
||||
page: 0,
|
||||
page_size: ACTION_LOG_PAGE_SIZE,
|
||||
has_more: false,
|
||||
};
|
||||
}
|
||||
const existingIdx = prev.items.findIndex((a) => a.id === event.id);
|
||||
if (existingIdx >= 0) {
|
||||
const merged = [...prev.items];
|
||||
const existing = merged[existingIdx];
|
||||
if (existing) {
|
||||
merged[existingIdx] = {
|
||||
...existing,
|
||||
reversible: event.reversible,
|
||||
tool_call_id: event.lc_tool_call_id ?? existing.tool_call_id,
|
||||
chat_turn_id: event.chat_turn_id ?? existing.chat_turn_id,
|
||||
};
|
||||
}
|
||||
dbg("applyActionLogSse: merged into existing entry", {
|
||||
id: event.id,
|
||||
tool_call_id: merged[existingIdx]?.tool_call_id,
|
||||
reversible: merged[existingIdx]?.reversible,
|
||||
});
|
||||
return { ...prev, items: merged };
|
||||
}
|
||||
dbg("applyActionLogSse: appended new placeholder", {
|
||||
id: event.id,
|
||||
tool_call_id: placeholder.tool_call_id,
|
||||
tool_name: placeholder.tool_name,
|
||||
reversible: placeholder.reversible,
|
||||
cacheSizeAfter: prev.items.length + 1,
|
||||
});
|
||||
// REST returns newest-first — keep that ordering when
|
||||
// the server eventually refetches by prepending.
|
||||
return {
|
||||
...prev,
|
||||
items: [placeholder, ...prev.items],
|
||||
total: prev.total + 1,
|
||||
};
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Apply a post-SAVEPOINT reversibility flip
|
||||
* (``data-action-log-updated`` SSE event) to the cache.
|
||||
*/
|
||||
export function applyActionLogUpdatedSse(
|
||||
queryClient: QueryClient,
|
||||
threadId: number,
|
||||
id: number,
|
||||
reversible: boolean
|
||||
): void {
|
||||
dbg("applyActionLogUpdatedSse: reversibility flip", {
|
||||
threadId,
|
||||
id,
|
||||
reversible,
|
||||
});
|
||||
queryClient.setQueryData<AgentActionListResponse>(
|
||||
agentActionsQueryKey(threadId),
|
||||
(prev) => {
|
||||
if (!prev) {
|
||||
dbg("applyActionLogUpdatedSse: NO prev cache for thread; flip dropped", {
|
||||
threadId,
|
||||
id,
|
||||
});
|
||||
return prev;
|
||||
}
|
||||
let mutated = false;
|
||||
const items = prev.items.map((a) => {
|
||||
if (a.id !== id) return a;
|
||||
mutated = true;
|
||||
return { ...a, reversible };
|
||||
});
|
||||
if (!mutated) {
|
||||
dbg("applyActionLogUpdatedSse: id not in cache; flip dropped", {
|
||||
threadId,
|
||||
id,
|
||||
cacheSize: prev.items.length,
|
||||
cacheIds: prev.items.map((a) => a.id),
|
||||
});
|
||||
}
|
||||
return mutated ? { ...prev, items } : prev;
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Optimistically mark ``id`` as reverted.
|
||||
*
|
||||
* Used by the inline / per-turn Revert button immediately after the
|
||||
* server returns success so the UI flips to "Reverted" without
|
||||
* waiting for a refetch. ``newActionId`` is the id of the new
|
||||
* ``is_revert_action`` row the server inserted; pass ``null`` if the
|
||||
* server didn't return it.
|
||||
*/
|
||||
export function markActionRevertedInCache(
|
||||
queryClient: QueryClient,
|
||||
threadId: number,
|
||||
id: number,
|
||||
newActionId: number | null
|
||||
): void {
|
||||
queryClient.setQueryData<AgentActionListResponse>(
|
||||
agentActionsQueryKey(threadId),
|
||||
(prev) => {
|
||||
if (!prev) return prev;
|
||||
let mutated = false;
|
||||
const items = prev.items.map((a) => {
|
||||
if (a.id !== id) return a;
|
||||
mutated = true;
|
||||
// ``-1`` is a sentinel meaning "we know it was reverted
|
||||
// but the server didn't tell us the new row's id".
|
||||
return {
|
||||
...a,
|
||||
reverted_by_action_id: newActionId ?? -1,
|
||||
};
|
||||
});
|
||||
return mutated ? { ...prev, items } : prev;
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Apply a batch of revert results (per-turn revert response) to the
|
||||
* cache. Anything in the ``reverted`` / ``already_reverted`` buckets
|
||||
* gets its ``reverted_by_action_id`` set; other rows are left alone.
|
||||
*/
|
||||
export function applyRevertTurnResultsToCache(
|
||||
queryClient: QueryClient,
|
||||
threadId: number,
|
||||
entries: Array<{ id: number; newActionId: number | null }>
|
||||
): void {
|
||||
if (entries.length === 0) return;
|
||||
queryClient.setQueryData<AgentActionListResponse>(
|
||||
agentActionsQueryKey(threadId),
|
||||
(prev) => {
|
||||
if (!prev) return prev;
|
||||
const lookup = new Map(entries.map((e) => [e.id, e.newActionId]));
|
||||
let mutated = false;
|
||||
const items = prev.items.map((a) => {
|
||||
if (!lookup.has(a.id)) return a;
|
||||
mutated = true;
|
||||
const newActionId = lookup.get(a.id) ?? null;
|
||||
return { ...a, reverted_by_action_id: newActionId ?? -1 };
|
||||
});
|
||||
return mutated ? { ...prev, items } : prev;
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Read-side hook used by the card, the turn button, the sheet, and
|
||||
* the edit-from-position pre-flight.
|
||||
*
|
||||
* Returns the raw query state plus convenience selectors so consumers
|
||||
* don't reach into ``data.items`` directly. ``enabled`` is the only
|
||||
* knob — pass ``false`` to keep the query dormant when the consumer
|
||||
* doesn't yet have a thread id.
|
||||
*/
|
||||
export function useAgentActionsQuery(
|
||||
threadId: number | null,
|
||||
options: { enabled?: boolean } = {}
|
||||
) {
|
||||
const enabled = (options.enabled ?? true) && threadId !== null;
|
||||
const query = useQuery({
|
||||
queryKey: agentActionsQueryKey(threadId),
|
||||
queryFn: async () => {
|
||||
dbg("useAgentActionsQuery: REST fetch START", {
|
||||
threadId,
|
||||
pageSize: ACTION_LOG_PAGE_SIZE,
|
||||
});
|
||||
const res = await agentActionsApiService.listForThread(threadId as number, {
|
||||
page: 0,
|
||||
pageSize: ACTION_LOG_PAGE_SIZE,
|
||||
});
|
||||
dbg("useAgentActionsQuery: REST fetch DONE", {
|
||||
threadId,
|
||||
total: res.total,
|
||||
returned: res.items.length,
|
||||
items: res.items.map((a) => ({
|
||||
id: a.id,
|
||||
tool_name: a.tool_name,
|
||||
tool_call_id: a.tool_call_id,
|
||||
reversible: a.reversible,
|
||||
reverted_by_action_id: a.reverted_by_action_id,
|
||||
is_revert_action: a.is_revert_action,
|
||||
})),
|
||||
});
|
||||
return res;
|
||||
},
|
||||
enabled,
|
||||
staleTime: 15 * 1000,
|
||||
});
|
||||
|
||||
const items = useMemo(() => query.data?.items ?? [], [query.data]);
|
||||
|
||||
// Index ``items`` once per change so the lookups below are O(1)
|
||||
// instead of O(N) per card per render. With the cache sized to 200
|
||||
// rows and many tool cards visible at once, the unindexed scan was
|
||||
// the hottest path on every assistant text-delta. (Vercel React
|
||||
// rule ``js-index-maps`` / ``js-set-map-lookups``.)
|
||||
const byToolCallId = useMemo(() => {
|
||||
const m = new Map<string, AgentAction>();
|
||||
for (const a of items) {
|
||||
if (a.tool_call_id) m.set(a.tool_call_id, a);
|
||||
}
|
||||
return m;
|
||||
}, [items]);
|
||||
|
||||
// Pre-grouped + pre-sorted (oldest-first, the order the agent
|
||||
// actually executed them in) so the (chat_turn_id, tool_name,
|
||||
// position) fallback in ``tool-fallback.tsx`` is also O(1) per
|
||||
// card. Excludes ``is_revert_action`` rows so the position index
|
||||
// matches the agent's original execution order.
|
||||
const byTurnAndTool = useMemo(() => {
|
||||
const m = new Map<string, AgentAction[]>();
|
||||
for (const a of items) {
|
||||
if (!a.chat_turn_id || a.is_revert_action) continue;
|
||||
const key = `${a.chat_turn_id}::${a.tool_name}`;
|
||||
const bucket = m.get(key);
|
||||
if (bucket) bucket.push(a);
|
||||
else m.set(key, [a]);
|
||||
}
|
||||
for (const bucket of m.values()) {
|
||||
bucket.sort(
|
||||
(a, b) =>
|
||||
new Date(a.created_at).getTime() - new Date(b.created_at).getTime()
|
||||
);
|
||||
}
|
||||
return m;
|
||||
}, [items]);
|
||||
|
||||
// Snapshot the cache shape when its size changes — easiest way to
|
||||
// spot when the cache is empty or stale at the moment a card
|
||||
// mounts. Tracked on a ref so we don't re-run the diff on
|
||||
// reference-equal cache reads.
|
||||
const lastSnapshotRef = useRef<{ threadId: number | null; size: number } | null>(null);
|
||||
useEffect(() => {
|
||||
const last = lastSnapshotRef.current;
|
||||
if (!last || last.threadId !== threadId || last.size !== items.length) {
|
||||
dbg("useAgentActionsQuery: cache snapshot", {
|
||||
threadId,
|
||||
enabled,
|
||||
itemCount: items.length,
|
||||
itemKeys: items.slice(0, 8).map((a) => ({
|
||||
id: a.id,
|
||||
tool_name: a.tool_name,
|
||||
tool_call_id: a.tool_call_id,
|
||||
chat_turn_id: a.chat_turn_id,
|
||||
reversible: a.reversible,
|
||||
})),
|
||||
});
|
||||
lastSnapshotRef.current = { threadId, size: items.length };
|
||||
}
|
||||
}, [threadId, enabled, items]);
|
||||
|
||||
const findByToolCallId = useCallback(
|
||||
(toolCallId: string | null | undefined): AgentAction | null => {
|
||||
if (!toolCallId) return null;
|
||||
const found = byToolCallId.get(toolCallId) ?? null;
|
||||
if (!found && items.length > 0) {
|
||||
dbg("findByToolCallId: MISS", {
|
||||
queriedToolCallId: toolCallId,
|
||||
itemCount: items.length,
|
||||
availableToolCallIds: Array.from(byToolCallId.keys()),
|
||||
});
|
||||
}
|
||||
return found;
|
||||
},
|
||||
[byToolCallId, items.length]
|
||||
);
|
||||
|
||||
const findByChatTurnId = useCallback(
|
||||
(chatTurnId: string | null | undefined): AgentAction[] => {
|
||||
if (!chatTurnId) return [];
|
||||
// Per-turn aggregation is uncommon enough (only the
|
||||
// "Revert turn" button uses it) that re-scanning is fine;
|
||||
// indexing it would just bloat memory.
|
||||
return items.filter((a) => a.chat_turn_id === chatTurnId);
|
||||
},
|
||||
[items]
|
||||
);
|
||||
|
||||
const findByChatTurnAndTool = useCallback(
|
||||
(
|
||||
chatTurnId: string | null | undefined,
|
||||
toolName: string | null | undefined
|
||||
): AgentAction[] => {
|
||||
if (!chatTurnId || !toolName) return [];
|
||||
return byTurnAndTool.get(`${chatTurnId}::${toolName}`) ?? [];
|
||||
},
|
||||
[byTurnAndTool]
|
||||
);
|
||||
|
||||
return {
|
||||
...query,
|
||||
items,
|
||||
findByToolCallId,
|
||||
findByChatTurnId,
|
||||
findByChatTurnAndTool,
|
||||
};
|
||||
}
|
||||
|
|
@ -31,6 +31,14 @@ export function useMessagesSync(
|
|||
content: msg.content,
|
||||
author_id: msg.authorId ?? null,
|
||||
created_at: new Date(msg.createdAt).toISOString(),
|
||||
// Forward the per-turn correlation id so post-stream Zero
|
||||
// re-syncs preserve ``metadata.custom.chatTurnId`` on the
|
||||
// converted ``ThreadMessageLike``. Without this the inline
|
||||
// Revert button's ``(chat_turn_id, tool_name, position)``
|
||||
// fallback breaks the moment Zero overwrites the messages
|
||||
// state after a live stream completes (see
|
||||
// ``handleSyncedMessagesUpdate`` in the chat page).
|
||||
turn_id: msg.turnId ?? null,
|
||||
}));
|
||||
|
||||
onMessagesUpdateRef.current(mapped);
|
||||
|
|
|
|||
|
|
@ -16,6 +16,23 @@ export type ContentPart =
|
|||
toolName: string;
|
||||
args: Record<string, unknown>;
|
||||
result?: unknown;
|
||||
/**
|
||||
* Live / finalized JSON text for the tool's input arguments.
|
||||
*
|
||||
* - During streaming: accumulated partial JSON text from
|
||||
* ``tool-input-delta`` events (may be invalid JSON
|
||||
* mid-stream). assistant-ui's argsText parser tolerates
|
||||
* invalid JSON gracefully (changelog 0.7.32 / 0.7.78).
|
||||
* - On completion (``tool-input-available``): replaced with
|
||||
* ``JSON.stringify(input, null, 2)`` so the post-stream
|
||||
* card renders pretty-printed JSON instead of the
|
||||
* model's possibly-fragmented formatting.
|
||||
*
|
||||
* Per assistant-ui ``ThreadMessageLike`` precedence
|
||||
* (changelog 0.11.6 ``d318c83``), when ``argsText`` is
|
||||
* supplied it wins over ``JSON.stringify(args)``.
|
||||
*/
|
||||
argsText?: string;
|
||||
/**
|
||||
* Authoritative LangChain ``tool_call.id`` propagated by the backend
|
||||
* via ``langchainToolCallId`` on tool-input-start/available and
|
||||
|
|
@ -282,12 +299,22 @@ export function findToolCallIdByLcId(
|
|||
export function updateToolCall(
|
||||
state: ContentPartsState,
|
||||
toolCallId: string,
|
||||
update: { args?: Record<string, unknown>; result?: unknown; langchainToolCallId?: string }
|
||||
update: {
|
||||
args?: Record<string, unknown>;
|
||||
argsText?: string;
|
||||
result?: unknown;
|
||||
langchainToolCallId?: string;
|
||||
}
|
||||
): void {
|
||||
const index = state.toolCallIndices.get(toolCallId);
|
||||
if (index !== undefined && state.contentParts[index]?.type === "tool-call") {
|
||||
const tc = state.contentParts[index] as ContentPart & { type: "tool-call" };
|
||||
if (update.args) tc.args = update.args;
|
||||
// ``!== undefined`` (NOT a truthy check): an explicit empty
|
||||
// string CAN clear, and a finalization with
|
||||
// ``JSON.stringify({}, null, 2) === "{}"`` (truthy but
|
||||
// represents an empty-input call) still applies.
|
||||
if (update.argsText !== undefined) tc.argsText = update.argsText;
|
||||
if (update.result !== undefined) tc.result = update.result;
|
||||
// Only backfill langchainToolCallId if not already set — the
|
||||
// authoritative ``on_tool_end`` value should override an earlier
|
||||
|
|
@ -299,6 +326,25 @@ export function updateToolCall(
|
|||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Append a streamed args-delta chunk to the active tool call's
|
||||
* ``argsText``. No-ops when no card has been registered yet for the
|
||||
* given ``toolCallId`` (the matching ``tool-input-start`` either lost
|
||||
* the wire race or this id never had a card — either way the deltas
|
||||
* have nowhere safe to land).
|
||||
*/
|
||||
export function appendToolInputDelta(
|
||||
state: ContentPartsState,
|
||||
toolCallId: string,
|
||||
delta: string
|
||||
): void {
|
||||
const idx = state.toolCallIndices.get(toolCallId);
|
||||
if (idx === undefined) return;
|
||||
const tc = state.contentParts[idx];
|
||||
if (tc?.type !== "tool-call") return;
|
||||
tc.argsText = (tc.argsText ?? "") + delta;
|
||||
}
|
||||
|
||||
function _hasInterruptResult(part: ContentPart): boolean {
|
||||
if (part.type !== "tool-call") return false;
|
||||
const r = (part as { result?: unknown }).result;
|
||||
|
|
@ -371,6 +417,18 @@ export type SSEEvent =
|
|||
/** Authoritative LangChain ``tool_call.id``. Optional. */
|
||||
langchainToolCallId?: string;
|
||||
}
|
||||
| {
|
||||
/**
|
||||
* Live tool-call argument delta. Concatenated into
|
||||
* ``argsText`` on the matching ``tool-call`` content part
|
||||
* by ``appendToolInputDelta``. parity_v2 only — the legacy
|
||||
* code path emits ``tool-input-available`` without prior
|
||||
* deltas.
|
||||
*/
|
||||
type: "tool-input-delta";
|
||||
toolCallId: string;
|
||||
inputTextDelta: string;
|
||||
}
|
||||
| {
|
||||
type: "tool-input-available";
|
||||
toolCallId: string;
|
||||
|
|
|
|||
|
|
@ -8,6 +8,13 @@ export const newChatMessageTable = table("new_chat_messages")
|
|||
threadId: number().from("thread_id"),
|
||||
authorId: string().optional().from("author_id"),
|
||||
createdAt: number().from("created_at"),
|
||||
// Per-turn correlation id sourced from ``configurable.turn_id``
|
||||
// at streaming time. Required by the inline Revert button's
|
||||
// (chat_turn_id, tool_name, position) fallback in tool-fallback.tsx
|
||||
// — without it the live-collab Zero sync would clobber the
|
||||
// metadata we set during streaming and the button would vanish
|
||||
// the moment Zero re-syncs after the stream finishes.
|
||||
turnId: string().optional().from("turn_id"),
|
||||
})
|
||||
.primaryKey("id");
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue