mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-12 09:12:40 +02:00
433 lines
14 KiB
TypeScript
433 lines
14 KiB
TypeScript
import {
|
|
makeAssistantDataUI,
|
|
type ToolCallMessagePartComponent,
|
|
useAuiState,
|
|
} from "@assistant-ui/react";
|
|
import { ChevronRightIcon } from "lucide-react";
|
|
import type { FC } from "react";
|
|
import { useCallback, useEffect, useMemo, useState } from "react";
|
|
import { TOOLS_BY_NAME, TOOLS_FALLBACK } from "@/components/assistant-ui/assistant-message";
|
|
import { ChainOfThoughtItem } from "@/components/prompt-kit/chain-of-thought";
|
|
import { TextShimmerLoader } from "@/components/prompt-kit/loader";
|
|
import { HitlRenderTargetProvider, isInterruptResult } from "@/lib/hitl";
|
|
import { cn } from "@/lib/utils";
|
|
|
|
export interface ThinkingStep {
|
|
id: string;
|
|
title: string;
|
|
items: string[];
|
|
status: "pending" | "in_progress" | "completed";
|
|
/**
|
|
* Optional relay metadata forwarded from ``data-thinking-step`` SSE
|
|
* (e.g. ``spanId`` set by ``AgentEventRelayState.span_metadata_if_active``).
|
|
* Steps under an open delegating ``task`` carry ``metadata.spanId`` and are
|
|
* grouped under the preceding parent (``task`` step) as indented children.
|
|
*/
|
|
metadata?: Record<string, unknown>;
|
|
}
|
|
|
|
/**
|
|
* Per-step info joined from the assistant message ``tool-call`` parts via
|
|
* the shared ``metadata.thinkingStepId`` correlation
|
|
* (set on the server in ``AgentEventRelayState.tool_activity_metadata``).
|
|
*
|
|
* Carries enough of the part to:
|
|
* - identify the opening ``task`` step and substitute the subagent display
|
|
* name on the parent header (uses ``toolName`` and ``args``);
|
|
* - render the matching tool component inline under the step row when the
|
|
* card's result is an HITL interrupt (uses ``toolCallId``, ``argsText``,
|
|
* ``result``, ``langchainToolCallId``).
|
|
*/
|
|
interface StepToolInfo {
|
|
toolCallId: string;
|
|
toolName: string;
|
|
args: Record<string, unknown>;
|
|
argsText?: string;
|
|
result?: unknown;
|
|
langchainToolCallId?: string;
|
|
}
|
|
|
|
export type ThinkingStepToolInfoMap = ReadonlyMap<string, StepToolInfo>;
|
|
|
|
/**
|
|
* Build ``thinkingStepId → StepToolInfo`` from message content. Used to
|
|
* - identify the opening ``task`` step (parent header, never indents) without
|
|
* relying on the human-readable title;
|
|
* - render the parent's display title from ``args.subagent_type`` instead of
|
|
* the generic "Task" copy;
|
|
* - mount the matching tool-call card inline under a step row when the
|
|
* result is an HITL interrupt (see ``TimelineHitlCard``).
|
|
*/
|
|
export function buildThinkingStepToolInfo(
|
|
content: readonly unknown[] | undefined
|
|
): ThinkingStepToolInfoMap {
|
|
const m = new Map<string, StepToolInfo>();
|
|
if (!content) return m;
|
|
for (const part of content) {
|
|
if (!part || typeof part !== "object") continue;
|
|
const o = part as {
|
|
type?: string;
|
|
toolCallId?: string;
|
|
toolName?: string;
|
|
args?: Record<string, unknown>;
|
|
argsText?: string;
|
|
result?: unknown;
|
|
langchainToolCallId?: string;
|
|
metadata?: Record<string, unknown>;
|
|
};
|
|
if (o.type !== "tool-call" || !o.toolName || !o.toolCallId) continue;
|
|
const tid = o.metadata?.thinkingStepId;
|
|
if (typeof tid === "string" && tid.trim().length > 0) {
|
|
m.set(tid, {
|
|
toolCallId: o.toolCallId,
|
|
toolName: o.toolName,
|
|
args: o.args ?? {},
|
|
argsText: o.argsText,
|
|
result: o.result,
|
|
langchainToolCallId: o.langchainToolCallId,
|
|
});
|
|
}
|
|
}
|
|
return m;
|
|
}
|
|
|
|
function asNonEmptyString(v: unknown): string | undefined {
|
|
return typeof v === "string" && v.trim().length > 0 ? v.trim() : undefined;
|
|
}
|
|
|
|
function titleCaseSubagent(raw: string): string {
|
|
// "notion" → "Notion", "doc_research" → "Doc Research".
|
|
return raw
|
|
.split(/[\s_-]+/)
|
|
.filter(Boolean)
|
|
.map((part) => part.charAt(0).toUpperCase() + part.slice(1))
|
|
.join(" ");
|
|
}
|
|
|
|
/**
|
|
* Display title for a step. For the opening ``task`` step we substitute the
|
|
* subagent type from the matching tool-call args (e.g. ``"Notion"`` instead of
|
|
* the generic ``"Task"``). Falls back to the step's own title if the tool-call
|
|
* hasn't streamed in yet.
|
|
*/
|
|
function resolveDisplayTitle(step: ThinkingStep, info: StepToolInfo | undefined): string {
|
|
if (info?.toolName === "task") {
|
|
const subagent = asNonEmptyString(info.args?.subagent_type);
|
|
if (subagent) return titleCaseSubagent(subagent);
|
|
}
|
|
return step.title;
|
|
}
|
|
|
|
function isDelegatedChild(step: ThinkingStep, info: StepToolInfo | undefined): boolean {
|
|
const sid = asNonEmptyString(step.metadata?.spanId);
|
|
if (!sid) return false;
|
|
// The opening ``task`` step also carries ``spanId`` (it owns the span) but
|
|
// must render as the parent header. Prefer the joined ``toolName`` (set by
|
|
// ``buildThinkingStepToolInfo`` from ``tool-call.metadata.thinkingStepId``).
|
|
// Fall back to the title heuristic when no tool-call is matched — happens
|
|
// for messages persisted before ``thinkingStepId`` shipped, and briefly
|
|
// during streaming if the ``tool-input-start`` frame hasn't been processed
|
|
// yet for some reason.
|
|
if (info) return info.toolName !== "task";
|
|
return step.title !== "Task";
|
|
}
|
|
|
|
interface StepGroup {
|
|
parent: ThinkingStep;
|
|
children: ThinkingStep[];
|
|
}
|
|
|
|
/**
|
|
* Group consecutive delegated child steps under the preceding parent step.
|
|
* If the very first step is a child (no parent yet seen), it's promoted to a
|
|
* parent so it still renders — defensive only, real flows always start with a
|
|
* parent step.
|
|
*/
|
|
const EMPTY_STEP_TOOL_INFO: ThinkingStepToolInfoMap = new Map();
|
|
|
|
function groupSteps(
|
|
steps: readonly ThinkingStep[],
|
|
stepToolInfo: ThinkingStepToolInfoMap
|
|
): StepGroup[] {
|
|
const groups: StepGroup[] = [];
|
|
for (const step of steps) {
|
|
if (isDelegatedChild(step, stepToolInfo.get(step.id)) && groups.length > 0) {
|
|
groups[groups.length - 1].children.push(step);
|
|
} else {
|
|
groups.push({ parent: step, children: [] });
|
|
}
|
|
}
|
|
return groups;
|
|
}
|
|
|
|
const StepBody: FC<{
|
|
step: ThinkingStep;
|
|
status: "pending" | "in_progress" | "completed";
|
|
displayTitle: string;
|
|
}> = ({ step, status, displayTitle }) => (
|
|
<div className="min-w-0">
|
|
<div
|
|
className={cn(
|
|
"text-sm leading-5",
|
|
status === "in_progress" && "text-foreground font-medium",
|
|
status === "completed" && "text-muted-foreground",
|
|
status === "pending" && "text-muted-foreground/60"
|
|
)}
|
|
>
|
|
{displayTitle}
|
|
</div>
|
|
|
|
{step.items && step.items.length > 0 && (
|
|
<div className="mt-1 space-y-0.5">
|
|
{step.items.map((item) => (
|
|
<ChainOfThoughtItem key={`${step.id}-${item}`} className="text-xs">
|
|
{item}
|
|
</ChainOfThoughtItem>
|
|
))}
|
|
</div>
|
|
)}
|
|
</div>
|
|
);
|
|
|
|
/**
|
|
* Mount the same tool-call UI used in the message body, but inside the
|
|
* chain-of-thought timeline. The body copy returns ``null`` (see
|
|
* ``withHitlInTimeline`` in ``lib/hitl/render-target``), so the card
|
|
* effectively moves from the body to the timeline for the lifetime of the
|
|
* interrupt (pending → processing → complete / rejected).
|
|
*
|
|
* ``metadata`` is intentionally omitted from the props we forward — the
|
|
* step row already provides any indentation it needs, so we don't want
|
|
* ``withDelegationSpanIndent`` to add a second indent + border on top.
|
|
*
|
|
* ``status`` is a placeholder (HITL UIs read only ``args`` + ``result``)
|
|
* so we don't need to mirror assistant-ui's runtime status object here.
|
|
*/
|
|
const TimelineHitlCard: FC<{ info: StepToolInfo }> = ({ info }) => {
|
|
const Comp =
|
|
(TOOLS_BY_NAME as Record<string, ToolCallMessagePartComponent | undefined>)[info.toolName] ??
|
|
TOOLS_FALLBACK;
|
|
const props = {
|
|
toolCallId: info.toolCallId,
|
|
toolName: info.toolName,
|
|
args: info.args,
|
|
argsText: info.argsText,
|
|
result: info.result,
|
|
langchainToolCallId: info.langchainToolCallId,
|
|
status: { type: "complete" } as const,
|
|
};
|
|
return (
|
|
<HitlRenderTargetProvider value="timeline">
|
|
{/* biome-ignore lint/suspicious/noExplicitAny: ToolCallMessagePartProps requires
|
|
runtime-only fields (addResult, resume, MessagePartState) we don't have when
|
|
re-rendering manually; HITL components only read args + result. */}
|
|
<Comp {...(props as any)} />
|
|
</HitlRenderTargetProvider>
|
|
);
|
|
};
|
|
|
|
function hitlInterruptInfo(info: StepToolInfo | undefined): StepToolInfo | undefined {
|
|
return info && isInterruptResult(info.result) ? info : undefined;
|
|
}
|
|
|
|
/**
|
|
* Chain of thought display component - single collapsible dropdown design.
|
|
*
|
|
* ``stepToolInfo`` joins each step (by ``thinkingStepId``) to its ``tool-call``
|
|
* part so we can:
|
|
* - replace the generic ``"Task"`` title with the real subagent name
|
|
* (``args.subagent_type``) on the parent header;
|
|
* - decide parent-vs-child purely from the matched ``toolName`` instead of
|
|
* relying on the displayed title.
|
|
*/
|
|
export const ThinkingStepsDisplay: FC<{
|
|
steps: ThinkingStep[];
|
|
isThreadRunning?: boolean;
|
|
stepToolInfo?: ThinkingStepToolInfoMap;
|
|
}> = ({ steps, isThreadRunning = true, stepToolInfo }) => {
|
|
const toolInfo = stepToolInfo ?? EMPTY_STEP_TOOL_INFO;
|
|
const getEffectiveStatus = useCallback(
|
|
(step: ThinkingStep): "pending" | "in_progress" | "completed" => {
|
|
if (step.status === "in_progress" && !isThreadRunning) {
|
|
return "completed";
|
|
}
|
|
return step.status;
|
|
},
|
|
[isThreadRunning]
|
|
);
|
|
|
|
const inProgressStep = steps.find((s) => getEffectiveStatus(s) === "in_progress");
|
|
const inProgressDisplayTitle = inProgressStep
|
|
? resolveDisplayTitle(inProgressStep, toolInfo.get(inProgressStep.id))
|
|
: undefined;
|
|
const allCompleted =
|
|
steps.length > 0 &&
|
|
!isThreadRunning &&
|
|
steps.every((s) => getEffectiveStatus(s) === "completed");
|
|
const isProcessing = isThreadRunning && !allCompleted;
|
|
const [isOpen, setIsOpen] = useState(() => isProcessing);
|
|
|
|
useEffect(() => {
|
|
if (isProcessing) {
|
|
setIsOpen(true);
|
|
return;
|
|
}
|
|
|
|
if (allCompleted) {
|
|
setIsOpen(false);
|
|
}
|
|
}, [allCompleted, isProcessing]);
|
|
|
|
const groups = useMemo(() => groupSteps(steps, toolInfo), [steps, toolInfo]);
|
|
|
|
if (steps.length === 0) return null;
|
|
|
|
const getHeaderText = () => {
|
|
if (allCompleted) {
|
|
return "Reviewed";
|
|
}
|
|
if (inProgressDisplayTitle) {
|
|
return inProgressDisplayTitle;
|
|
}
|
|
if (isProcessing) {
|
|
return "Processing";
|
|
}
|
|
return "Reviewed";
|
|
};
|
|
|
|
return (
|
|
<div className="mx-auto w-full max-w-(--thread-max-width) px-2 py-2">
|
|
<div className="rounded-lg">
|
|
<button
|
|
type="button"
|
|
onClick={() => setIsOpen((prev) => !prev)}
|
|
className={cn(
|
|
"flex w-full items-center gap-1.5 text-left text-sm transition-colors",
|
|
"text-muted-foreground hover:text-foreground"
|
|
)}
|
|
>
|
|
{isProcessing ? (
|
|
<TextShimmerLoader text={getHeaderText()} size="sm" />
|
|
) : (
|
|
<span>{getHeaderText()}</span>
|
|
)}
|
|
|
|
<ChevronRightIcon
|
|
className={cn("size-4 transition-transform duration-200", isOpen && "rotate-90")}
|
|
/>
|
|
</button>
|
|
|
|
<div
|
|
className={cn(
|
|
"grid transition-[grid-template-rows] duration-300 ease-out",
|
|
isOpen ? "grid-rows-[1fr]" : "grid-rows-[0fr]"
|
|
)}
|
|
>
|
|
<div className="overflow-hidden">
|
|
<div className="mt-3 pl-1">
|
|
{groups.map((group, groupIndex) => {
|
|
const isLastGroup = groupIndex === groups.length - 1;
|
|
const parentStatus = getEffectiveStatus(group.parent);
|
|
const parentInfo = toolInfo.get(group.parent.id);
|
|
const parentTitle = resolveDisplayTitle(group.parent, parentInfo);
|
|
const hasChildren = group.children.length > 0;
|
|
// Parent dots are connected by a vertical line that runs through
|
|
// any indented children (their column has no dot, so the line
|
|
// passes cleanly behind them) and overshoots by ~15px to reach
|
|
// the next group's dot center (top-[15px]).
|
|
const showParentLine = !isLastGroup;
|
|
|
|
return (
|
|
<div key={group.parent.id} className="relative flex gap-3">
|
|
<div className="relative flex flex-col items-center w-2 self-stretch">
|
|
{showParentLine && (
|
|
<div className="absolute left-1/2 top-[15px] -bottom-[15px] w-px -translate-x-1/2 bg-muted-foreground/30" />
|
|
)}
|
|
<div className="relative z-10 mt-[7px] flex shrink-0 items-center justify-center">
|
|
{parentStatus === "in_progress" ? (
|
|
<span className="relative flex size-2">
|
|
<span className="absolute inline-flex size-full animate-ping rounded-full bg-primary/60" />
|
|
<span className="relative inline-flex size-2 rounded-full bg-primary" />
|
|
</span>
|
|
) : (
|
|
<span className="size-2 rounded-full bg-muted-foreground/30" />
|
|
)}
|
|
</div>
|
|
</div>
|
|
|
|
<div className="flex-1 min-w-0 pb-4">
|
|
<StepBody
|
|
step={group.parent}
|
|
status={parentStatus}
|
|
displayTitle={parentTitle}
|
|
/>
|
|
|
|
{(() => {
|
|
const hitl = hitlInterruptInfo(parentInfo);
|
|
return hitl ? <TimelineHitlCard info={hitl} /> : null;
|
|
})()}
|
|
|
|
{hasChildren && (
|
|
<div className="mt-2 ml-3 space-y-2">
|
|
{group.children.map((child) => {
|
|
const childInfo = toolInfo.get(child.id);
|
|
const childHitl = hitlInterruptInfo(childInfo);
|
|
return (
|
|
<div key={child.id}>
|
|
<StepBody
|
|
step={child}
|
|
status={getEffectiveStatus(child)}
|
|
displayTitle={resolveDisplayTitle(child, childInfo)}
|
|
/>
|
|
{childHitl && <TimelineHitlCard info={childHitl} />}
|
|
</div>
|
|
);
|
|
})}
|
|
</div>
|
|
)}
|
|
</div>
|
|
</div>
|
|
);
|
|
})}
|
|
</div>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
);
|
|
};
|
|
|
|
/**
|
|
* assistant-ui data UI component that renders thinking steps from message content.
|
|
* Registered globally via makeAssistantDataUI — renders inside MessagePrimitive.Parts
|
|
* at the position of the data part in the content array.
|
|
*/
|
|
function ThinkingStepsDataRenderer({ data }: { name: string; data: unknown }) {
|
|
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
|
|
const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false);
|
|
const isMessageStreaming = isThreadRunning && isLastMessage;
|
|
const content = useAuiState(({ message }) => message?.content);
|
|
|
|
const stepToolInfo = useMemo(
|
|
() => buildThinkingStepToolInfo(Array.isArray(content) ? content : undefined),
|
|
[content]
|
|
);
|
|
|
|
const steps = (data as { steps: ThinkingStep[] } | null)?.steps ?? [];
|
|
if (steps.length === 0) return null;
|
|
|
|
return (
|
|
<div className="mb-3 -mx-2 leading-normal">
|
|
<ThinkingStepsDisplay
|
|
steps={steps}
|
|
isThreadRunning={isMessageStreaming}
|
|
stepToolInfo={stepToolInfo}
|
|
/>
|
|
</div>
|
|
);
|
|
}
|
|
|
|
export const ThinkingStepsDataUI = makeAssistantDataUI({
|
|
name: "thinking-steps",
|
|
render: ThinkingStepsDataRenderer,
|
|
});
|