[pull] main from LibreChat-AI:main

packages/code/src/native-process-child.ts

@@ -101,6 +101,10 @@ process.on('message', async (raw: unknown) => {
         error instanceof WorkspaceToolError
           ? error.mutationMayHaveCommitted
           : true,
+      requiresQuarantine:
+        error instanceof WorkspaceToolError
+          ? error.requiresQuarantine
+          : true,
     });
   } finally {
     active = undefined;

packages/code/src/native-process.test.ts

@@ -179,13 +179,47 @@ test('executor cancellation targets the active request and preserves mutation ce
     ok: false,
     code: 'EXECUTION_ABORTED',
     mutation: true,
+    requiresQuarantine: false,
   });
   await assert.rejects(
     execution,
     (error: unknown) =>
       error instanceof WorkspaceToolError &&
       error.code === 'EXECUTION_ABORTED' &&
-      error.mutationMayHaveCommitted,
+      error.mutationMayHaveCommitted &&
+      !error.requiresQuarantine,
+  );
+  await sandbox.close();
+});
+
+test('executor ignores a cleanup exemption on non-cancellation failures', async () => {
+  let dispatched!: () => void;
+  const dispatch = new Promise<void>((resolve) => {
+    dispatched = resolve;
+  });
+  const fake = fixture(() => dispatched());
+  const sandbox = new NativeProcessWorkspaceCommandSandbox(
+    { workspaceRoot: '/workspace' },
+    fake.fork,
+  );
+  const execution = sandbox.execute(request);
+  await dispatch;
+  const command = fake.messages.find((message) => message.type === 'execute')!;
+  fake.child.emit('message', {
+    id: command.id,
+    ok: false,
+    code: 'COMMAND_UNAVAILABLE',
+    mutation: true,
+    requiresQuarantine: false,
+  });
+
+  await assert.rejects(
+    execution,
+    (error: unknown) =>
+      error instanceof WorkspaceToolError &&
+      error.code === 'COMMAND_UNAVAILABLE' &&
+      error.mutationMayHaveCommitted &&
+      error.requiresQuarantine,
   );
   await sandbox.close();
 });

packages/code/src/native-process.ts

@@ -134,6 +134,7 @@ export class NativeProcessWorkspaceCommandSandbox
         ok?: unknown;
         result?: unknown;
         mutation?: unknown;
+        requiresQuarantine?: unknown;
         code?: unknown;
         errorMessage?: unknown;
         fatal?: unknown;
@@ -156,6 +157,9 @@ export class NativeProcessWorkspaceCommandSandbox
           message.code === 'REGISTRATION_INVALID'
             ? message.code
             : 'COMMAND_UNAVAILABLE';
+        const processTerminationConfirmed =
+          code === 'EXECUTION_ABORTED' &&
+          message.requiresQuarantine === false;
         pending.reject(
           new WorkspaceToolError(
             typeof message.errorMessage === 'string' &&
@@ -164,6 +168,7 @@ export class NativeProcessWorkspaceCommandSandbox
               : 'Native executor request failed',
             code,
             pending.mutation && message.mutation !== false,
+            pending.mutation && !processTerminationConfirmed,
           ),
         );
       }

packages/code/src/native-sandbox.test.ts

@@ -1030,7 +1030,8 @@ test('reports cancellation after command start as a potentially committed mutati
     (error: unknown) =>
       error instanceof WorkspaceToolError &&
       error.code === 'EXECUTION_ABORTED' &&
-      error.mutationMayHaveCommitted === true,
+      error.mutationMayHaveCommitted === true &&
+      error.requiresQuarantine === false,
   );
 });
 
@@ -1157,7 +1158,9 @@ test('cleans allocated command state exactly once on every execution exit', asyn
             error.code === (outcome.startsWith('abort')
               ? 'EXECUTION_ABORTED'
               : 'COMMAND_UNAVAILABLE') &&
-            error.mutationMayHaveCommitted === (outcome === 'abort-after-spawn'),
+            error.mutationMayHaveCommitted === (outcome === 'abort-after-spawn') &&
+            error.requiresQuarantine ===
+              (outcome === 'abort-after-spawn' && process.platform === 'win32'),
           );
         }
         assert.equal(

packages/code/src/native-sandbox.ts

@@ -681,6 +681,10 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox
                 'Workspace command execution aborted',
                 'EXECUTION_ABORTED',
                 true,
+                // POSIX commands run in a detached process group, so its
+                // observed close follows a group-wide SIGKILL. The Windows
+                // fallback cannot yet prove descendant termination.
+                this.platform === 'win32',
               ),
             );
             return;

packages/code/src/worker.ts

@@ -1562,7 +1562,7 @@ export class BridgeWorker {
           !(
             error instanceof WorkspaceToolError &&
             this.options.workspaceTools?.mutationFailuresAreAtomic === true &&
-            !error.mutationMayHaveCommitted
+            !error.requiresQuarantine
           ))
       ) {
         ambiguousWorkspaceMutationError = error;

packages/code/src/workspace-worker.test.ts

@@ -1410,6 +1410,76 @@ test('worker clears quarantine after a composed command is cleanly rejected', as
   assert.deepEqual(lifecycle, ['arm', 'execute', 'settle', 'clear']);
 });
 
