diff --git a/components/VercelChat/ChatInput.tsx b/components/VercelChat/ChatInput.tsx index 1e9eb99f0..44dcc2bba 100644 --- a/components/VercelChat/ChatInput.tsx +++ b/components/VercelChat/ChatInput.tsx @@ -23,6 +23,7 @@ export function ChatInput() { isLoadingSignedUrls, handleSendMessage, isGeneratingResponse, + isStopping, workspaceStatus, stop, setInput, @@ -42,6 +43,10 @@ export function ChatInput() { const handleSend = (event: React.FormEvent) => { event.preventDefault(); + // Already cancelling — ignore further clicks until the backend + // round-trip resolves and the SSE watcher closes the stream. + if (isStopping) return; + // Allow stop action regardless of input state if (isGeneratingResponse) { stop(); @@ -94,12 +99,15 @@ export function ChatInput() { diff --git a/components/VercelChat/MessageParts.tsx b/components/VercelChat/MessageParts.tsx index 1e522c846..1383ade1c 100644 --- a/components/VercelChat/MessageParts.tsx +++ b/components/VercelChat/MessageParts.tsx @@ -10,7 +10,11 @@ import { Dispatch, SetStateAction } from "react"; import { cn } from "@/lib/utils"; import ViewingMessage from "./ViewingMessage"; import EditingMessage from "./EditingMessage"; -import { getToolCallComponent, getToolResultComponent } from "./ToolComponents"; +import { + getCancelledToolComponent, + getToolCallComponent, + getToolResultComponent, +} from "./ToolComponents"; import MessageFileViewer from "./message-file-viewer"; import { EnhancedReasoning } from "@/components/reasoning/EnhancedReasoning"; import { Actions, Action } from "@/components/actions"; @@ -107,11 +111,17 @@ export function MessageParts({ message, mode, setMode }: MessagePartsProps) { if (isToolOrDynamicToolUIPart(part)) { const { state } = part as ToolUIPart; + // Terminal-error states (user-cancel via stop, approval-denied, + // tool-thrown errors) need their own renderer — falling through + // to getToolCallComponent paints the spinning skeleton, which + // makes a cancelled tool look like it's still running. + if (state === "output-error" || state === "output-denied") { + return getCancelledToolComponent(part as ToolUIPart); + } if (state !== "output-available") { return getToolCallComponent(part as ToolUIPart); - } else { - return getToolResultComponent(part as ToolUIPart); } + return getToolResultComponent(part as ToolUIPart); } } )} diff --git a/components/VercelChat/ToolComponents.tsx b/components/VercelChat/ToolComponents.tsx index 763a69408..9267072b2 100644 --- a/components/VercelChat/ToolComponents.tsx +++ b/components/VercelChat/ToolComponents.tsx @@ -23,7 +23,7 @@ import UpdateArtistSocialsSuccess from "./tools/UpdateArtistSocialsSuccess"; import { UpdateArtistSocialsResult } from "./tools/UpdateArtistSocialsSuccess"; import { TxtFileResult } from "@/components/ui/TxtFileResult"; import { TxtFileGenerationResult } from "@/components/ui/TxtFileResult"; -import { Loader } from "lucide-react"; +import { Loader, OctagonX } from "lucide-react"; import GenericSuccess from "./tools/GenericSuccess"; import getToolInfo from "@/lib/utils/getToolsInfo"; import { isSearchProgressUpdate } from "@/lib/search/searchProgressUtils"; @@ -220,6 +220,30 @@ export function getToolCallComponent(part: ToolUIPart) { ); } +/** + * Render a tool-call that terminated in `output-error` or `output-denied` + * (cancelled by the user, denied by approval flow, or errored). Without + * this branch the renderer falls through to the spinning skeleton and + * the pill keeps "running" forever. + */ +export function getCancelledToolComponent(part: ToolUIPart | DynamicToolUIPart) { + const { toolCallId } = part; + const toolName = getToolOrDynamicToolName(part); + const label = + (part as { errorText?: string }).errorText === "Cancelled" + ? `Cancelled ${toolName}` + : `Failed ${toolName}`; + return ( +
+ + {label} +
+ ); +} + export function getToolResultComponent(part: ToolUIPart | DynamicToolUIPart) { const { toolCallId, output, type } = part; const isMcp = type === "dynamic-tool"; diff --git a/hooks/useStopChatWorkflow.ts b/hooks/useStopChatWorkflow.ts new file mode 100644 index 000000000..21c02ba0e --- /dev/null +++ b/hooks/useStopChatWorkflow.ts @@ -0,0 +1,29 @@ +import { useMutation } from "@tanstack/react-query"; +import { usePrivy } from "@privy-io/react-auth"; +import { stopChatWorkflow } from "@/lib/chat/stopChatWorkflow"; + +/** + * Wraps the `POST /api/chat/{chatId}/stop` round-trip with a React Query + * mutation. Consumers read `isStopping` (`mutation.isPending`) to flip + * the submit button to a "stopping" state the instant the user clicks, + * so the UI doesn't sit dead for the 1–2s while the backend cancels the + * workflow and the SSE watcher closes the stream. + * + * Workflow-chat-only: legacy `/api/chat` aborts locally via the AI SDK's + * `stop()` and never hits this hook. + */ +export function useStopChatWorkflow(chatId: string) { + const { getAccessToken } = usePrivy(); + + const mutation = useMutation({ + mutationFn: async () => { + const token = await getAccessToken().catch(() => null); + await stopChatWorkflow(chatId, token); + }, + }); + + return { + stop: mutation.mutateAsync, + isStopping: mutation.isPending, + }; +} diff --git a/hooks/useVercelChat.ts b/hooks/useVercelChat.ts index ab9d202c2..1aad4c367 100644 --- a/hooks/useVercelChat.ts +++ b/hooks/useVercelChat.ts @@ -23,6 +23,7 @@ import { formatTextAttachments } from "@/lib/chat/formatTextAttachments"; import { useDeleteTrailingMessages } from "./useDeleteTrailingMessages"; import { getFileContents } from "@/lib/sandboxes/getFileContents"; import getMimeFromPath from "@/lib/files/getMimeFromPath"; +import { useStopChatWorkflow } from "./useStopChatWorkflow"; import { getChatPath } from "@/lib/chat/getChatPath"; import { usePersistSelectedModel } from "./usePersistSelectedModel"; @@ -205,7 +206,7 @@ export function useVercelChat({ [id, artistId, organizationId, accountIdOverride, model], ); - const { messages, status, stop, sendMessage, setMessages, regenerate } = + const { messages, status, stop: aiStop, sendMessage, setMessages, regenerate } = useChat({ id, transport, @@ -221,6 +222,19 @@ export function useVercelChat({ }, }); + // Workflow chats: await the backend cancel and let SSE close naturally — + // calling aiStop here would tear down SSE before in-flight chunks reach + // the UI and frontend/DB would disagree on reload. Legacy chats abort + // locally via the AI SDK's `stop()`. + const { stop: stopWorkflow, isStopping } = useStopChatWorkflow(transportChatId); + const stop = useCallback(async () => { + if (sessionId) { + await stopWorkflow(); + return; + } + await aiStop(); + }, [aiStop, sessionId, stopWorkflow]); + const earliestFailedUserMessageId = useMemo( () => getEarliestFailedUserMessageId(messages), [messages], @@ -408,6 +422,7 @@ export function useVercelChat({ isLoading, hasError, isGeneratingResponse, + isStopping, model, isLoadingSignedUrls, diff --git a/lib/chat/__tests__/stopChatWorkflow.test.ts b/lib/chat/__tests__/stopChatWorkflow.test.ts new file mode 100644 index 000000000..fedf56d89 --- /dev/null +++ b/lib/chat/__tests__/stopChatWorkflow.test.ts @@ -0,0 +1,54 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { stopChatWorkflow } from "../stopChatWorkflow"; +import { NEW_API_BASE_URL } from "../../consts"; + +const mockFetch = vi.fn(); +global.fetch = mockFetch; + +describe("stopChatWorkflow", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("POSTs to the chat stop endpoint with a bearer token", async () => { + mockFetch.mockResolvedValueOnce({ ok: true, json: async () => ({}) }); + + await stopChatWorkflow("chat-1", "tok-123"); + + expect(mockFetch).toHaveBeenCalledWith( + `${NEW_API_BASE_URL}/api/chat/chat-1/stop`, + { + method: "POST", + headers: { Authorization: "Bearer tok-123" }, + }, + ); + }); + + it("omits the Authorization header when unauthenticated", async () => { + mockFetch.mockResolvedValueOnce({ ok: true, json: async () => ({}) }); + + await stopChatWorkflow("chat-1", null); + + expect(mockFetch).toHaveBeenCalledWith( + `${NEW_API_BASE_URL}/api/chat/chat-1/stop`, + { method: "POST", headers: {} }, + ); + }); + + it("url-encodes the chat id", async () => { + mockFetch.mockResolvedValueOnce({ ok: true, json: async () => ({}) }); + + await stopChatWorkflow("a/b c", "tok"); + + expect(mockFetch).toHaveBeenCalledWith( + `${NEW_API_BASE_URL}/api/chat/a%2Fb%20c/stop`, + expect.any(Object), + ); + }); + + it("never throws when the request fails", async () => { + mockFetch.mockRejectedValueOnce(new Error("network down")); + + await expect(stopChatWorkflow("chat-1", "tok")).resolves.toBeUndefined(); + }); +}); diff --git a/lib/chat/stopChatWorkflow.ts b/lib/chat/stopChatWorkflow.ts new file mode 100644 index 000000000..970880d04 --- /dev/null +++ b/lib/chat/stopChatWorkflow.ts @@ -0,0 +1,30 @@ +import { getClientApiBaseUrl } from "@/lib/api/getClientApiBaseUrl"; + +/** + * Cancels the in-flight workflow run for a chat via recoup-api + * `POST /api/chat/{chatId}/stop`. + * + * The client AI SDK `stop()` only aborts the local fetch; the durable + * workflow run keeps streaming (and billing) server-side until it's + * cancelled here. Callers should fire this without blocking the UI stop. + * + * @param chatId - Chat row id (the workflow run is keyed off it). + * @param accessToken - Privy access token; omitted when unauthenticated. + * @returns Resolves once the request settles; never throws. + */ +export async function stopChatWorkflow( + chatId: string, + accessToken: string | null, +): Promise { + try { + await fetch( + `${getClientApiBaseUrl()}/api/chat/${encodeURIComponent(chatId)}/stop`, + { + method: "POST", + headers: accessToken ? { Authorization: `Bearer ${accessToken}` } : {}, + }, + ); + } catch { + // Best-effort: the run also self-cancels when its slot is cleared. + } +} diff --git a/providers/VercelChatProvider.tsx b/providers/VercelChatProvider.tsx index d085a8f4c..6e3370493 100644 --- a/providers/VercelChatProvider.tsx +++ b/providers/VercelChatProvider.tsx @@ -25,6 +25,7 @@ interface VercelChatContextType { isLoading: boolean; hasError: boolean; isGeneratingResponse: boolean; + isStopping: boolean; isLoadingSignedUrls: boolean; handleSendMessage: (event: React.FormEvent) => Promise; stop: UseChatHelpers["stop"]; @@ -118,6 +119,7 @@ export function VercelChatProvider({ isLoading, hasError, isGeneratingResponse, + isStopping, isLoadingSignedUrls, handleSendMessage, stop, @@ -163,6 +165,7 @@ export function VercelChatProvider({ isLoading, hasError, isGeneratingResponse, + isStopping, isLoadingSignedUrls, handleSendMessage: handleSendMessageWithClear, stop,