diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 6d4a7e2a31..236dd31c39 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -4028,6 +4028,23 @@ const useOwnsResearchMessage = () => { ); }; +// Whether the active thread has a non-terminal durable research run. After a +// reload the run is followed by the research store rather than an assistant-ui +// run, so `thread.isRunning` is false while research is still active; message +// edit/reload/branch actions must also gate on this to preserve one-run-per-chat. +const useThreadResearchActive = (): boolean => { + const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); + return useResearchRunStore((state) => { + const runId = activeThreadId + ? state.latestRunByThreadId[activeThreadId] + : undefined; + const run = runId ? state.sessions[runId]?.run : undefined; + return Boolean( + run && !["completed", "failed", "cancelled"].includes(run.status), + ); + }); +}; + const DeleteMessageButton: FC = () => { const aui = useAui(); const messageId = useAuiState(({ message }) => message.id); @@ -4132,6 +4149,7 @@ const EditAssistantMessageButton: FC = () => { const messageId = useAuiState(({ message }) => message.id); const researchRunId = useResearchMessageRunId(); const isRunning = useAuiState(({ thread }) => thread.isRunning); + const researchActive = useThreadResearchActive(); const setEditingId = useChatRuntimeStore((s) => s.setEditingMessageId); if (researchRunId) return null; @@ -4139,7 +4157,7 @@ const EditAssistantMessageButton: FC = () => { return ( setEditingId(messageId)} > { const AssistantActionBar: FC = () => { const { forkMessage, forkDisabled } = useForkMessageAction(); const researchRunId = useResearchMessageRunId(); + const researchActive = useThreadResearchActive(); const [detailsOpen, setDetailsOpen] = useState(false); const ttsEnabled = useVoiceSettingsStore((s) => s.ttsEnabled); // hideWhenRunning is thread-level, so a new run would hide this bar and its @@ -4168,7 +4187,7 @@ const AssistantActionBar: FC = () => { > - {!researchRunId && ( + {!researchRunId && !researchActive && ( @@ -4296,13 +4315,14 @@ const UserMessage: FC = () => { const UserActionBar: FC = () => { const ownsResearchMessage = useOwnsResearchMessage(); + const researchActive = useThreadResearchActive(); return ( - {!ownsResearchMessage && ( + {!ownsResearchMessage && !researchActive && ( { const aui = useAui(); const { inputProps, isComposingRef } = useImeComposerInputHandlers(); const resendAfterCancelRef = useRef(false); + const researchActive = useThreadResearchActive(); useAuiEvent("thread.runEnd", () => { if (!resendAfterCancelRef.current) { @@ -4352,6 +4373,7 @@ const EditComposer: FC = () => {