Fixes #397: Abort sandbox exec on turn cancellation

.changeset/tidy-planes-stop.md

@@ -0,0 +1,6 @@
+---
+"@truefoundry/trueforge-core": patch
+"@truefoundry/trueforge": patch
+---
+
+Abort sandbox exec work when a turn cancellation signal aborts.

packages/trueforge-core/src/core/index.ts

@@ -155,7 +155,7 @@ export type { CodeModeErrorSource, CodeModeReply, CodeModeRequest } from './sand
 export { DaytonaSandboxProvider } from './sandbox/provider/DaytonaProvider';
 export type { DaytonaSandboxProviderOptions } from './sandbox/provider/DaytonaProvider';
 export { absolutizeRelativeExecEnv } from './sandbox/provider/execEnv';
-export { ensureExecSuccess, shellEscape } from './sandbox/provider/Provider';
+export { SANDBOX_EXEC_ABORTED, ensureExecSuccess, shellEscape } from './sandbox/provider/Provider';
 export type {
   ExecErrorResult,
   ExecResult,
@@ -182,6 +182,7 @@ export {
 export { SANDBOX_IMAGE_URI } from './sandbox/sandboxImage';
 export { existingSandboxIdForProvider, formatSandboxId, parseSandboxId, rawSandboxId } from './sandbox/sandboxRef';
 export type { SandboxRefParts } from './sandbox/sandboxRef';
+export { onSignalAbort } from './util/abort';
 
 // Web search
 export { ParallelWebSearchProvider } from './web-search/ParallelWebSearchProvider';

packages/trueforge-core/src/core/sandbox/Sandbox.ts

@@ -87,6 +87,7 @@ export interface SandboxStoredFile {
 export interface SandboxOptions {
   provider: SandboxProvider;
   existingSandboxId?: string | undefined;
+  signal?: AbortSignal | undefined;
   skillMounter?: ISkillMounter | undefined;
   fileDownloadEnabled?: boolean | undefined;
   /** Pre-resolved credential-store file content (null = clear / no git auth). */
@@ -198,6 +199,7 @@ export class Sandbox extends LocalToolMCP {
 
   private readonly provider: SandboxProvider;
   private readonly existingSandboxId?: string | undefined;
+  private readonly signal?: AbortSignal | undefined;
   private existingSandboxInfo: SandboxInfo | undefined;
   // Cached promise to prevent concurrent sub-agents from creating duplicate sandboxes.
   private sandboxCreationPromise?: Promise<SandboxInfo> | undefined;
@@ -228,6 +230,7 @@ export class Sandbox extends LocalToolMCP {
     super({ tracing: options.tracing });
     this.provider = options.provider;
     this.existingSandboxId = options.existingSandboxId;
+    this.signal = options.signal;
     this.skillMounter = options.skillMounter;
     this.fileDownloadEnabled = options.fileDownloadEnabled ?? false;
     const mcpBoundTimeoutMs = options.mcpRequestTimeoutMs + options.mcpConnectTimeoutMs;
@@ -560,6 +563,7 @@ export class Sandbox extends LocalToolMCP {
         command: input.command,
         cwd: input.cwd,
         env,
+        signal: this.signal,
       });
 
       return {
@@ -595,6 +599,7 @@ export class Sandbox extends LocalToolMCP {
         command: input.command,
         cwd: input.cwd,
         env: retryEnv,
+        signal: this.signal,
       });
       return {
         result: {

packages/trueforge-core/src/core/sandbox/provider/DaytonaProvider.ts

@@ -5,6 +5,7 @@ import { suppressTracing } from '@opentelemetry/core';
 import { randomUUID } from 'node:crypto';
 import { join } from 'node:path/posix';
 import type { Logger } from 'winston';
+import { isSignalAborted, onSignalAbort } from '../../util/abort';
 import { extractErrorLogFields } from '../../util/errorLogFields';
 import {
   SandboxFileNotFoundError,
@@ -16,7 +17,14 @@ import {
 import type { CodeModeTransport } from '../codeMode/CodeModeTransport';
 import { CodeModeNatsTransport } from '../codeMode/nats/CodeModeNatsTransport';
 import { DEFAULT_PREVIEW_URL_EXPIRY_SECONDS, DEFAULT_SANDBOX_NATS_WS_PORT } from '../constants';
-import type { ExecResult, SandboxBuild, SandboxExecParams, SandboxFileInfo, SandboxProvider } from './Provider';
+import {
+  SANDBOX_EXEC_ABORTED,
+  type ExecResult,
+  type SandboxBuild,
+  type SandboxExecParams,
+  type SandboxFileInfo,
+  type SandboxProvider,
+} from './Provider';
 
 const SANDBOX_NOT_FOUND_STATUS = 404;
 /** Another replica already registered this build name; its create is the one that counts. */
@@ -193,18 +201,22 @@ export class DaytonaSandboxProvider implements SandboxProvider {
     return recovery;
   }
 
-  private async executeWithSandboxRecovery<T>(sandboxId: string, operation: () => Promise<T>): Promise<T> {
+  private async executeWithSandboxRecovery<T>(params: {
+    sandboxId: string;
+    signal: AbortSignal | undefined;
+    operation: () => Promise<T>;
+  }): Promise<T> {
     try {
-      return await operation();
+      return await params.operation();
     } catch (originalError) {
       // TODO: Narrow to a specific Daytona error code once @daytona/sdk exposes one for "sandbox not running".
-      if (!(originalError instanceof DaytonaError)) {
+      if (!(originalError instanceof DaytonaError) || isSignalAborted(params.signal)) {
         throw originalError;
       }
 
       let recovered: boolean;
       try {
-        recovered = await DaytonaSandboxProvider.recoverSandboxIfStopped(sandboxId);
+        recovered = await DaytonaSandboxProvider.recoverSandboxIfStopped(params.sandboxId);
       } catch (recoveryError) {
         this.logger.error('Sandbox recovery failed', {
           ...extractErrorLogFields(recoveryError),
@@ -218,7 +230,7 @@ export class DaytonaSandboxProvider implements SandboxProvider {
       }
 
       try {
-        return await operation();
+        return await params.operation();
       } catch (retryError) {
         this.logger.error('Sandbox operation failed after successful recovery', {
           ...extractErrorLogFields(retryError),
@@ -384,22 +396,56 @@ export class DaytonaSandboxProvider implements SandboxProvider {
 
   async exec(params: SandboxExecParams): Promise<ExecResult> {
     return context.with(suppressTracing(context.active()), async (): Promise<ExecResult> => {
+      if (isSignalAborted(params.signal)) {
+        return { success: false, error: SANDBOX_EXEC_ABORTED };
+      }
+      // Map the catch to abort only if this exec's handler called stop(); a later signal abort must not hide a Daytona error.
+      const execAbortRequested = { value: false };
       try {
-        return await this.executeWithSandboxRecovery(params.sandboxId, async () => {
-          const { sandbox, defaultTimeoutMs } = await this.getOrCreateSandbox(params.sandboxId);
-          const response = await sandbox.process.executeCommand(
-            params.command,
-            params.cwd,
-            params.env ?? {},
-            params.timeoutSeconds ?? defaultTimeoutMs / 1000,
-          );
-          return {
-            success: true,
-            response: { exitCode: response.exitCode, result: response.result },
-          };
+        return await this.executeWithSandboxRecovery({
+          sandboxId: params.sandboxId,
+          signal: params.signal,
+          operation: async () => {
+            const { sandbox, defaultTimeoutMs } = await this.getOrCreateSandbox(params.sandboxId);
+            const timeoutSeconds = params.timeoutSeconds ?? defaultTimeoutMs / 1000;
+            if (isSignalAborted(params.signal)) {
+              return { success: false, error: SANDBOX_EXEC_ABORTED };
+            }
+            const cleanupAbort = onSignalAbort(params.signal, () => {
+              execAbortRequested.value = true;
+              void sandbox.stop(timeoutSeconds, true).catch((error: unknown) => {
+                this.logger.error('Failed to stop Daytona sandbox after exec abort', {
+                  ...extractErrorLogFields(error),
+                  sandboxId: params.sandboxId,
+                });
+              });
+            });
+            try {
+              const response = await sandbox.process.executeCommand(
+                params.command,
+                params.cwd,
+                params.env ?? {},
+                timeoutSeconds,
+              );
+              if (execAbortRequested.value) {
+                // Drop only the stale SDK object. The remote sandbox and filesystem stay under the same id.
+                DaytonaSandboxProvider.cachedSandboxes.delete(params.sandboxId);
+                return { success: false, error: SANDBOX_EXEC_ABORTED };
+              }
+              return {
+                success: true,
+                response: { exitCode: response.exitCode, result: response.result },
+              };
+            } finally {
+              cleanupAbort();
+            }
+          },
         });
       } catch (e: unknown) {
         DaytonaSandboxProvider.cachedSandboxes.delete(params.sandboxId);
+        if (execAbortRequested.value) {
+          return { success: false, error: SANDBOX_EXEC_ABORTED };
+        }
         if (e instanceof SandboxNotAvailableError) {
           throw e;
         }
@@ -418,18 +464,22 @@ export class DaytonaSandboxProvider implements SandboxProvider {
   async downloadFile(params: { sandboxId: string; path: string }): Promise<Buffer> {
     return context.with(suppressTracing(context.active()), async () => {
       try {
-        return await this.executeWithSandboxRecovery(params.sandboxId, async () => {
-          const { sandbox } = await this.getOrCreateSandbox(params.sandboxId);
-
-          const info = await this.getFileInfo(sandbox, params.path);
-          if (info.isDir) {
-            throw new SandboxPathIsDirectoryError(params.path);
-          }
-          if (info.size > this.fileMaxBytesForDownload) {
-            throw new SandboxFileTooLargeError(params.path, info.size, this.fileMaxBytesForDownload);
-          }
-
-          return await sandbox.fs.downloadFile(params.path);
+        return await this.executeWithSandboxRecovery({
+          sandboxId: params.sandboxId,
+          signal: undefined,
+          operation: async () => {
+            const { sandbox } = await this.getOrCreateSandbox(params.sandboxId);
+
+            const info = await this.getFileInfo(sandbox, params.path);
+            if (info.isDir) {
+              throw new SandboxPathIsDirectoryError(params.path);
+            }
+            if (info.size > this.fileMaxBytesForDownload) {
+              throw new SandboxFileTooLargeError(params.path, info.size, this.fileMaxBytesForDownload);
+            }
+
+            return await sandbox.fs.downloadFile(params.path);
+          },
         });
       } catch (e: unknown) {
         if (e instanceof SandboxPathIsDirectoryError || e instanceof SandboxFileTooLargeError) {
@@ -447,9 +497,13 @@ export class DaytonaSandboxProvider implements SandboxProvider {
   async uploadFile(params: { sandboxId: string; remotePath: string; content: Buffer }): Promise<void> {
     return context.with(suppressTracing(context.active()), async () => {
       try {
-        await this.executeWithSandboxRecovery(params.sandboxId, async () => {
-          const { sandbox } = await this.getOrCreateSandbox(params.sandboxId);
-          await sandbox.fs.uploadFile(params.content, params.remotePath);
+        await this.executeWithSandboxRecovery({
+          sandboxId: params.sandboxId,
+          signal: undefined,
+          operation: async () => {
+            const { sandbox } = await this.getOrCreateSandbox(params.sandboxId);
+            await sandbox.fs.uploadFile(params.content, params.remotePath);
+          },
         });
       } catch (e: unknown) {
         DaytonaSandboxProvider.cachedSandboxes.delete(params.sandboxId);
@@ -462,10 +516,14 @@ export class DaytonaSandboxProvider implements SandboxProvider {
   private async getPreviewUrl(params: { sandboxId: string; port: number; expiresInSeconds: number }): Promise<string> {
     return context.with(suppressTracing(context.active()), async () => {
       try {
-        return await this.executeWithSandboxRecovery(params.sandboxId, async () => {
-          const { sandbox } = await this.getOrCreateSandbox(params.sandboxId);
-          const signed = await sandbox.getSignedPreviewUrl(params.port, params.expiresInSeconds);
-          return signed.url;
+        return await this.executeWithSandboxRecovery({
+          sandboxId: params.sandboxId,
+          signal: undefined,
+          operation: async () => {
+            const { sandbox } = await this.getOrCreateSandbox(params.sandboxId);
+            const signed = await sandbox.getSignedPreviewUrl(params.port, params.expiresInSeconds);
+            return signed.url;
+          },
         });
       } catch (e: unknown) {
         DaytonaSandboxProvider.cachedSandboxes.delete(params.sandboxId);

packages/trueforge-core/src/core/sandbox/provider/Provider.ts

@@ -14,6 +14,8 @@ export interface ExecErrorResult {
 
 export type ExecResult = ExecSuccessResult | ExecErrorResult;
 
+export const SANDBOX_EXEC_ABORTED = 'Sandbox exec aborted';
+
 /**
  * Wraps a value in single quotes for safe use in a shell command. Inner single quotes are
  * escaped via the standard shell idiom: ' -> '\''. Use for any user-controlled value that is
@@ -44,6 +46,7 @@ export interface SandboxExecParams {
   command: string;
   cwd?: string | undefined;
   env?: Record<string, string> | undefined;
+  signal?: AbortSignal | undefined;
   /** Overrides the provider's default exec timeout (e.g. for long skill downloads). */
   timeoutSeconds?: number | undefined;
 }

packages/trueforge-core/src/core/sandbox/provider/TFYSandboxProvider.ts

@@ -4,6 +4,7 @@ import dedent from 'dedent';
 import { randomUUID } from 'node:crypto';
 import { join } from 'node:path/posix';
 import type { Logger } from 'winston';
+import { isSignalAborted, onSignalAbort } from '../../util/abort';
 import { extractErrorLogFields } from '../../util/errorLogFields';
 import type { CodeModeTransport } from '../codeMode/CodeModeTransport';
 import { CodeModeNatsTransport } from '../codeMode/nats/CodeModeNatsTransport';
@@ -17,6 +18,7 @@ import {
 import { absolutizeRelativeExecEnv } from './execEnv';
 import {
   ensureExecSuccess,
+  SANDBOX_EXEC_ABORTED,
   shellEscape,
   type ExecResult,
   type SandboxBuild,
@@ -152,9 +154,15 @@ export class TFYSandboxProvider implements SandboxProvider {
 
       const controller = new AbortController();
       const clientTimeoutMs = (timeoutSeconds + CLIENT_TIMEOUT_BUFFER_SECONDS) * 1000;
+      const abortState: { source: 'timeout' | 'turn' | undefined } = { source: undefined };
       const timer = setTimeout(() => {
+        abortState.source = 'timeout';
         controller.abort();
       }, clientTimeoutMs);
+      const cleanupAbort = onSignalAbort(params.signal, () => {
+        abortState.source = 'turn';
+        controller.abort();
+      });
 
       try {
         const response = await fetch(`${this.serverUrl}/exec`, {
@@ -174,6 +182,9 @@ export class TFYSandboxProvider implements SandboxProvider {
         return result;
       } catch (e: unknown) {
         if (e instanceof Error && e.name === 'AbortError') {
+          if (abortState.source === 'turn' && isSignalAborted(params.signal)) {
+            return { success: false, error: SANDBOX_EXEC_ABORTED };
+          }
           this.logger.error(`Sandbox exec timed out after ${String(timeoutSeconds)}s`, extractErrorLogFields(e));
           return { success: false, error: `Sandbox exec timed out after ${String(timeoutSeconds)}s` };
         }
@@ -182,6 +193,7 @@ export class TFYSandboxProvider implements SandboxProvider {
         return { success: false, error: message };
       } finally {
         clearTimeout(timer);
+        cleanupAbort();
       }
     });
   }

packages/trueforge-core/src/core/util/abort.ts

@@ -1,10 +1,17 @@
-export function onSignalAbort(signal: AbortSignal | undefined, callback: () => void): void {
+export function isSignalAborted(signal: AbortSignal | undefined): boolean {
+  return signal?.aborted === true;
+}
+
+export function onSignalAbort(signal: AbortSignal | undefined, callback: () => void): () => void {
   if (!signal) {
-    return;
+    return () => undefined;
   }
   if (signal.aborted) {
     callback();
-    return;
+    return () => undefined;
   }
   signal.addEventListener('abort', callback, { once: true });
+  return () => {
+    signal.removeEventListener('abort', callback);
+  };
 }