diff --git a/sdk/src/__tests__/client-cancellable-run.test.ts b/sdk/src/__tests__/client-cancellable-run.test.ts new file mode 100644 index 0000000000..506d228241 --- /dev/null +++ b/sdk/src/__tests__/client-cancellable-run.test.ts @@ -0,0 +1,116 @@ +import * as mainPromptModule from '@codebuff/agent-runtime/main-prompt' +import { afterEach, describe, expect, it, mock, spyOn } from 'bun:test' + +import { CodebuffClient } from '../client' +import { createRunController } from '../run-controller' +import * as databaseModule from '../impl/database' + +function mockDatabase() { + spyOn(databaseModule, 'getUserInfoFromApiKey').mockResolvedValue({ + id: 'user-123', + email: 'test@example.com', + discord_id: null, + stripe_customer_id: null, + banned: false, + created_at: new Date('2024-01-01T00:00:00Z'), + }) + spyOn(databaseModule, 'fetchAgentFromDatabase').mockResolvedValue(null) + spyOn(databaseModule, 'startAgentRun').mockResolvedValue('run-1') + spyOn(databaseModule, 'finishAgentRun').mockResolvedValue(undefined) + spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1') +} + +describe('CodebuffClient runCancellable', () => { + afterEach(() => { + mock.restore() + }) + + it('aborts the active run signal when cancel is called', async () => { + mockDatabase() + + let runtimeSignal: AbortSignal | undefined + let markRuntimeStarted: () => void = () => {} + const runtimeStarted = new Promise((resolve) => { + markRuntimeStarted = resolve + }) + spyOn(mainPromptModule, 'callMainPrompt').mockImplementation( + async (params: Parameters[0]) => { + runtimeSignal = params.signal + markRuntimeStarted() + return await new Promise((_, reject) => { + params.signal.addEventListener( + 'abort', + () => reject(params.signal.reason), + { once: true }, + ) + }) + }, + ) + + const client = new CodebuffClient({ apiKey: 'test-key' }) + const activeRun = client.runCancellable({ + agent: 'base2', + prompt: 'set up preview', + }) + + await runtimeStarted + activeRun.cancel('Stopped from UI') + const result = await activeRun.result + + expect(activeRun.controller.cancelled).toBe(true) + expect(activeRun.signal.aborted).toBe(true) + expect(runtimeSignal?.aborted).toBe(true) + expect(result.output.type).toBe('error') + if (result.output.type === 'error') { + expect(result.output.message).toBe('Stopped from UI') + } + }) + + it('combines caller-provided signals with the run controller', async () => { + mockDatabase() + + let runtimeSignal: AbortSignal | undefined + let markRuntimeStarted: () => void = () => {} + const runtimeStarted = new Promise((resolve) => { + markRuntimeStarted = resolve + }) + spyOn(mainPromptModule, 'callMainPrompt').mockImplementation( + async (params: Parameters[0]) => { + runtimeSignal = params.signal + markRuntimeStarted() + return await new Promise((_, reject) => { + params.signal.addEventListener( + 'abort', + () => reject(params.signal.reason), + { once: true }, + ) + }) + }, + ) + + const externalController = new AbortController() + const runController = createRunController('cloud-run-1') + const client = new CodebuffClient({ apiKey: 'test-key' }) + const activeRun = client.runCancellable( + { + agent: 'base2', + prompt: 'build mobile preview', + signal: externalController.signal, + }, + runController, + ) + + await runtimeStarted + externalController.abort(new Error('Request disconnected')) + const result = await activeRun.result + + expect(activeRun.id).toBe('cloud-run-1') + expect(runController.cancelled).toBe(false) + expect(activeRun.signal.aborted).toBe(true) + expect(runtimeSignal?.aborted).toBe(true) + expect(result.output.type).toBe('error') + if (result.output.type === 'error') { + expect(result.output.message).toBe('Request disconnected') + } + }) +}) diff --git a/sdk/src/client.ts b/sdk/src/client.ts index b2f5d6bbe5..e1aa1b26c3 100644 --- a/sdk/src/client.ts +++ b/sdk/src/client.ts @@ -3,8 +3,14 @@ import { API_KEY_ENV_VAR } from '@codebuff/common/constants/paths' import { getWebsiteUrl } from './constants' import { getCodebuffApiKeyFromEnv } from './env' import { run } from './run' +import { + anySignal, + CodebuffRunController, + createRunController, +} from './run-controller' import type { RunOptions, CodebuffClientOptions } from './run' +import type { CancellableRun } from './run-controller' import type { RunState } from './run-state' export class CodebuffClient { @@ -59,6 +65,30 @@ export class CodebuffClient { return run({ ...this.options, ...options }) } + /** + * Start a run and return an explicit controller that can terminate it. + * + * Hosts with their own UI or API surface can keep `controller` (or the + * returned `cancel` function) in their in-flight run registry and call it + * from a Stop/Terminate action. Cancellation flows through the same + * AbortSignal used by terminal commands, LLM streams, and subagents. + */ + public runCancellable( + options: RunOptions & CodebuffClientOptions, + controller: CodebuffRunController = createRunController(), + ): CancellableRun { + const signal = anySignal([options.signal, controller.signal]) + const result = run({ ...this.options, ...options, signal }) + + return { + id: controller.id, + controller, + signal, + result, + cancel: (reason?: string | Error) => controller.cancel(reason), + } + } + /** * Check connection to the Codebuff backend by hitting the /healthz endpoint. * diff --git a/sdk/src/index.ts b/sdk/src/index.ts index 08ffe2a6a6..96a4bf9a00 100644 --- a/sdk/src/index.ts +++ b/sdk/src/index.ts @@ -7,6 +7,11 @@ export type { ImagePart, } from '@codebuff/common/types/messages/content-part' export { run, STATE_SNAPSHOT_INTERRUPTION_MESSAGE } from './run' +export { + CodebuffRunController, + createRunController, +} from './run-controller' +export type { CancellableRun } from './run-controller' export { getFiles } from './tools/read-files' export type { FileFilter, FileFilterResult } from './tools/read-files' export type { diff --git a/sdk/src/run-controller.ts b/sdk/src/run-controller.ts new file mode 100644 index 0000000000..10247d9d29 --- /dev/null +++ b/sdk/src/run-controller.ts @@ -0,0 +1,69 @@ +import type { RunState } from './run-state' + +const DEFAULT_CANCEL_REASON = 'Run cancelled by user.' + +export class CodebuffRunController { + public readonly id: string + private readonly controller: AbortController + + constructor(id: string = crypto.randomUUID()) { + this.id = id + this.controller = new AbortController() + } + + public get signal(): AbortSignal { + return this.controller.signal + } + + public get cancelled(): boolean { + return this.controller.signal.aborted + } + + public cancel(reason: string | Error = DEFAULT_CANCEL_REASON): void { + if (this.cancelled) return + this.controller.abort( + reason instanceof Error ? reason : new Error(reason), + ) + } +} + +export type CancellableRun = { + id: string + controller: CodebuffRunController + signal: AbortSignal + result: Promise + cancel: (reason?: string | Error) => void +} + +export function createRunController(id?: string): CodebuffRunController { + return new CodebuffRunController(id) +} + +export function anySignal(signals: Array): AbortSignal { + const activeSignals = signals.filter((signal): signal is AbortSignal => + Boolean(signal), + ) + if (activeSignals.length === 0) { + return new AbortController().signal + } + if (activeSignals.length === 1) { + return activeSignals[0] + } + if (typeof AbortSignal.any === 'function') { + return AbortSignal.any(activeSignals) + } + + const controller = new AbortController() + const abort = (signal: AbortSignal) => { + if (controller.signal.aborted) return + controller.abort(signal.reason) + } + for (const signal of activeSignals) { + if (signal.aborted) { + abort(signal) + break + } + signal.addEventListener('abort', () => abort(signal), { once: true }) + } + return controller.signal +}