diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index d465e23f6..a201ece22 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -217,24 +217,11 @@ async def stream_new_chat( yield completion_event if just_finished_tool: - # We just finished a tool - don't create a step here, - # text will flow silently after tools. - # Clear the active step tracking. + # Clear the active step tracking - text flows without a dedicated step last_active_step_id = None last_active_step_title = "" last_active_step_items = [] just_finished_tool = False - else: - # Normal text generation (not after a tool) - gen_step_id = next_thinking_step_id() - last_active_step_id = gen_step_id - last_active_step_title = "Generating response" - last_active_step_items = [] - yield streaming_service.format_thinking_step( - step_id=gen_step_id, - title="Generating response", - status="in_progress", - ) current_text_id = streaming_service.generate_text_id() yield streaming_service.format_text_start(current_text_id) diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 57408391d..e3ee85dce 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -30,7 +30,7 @@ import { } from "lucide-react"; import Image from "next/image"; import Link from "next/link"; -import { type FC, useState, useRef, useCallback } from "react"; +import { type FC, useState, useRef, useCallback, useEffect } from "react"; import { useAtomValue } from "jotai"; import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms"; import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors"; @@ -94,35 +94,102 @@ function getStepIcon(status: "pending" | "in_progress" | "completed", title: str } /** - * Chain of thought display component + * Chain of thought display component with smart expand/collapse behavior */ -const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[] }> = ({ steps }) => { +const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolean }> = ({ steps, isThreadRunning = true }) => { + // Track which steps the user has manually toggled (overrides auto behavior) + const [manualOverrides, setManualOverrides] = useState>({}); + // Track previous step statuses to detect changes + const prevStatusesRef = useRef>({}); + + // Derive effective status: if thread stopped and step is in_progress, treat as completed + const getEffectiveStatus = (step: ThinkingStep): "pending" | "in_progress" | "completed" => { + if (step.status === "in_progress" && !isThreadRunning) { + return "completed"; // Thread was stopped, so mark as completed + } + return step.status; + }; + + // Check if any step is effectively in progress + const hasInProgressStep = steps.some(step => getEffectiveStatus(step) === "in_progress"); + + // Find the last completed step index (using effective status) + const lastCompletedIndex = steps + .map((s, i) => getEffectiveStatus(s) === "completed" ? i : -1) + .filter(i => i !== -1) + .pop(); + + // Clear manual overrides when a step's status changes + useEffect(() => { + const currentStatuses: Record = {}; + steps.forEach(step => { + currentStatuses[step.id] = step.status; + // If status changed, clear any manual override for this step + if (prevStatusesRef.current[step.id] && prevStatusesRef.current[step.id] !== step.status) { + setManualOverrides(prev => { + const next = { ...prev }; + delete next[step.id]; + return next; + }); + } + }); + prevStatusesRef.current = currentStatuses; + }, [steps]); + if (steps.length === 0) return null; + const getStepOpenState = (step: ThinkingStep, index: number): boolean => { + const effectiveStatus = getEffectiveStatus(step); + // If user has manually toggled, respect that + if (manualOverrides[step.id] !== undefined) { + return manualOverrides[step.id]; + } + // Auto behavior: open if in progress + if (effectiveStatus === "in_progress") { + return true; + } + // Auto behavior: keep last completed step open if no in-progress step + if (!hasInProgressStep && index === lastCompletedIndex) { + return true; + } + // Default: collapsed + return false; + }; + + const handleToggle = (stepId: string, currentOpen: boolean) => { + setManualOverrides(prev => ({ + ...prev, + [stepId]: !currentOpen, + })); + }; + return (
- {steps.map((step) => { - const icon = getStepIcon(step.status, step.title); + {steps.map((step, index) => { + const effectiveStatus = getEffectiveStatus(step); + const icon = getStepIcon(effectiveStatus, step.title); + const isOpen = getStepOpenState(step, index); return ( handleToggle(step.id, isOpen)} > {step.title} {step.items && step.items.length > 0 && ( - {step.items.map((item, index) => ( - + {step.items.map((item, idx) => ( + {item} ))} @@ -263,8 +330,8 @@ const ThreadWelcome: FC = () => { return (
- {/* Greeting positioned near the composer */} -
+ {/* Greeting positioned above the composer - fixed position */} +

{/** biome-ignore lint/a11y/noStaticElementInteractions: wrong lint error, this is a workaround to fix the lint error */}
{ {getTimeBasedGreeting(user?.email)}

- {/* Composer centered in the middle of the screen */} -
+ {/* Composer - top edge fixed, expands downward only */} +
@@ -525,12 +592,15 @@ const AssistantMessageInner: FC = () => { const messageId = useMessage((m) => m.id); const thinkingSteps = thinkingStepsMap.get(messageId) || []; + // Check if thread is still running (for stopping the spinner when cancelled) + const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning); + return ( <> {/* Show thinking steps BEFORE the text response */} {thinkingSteps.length > 0 && (
- +
)} diff --git a/surfsense_web/components/tool-ui/deepagent-thinking.tsx b/surfsense_web/components/tool-ui/deepagent-thinking.tsx index f0c03e570..8cd901e69 100644 --- a/surfsense_web/components/tool-ui/deepagent-thinking.tsx +++ b/surfsense_web/components/tool-ui/deepagent-thinking.tsx @@ -2,7 +2,7 @@ import { makeAssistantToolUI } from "@assistant-ui/react"; import { Brain, CheckCircle2, Loader2, Search, Sparkles } from "lucide-react"; -import { useMemo } from "react"; +import { useMemo, useState, useEffect, useRef } from "react"; import { ChainOfThought, ChainOfThoughtContent, @@ -61,13 +61,21 @@ function getStepIcon(status: "pending" | "in_progress" | "completed", title: str } /** - * Component to display a single thinking step + * Component to display a single thinking step with controlled open state */ -function ThinkingStepDisplay({ step }: { step: ThinkingStep }) { +function ThinkingStepDisplay({ + step, + isOpen, + onToggle +}: { + step: ThinkingStep; + isOpen: boolean; + onToggle: () => void; +}) { const icon = useMemo(() => getStepIcon(step.status, step.title), [step.status, step.title]); return ( - + >({}); + // Track previous step statuses to detect changes + const prevStatusesRef = useRef>({}); + + // Check if any step is currently in progress + const hasInProgressStep = steps.some(step => step.status === "in_progress"); + + // Find the last completed step index + const lastCompletedIndex = steps + .map((s, i) => s.status === "completed" ? i : -1) + .filter(i => i !== -1) + .pop(); + + // Clear manual overrides when a step's status changes + useEffect(() => { + const currentStatuses: Record = {}; + steps.forEach(step => { + currentStatuses[step.id] = step.status; + // If status changed, clear any manual override for this step + if (prevStatusesRef.current[step.id] && prevStatusesRef.current[step.id] !== step.status) { + setManualOverrides(prev => { + const next = { ...prev }; + delete next[step.id]; + return next; + }); + } + }); + prevStatusesRef.current = currentStatuses; + }, [steps]); + + const getStepOpenState = (step: ThinkingStep, index: number): boolean => { + // If user has manually toggled, respect that + if (manualOverrides[step.id] !== undefined) { + return manualOverrides[step.id]; + } + // Auto behavior: open if in progress + if (step.status === "in_progress") { + return true; + } + // Auto behavior: keep last completed step open if no in-progress step + if (!hasInProgressStep && index === lastCompletedIndex) { + return true; + } + // Default: collapsed + return false; + }; + + const handleToggle = (stepId: string, currentOpen: boolean) => { + setManualOverrides(prev => ({ + ...prev, + [stepId]: !currentOpen, + })); + }; + + return ( + + {steps.map((step, index) => { + const isOpen = getStepOpenState(step, index); + return ( + handleToggle(step.id, isOpen)} + /> + ); + })} + + ); +} + /** * DeepAgent Thinking Tool UI Component * @@ -131,7 +215,7 @@ export const DeepAgentThinkingToolUI = makeAssistantToolUI< DeepAgentThinkingResult >({ toolName: "deepagent_thinking", - render: function DeepAgentThinkingUI({ args, result, status }) { + render: function DeepAgentThinkingUI({ result, status }) { // Loading state - tool is still running if (status.type === "running" || status.type === "requires-action") { return ; @@ -155,11 +239,7 @@ export const DeepAgentThinkingToolUI = makeAssistantToolUI< // Render the chain of thought return (
- - {result.steps.map((step) => ( - - ))} - +
); }, @@ -189,11 +269,7 @@ export function InlineThinkingDisplay({ {isStreaming && steps.length === 0 ? ( ) : ( - - {steps.map((step) => ( - - ))} - + )}
);