mirror of
https://github.com/rowboatlabs/rowboat.git
synced 2026-07-15 21:11:08 +02:00
feat: add stop execution with hybrid graceful/force abort
Implement a stop execution feature that allows users to abort ongoing LLM streaming, kill running tool calls, and clear pending permission/human input requests. Uses a hybrid approach: first click sends graceful SIGTERM, second click within 2s sends SIGKILL and force-closes MCP clients. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
parent
9828ea8301
commit
a3e681a7c4
13 changed files with 642 additions and 54 deletions
|
|
@ -295,7 +295,7 @@ export function setupIpcHandlers() {
|
||||||
return { success: true };
|
return { success: true };
|
||||||
},
|
},
|
||||||
'runs:stop': async (_event, args) => {
|
'runs:stop': async (_event, args) => {
|
||||||
await runsCore.stop(args.runId);
|
await runsCore.stop(args.runId, args.force);
|
||||||
return { success: true };
|
return { success: true };
|
||||||
},
|
},
|
||||||
'runs:fetch': async (_event, args) => {
|
'runs:fetch': async (_event, args) => {
|
||||||
|
|
|
||||||
|
|
@ -6,7 +6,7 @@ import type { LanguageModelUsage, ToolUIPart } from 'ai';
|
||||||
import './App.css'
|
import './App.css'
|
||||||
import z from 'zod';
|
import z from 'zod';
|
||||||
import { Button } from './components/ui/button';
|
import { Button } from './components/ui/button';
|
||||||
import { CheckIcon, LoaderIcon, ArrowUp, PanelRightIcon, SquarePen } from 'lucide-react';
|
import { CheckIcon, LoaderIcon, ArrowUp, PanelRightIcon, SquarePen, Square } from 'lucide-react';
|
||||||
import { cn } from '@/lib/utils';
|
import { cn } from '@/lib/utils';
|
||||||
import { MarkdownEditor } from './components/markdown-editor';
|
import { MarkdownEditor } from './components/markdown-editor';
|
||||||
import { ChatInputBar } from './components/chat-button';
|
import { ChatInputBar } from './components/chat-button';
|
||||||
|
|
@ -279,7 +279,9 @@ const collectFilePaths = (nodes: TreeNode[]): string[] =>
|
||||||
// Inner component that uses the controller to access mentions
|
// Inner component that uses the controller to access mentions
|
||||||
interface ChatInputInnerProps {
|
interface ChatInputInnerProps {
|
||||||
onSubmit: (message: PromptInputMessage, mentions?: FileMention[]) => void
|
onSubmit: (message: PromptInputMessage, mentions?: FileMention[]) => void
|
||||||
|
onStop?: () => void
|
||||||
isProcessing: boolean
|
isProcessing: boolean
|
||||||
|
isStopping?: boolean
|
||||||
presetMessage?: string
|
presetMessage?: string
|
||||||
onPresetMessageConsumed?: () => void
|
onPresetMessageConsumed?: () => void
|
||||||
runId?: string | null
|
runId?: string | null
|
||||||
|
|
@ -287,7 +289,9 @@ interface ChatInputInnerProps {
|
||||||
|
|
||||||
function ChatInputInner({
|
function ChatInputInner({
|
||||||
onSubmit,
|
onSubmit,
|
||||||
|
onStop,
|
||||||
isProcessing,
|
isProcessing,
|
||||||
|
isStopping,
|
||||||
presetMessage,
|
presetMessage,
|
||||||
onPresetMessageConsumed,
|
onPresetMessageConsumed,
|
||||||
runId,
|
runId,
|
||||||
|
|
@ -327,19 +331,39 @@ function ChatInputInner({
|
||||||
focusTrigger={runId}
|
focusTrigger={runId}
|
||||||
className="min-h-6 py-0 border-0 shadow-none focus-visible:ring-0 rounded-none"
|
className="min-h-6 py-0 border-0 shadow-none focus-visible:ring-0 rounded-none"
|
||||||
/>
|
/>
|
||||||
<Button
|
{isProcessing ? (
|
||||||
size="icon"
|
<Button
|
||||||
onClick={handleSubmit}
|
size="icon"
|
||||||
disabled={!canSubmit}
|
onClick={onStop}
|
||||||
className={cn(
|
title={isStopping ? "Click again to force stop" : "Stop generation"}
|
||||||
"h-7 w-7 rounded-full shrink-0 transition-all",
|
className={cn(
|
||||||
canSubmit
|
"h-7 w-7 rounded-full shrink-0 transition-all",
|
||||||
? "bg-primary text-primary-foreground hover:bg-primary/90"
|
isStopping
|
||||||
: "bg-muted text-muted-foreground"
|
? "bg-destructive text-destructive-foreground hover:bg-destructive/90"
|
||||||
)}
|
: "bg-primary text-primary-foreground hover:bg-primary/90"
|
||||||
>
|
)}
|
||||||
<ArrowUp className="h-4 w-4" />
|
>
|
||||||
</Button>
|
{isStopping ? (
|
||||||
|
<LoaderIcon className="h-4 w-4 animate-spin" />
|
||||||
|
) : (
|
||||||
|
<Square className="h-3 w-3 fill-current" />
|
||||||
|
)}
|
||||||
|
</Button>
|
||||||
|
) : (
|
||||||
|
<Button
|
||||||
|
size="icon"
|
||||||
|
onClick={handleSubmit}
|
||||||
|
disabled={!canSubmit}
|
||||||
|
className={cn(
|
||||||
|
"h-7 w-7 rounded-full shrink-0 transition-all",
|
||||||
|
canSubmit
|
||||||
|
? "bg-primary text-primary-foreground hover:bg-primary/90"
|
||||||
|
: "bg-muted text-muted-foreground"
|
||||||
|
)}
|
||||||
|
>
|
||||||
|
<ArrowUp className="h-4 w-4" />
|
||||||
|
</Button>
|
||||||
|
)}
|
||||||
</div>
|
</div>
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
@ -350,7 +374,9 @@ interface ChatInputWithMentionsProps {
|
||||||
recentFiles: string[]
|
recentFiles: string[]
|
||||||
visibleFiles: string[]
|
visibleFiles: string[]
|
||||||
onSubmit: (message: PromptInputMessage, mentions?: FileMention[]) => void
|
onSubmit: (message: PromptInputMessage, mentions?: FileMention[]) => void
|
||||||
|
onStop?: () => void
|
||||||
isProcessing: boolean
|
isProcessing: boolean
|
||||||
|
isStopping?: boolean
|
||||||
presetMessage?: string
|
presetMessage?: string
|
||||||
onPresetMessageConsumed?: () => void
|
onPresetMessageConsumed?: () => void
|
||||||
runId?: string | null
|
runId?: string | null
|
||||||
|
|
@ -361,7 +387,9 @@ function ChatInputWithMentions({
|
||||||
recentFiles,
|
recentFiles,
|
||||||
visibleFiles,
|
visibleFiles,
|
||||||
onSubmit,
|
onSubmit,
|
||||||
|
onStop,
|
||||||
isProcessing,
|
isProcessing,
|
||||||
|
isStopping,
|
||||||
presetMessage,
|
presetMessage,
|
||||||
onPresetMessageConsumed,
|
onPresetMessageConsumed,
|
||||||
runId,
|
runId,
|
||||||
|
|
@ -370,7 +398,9 @@ function ChatInputWithMentions({
|
||||||
<PromptInputProvider knowledgeFiles={knowledgeFiles} recentFiles={recentFiles} visibleFiles={visibleFiles}>
|
<PromptInputProvider knowledgeFiles={knowledgeFiles} recentFiles={recentFiles} visibleFiles={visibleFiles}>
|
||||||
<ChatInputInner
|
<ChatInputInner
|
||||||
onSubmit={onSubmit}
|
onSubmit={onSubmit}
|
||||||
|
onStop={onStop}
|
||||||
isProcessing={isProcessing}
|
isProcessing={isProcessing}
|
||||||
|
isStopping={isStopping}
|
||||||
presetMessage={presetMessage}
|
presetMessage={presetMessage}
|
||||||
onPresetMessageConsumed={onPresetMessageConsumed}
|
onPresetMessageConsumed={onPresetMessageConsumed}
|
||||||
runId={runId}
|
runId={runId}
|
||||||
|
|
@ -414,6 +444,8 @@ function App() {
|
||||||
const [runId, setRunId] = useState<string | null>(null)
|
const [runId, setRunId] = useState<string | null>(null)
|
||||||
const runIdRef = useRef<string | null>(null)
|
const runIdRef = useRef<string | null>(null)
|
||||||
const [isProcessing, setIsProcessing] = useState(false)
|
const [isProcessing, setIsProcessing] = useState(false)
|
||||||
|
const [isStopping, setIsStopping] = useState(false)
|
||||||
|
const [stopClickedAt, setStopClickedAt] = useState<number | null>(null)
|
||||||
const [agentId] = useState<string>('copilot')
|
const [agentId] = useState<string>('copilot')
|
||||||
const [presetMessage, setPresetMessage] = useState<string | undefined>(undefined)
|
const [presetMessage, setPresetMessage] = useState<string | undefined>(undefined)
|
||||||
|
|
||||||
|
|
@ -758,6 +790,8 @@ function App() {
|
||||||
|
|
||||||
case 'run-processing-end':
|
case 'run-processing-end':
|
||||||
setIsProcessing(false)
|
setIsProcessing(false)
|
||||||
|
setIsStopping(false)
|
||||||
|
setStopClickedAt(null)
|
||||||
break
|
break
|
||||||
|
|
||||||
case 'start':
|
case 'start':
|
||||||
|
|
@ -936,8 +970,32 @@ function App() {
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
|
|
||||||
|
case 'run-stopped':
|
||||||
|
setIsProcessing(false)
|
||||||
|
setIsStopping(false)
|
||||||
|
setStopClickedAt(null)
|
||||||
|
// Clear pending requests since they've been aborted
|
||||||
|
setPendingPermissionRequests(new Map())
|
||||||
|
setPendingAskHumanRequests(new Map())
|
||||||
|
// Flush any streaming content as a message
|
||||||
|
setCurrentAssistantMessage(currentMsg => {
|
||||||
|
if (currentMsg) {
|
||||||
|
setConversation(prev => [...prev, {
|
||||||
|
id: `assistant-stopped-${Date.now()}`,
|
||||||
|
role: 'assistant',
|
||||||
|
content: currentMsg,
|
||||||
|
timestamp: Date.now(),
|
||||||
|
}])
|
||||||
|
}
|
||||||
|
return ''
|
||||||
|
})
|
||||||
|
setCurrentReasoning('')
|
||||||
|
break
|
||||||
|
|
||||||
case 'error':
|
case 'error':
|
||||||
setIsProcessing(false)
|
setIsProcessing(false)
|
||||||
|
setIsStopping(false)
|
||||||
|
setStopClickedAt(null)
|
||||||
console.error('Run error:', event.error)
|
console.error('Run error:', event.error)
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
|
|
@ -1009,6 +1067,21 @@ function App() {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const handleStop = useCallback(async () => {
|
||||||
|
if (!runId) return
|
||||||
|
const now = Date.now()
|
||||||
|
const isForce = isStopping && stopClickedAt !== null && (now - stopClickedAt) < 2000
|
||||||
|
|
||||||
|
setStopClickedAt(now)
|
||||||
|
setIsStopping(true)
|
||||||
|
|
||||||
|
try {
|
||||||
|
await window.ipc.invoke('runs:stop', { runId, force: isForce })
|
||||||
|
} catch (error) {
|
||||||
|
console.error('Failed to stop run:', error)
|
||||||
|
}
|
||||||
|
}, [runId, isStopping, stopClickedAt])
|
||||||
|
|
||||||
const handlePermissionResponse = useCallback(async (toolCallId: string, subflow: string[], response: 'approve' | 'deny') => {
|
const handlePermissionResponse = useCallback(async (toolCallId: string, subflow: string[], response: 'approve' | 'deny') => {
|
||||||
if (!runId) return
|
if (!runId) return
|
||||||
|
|
||||||
|
|
@ -1779,7 +1852,9 @@ function App() {
|
||||||
recentFiles={recentWikiFiles}
|
recentFiles={recentWikiFiles}
|
||||||
visibleFiles={visibleKnowledgeFiles}
|
visibleFiles={visibleKnowledgeFiles}
|
||||||
onSubmit={handlePromptSubmit}
|
onSubmit={handlePromptSubmit}
|
||||||
|
onStop={handleStop}
|
||||||
isProcessing={isProcessing}
|
isProcessing={isProcessing}
|
||||||
|
isStopping={isStopping}
|
||||||
presetMessage={presetMessage}
|
presetMessage={presetMessage}
|
||||||
onPresetMessageConsumed={() => setPresetMessage(undefined)}
|
onPresetMessageConsumed={() => setPresetMessage(undefined)}
|
||||||
runId={runId}
|
runId={runId}
|
||||||
|
|
@ -1801,6 +1876,8 @@ function App() {
|
||||||
currentAssistantMessage={currentAssistantMessage}
|
currentAssistantMessage={currentAssistantMessage}
|
||||||
currentReasoning={currentReasoning}
|
currentReasoning={currentReasoning}
|
||||||
isProcessing={isProcessing}
|
isProcessing={isProcessing}
|
||||||
|
isStopping={isStopping}
|
||||||
|
onStop={handleStop}
|
||||||
message={message}
|
message={message}
|
||||||
onMessageChange={setMessage}
|
onMessageChange={setMessage}
|
||||||
onSubmit={handlePromptSubmit}
|
onSubmit={handlePromptSubmit}
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,5 @@
|
||||||
import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
|
import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
|
||||||
import { ArrowUp, Expand, Plus } from 'lucide-react'
|
import { ArrowUp, Expand, LoaderIcon, Plus, Square } from 'lucide-react'
|
||||||
import type { ToolUIPart } from 'ai'
|
import type { ToolUIPart } from 'ai'
|
||||||
import { Button } from '@/components/ui/button'
|
import { Button } from '@/components/ui/button'
|
||||||
import { cn } from '@/lib/utils'
|
import { cn } from '@/lib/utils'
|
||||||
|
|
@ -115,6 +115,8 @@ interface ChatSidebarProps {
|
||||||
currentAssistantMessage: string
|
currentAssistantMessage: string
|
||||||
currentReasoning: string
|
currentReasoning: string
|
||||||
isProcessing: boolean
|
isProcessing: boolean
|
||||||
|
isStopping?: boolean
|
||||||
|
onStop?: () => void
|
||||||
message: string
|
message: string
|
||||||
onMessageChange: (message: string) => void
|
onMessageChange: (message: string) => void
|
||||||
onSubmit: (message: PromptInputMessage, mentions?: FileMention[]) => void
|
onSubmit: (message: PromptInputMessage, mentions?: FileMention[]) => void
|
||||||
|
|
@ -139,6 +141,8 @@ export function ChatSidebar({
|
||||||
currentAssistantMessage,
|
currentAssistantMessage,
|
||||||
currentReasoning,
|
currentReasoning,
|
||||||
isProcessing,
|
isProcessing,
|
||||||
|
isStopping,
|
||||||
|
onStop,
|
||||||
message,
|
message,
|
||||||
onMessageChange,
|
onMessageChange,
|
||||||
onSubmit,
|
onSubmit,
|
||||||
|
|
@ -595,19 +599,39 @@ export function ChatSidebar({
|
||||||
style={{ fieldSizing: 'content' } as React.CSSProperties}
|
style={{ fieldSizing: 'content' } as React.CSSProperties}
|
||||||
/>
|
/>
|
||||||
</div>
|
</div>
|
||||||
<Button
|
{isProcessing ? (
|
||||||
size="icon"
|
<Button
|
||||||
onClick={handleSubmit}
|
size="icon"
|
||||||
disabled={!canSubmit}
|
onClick={onStop}
|
||||||
className={cn(
|
title={isStopping ? "Click again to force stop" : "Stop generation"}
|
||||||
"h-7 w-7 rounded-full shrink-0 transition-all",
|
className={cn(
|
||||||
canSubmit
|
"h-7 w-7 rounded-full shrink-0 transition-all",
|
||||||
? "bg-primary text-primary-foreground hover:bg-primary/90"
|
isStopping
|
||||||
: "bg-muted text-muted-foreground"
|
? "bg-destructive text-destructive-foreground hover:bg-destructive/90"
|
||||||
)}
|
: "bg-primary text-primary-foreground hover:bg-primary/90"
|
||||||
>
|
)}
|
||||||
<ArrowUp className="h-4 w-4" />
|
>
|
||||||
</Button>
|
{isStopping ? (
|
||||||
|
<LoaderIcon className="h-4 w-4 animate-spin" />
|
||||||
|
) : (
|
||||||
|
<Square className="h-3 w-3 fill-current" />
|
||||||
|
)}
|
||||||
|
</Button>
|
||||||
|
) : (
|
||||||
|
<Button
|
||||||
|
size="icon"
|
||||||
|
onClick={handleSubmit}
|
||||||
|
disabled={!canSubmit}
|
||||||
|
className={cn(
|
||||||
|
"h-7 w-7 rounded-full shrink-0 transition-all",
|
||||||
|
canSubmit
|
||||||
|
? "bg-primary text-primary-foreground hover:bg-primary/90"
|
||||||
|
: "bg-muted text-muted-foreground"
|
||||||
|
)}
|
||||||
|
>
|
||||||
|
<ArrowUp className="h-4 w-4" />
|
||||||
|
</Button>
|
||||||
|
)}
|
||||||
</div>
|
</div>
|
||||||
{knowledgeFiles.length > 0 && (
|
{knowledgeFiles.length > 0 && (
|
||||||
<MentionPopover
|
<MentionPopover
|
||||||
|
|
|
||||||
|
|
@ -22,6 +22,7 @@ import { IBus } from "../application/lib/bus.js";
|
||||||
import { IMessageQueue } from "../application/lib/message-queue.js";
|
import { IMessageQueue } from "../application/lib/message-queue.js";
|
||||||
import { IRunsRepo } from "../runs/repo.js";
|
import { IRunsRepo } from "../runs/repo.js";
|
||||||
import { IRunsLock } from "../runs/lock.js";
|
import { IRunsLock } from "../runs/lock.js";
|
||||||
|
import { IAbortRegistry } from "../runs/abort-registry.js";
|
||||||
import { PrefixLogger } from "@x/shared";
|
import { PrefixLogger } from "@x/shared";
|
||||||
import { parse } from "yaml";
|
import { parse } from "yaml";
|
||||||
import { raw as noteCreationMediumRaw } from "../knowledge/note_creation_medium.js";
|
import { raw as noteCreationMediumRaw } from "../knowledge/note_creation_medium.js";
|
||||||
|
|
@ -39,6 +40,7 @@ export class AgentRuntime implements IAgentRuntime {
|
||||||
private messageQueue: IMessageQueue;
|
private messageQueue: IMessageQueue;
|
||||||
private modelConfigRepo: IModelConfigRepo;
|
private modelConfigRepo: IModelConfigRepo;
|
||||||
private runsLock: IRunsLock;
|
private runsLock: IRunsLock;
|
||||||
|
private abortRegistry: IAbortRegistry;
|
||||||
|
|
||||||
constructor({
|
constructor({
|
||||||
runsRepo,
|
runsRepo,
|
||||||
|
|
@ -47,6 +49,7 @@ export class AgentRuntime implements IAgentRuntime {
|
||||||
messageQueue,
|
messageQueue,
|
||||||
modelConfigRepo,
|
modelConfigRepo,
|
||||||
runsLock,
|
runsLock,
|
||||||
|
abortRegistry,
|
||||||
}: {
|
}: {
|
||||||
runsRepo: IRunsRepo;
|
runsRepo: IRunsRepo;
|
||||||
idGenerator: IMonotonicallyIncreasingIdGenerator;
|
idGenerator: IMonotonicallyIncreasingIdGenerator;
|
||||||
|
|
@ -54,6 +57,7 @@ export class AgentRuntime implements IAgentRuntime {
|
||||||
messageQueue: IMessageQueue;
|
messageQueue: IMessageQueue;
|
||||||
modelConfigRepo: IModelConfigRepo;
|
modelConfigRepo: IModelConfigRepo;
|
||||||
runsLock: IRunsLock;
|
runsLock: IRunsLock;
|
||||||
|
abortRegistry: IAbortRegistry;
|
||||||
}) {
|
}) {
|
||||||
this.runsRepo = runsRepo;
|
this.runsRepo = runsRepo;
|
||||||
this.idGenerator = idGenerator;
|
this.idGenerator = idGenerator;
|
||||||
|
|
@ -61,6 +65,7 @@ export class AgentRuntime implements IAgentRuntime {
|
||||||
this.messageQueue = messageQueue;
|
this.messageQueue = messageQueue;
|
||||||
this.modelConfigRepo = modelConfigRepo;
|
this.modelConfigRepo = modelConfigRepo;
|
||||||
this.runsLock = runsLock;
|
this.runsLock = runsLock;
|
||||||
|
this.abortRegistry = abortRegistry;
|
||||||
}
|
}
|
||||||
|
|
||||||
async trigger(runId: string): Promise<void> {
|
async trigger(runId: string): Promise<void> {
|
||||||
|
|
@ -68,6 +73,7 @@ export class AgentRuntime implements IAgentRuntime {
|
||||||
console.log(`unable to acquire lock on run ${runId}`);
|
console.log(`unable to acquire lock on run ${runId}`);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
const signal = this.abortRegistry.createForRun(runId);
|
||||||
try {
|
try {
|
||||||
await this.bus.publish({
|
await this.bus.publish({
|
||||||
runId,
|
runId,
|
||||||
|
|
@ -75,6 +81,11 @@ export class AgentRuntime implements IAgentRuntime {
|
||||||
subflow: [],
|
subflow: [],
|
||||||
});
|
});
|
||||||
while (true) {
|
while (true) {
|
||||||
|
// Check for abort before each iteration
|
||||||
|
if (signal.aborted) {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
let eventCount = 0;
|
let eventCount = 0;
|
||||||
const run = await this.runsRepo.fetch(runId);
|
const run = await this.runsRepo.fetch(runId);
|
||||||
if (!run) {
|
if (!run) {
|
||||||
|
|
@ -84,18 +95,28 @@ export class AgentRuntime implements IAgentRuntime {
|
||||||
for (const event of run.log) {
|
for (const event of run.log) {
|
||||||
state.ingest(event);
|
state.ingest(event);
|
||||||
}
|
}
|
||||||
for await (const event of streamAgent({
|
try {
|
||||||
state,
|
for await (const event of streamAgent({
|
||||||
idGenerator: this.idGenerator,
|
state,
|
||||||
runId,
|
idGenerator: this.idGenerator,
|
||||||
messageQueue: this.messageQueue,
|
runId,
|
||||||
modelConfigRepo: this.modelConfigRepo,
|
messageQueue: this.messageQueue,
|
||||||
})) {
|
modelConfigRepo: this.modelConfigRepo,
|
||||||
eventCount++;
|
signal,
|
||||||
if (event.type !== "llm-stream-event") {
|
abortRegistry: this.abortRegistry,
|
||||||
await this.runsRepo.appendEvents(runId, [event]);
|
})) {
|
||||||
|
eventCount++;
|
||||||
|
if (event.type !== "llm-stream-event") {
|
||||||
|
await this.runsRepo.appendEvents(runId, [event]);
|
||||||
|
}
|
||||||
|
await this.bus.publish(event);
|
||||||
}
|
}
|
||||||
await this.bus.publish(event);
|
} catch (error) {
|
||||||
|
if (error instanceof Error && error.name === "AbortError") {
|
||||||
|
// Abort detected — exit cleanly
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
throw error;
|
||||||
}
|
}
|
||||||
|
|
||||||
// if no events, break
|
// if no events, break
|
||||||
|
|
@ -103,7 +124,20 @@ export class AgentRuntime implements IAgentRuntime {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Emit run-stopped event if aborted
|
||||||
|
if (signal.aborted) {
|
||||||
|
const stoppedEvent: z.infer<typeof RunEvent> = {
|
||||||
|
runId,
|
||||||
|
type: "run-stopped",
|
||||||
|
reason: "user-requested",
|
||||||
|
subflow: [],
|
||||||
|
};
|
||||||
|
await this.runsRepo.appendEvents(runId, [stoppedEvent]);
|
||||||
|
await this.bus.publish(stoppedEvent);
|
||||||
|
}
|
||||||
} finally {
|
} finally {
|
||||||
|
this.abortRegistry.cleanup(runId);
|
||||||
await this.runsLock.release(runId);
|
await this.runsLock.release(runId);
|
||||||
await this.bus.publish({
|
await this.bus.publish({
|
||||||
runId,
|
runId,
|
||||||
|
|
@ -428,6 +462,39 @@ export class AgentState {
|
||||||
return response;
|
return response;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Returns tool-result messages for all pending tool calls, marking them as aborted.
|
||||||
|
* This is called when a run is stopped so the LLM knows what happened to its tool requests.
|
||||||
|
*/
|
||||||
|
getAbortedToolResults(): z.infer<typeof ToolMessage>[] {
|
||||||
|
const results: z.infer<typeof ToolMessage>[] = [];
|
||||||
|
for (const toolCallId of Object.keys(this.pendingToolCalls)) {
|
||||||
|
const toolCall = this.toolCallIdMap[toolCallId];
|
||||||
|
if (toolCall) {
|
||||||
|
results.push({
|
||||||
|
role: "tool",
|
||||||
|
content: JSON.stringify({ error: "Tool execution aborted" }),
|
||||||
|
toolCallId,
|
||||||
|
toolName: toolCall.toolName,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return results;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Clear all pending state (permissions, ask-human, tool calls).
|
||||||
|
* Used when a run is stopped.
|
||||||
|
*/
|
||||||
|
clearAllPending(): void {
|
||||||
|
this.pendingToolPermissionRequests = {};
|
||||||
|
this.pendingAskHumanRequests = {};
|
||||||
|
// Recursively clear subflows
|
||||||
|
for (const subflow of Object.values(this.subflowStates)) {
|
||||||
|
subflow.clearAllPending();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
finalResponse(): string {
|
finalResponse(): string {
|
||||||
if (!this.lastAssistantMsg) {
|
if (!this.lastAssistantMsg) {
|
||||||
return '';
|
return '';
|
||||||
|
|
@ -526,12 +593,16 @@ export async function* streamAgent({
|
||||||
runId,
|
runId,
|
||||||
messageQueue,
|
messageQueue,
|
||||||
modelConfigRepo,
|
modelConfigRepo,
|
||||||
|
signal,
|
||||||
|
abortRegistry,
|
||||||
}: {
|
}: {
|
||||||
state: AgentState,
|
state: AgentState,
|
||||||
idGenerator: IMonotonicallyIncreasingIdGenerator;
|
idGenerator: IMonotonicallyIncreasingIdGenerator;
|
||||||
runId: string;
|
runId: string;
|
||||||
messageQueue: IMessageQueue;
|
messageQueue: IMessageQueue;
|
||||||
modelConfigRepo: IModelConfigRepo;
|
modelConfigRepo: IModelConfigRepo;
|
||||||
|
signal: AbortSignal;
|
||||||
|
abortRegistry: IAbortRegistry;
|
||||||
}): AsyncGenerator<z.infer<typeof RunEvent>, void, unknown> {
|
}): AsyncGenerator<z.infer<typeof RunEvent>, void, unknown> {
|
||||||
const logger = new PrefixLogger(`run-${runId}-${state.agentName}`);
|
const logger = new PrefixLogger(`run-${runId}-${state.agentName}`);
|
||||||
|
|
||||||
|
|
@ -557,6 +628,9 @@ export async function* streamAgent({
|
||||||
|
|
||||||
let loopCounter = 0;
|
let loopCounter = 0;
|
||||||
while (true) {
|
while (true) {
|
||||||
|
// Check abort at the top of each iteration
|
||||||
|
signal.throwIfAborted();
|
||||||
|
|
||||||
loopCounter++;
|
loopCounter++;
|
||||||
const loopLogger = logger.child(`iter-${loopCounter}`);
|
const loopLogger = logger.child(`iter-${loopCounter}`);
|
||||||
loopLogger.log('starting loop iteration');
|
loopLogger.log('starting loop iteration');
|
||||||
|
|
@ -598,6 +672,11 @@ export async function* streamAgent({
|
||||||
}
|
}
|
||||||
|
|
||||||
// execute approved tool
|
// execute approved tool
|
||||||
|
// Check abort before starting tool execution
|
||||||
|
if (signal.aborted) {
|
||||||
|
_logger.log('skipping, reason: aborted');
|
||||||
|
break;
|
||||||
|
}
|
||||||
_logger.log('executing tool');
|
_logger.log('executing tool');
|
||||||
yield* processEvent({
|
yield* processEvent({
|
||||||
runId,
|
runId,
|
||||||
|
|
@ -616,6 +695,8 @@ export async function* streamAgent({
|
||||||
runId,
|
runId,
|
||||||
messageQueue,
|
messageQueue,
|
||||||
modelConfigRepo,
|
modelConfigRepo,
|
||||||
|
signal,
|
||||||
|
abortRegistry,
|
||||||
})) {
|
})) {
|
||||||
yield* processEvent({
|
yield* processEvent({
|
||||||
...event,
|
...event,
|
||||||
|
|
@ -626,7 +707,7 @@ export async function* streamAgent({
|
||||||
result = subflowState.finalResponse();
|
result = subflowState.finalResponse();
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
result = await execTool(agent.tools![toolCall.toolName], toolCall.arguments);
|
result = await execTool(agent.tools![toolCall.toolName], toolCall.arguments, { runId, signal, abortRegistry });
|
||||||
}
|
}
|
||||||
const resultPayload = result === undefined ? null : result;
|
const resultPayload = result === undefined ? null : result;
|
||||||
const resultMsg: z.infer<typeof ToolMessage> = {
|
const resultMsg: z.infer<typeof ToolMessage> = {
|
||||||
|
|
@ -709,6 +790,7 @@ export async function* streamAgent({
|
||||||
state.messages,
|
state.messages,
|
||||||
instructionsWithDateTime,
|
instructionsWithDateTime,
|
||||||
tools,
|
tools,
|
||||||
|
signal,
|
||||||
)) {
|
)) {
|
||||||
// Only log significant events (not text-delta to reduce noise)
|
// Only log significant events (not text-delta to reduce noise)
|
||||||
if (event.type !== 'text-delta') {
|
if (event.type !== 'text-delta') {
|
||||||
|
|
@ -791,6 +873,7 @@ async function* streamLlm(
|
||||||
messages: z.infer<typeof MessageList>,
|
messages: z.infer<typeof MessageList>,
|
||||||
instructions: string,
|
instructions: string,
|
||||||
tools: ToolSet,
|
tools: ToolSet,
|
||||||
|
signal?: AbortSignal,
|
||||||
): AsyncGenerator<z.infer<typeof LlmStepStreamEvent>, void, unknown> {
|
): AsyncGenerator<z.infer<typeof LlmStepStreamEvent>, void, unknown> {
|
||||||
const { fullStream } = streamText({
|
const { fullStream } = streamText({
|
||||||
model,
|
model,
|
||||||
|
|
@ -798,8 +881,11 @@ async function* streamLlm(
|
||||||
system: instructions,
|
system: instructions,
|
||||||
tools,
|
tools,
|
||||||
stopWhen: stepCountIs(1),
|
stopWhen: stepCountIs(1),
|
||||||
|
abortSignal: signal,
|
||||||
});
|
});
|
||||||
for await (const event of fullStream) {
|
for await (const event of fullStream) {
|
||||||
|
// Check abort on every chunk for responsiveness
|
||||||
|
signal?.throwIfAborted();
|
||||||
// console.log("\n\n\t>>>>\t\tstream event", JSON.stringify(event));
|
// console.log("\n\n\t>>>>\t\tstream event", JSON.stringify(event));
|
||||||
switch (event.type) {
|
switch (event.type) {
|
||||||
case "reasoning-start":
|
case "reasoning-start":
|
||||||
|
|
|
||||||
|
|
@ -2,7 +2,7 @@ import { z, ZodType } from "zod";
|
||||||
import * as path from "path";
|
import * as path from "path";
|
||||||
import { execSync } from "child_process";
|
import { execSync } from "child_process";
|
||||||
import { glob } from "glob";
|
import { glob } from "glob";
|
||||||
import { executeCommand } from "./command-executor.js";
|
import { executeCommand, executeCommandAbortable } from "./command-executor.js";
|
||||||
import { resolveSkill, availableSkills } from "../assistant/skills/index.js";
|
import { resolveSkill, availableSkills } from "../assistant/skills/index.js";
|
||||||
import { executeTool, listServers, listTools } from "../../mcp/mcp.js";
|
import { executeTool, listServers, listTools } from "../../mcp/mcp.js";
|
||||||
import container from "../../di/container.js";
|
import container from "../../di/container.js";
|
||||||
|
|
@ -11,13 +11,14 @@ import { McpServerDefinition } from "@x/shared/dist/mcp.js";
|
||||||
import * as workspace from "../../workspace/workspace.js";
|
import * as workspace from "../../workspace/workspace.js";
|
||||||
import { IAgentsRepo } from "../../agents/repo.js";
|
import { IAgentsRepo } from "../../agents/repo.js";
|
||||||
import { WorkDir } from "../../config/config.js";
|
import { WorkDir } from "../../config/config.js";
|
||||||
|
import type { ToolContext } from "./exec-tool.js";
|
||||||
|
|
||||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||||
const BuiltinToolsSchema = z.record(z.string(), z.object({
|
const BuiltinToolsSchema = z.record(z.string(), z.object({
|
||||||
description: z.string(),
|
description: z.string(),
|
||||||
inputSchema: z.custom<ZodType>(),
|
inputSchema: z.custom<ZodType>(),
|
||||||
execute: z.function({
|
execute: z.function({
|
||||||
input: z.any(),
|
input: z.any(), // (input, ctx?) => Promise<any>
|
||||||
output: z.promise(z.any()),
|
output: z.promise(z.any()),
|
||||||
}),
|
}),
|
||||||
}));
|
}));
|
||||||
|
|
@ -611,15 +612,15 @@ export const BuiltinTools: z.infer<typeof BuiltinToolsSchema> = {
|
||||||
command: z.string().describe('The shell command to execute (e.g., "ls -la", "cat file.txt")'),
|
command: z.string().describe('The shell command to execute (e.g., "ls -la", "cat file.txt")'),
|
||||||
cwd: z.string().optional().describe('Working directory to execute the command in (defaults to workspace root)'),
|
cwd: z.string().optional().describe('Working directory to execute the command in (defaults to workspace root)'),
|
||||||
}),
|
}),
|
||||||
execute: async ({ command, cwd }: { command: string, cwd?: string }) => {
|
execute: async ({ command, cwd }: { command: string, cwd?: string }, ctx?: ToolContext) => {
|
||||||
try {
|
try {
|
||||||
const rootDir = path.resolve(WorkDir);
|
const rootDir = path.resolve(WorkDir);
|
||||||
const workingDir = cwd ? path.resolve(rootDir, cwd) : rootDir;
|
const workingDir = cwd ? path.resolve(rootDir, cwd) : rootDir;
|
||||||
const rootPrefix = rootDir.endsWith(path.sep)
|
|
||||||
? rootDir
|
|
||||||
: `${rootDir}${path.sep}`;
|
|
||||||
|
|
||||||
// TODO: Re-enable this check
|
// TODO: Re-enable this check
|
||||||
|
// const rootPrefix = rootDir.endsWith(path.sep)
|
||||||
|
// ? rootDir
|
||||||
|
// : `${rootDir}${path.sep}`;
|
||||||
// if (workingDir !== rootDir && !workingDir.startsWith(rootPrefix)) {
|
// if (workingDir !== rootDir && !workingDir.startsWith(rootPrefix)) {
|
||||||
// return {
|
// return {
|
||||||
// success: false,
|
// success: false,
|
||||||
|
|
@ -629,8 +630,32 @@ export const BuiltinTools: z.infer<typeof BuiltinToolsSchema> = {
|
||||||
// };
|
// };
|
||||||
// }
|
// }
|
||||||
|
|
||||||
|
// Use abortable version when we have a signal
|
||||||
|
if (ctx?.signal) {
|
||||||
|
const { promise, process: proc } = executeCommandAbortable(command, {
|
||||||
|
cwd: workingDir,
|
||||||
|
signal: ctx.signal,
|
||||||
|
});
|
||||||
|
|
||||||
|
// Register process with abort registry for force-kill
|
||||||
|
ctx.abortRegistry.registerProcess(ctx.runId, proc);
|
||||||
|
|
||||||
|
const result = await promise;
|
||||||
|
|
||||||
|
return {
|
||||||
|
success: result.exitCode === 0 && !result.wasAborted,
|
||||||
|
stdout: result.stdout,
|
||||||
|
stderr: result.stderr,
|
||||||
|
exitCode: result.exitCode,
|
||||||
|
wasAborted: result.wasAborted,
|
||||||
|
command,
|
||||||
|
workingDir,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
// Fallback to original for backward compatibility
|
||||||
const result = await executeCommand(command, { cwd: workingDir });
|
const result = await executeCommand(command, { cwd: workingDir });
|
||||||
|
|
||||||
return {
|
return {
|
||||||
success: result.exitCode === 0,
|
success: result.exitCode === 0,
|
||||||
stdout: result.stdout,
|
stdout: result.stdout,
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,4 @@
|
||||||
import { exec, execSync } from 'child_process';
|
import { exec, execSync, spawn, ChildProcess } from 'child_process';
|
||||||
import { promisify } from 'util';
|
import { promisify } from 'util';
|
||||||
import { getSecurityAllowList } from '../../config/security.js';
|
import { getSecurityAllowList } from '../../config/security.js';
|
||||||
|
|
||||||
|
|
@ -110,6 +110,159 @@ export async function executeCommand(
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export interface AbortableCommandResult extends CommandResult {
|
||||||
|
wasAborted: boolean;
|
||||||
|
}
|
||||||
|
|
||||||
|
const SIGKILL_GRACE_MS = 200;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Kill a process tree using negative PID (process group kill on Unix).
|
||||||
|
* Falls back to direct kill if group kill fails.
|
||||||
|
*/
|
||||||
|
function killProcessTree(proc: ChildProcess, signal: NodeJS.Signals): void {
|
||||||
|
if (!proc.pid || proc.killed) return;
|
||||||
|
try {
|
||||||
|
// Negative PID kills the entire process group (Unix)
|
||||||
|
process.kill(-proc.pid, signal);
|
||||||
|
} catch {
|
||||||
|
try {
|
||||||
|
proc.kill(signal);
|
||||||
|
} catch {
|
||||||
|
// Process may already be dead
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Executes a shell command with abort support.
|
||||||
|
* Uses spawn with detached=true to create a process group for proper tree killing.
|
||||||
|
* Returns both the promise and the child process handle.
|
||||||
|
*/
|
||||||
|
export function executeCommandAbortable(
|
||||||
|
command: string,
|
||||||
|
options?: {
|
||||||
|
cwd?: string;
|
||||||
|
timeout?: number;
|
||||||
|
maxBuffer?: number;
|
||||||
|
signal?: AbortSignal;
|
||||||
|
}
|
||||||
|
): { promise: Promise<AbortableCommandResult>; process: ChildProcess } {
|
||||||
|
// Check if already aborted before spawning
|
||||||
|
if (options?.signal?.aborted) {
|
||||||
|
// Return a dummy process and a resolved result
|
||||||
|
const dummyProc = spawn('true', { shell: true });
|
||||||
|
dummyProc.kill();
|
||||||
|
return {
|
||||||
|
process: dummyProc,
|
||||||
|
promise: Promise.resolve({
|
||||||
|
stdout: '',
|
||||||
|
stderr: '',
|
||||||
|
exitCode: 130,
|
||||||
|
wasAborted: true,
|
||||||
|
}),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
const proc = spawn(command, [], {
|
||||||
|
shell: '/bin/sh',
|
||||||
|
cwd: options?.cwd,
|
||||||
|
detached: process.platform !== 'win32', // Create process group on Unix
|
||||||
|
stdio: ['ignore', 'pipe', 'pipe'],
|
||||||
|
});
|
||||||
|
|
||||||
|
const promise = new Promise<AbortableCommandResult>((resolve) => {
|
||||||
|
let stdout = '';
|
||||||
|
let stderr = '';
|
||||||
|
let wasAborted = false;
|
||||||
|
let exited = false;
|
||||||
|
|
||||||
|
// Collect output
|
||||||
|
proc.stdout?.on('data', (chunk: Buffer) => {
|
||||||
|
const maxBuffer = options?.maxBuffer || 1024 * 1024;
|
||||||
|
if (stdout.length < maxBuffer) {
|
||||||
|
stdout += chunk.toString();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
proc.stderr?.on('data', (chunk: Buffer) => {
|
||||||
|
const maxBuffer = options?.maxBuffer || 1024 * 1024;
|
||||||
|
if (stderr.length < maxBuffer) {
|
||||||
|
stderr += chunk.toString();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// Abort handler
|
||||||
|
const abortHandler = () => {
|
||||||
|
wasAborted = true;
|
||||||
|
killProcessTree(proc, 'SIGTERM');
|
||||||
|
// Force kill after grace period
|
||||||
|
setTimeout(() => {
|
||||||
|
if (!exited) {
|
||||||
|
killProcessTree(proc, 'SIGKILL');
|
||||||
|
}
|
||||||
|
}, SIGKILL_GRACE_MS);
|
||||||
|
};
|
||||||
|
|
||||||
|
if (options?.signal) {
|
||||||
|
options.signal.addEventListener('abort', abortHandler, { once: true });
|
||||||
|
}
|
||||||
|
|
||||||
|
// Timeout handler
|
||||||
|
let timeoutId: ReturnType<typeof setTimeout> | undefined;
|
||||||
|
if (options?.timeout) {
|
||||||
|
timeoutId = setTimeout(() => {
|
||||||
|
wasAborted = true;
|
||||||
|
killProcessTree(proc, 'SIGTERM');
|
||||||
|
setTimeout(() => {
|
||||||
|
if (!exited) {
|
||||||
|
killProcessTree(proc, 'SIGKILL');
|
||||||
|
}
|
||||||
|
}, SIGKILL_GRACE_MS);
|
||||||
|
}, options.timeout);
|
||||||
|
}
|
||||||
|
|
||||||
|
proc.once('exit', (code) => {
|
||||||
|
exited = true;
|
||||||
|
// Cleanup listeners
|
||||||
|
if (options?.signal) {
|
||||||
|
options.signal.removeEventListener('abort', abortHandler);
|
||||||
|
}
|
||||||
|
if (timeoutId) {
|
||||||
|
clearTimeout(timeoutId);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (wasAborted) {
|
||||||
|
stdout += '\n\n(Command was aborted)';
|
||||||
|
}
|
||||||
|
|
||||||
|
resolve({
|
||||||
|
stdout: stdout.trim(),
|
||||||
|
stderr: stderr.trim(),
|
||||||
|
exitCode: code ?? 1,
|
||||||
|
wasAborted,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
proc.once('error', (err) => {
|
||||||
|
exited = true;
|
||||||
|
if (options?.signal) {
|
||||||
|
options.signal.removeEventListener('abort', abortHandler);
|
||||||
|
}
|
||||||
|
if (timeoutId) {
|
||||||
|
clearTimeout(timeoutId);
|
||||||
|
}
|
||||||
|
resolve({
|
||||||
|
stdout: '',
|
||||||
|
stderr: err.message,
|
||||||
|
exitCode: 1,
|
||||||
|
wasAborted,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
return { promise, process: proc };
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Executes a command synchronously (blocking)
|
* Executes a command synchronously (blocking)
|
||||||
* Use with caution - prefer executeCommand for async execution
|
* Use with caution - prefer executeCommand for async execution
|
||||||
|
|
|
||||||
|
|
@ -2,22 +2,36 @@ import { ToolAttachment } from "@x/shared/dist/agent.js";
|
||||||
import { z } from "zod";
|
import { z } from "zod";
|
||||||
import { BuiltinTools } from "./builtin-tools.js";
|
import { BuiltinTools } from "./builtin-tools.js";
|
||||||
import { executeTool } from "../../mcp/mcp.js";
|
import { executeTool } from "../../mcp/mcp.js";
|
||||||
|
import { IAbortRegistry } from "../../runs/abort-registry.js";
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Context passed to every tool execution, providing abort signal and run metadata.
|
||||||
|
*/
|
||||||
|
export interface ToolContext {
|
||||||
|
runId: string;
|
||||||
|
signal: AbortSignal;
|
||||||
|
abortRegistry: IAbortRegistry;
|
||||||
|
}
|
||||||
|
|
||||||
async function execMcpTool(agentTool: z.infer<typeof ToolAttachment> & { type: "mcp" }, input: Record<string, unknown>): Promise<unknown> {
|
async function execMcpTool(agentTool: z.infer<typeof ToolAttachment> & { type: "mcp" }, input: Record<string, unknown>): Promise<unknown> {
|
||||||
const result = await executeTool(agentTool.mcpServerName, agentTool.name, input);
|
const result = await executeTool(agentTool.mcpServerName, agentTool.name, input);
|
||||||
return result;
|
return result;
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function execTool(agentTool: z.infer<typeof ToolAttachment>, input: Record<string, unknown>): Promise<unknown> {
|
export async function execTool(agentTool: z.infer<typeof ToolAttachment>, input: Record<string, unknown>, ctx?: ToolContext): Promise<unknown> {
|
||||||
|
// Check abort before starting any tool
|
||||||
|
ctx?.signal.throwIfAborted();
|
||||||
|
|
||||||
switch (agentTool.type) {
|
switch (agentTool.type) {
|
||||||
case "mcp":
|
case "mcp":
|
||||||
|
// MCP tools: let complete on graceful stop (most are fast)
|
||||||
return execMcpTool(agentTool, input);
|
return execMcpTool(agentTool, input);
|
||||||
case "builtin": {
|
case "builtin": {
|
||||||
const builtinTool = BuiltinTools[agentTool.name];
|
const builtinTool = BuiltinTools[agentTool.name];
|
||||||
if (!builtinTool || !builtinTool.execute) {
|
if (!builtinTool || !builtinTool.execute) {
|
||||||
throw new Error(`Unsupported builtin tool: ${agentTool.name}`);
|
throw new Error(`Unsupported builtin tool: ${agentTool.name}`);
|
||||||
}
|
}
|
||||||
return builtinTool.execute(input);
|
return builtinTool.execute(input, ctx);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -11,6 +11,7 @@ import { IAgentRuntime, AgentRuntime } from "../agents/runtime.js";
|
||||||
import { FSOAuthRepo, IOAuthRepo } from "../auth/repo.js";
|
import { FSOAuthRepo, IOAuthRepo } from "../auth/repo.js";
|
||||||
import { FSClientRegistrationRepo, IClientRegistrationRepo } from "../auth/client-repo.js";
|
import { FSClientRegistrationRepo, IClientRegistrationRepo } from "../auth/client-repo.js";
|
||||||
import { FSGranolaConfigRepo, IGranolaConfigRepo } from "../knowledge/granola/repo.js";
|
import { FSGranolaConfigRepo, IGranolaConfigRepo } from "../knowledge/granola/repo.js";
|
||||||
|
import { IAbortRegistry, InMemoryAbortRegistry } from "../runs/abort-registry.js";
|
||||||
|
|
||||||
const container = createContainer({
|
const container = createContainer({
|
||||||
injectionMode: InjectionMode.PROXY,
|
injectionMode: InjectionMode.PROXY,
|
||||||
|
|
@ -22,6 +23,7 @@ container.register({
|
||||||
messageQueue: asClass<IMessageQueue>(InMemoryMessageQueue).singleton(),
|
messageQueue: asClass<IMessageQueue>(InMemoryMessageQueue).singleton(),
|
||||||
bus: asClass<IBus>(InMemoryBus).singleton(),
|
bus: asClass<IBus>(InMemoryBus).singleton(),
|
||||||
runsLock: asClass<IRunsLock>(InMemoryRunsLock).singleton(),
|
runsLock: asClass<IRunsLock>(InMemoryRunsLock).singleton(),
|
||||||
|
abortRegistry: asClass<IAbortRegistry>(InMemoryAbortRegistry).singleton(),
|
||||||
agentRuntime: asClass<IAgentRuntime>(AgentRuntime).singleton(),
|
agentRuntime: asClass<IAgentRuntime>(AgentRuntime).singleton(),
|
||||||
|
|
||||||
mcpConfigRepo: asClass<IMcpConfigRepo>(FSMcpConfigRepo).singleton(),
|
mcpConfigRepo: asClass<IMcpConfigRepo>(FSMcpConfigRepo).singleton(),
|
||||||
|
|
|
||||||
|
|
@ -84,6 +84,22 @@ export async function cleanup() {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Force-close all MCP client connections.
|
||||||
|
* Used during force abort to immediately reject any pending MCP tool calls.
|
||||||
|
* Clients will be lazily reconnected on next use.
|
||||||
|
*/
|
||||||
|
export async function forceCloseAllMcpClients(): Promise<void> {
|
||||||
|
for (const [serverName, { client }] of Object.entries(clients)) {
|
||||||
|
try {
|
||||||
|
await client?.close();
|
||||||
|
} catch {
|
||||||
|
// Ignore errors during force close
|
||||||
|
}
|
||||||
|
delete clients[serverName];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
export async function listServers(): Promise<z.infer<typeof McpServerList>> {
|
export async function listServers(): Promise<z.infer<typeof McpServerList>> {
|
||||||
const repo = container.resolve<IMcpConfigRepo>('mcpConfigRepo');
|
const repo = container.resolve<IMcpConfigRepo>('mcpConfigRepo');
|
||||||
const { mcpServers } = await repo.getConfig();
|
const { mcpServers } = await repo.getConfig();
|
||||||
|
|
|
||||||
170
apps/x/packages/core/src/runs/abort-registry.ts
Normal file
170
apps/x/packages/core/src/runs/abort-registry.ts
Normal file
|
|
@ -0,0 +1,170 @@
|
||||||
|
import { ChildProcess } from "child_process";
|
||||||
|
|
||||||
|
export interface IAbortRegistry {
|
||||||
|
/**
|
||||||
|
* Create and track an AbortController for a run.
|
||||||
|
* Returns the AbortSignal to thread through all operations.
|
||||||
|
*/
|
||||||
|
createForRun(runId: string): AbortSignal;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Track a child process for a run (so we can kill it on abort).
|
||||||
|
*/
|
||||||
|
registerProcess(runId: string, process: ChildProcess): void;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Untrack a child process after it exits.
|
||||||
|
*/
|
||||||
|
unregisterProcess(runId: string, process: ChildProcess): void;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Graceful abort:
|
||||||
|
* 1. Fires the AbortSignal (cancels LLM streaming, etc.)
|
||||||
|
* 2. Sends SIGTERM to all tracked process groups
|
||||||
|
* 3. Schedules SIGKILL fallback after grace period
|
||||||
|
*/
|
||||||
|
abort(runId: string): void;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Force abort:
|
||||||
|
* 1. Fires AbortSignal if not already fired
|
||||||
|
* 2. Sends SIGKILL to all tracked process groups immediately
|
||||||
|
*/
|
||||||
|
forceAbort(runId: string): void;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Check if a run has been aborted.
|
||||||
|
*/
|
||||||
|
isAborted(runId: string): boolean;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Clean up tracking state after a run completes or is fully stopped.
|
||||||
|
*/
|
||||||
|
cleanup(runId: string): void;
|
||||||
|
}
|
||||||
|
|
||||||
|
interface RunAbortState {
|
||||||
|
controller: AbortController;
|
||||||
|
processes: Set<ChildProcess>;
|
||||||
|
killTimers: Set<ReturnType<typeof setTimeout>>;
|
||||||
|
}
|
||||||
|
|
||||||
|
const SIGKILL_GRACE_MS = 200;
|
||||||
|
|
||||||
|
export class InMemoryAbortRegistry implements IAbortRegistry {
|
||||||
|
private runs: Map<string, RunAbortState> = new Map();
|
||||||
|
|
||||||
|
createForRun(runId: string): AbortSignal {
|
||||||
|
// If a previous run state exists, clean it up first
|
||||||
|
this.cleanup(runId);
|
||||||
|
|
||||||
|
const state: RunAbortState = {
|
||||||
|
controller: new AbortController(),
|
||||||
|
processes: new Set(),
|
||||||
|
killTimers: new Set(),
|
||||||
|
};
|
||||||
|
this.runs.set(runId, state);
|
||||||
|
return state.controller.signal;
|
||||||
|
}
|
||||||
|
|
||||||
|
registerProcess(runId: string, process: ChildProcess): void {
|
||||||
|
const state = this.runs.get(runId);
|
||||||
|
if (!state) return;
|
||||||
|
state.processes.add(process);
|
||||||
|
|
||||||
|
// Auto-unregister when process exits
|
||||||
|
const onExit = () => {
|
||||||
|
state.processes.delete(process);
|
||||||
|
};
|
||||||
|
process.once("exit", onExit);
|
||||||
|
process.once("error", onExit);
|
||||||
|
}
|
||||||
|
|
||||||
|
unregisterProcess(runId: string, process: ChildProcess): void {
|
||||||
|
const state = this.runs.get(runId);
|
||||||
|
if (!state) return;
|
||||||
|
state.processes.delete(process);
|
||||||
|
}
|
||||||
|
|
||||||
|
abort(runId: string): void {
|
||||||
|
const state = this.runs.get(runId);
|
||||||
|
if (!state) return;
|
||||||
|
|
||||||
|
// 1. Fire the abort signal
|
||||||
|
if (!state.controller.signal.aborted) {
|
||||||
|
state.controller.abort();
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. SIGTERM all tracked process groups
|
||||||
|
for (const proc of state.processes) {
|
||||||
|
this.killProcessTree(proc, "SIGTERM");
|
||||||
|
|
||||||
|
// 3. Schedule SIGKILL fallback
|
||||||
|
const timer = setTimeout(() => {
|
||||||
|
if (!proc.killed) {
|
||||||
|
this.killProcessTree(proc, "SIGKILL");
|
||||||
|
}
|
||||||
|
state.killTimers.delete(timer);
|
||||||
|
}, SIGKILL_GRACE_MS);
|
||||||
|
state.killTimers.add(timer);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
forceAbort(runId: string): void {
|
||||||
|
const state = this.runs.get(runId);
|
||||||
|
if (!state) return;
|
||||||
|
|
||||||
|
// 1. Fire abort signal if not already
|
||||||
|
if (!state.controller.signal.aborted) {
|
||||||
|
state.controller.abort();
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Clear any pending graceful kill timers
|
||||||
|
for (const timer of state.killTimers) {
|
||||||
|
clearTimeout(timer);
|
||||||
|
}
|
||||||
|
state.killTimers.clear();
|
||||||
|
|
||||||
|
// 3. SIGKILL all tracked process groups immediately
|
||||||
|
for (const proc of state.processes) {
|
||||||
|
this.killProcessTree(proc, "SIGKILL");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
isAborted(runId: string): boolean {
|
||||||
|
const state = this.runs.get(runId);
|
||||||
|
return state?.controller.signal.aborted ?? false;
|
||||||
|
}
|
||||||
|
|
||||||
|
cleanup(runId: string): void {
|
||||||
|
const state = this.runs.get(runId);
|
||||||
|
if (!state) return;
|
||||||
|
|
||||||
|
// Clear any pending kill timers
|
||||||
|
for (const timer of state.killTimers) {
|
||||||
|
clearTimeout(timer);
|
||||||
|
}
|
||||||
|
|
||||||
|
this.runs.delete(runId);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Kill a process tree using negative PID (process group kill on Unix).
|
||||||
|
* Falls back to direct kill if group kill fails.
|
||||||
|
*/
|
||||||
|
private killProcessTree(proc: ChildProcess, signal: NodeJS.Signals): void {
|
||||||
|
if (!proc.pid || proc.killed) return;
|
||||||
|
|
||||||
|
try {
|
||||||
|
// Negative PID kills the entire process group (Unix)
|
||||||
|
process.kill(-proc.pid, signal);
|
||||||
|
} catch {
|
||||||
|
// Fallback: kill just the process directly
|
||||||
|
try {
|
||||||
|
proc.kill(signal);
|
||||||
|
} catch {
|
||||||
|
// Process may already be dead
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -5,6 +5,8 @@ import { AskHumanResponseEvent, ToolPermissionResponseEvent, CreateRunOptions, R
|
||||||
import { IRunsRepo } from "./repo.js";
|
import { IRunsRepo } from "./repo.js";
|
||||||
import { IAgentRuntime } from "../agents/runtime.js";
|
import { IAgentRuntime } from "../agents/runtime.js";
|
||||||
import { IBus } from "../application/lib/bus.js";
|
import { IBus } from "../application/lib/bus.js";
|
||||||
|
import { IAbortRegistry } from "./abort-registry.js";
|
||||||
|
import { forceCloseAllMcpClients } from "../mcp/mcp.js";
|
||||||
|
|
||||||
export async function createRun(opts: z.infer<typeof CreateRunOptions>): Promise<z.infer<typeof Run>> {
|
export async function createRun(opts: z.infer<typeof CreateRunOptions>): Promise<z.infer<typeof Run>> {
|
||||||
const repo = container.resolve<IRunsRepo>('runsRepo');
|
const repo = container.resolve<IRunsRepo>('runsRepo');
|
||||||
|
|
@ -46,9 +48,21 @@ export async function replyToHumanInputRequest(runId: string, ev: z.infer<typeof
|
||||||
runtime.trigger(runId);
|
runtime.trigger(runId);
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function stop(runId: string): Promise<void> {
|
export async function stop(runId: string, force: boolean = false): Promise<void> {
|
||||||
console.log(`Stopping run ${runId}`);
|
const abortRegistry = container.resolve<IAbortRegistry>('abortRegistry');
|
||||||
throw new Error('Not implemented');
|
|
||||||
|
if (force && abortRegistry.isAborted(runId)) {
|
||||||
|
// Second click: aggressive cleanup — SIGKILL + force close MCP clients
|
||||||
|
console.log(`Force stopping run ${runId}`);
|
||||||
|
abortRegistry.forceAbort(runId);
|
||||||
|
await forceCloseAllMcpClients();
|
||||||
|
} else {
|
||||||
|
// First click: graceful — fires AbortSignal + SIGTERM
|
||||||
|
console.log(`Gracefully stopping run ${runId}`);
|
||||||
|
abortRegistry.abort(runId);
|
||||||
|
}
|
||||||
|
// Note: The run-stopped event is emitted by AgentRuntime.trigger() when it detects the abort.
|
||||||
|
// This avoids duplicate events and ensures proper sequencing.
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function fetchRun(runId: string): Promise<z.infer<typeof Run>> {
|
export async function fetchRun(runId: string): Promise<z.infer<typeof Run>> {
|
||||||
|
|
|
||||||
|
|
@ -151,6 +151,7 @@ const ipcSchemas = {
|
||||||
'runs:stop': {
|
'runs:stop': {
|
||||||
req: z.object({
|
req: z.object({
|
||||||
runId: z.string(),
|
runId: z.string(),
|
||||||
|
force: z.boolean().optional().default(false),
|
||||||
}),
|
}),
|
||||||
res: z.object({
|
res: z.object({
|
||||||
success: z.literal(true),
|
success: z.literal(true),
|
||||||
|
|
|
||||||
|
|
@ -80,6 +80,11 @@ export const RunErrorEvent = BaseRunEvent.extend({
|
||||||
error: z.string(),
|
error: z.string(),
|
||||||
});
|
});
|
||||||
|
|
||||||
|
export const RunStoppedEvent = BaseRunEvent.extend({
|
||||||
|
type: z.literal("run-stopped"),
|
||||||
|
reason: z.enum(["user-requested", "force-stopped"]).optional(),
|
||||||
|
});
|
||||||
|
|
||||||
export const RunEvent = z.union([
|
export const RunEvent = z.union([
|
||||||
RunProcessingStartEvent,
|
RunProcessingStartEvent,
|
||||||
RunProcessingEndEvent,
|
RunProcessingEndEvent,
|
||||||
|
|
@ -94,6 +99,7 @@ export const RunEvent = z.union([
|
||||||
ToolPermissionRequestEvent,
|
ToolPermissionRequestEvent,
|
||||||
ToolPermissionResponseEvent,
|
ToolPermissionResponseEvent,
|
||||||
RunErrorEvent,
|
RunErrorEvent,
|
||||||
|
RunStoppedEvent,
|
||||||
]);
|
]);
|
||||||
|
|
||||||
export const ToolPermissionAuthorizePayload = ToolPermissionResponseEvent.pick({
|
export const ToolPermissionAuthorizePayload = ToolPermissionResponseEvent.pick({
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue