Skip to content

feat/cancel-generation - stop an ongoing generation query #1019

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
udaykumar-dhokia wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat/cancel-generation - stop an ongoing generation query #1019

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10,821 changes: 10,821 additions & 0 deletions package-lock.json
View file
Open in desktop

Large diffs are not rendered by default.

3 changes: 2 additions & 1 deletion src/app/api/chat/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,8 +233,9 @@ export const POST = async (req: Request) => {
});

req.signal.addEventListener('abort', () => {
session.abort();
disconnect();
writer.close();
writer.close().catch(() => {});
});

return new Response(responseStream.readable, {
Expand Down
47 changes: 33 additions & 14 deletions src/components/MessageInput.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
import { cn } from '@/lib/utils';
import { ArrowUp } from 'lucide-react';
import { ArrowUp, Square } from 'lucide-react';
import { useEffect, useRef, useState } from 'react';
import TextareaAutosize from 'react-textarea-autosize';
import AttachSmall from './MessageInputActions/AttachSmall';
import { useChat } from '@/lib/hooks/useChat';

const MessageInput = () => {
const { loading, sendMessage } = useChat();
const { loading, sendMessage, cancelMessage } = useChat();

const [copilotEnabled, setCopilotEnabled] = useState(false);
const [message, setMessage] = useState('');
Expand Down Expand Up @@ -76,23 +76,42 @@ const MessageInput = () => {
className="transition bg-transparent dark:placeholder:text-white/50 placeholder:text-sm text-sm dark:text-white resize-none focus:outline-none w-full px-2 max-h-24 lg:max-h-36 xl:max-h-48 flex-grow flex-shrink"
placeholder="Ask a follow-up"
/>
{mode === 'single' && (
<button
disabled={message.trim().length === 0 || loading}
className="bg-[#24A0ED] text-white disabled:text-black/50 dark:disabled:text-white/50 hover:bg-opacity-85 transition duration-100 disabled:bg-[#e0e0dc79] dark:disabled:bg-[#ececec21] rounded-full p-2"
>
<ArrowUp className="bg-background" size={17} />
</button>
)}
{mode === 'multi' && (
<div className="flex flex-row items-center justify-between w-full pt-2">
<AttachSmall />
{mode === 'single' &&
(loading ? (
<button
type="button"
onClick={cancelMessage}
className="bg-red-500 hover:bg-red-600 text-white transition duration-100 rounded-full p-2"
>
<Square size={15} fill="currentColor" />
</button>
) : (
<button
disabled={message.trim().length === 0 || loading}
disabled={message.trim().length === 0}
className="bg-[#24A0ED] text-white disabled:text-black/50 dark:disabled:text-white/50 hover:bg-opacity-85 transition duration-100 disabled:bg-[#e0e0dc79] dark:disabled:bg-[#ececec21] rounded-full p-2"
>
<ArrowUp className="bg-background" size={17} />
</button>
))}
{mode === 'multi' && (
<div className="flex flex-row items-center justify-between w-full pt-2">
<AttachSmall />
{loading ? (
<button
type="button"
onClick={cancelMessage}
className="bg-red-500 hover:bg-red-600 text-white transition duration-100 rounded-full p-2"
>
<Square size={15} fill="currentColor" />
</button>
) : (
<button
disabled={message.trim().length === 0}
className="bg-[#24A0ED] text-white disabled:text-black/50 dark:disabled:text-white/50 hover:bg-opacity-85 transition duration-100 disabled:bg-[#e0e0dc79] dark:disabled:bg-[#ececec21] rounded-full p-2"
>
<ArrowUp className="bg-background" size={17} />
</button>
)}
</div>
)}
</form>
Expand Down
75 changes: 52 additions & 23 deletions src/lib/agents/search/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,23 @@ class SearchAgent {
searchPromise,
]);

if (session.signal.aborted) {
Copy link
Contributor

@cubic-dev-ai cubic-dev-ai bot Mar 5, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: Abort handling is checked too late: searchAsync still waits for non-cancelable widget execution to finish before returning on cancel.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At src/lib/agents/search/index.ts, line 97:

<comment>Abort handling is checked too late: `searchAsync` still waits for non-cancelable widget execution to finish before returning on cancel.</comment>

<file context>
@@ -94,6 +94,23 @@ class SearchAgent {
       searchPromise,
     ]);
 
+    if (session.signal.aborted) {
+      await db
+        .update(messages)
</file context>
Fix with Cubic

await db
.update(messages)
.set({
status: 'completed',
responseBlocks: session.getAllBlocks(),
})
.where(
and(
eq(messages.chatId, input.chatId),
eq(messages.messageId, input.messageId),
),
)
.execute();
return;
}

session.emit('data', {
type: 'researchComplete',
});
Expand Down Expand Up @@ -131,41 +148,53 @@ class SearchAgent {
content: input.followUp,
},
],
signal: session.signal,
});

let responseBlockId = '';

for await (const chunk of answerStream) {
if (!responseBlockId) {
const block: TextBlock = {
id: crypto.randomUUID(),
type: 'text',
data: chunk.contentChunk,
};
try {
for await (const chunk of answerStream) {
if (session.signal.aborted) break;

session.emitBlock(block);
if (!responseBlockId) {
const block: TextBlock = {
id: crypto.randomUUID(),
type: 'text',
data: chunk.contentChunk,
};

responseBlockId = block.id;
} else {
const block = session.getBlock(responseBlockId) as TextBlock | null;
session.emitBlock(block);

if (!block) {
continue;
}
responseBlockId = block.id;
} else {
const block = session.getBlock(responseBlockId) as TextBlock | null;

block.data += chunk.contentChunk;
if (!block) {
continue;
}

session.updateBlock(block.id, [
{
op: 'replace',
path: '/data',
value: block.data,
},
]);
block.data += chunk.contentChunk;

session.updateBlock(block.id, [
{
op: 'replace',
path: '/data',
value: block.data,
},
]);
}
}
} catch (err: any) {
// Abort errors are expected when the user cancels
if (!session.signal.aborted) {
throw err;
}
}

session.emit('end', {});
if (!session.signal.aborted) {
session.emit('end', {});
}

await db
.update(messages)
Expand Down
114 changes: 62 additions & 52 deletions src/lib/agents/search/researcher/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,8 @@ class Researcher {
];

for (let i = 0; i < maxIteration; i++) {
if (session.signal.aborted) break;

const researcherPrompt = getResearcherPrompt(
availableActionsDescription,
input.config.mode,
Expand All @@ -74,6 +76,7 @@ class Researcher {
...agentMessageHistory,
],
tools: availableTools,
signal: session.signal,
});

const block = session.getBlock(researchBlockId);
Expand All @@ -83,68 +86,75 @@ class Researcher {

let finalToolCalls: ToolCall[] = [];

for await (const partialRes of actionStream) {
if (partialRes.toolCallChunk.length > 0) {
partialRes.toolCallChunk.forEach((tc) => {
if (
tc.name === '__reasoning_preamble' &&
tc.arguments['plan'] &&
!reasoningEmitted &&
block &&
block.type === 'research'
) {
reasoningEmitted = true;

block.data.subSteps.push({
id: reasoningId,
type: 'reasoning',
reasoning: tc.arguments['plan'],
});

session.updateBlock(researchBlockId, [
{
op: 'replace',
path: '/data/subSteps',
value: block.data.subSteps,
},
]);
} else if (
tc.name === '__reasoning_preamble' &&
tc.arguments['plan'] &&
reasoningEmitted &&
block &&
block.type === 'research'
) {
const subStepIndex = block.data.subSteps.findIndex(
(step: any) => step.id === reasoningId,
);
try {
for await (const partialRes of actionStream) {
if (partialRes.toolCallChunk.length > 0) {
partialRes.toolCallChunk.forEach((tc) => {
if (
tc.name === '__reasoning_preamble' &&
tc.arguments['plan'] &&
!reasoningEmitted &&
block &&
block.type === 'research'
) {
reasoningEmitted = true;

block.data.subSteps.push({
id: reasoningId,
type: 'reasoning',
reasoning: tc.arguments['plan'],
});

if (subStepIndex !== -1) {
const subStep = block.data.subSteps[
subStepIndex
] as ReasoningResearchBlock;
subStep.reasoning = tc.arguments['plan'];
session.updateBlock(researchBlockId, [
{
op: 'replace',
path: '/data/subSteps',
value: block.data.subSteps,
},
]);
} else if (
tc.name === '__reasoning_preamble' &&
tc.arguments['plan'] &&
reasoningEmitted &&
block &&
block.type === 'research'
) {
const subStepIndex = block.data.subSteps.findIndex(
(step: any) => step.id === reasoningId,
);

if (subStepIndex !== -1) {
const subStep = block.data.subSteps[
subStepIndex
] as ReasoningResearchBlock;
subStep.reasoning = tc.arguments['plan'];
session.updateBlock(researchBlockId, [
{
op: 'replace',
path: '/data/subSteps',
value: block.data.subSteps,
},
]);
}
}

const existingIndex = finalToolCalls.findIndex(
(ftc) => ftc.id === tc.id,
);

if (existingIndex !== -1) {
finalToolCalls[existingIndex].arguments = tc.arguments;
} else {
finalToolCalls.push(tc);
}
}

const existingIndex = finalToolCalls.findIndex(
(ftc) => ftc.id === tc.id,
);

if (existingIndex !== -1) {
finalToolCalls[existingIndex].arguments = tc.arguments;
} else {
finalToolCalls.push(tc);
}
});
});
}
}
} catch (err: any) {
if (!session.signal.aborted) {
throw err;
}
break;
}

if (finalToolCalls.length === 0) {
Expand Down
Loading