From ebebbbfc20cc0873fb0d6ba63465e6c0b9914d77 Mon Sep 17 00:00:00 2001
From: Muhammad Ahsan Farooq <ahsanfarooq210@gmail.com>
Date: Tue, 21 Apr 2026 23:20:07 +0500
Subject: [PATCH] Fix/allow for session persistence (#25176)

---
 packages/core/src/scheduler/policy.test.ts | 51 +++++++++++++++++++++-
 packages/core/src/scheduler/policy.ts      |  4 +-
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/packages/core/src/scheduler/policy.test.ts b/packages/core/src/scheduler/policy.test.ts
index c228ead10d..053e3b3833 100644
--- a/packages/core/src/scheduler/policy.test.ts
+++ b/packages/core/src/scheduler/policy.test.ts
@@ -239,7 +239,7 @@ describe('policy.ts', () => {
   });
 
   describe('updatePolicy', () => {
-    it('should set AUTO_EDIT mode for auto-edit transition tools', async () => {
+    it('should set AUTO_EDIT mode for auto-edit transition tools and publish policy update', async () => {
       const mockConfig = {
         getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
         setApprovalMode: vi.fn(),
@@ -266,7 +266,54 @@ describe('policy.ts', () => {
       expect(mockConfig.setApprovalMode).toHaveBeenCalledWith(
         ApprovalMode.AUTO_EDIT,
       );
-      expect(mockMessageBus.publish).not.toHaveBeenCalled();
+      expect(mockMessageBus.publish).toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: MessageBusType.UPDATE_POLICY,
+          toolName: 'replace',
+          persist: false,
+        }),
+      );
+    });
+
+    it('should preserve the original mode set when a session allow triggers AUTO_EDIT', async () => {
+      let currentMode = ApprovalMode.DEFAULT;
+      const mockConfig = {
+        getApprovalMode: vi.fn(() => currentMode),
+        setApprovalMode: vi.fn((mode: ApprovalMode) => {
+          currentMode = mode;
+        }),
+        getSessionId: vi.fn().mockReturnValue('test-session-id'),
+      } as unknown as Mocked<Config>;
+      (mockConfig as unknown as { config: Config }).config =
+        mockConfig as Config;
+      const mockMessageBus = {
+        publish: vi.fn(),
+      } as unknown as Mocked<MessageBus>;
+      const tool = { name: 'replace' } as AnyDeclarativeTool;
+
+      await updatePolicy(
+        tool,
+        ToolConfirmationOutcome.ProceedAlways,
+        undefined,
+        mockConfig,
+        mockMessageBus,
+      );
+
+      expect(mockConfig.setApprovalMode).toHaveBeenCalledWith(
+        ApprovalMode.AUTO_EDIT,
+      );
+      expect(mockMessageBus.publish).toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: MessageBusType.UPDATE_POLICY,
+          toolName: 'replace',
+          persist: false,
+          modes: [
+            ApprovalMode.DEFAULT,
+            ApprovalMode.AUTO_EDIT,
+            ApprovalMode.YOLO,
+          ],
+        }),
+      );
     });
 
     it('should handle standard policy updates (persist=false)', async () => {
diff --git a/packages/core/src/scheduler/policy.ts b/packages/core/src/scheduler/policy.ts
index 69e2a69e6c..71c5640db9 100644
--- a/packages/core/src/scheduler/policy.ts
+++ b/packages/core/src/scheduler/policy.ts
@@ -119,16 +119,16 @@ export async function updatePolicy(
   messageBus: MessageBus,
   toolInvocation?: AnyToolInvocation,
 ): Promise<void> {
+  const currentMode = context.config.getApprovalMode();
+
   // Mode Transitions (AUTO_EDIT)
   if (isAutoEditTransition(tool, outcome)) {
     context.config.setApprovalMode(ApprovalMode.AUTO_EDIT);
-    return;
   }
 
   // Determine persist scope if we are persisting.
   let persistScope: 'workspace' | 'user' | undefined;
   let modes: ApprovalMode[] | undefined;
-  const currentMode = context.config.getApprovalMode();
 
   // If this is an 'Always Allow' selection, we restrict it to the current mode
   // and more permissive modes.