Add cancellable SDK run controller

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<void>((resolve) => {
+      markRuntimeStarted = resolve
+    })
+    spyOn(mainPromptModule, 'callMainPrompt').mockImplementation(
+      async (params: Parameters<typeof mainPromptModule.callMainPrompt>[0]) => {
+        runtimeSignal = params.signal
+        markRuntimeStarted()
+        return await new Promise<never>((_, 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<void>((resolve) => {
+      markRuntimeStarted = resolve
+    })
+    spyOn(mainPromptModule, 'callMainPrompt').mockImplementation(
+      async (params: Parameters<typeof mainPromptModule.callMainPrompt>[0]) => {
+        runtimeSignal = params.signal
+        markRuntimeStarted()
+        return await new Promise<never>((_, 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')
+    }
+  })
+})

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.
    *

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 {

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<RunState>
+  cancel: (reason?: string | Error) => void
+}
+
+export function createRunController(id?: string): CodebuffRunController {
+  return new CodebuffRunController(id)
+}
+
+export function anySignal(signals: Array<AbortSignal | undefined>): 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
+}