+test('worker clears quarantine after a command cancellation confirms process termination', async () => {
+  const lifecycle: string[] = [];
+  const baseCapabilities = {
+    protocolVersion: 1 as const,
+    operations: ['read_file' as const],
+    workspaces: [{ id: 'primary', operations: ['read_file' as const] }],
+  };
+  const workspaceTools = new SandboxWorkspaceTools({
+    workspaceTools: {
+      capabilities: baseCapabilities,
+      mutationFailuresAreAtomic: true,
+      async execute() { throw new Error('base executor must not run'); },
+    },
+    commandWorkspaces: ['primary'],
+    commandSandbox: {
+      mutationFailuresAreAtomic: true,
+      async execute() {
+        lifecycle.push('execute');
+        throw new WorkspaceToolError(
+          'Workspace command execution aborted',
+          'EXECUTION_ABORTED',
+          true,
+          false,
+        );
+      },
+    },
+  });
+  const worker = new BridgeWorker({
+    codeApiUrl: 'https://code.example/v1',
+    token: 'worker-secret',
+    workerId: 'vm-1',
+    incarnationId,
+    sandboxEndpoint: 'http://127.0.0.1:2000/api/v2',
+    capabilities: {
+      statefulWorkspace: true,
+      sandboxProfile: 'nsjail',
+      runtimes: ['bash'],
+      workspaceTools: workspaceTools.capabilities,
+    },
+    workspaceTools,
+    workspaceMutationQuarantine: mutationQuarantine(
+      () => lifecycle.push('quarantine'),
+      () => lifecycle.push('arm'),
+      () => lifecycle.push('clear'),
+    ),
+    fetchImpl: async () => {
+      lifecycle.push('settle');
+      return Response.json({ protocolVersion: 1, accepted: true });
+    },
+  });
+
+  await worker.executeAndSettle({
+    protocolVersion: 1,
+    assignmentId: 'assignment-command-cancelled-cleanly',
+    workerId: 'vm-1',
+    incarnationId,
+    generation: 4,
+    leaseToken: 'lease-token-that-is-long-enough-for-testing',
+    expiresAt: new Date(Date.now() + 5_000).toISOString(),
+    executionKind: 'workspace_tool',
+    request: {
+      protocolVersion: 1,
+      operation: 'execute_command',
+      workspaceId: 'primary',
+      command: 'sleep 30',
+    },
+  });
+  assert.deepEqual(lifecycle, ['arm', 'execute', 'settle', 'clear']);
+});
+
 test('worker retains quarantine when an atomic executor cannot confirm durability', async () => {
   const lifecycle: string[] = [];
   const workspaceCapabilities = {

packages/code/src/workspace.ts

@@ -176,6 +176,8 @@ export class WorkspaceToolError extends Error {
     message: string,
     public readonly code: WorkspaceToolErrorCode,
     public readonly mutationMayHaveCommitted = false,
+    /** Retain the durable mutation guard when process or write settlement is uncertain. */
+    public readonly requiresQuarantine = mutationMayHaveCommitted,
   ) {
     super(message);
     this.name = 'WorkspaceToolError';

service/src/bridge/store.ts

@@ -23,6 +23,7 @@ import { BridgeWorkspaceSlots } from './slots';
 
 const PREFIX = 'codeapi:bridge:v1';
 const POLL_INTERVAL_MS = 100;
+const CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS = 5_000;
 const DEFAULT_WORKER_TTL_SECONDS = 60;
 const DEFAULT_REDIS_COMMAND_TIMEOUT_MS = 1_000;
 
@@ -1901,6 +1902,55 @@ export class RedisBridgeStore {
       // settlement before returning an error, even when the caller is gone.
       pollError = error;
     }
+    const workspaceRequest =
+      assignment.executionKind === 'workspace_tool' &&
+      isWorkspaceToolRequest(assignment.request)
+        ? assignment.request
+        : undefined;
+    const cancelledMutation =
+      signal.aborted &&
+      workspaceRequest != null &&
+      (workspaceRequest.operation === 'write_file' ||
+        workspaceRequest.operation === 'edit_file' ||
+        workspaceRequest.operation === 'execute_command');
+    if (cancelledMutation) {
+      try {
+        // Keep the acknowledged assignment available long enough for the
+        // worker to terminate its process tree and commit a clean rejection.
+        // Closing it first makes that rejection impossible to acknowledge and
+        // leaves the worker's durable mutation guard armed.
+        await this.cancel(assignment.assignmentId, assignment);
+        // Rejected settlements remain valid after the execution deadline.
+        // Give Stop its own grace so a near-timeout cancellation is not
+        // misclassified as an ambiguous timeout.
+        const cancellationDeadlineAtMs =
+          Date.now() + CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS;
+        let cancellationPollMs = POLL_INTERVAL_MS;
+        while (Date.now() < cancellationDeadlineAtMs) {
+          const raw = await boundedCommand(
+            this.redis.get(settlementKey(assignment.assignmentId)),
+            Math.max(
+              1,
+              Math.min(
+                this.redisCommandTimeoutMs,
+                cancellationDeadlineAtMs - Date.now(),
+              ),
+            ),
+            'Bridge cancelled workspace settlement poll',
+          );
+          if (raw != null) return JSON.parse(raw) as CodeBridgeSettlement;
+          await delay(
+            Math.min(
+              cancellationPollMs,
+              Math.max(0, cancellationDeadlineAtMs - Date.now()),
+            ),
+          );
+          cancellationPollMs = Math.min(cancellationPollMs * 2, 500);
+        }
+      } catch (error) {
+        pollError ??= error;
+      }
+    }
     const closeKeys = [
       assignmentKey(assignment.assignmentId),
       settlementKey(assignment.assignmentId),

service/src/bridge/workspace-store.test.ts

@@ -77,6 +77,108 @@ test('dispatches a workspace tool only to a worker advertising its workspace and
   });
 });
 
+test('drains an acknowledged workspace mutation cancellation before releasing it', async () => {
+  await store.register({
+    protocolVersion: BRIDGE_PROTOCOL_VERSION,
+    workerId: 'workspace-worker',
+    incarnationId,
+    capabilities: {
+      statefulWorkspace: false,
+      sandboxProfile: 'native-srt',
+      runtimes: [],
+      workspaceTools: {
+        protocolVersion: BRIDGE_PROTOCOL_VERSION,
+        operations: ['execute_command'],
+        workspaces: [{ id: 'primary' }],
+      },
+    },
+  });
+  const controller = new AbortController();
+  const completion = store.dispatchWorkspaceTool({
+    workerId: 'workspace-worker',
+    request: {
+      protocolVersion: BRIDGE_PROTOCOL_VERSION,
+      operation: 'execute_command',
+      workspaceId: 'primary',
+      command: 'sleep 30; touch delayed.txt',
+    },
+    deadlineAtMs: Date.now() + 5_000,
+    executionTimeoutMs: 500,
+    signal: controller.signal,
+  });
+  const assignment = (await store.lease(
+    'workspace-worker',
+    incarnationId,
+    1_000,
+  ))!;
+  await store.acknowledgeLease(
+    'workspace-worker',
+    incarnationId,
+    assignment.assignmentId,
+    assignment.generation,
+    assignment.leaseToken,
+  );
+
+  await new Promise((resolve) => setTimeout(resolve, 250));
+  controller.abort();
+  for (
+    let attempt = 0;
+    attempt < 100 &&
+    !(await store.cancelled(
+      'workspace-worker',
+      incarnationId,
+      assignment.assignmentId,
+    ));
+    attempt += 1
+  ) {
+    await new Promise((resolve) => setTimeout(resolve, 5));
+  }
+  await new Promise((resolve) => setTimeout(resolve, 350));
+  expect(Date.now()).toBeGreaterThan(Date.parse(assignment.expiresAt));
+  await store.settle('workspace-worker', assignment.assignmentId, {
+    protocolVersion: BRIDGE_PROTOCOL_VERSION,
+    generation: assignment.generation,
+    leaseToken: assignment.leaseToken,
+    incarnationId,
+    status: 'rejected',
+    errorCode: 'EXECUTION_ABORTED',
+    error: 'Workspace command execution aborted',
+  });
+
+  await expect(completion).resolves.toMatchObject({
+    status: 'rejected',
+    errorCode: 'EXECUTION_ABORTED',
+  });
+
+  const reuse = store.dispatchWorkspaceTool({
+    workerId: 'workspace-worker',
+    request: {
+      protocolVersion: BRIDGE_PROTOCOL_VERSION,
+      operation: 'execute_command',
+      workspaceId: 'primary',
+      command: 'printf reused',
+    },
+    deadlineAtMs: Date.now() + 5_000,
+    executionTimeoutMs: 5_000,
+    signal: new AbortController().signal,
+  });
+  const next = (await store.lease(
+    'workspace-worker',
+    incarnationId,
+    1_000,
+  ))!;
+  expect(next.request).toMatchObject({ command: 'printf reused' });
+  await store.settle('workspace-worker', next.assignmentId, {
+    protocolVersion: BRIDGE_PROTOCOL_VERSION,
+    generation: next.generation,
+    leaseToken: next.leaseToken,
+    incarnationId,
+    status: 'rejected',
+    error: 'fixture completion',
+  });
+  await expect(reuse).resolves.toMatchObject({ status: 'rejected' });
+});
+
 test('rejects a workspace tool that the selected worker did not advertise', async () => {
   await store.register({
     protocolVersion: BRIDGE_PROTOCOL_VERSION,