diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index a0a19572..99a70940 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -235,23 +235,51 @@ export async function POST(request: Request) { const proto = host.startsWith('localhost') ? 'http' : 'https'; const baseUrl = `${proto}://${host}`; + // Honor client-side abort (Stop button). When the user clicks Stop + // the browser's AbortController fires `request.signal.aborted` and + // the fetch stream is closed; we use it as a polite checkpoint + // between rounds and tool calls so we (a) don't keep paying Gemini + // for tokens the user no longer wants and (b) persist whatever the + // assistant produced before the cancel. + const clientSignal = request.signal; + // Stream response const encoder = new TextEncoder(); const stream = new ReadableStream({ async start(controller) { + let streamClosed = false; function emit(chunk: object) { - controller.enqueue(encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`)); + if (streamClosed) return; + try { + controller.enqueue(encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`)); + } catch { + // controller may have been closed by the abort handler + streamClosed = true; + } + } + function safeClose() { + if (streamClosed) return; + streamClosed = true; + try { + controller.close(); + } catch {} } let messages = [...history]; let round = 0; let assistantText = ''; const assistantToolCalls: ToolCall[] = []; + let aborted = clientSignal.aborted; + const onAbort = () => { + aborted = true; + }; + clientSignal.addEventListener('abort', onAbort); try { // Tool-calling loop: use non-streaming so thought_signature is // always present in the complete response (required by thinking models). while (round < MAX_TOOL_ROUNDS) { + if (aborted) break; round++; const toolDefs = mcp_token ? VIBN_TOOL_DEFINITIONS : []; @@ -283,9 +311,11 @@ export async function POST(request: Request) { }); if (!resp.toolCalls.length) break; + if (aborted) break; // Execute tool calls and add results for (const tc of resp.toolCalls) { + if (aborted) break; const result = mcp_token ? await executeMcpTool(tc.name, tc.args, mcp_token, baseUrl) : JSON.stringify({ error: 'No MCP token — read-only mode.' }); @@ -302,6 +332,19 @@ export async function POST(request: Request) { } } + // If the user clicked Stop, surface the cancel marker so the + // client renders "(stopped by user)" inline with the partial + // assistant message, then skip the round-cap recovery summary + // (we shouldn't pay Gemini for a turn the user just canceled). + if (aborted) { + const stopMarker = assistantText + ? '\n\n_(stopped by user)_' + : '_(stopped by user before any response)_'; + assistantText += stopMarker; + emit({ type: 'text', text: stopMarker }); + emit({ type: 'aborted' }); + } + // If the loop exited because we hit MAX_TOOL_ROUNDS while the // model still wanted to call tools, the user has only seen a // tray of ✓ icons with no narrative. Force one final no-tools @@ -309,7 +352,7 @@ export async function POST(request: Request) { const lastTurnHadTools = messages.length > 0 && messages[messages.length - 1].role === 'tool'; - if (round >= MAX_TOOL_ROUNDS && lastTurnHadTools) { + if (!aborted && round >= MAX_TOOL_ROUNDS && lastTurnHadTools) { try { const summary = await callGeminiChat({ systemPrompt: @@ -340,12 +383,28 @@ export async function POST(request: Request) { ); emit({ type: 'done' }); - controller.close(); + safeClose(); } catch (e) { - emit({ type: 'error', error: e instanceof Error ? e.message : String(e) }); - controller.close(); + // AbortError is the expected shape when the client cancels + // mid-Gemini-call — don't surface it as a real error. + const isAbort = + aborted || + (e instanceof Error && (e.name === 'AbortError' || /aborted/i.test(e.message))); + if (!isAbort) { + emit({ type: 'error', error: e instanceof Error ? e.message : String(e) }); + } else { + emit({ type: 'aborted' }); + } + safeClose(); + } finally { + clientSignal.removeEventListener('abort', onAbort); } }, + cancel() { + // Browser disconnected (tab closed, navigated away). Nothing to + // do — the abort handler above already flipped the flag and the + // loop will bail at the next checkpoint. + }, }); return new Response(stream, { diff --git a/components/vibn-chat/chat-panel.tsx b/components/vibn-chat/chat-panel.tsx index 2d941f03..66514727 100644 --- a/components/vibn-chat/chat-panel.tsx +++ b/components/vibn-chat/chat-panel.tsx @@ -13,6 +13,7 @@ import { Wrench, ChevronDown, Trash2, + Square, } from "lucide-react"; // ── Types ───────────────────────────────────────────────────────────────────── @@ -156,6 +157,10 @@ export function ChatPanel() { const [mcpToken, setMcpToken] = useState(null); const messagesEndRef = useRef(null); const inputRef = useRef(null); + // AbortController for the in-flight /api/chat fetch. Lives in a ref + // so the Stop button can reach it without re-rendering on every + // streaming chunk. + const abortRef = useRef(null); const scrollToBottom = useCallback(() => { messagesEndRef.current?.scrollIntoView({ behavior: "smooth" }); @@ -287,11 +292,15 @@ export function ChatPanel() { const assistantMsg: Message = { role: "assistant", content: "" }; let msgIndex = -1; + const controller = new AbortController(); + abortRef.current = controller; + try { const res = await fetch("/api/chat", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ thread_id: activeThread, message: text, workspace, mcp_token: mcpToken }), + signal: controller.signal, }); if (!res.ok || !res.body) throw new Error("Stream failed"); @@ -363,22 +372,50 @@ export function ChatPanel() { loadThreads(); } catch (e) { - setMessages((prev) => { - const next = [...prev]; - if (msgIndex >= 0 && next[msgIndex]) { - next[msgIndex] = { ...next[msgIndex], content: "⚠️ Failed to get response. Please try again." }; - } - return next; - }); + const isAbort = + e instanceof DOMException && e.name === "AbortError"; + if (isAbort) { + // Server-side will have appended "(stopped by user)" to the + // partial response and persisted it. We just need to make + // sure the local UI reflects whatever streamed in before the + // user clicked Stop — which it already does, because we've + // been mutating `messages[msgIndex]` chunk-by-chunk above. + setMessages((prev) => { + const next = [...prev]; + if (msgIndex >= 0 && next[msgIndex] && !next[msgIndex].content.includes("(stopped by user)")) { + next[msgIndex] = { + ...next[msgIndex], + content: (next[msgIndex].content || "") + "\n\n_(stopped by user)_", + }; + } + return next; + }); + } else { + setMessages((prev) => { + const next = [...prev]; + if (msgIndex >= 0 && next[msgIndex]) { + next[msgIndex] = { ...next[msgIndex], content: "⚠️ Failed to get response. Please try again." }; + } + return next; + }); + } } finally { + abortRef.current = null; setSending(false); } }, [input, sending, activeThread, workspace, mcpToken, threads, loadThreads]); + const cancelMessage = useCallback(() => { + abortRef.current?.abort(); + }, []); + const handleKeyDown = (e: React.KeyboardEvent) => { if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); sendMessage(); + } else if (e.key === "Escape" && sending) { + e.preventDefault(); + cancelMessage(); } }; @@ -598,9 +635,9 @@ export function ChatPanel() { value={input} onChange={(e) => setInput(e.target.value)} onKeyDown={handleKeyDown} - placeholder="Ask Vibn AI anything…" + placeholder={sending ? "Esc to stop generating…" : "Ask Vibn AI anything…"} rows={1} - disabled={sending || !activeThread} + disabled={!activeThread} style={{ flex: 1, border: "none", outline: "none", background: "transparent", fontSize: "0.84rem", lineHeight: 1.5, resize: "none", @@ -613,23 +650,48 @@ export function ChatPanel() { el.style.height = Math.min(el.scrollHeight, 120) + "px"; }} /> - + {(() => { + // While the AI is streaming or running tools, the button + // turns into a Stop control. Click → AbortController fires, + // server bails between rounds, partial text gets persisted. + const isActive = sending; + const canSend = !sending && input.trim() && activeThread; + return ( + + ); + })()}
Powered by Gemini 3.1 Pro