From b78efcfbbae42b161e4e4f8e3ed4944e3afe4487 Mon Sep 17 00:00:00 2001
From: Muhammad Ahsan Farooq <ahsanfarooq210@gmail.com>
Date: Sat, 14 Feb 2026 17:47:43 +0500
Subject: [PATCH 1/2] fix(core): publish policy update on auto-edit transition

Removes an early return in `updatePolicy` to ensure the `UPDATE_POLICY` message is broadcast when entering `AUTO_EDIT` mode.
---
 packages/core/src/scheduler/policy.test.ts | 10 ++++++++--
 packages/core/src/scheduler/policy.ts      |  1 -
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/packages/core/src/scheduler/policy.test.ts b/packages/core/src/scheduler/policy.test.ts
index a076e4c44fb..e852bb7bd0b 100644
--- a/packages/core/src/scheduler/policy.test.ts
+++ b/packages/core/src/scheduler/policy.test.ts
@@ -145,7 +145,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 = {
         setApprovalMode: vi.fn(),
       } as unknown as Mocked<Config>;
@@ -165,7 +165,13 @@ 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 handle standard policy updates (persist=false)', async () => {
diff --git a/packages/core/src/scheduler/policy.ts b/packages/core/src/scheduler/policy.ts
index 247b696f229..7c4d67c3b81 100644
--- a/packages/core/src/scheduler/policy.ts
+++ b/packages/core/src/scheduler/policy.ts
@@ -93,7 +93,6 @@ export async function updatePolicy(
   // Mode Transitions (AUTO_EDIT)
   if (isAutoEditTransition(tool, outcome)) {
     deps.config.setApprovalMode(ApprovalMode.AUTO_EDIT);
-    return;
   }
 
   // Specialized Tools (MCP)

From 978a944b6b53d0070d683d0d24b8390bd1741458 Mon Sep 17 00:00:00 2001
From: Muhammad Ahsan Farooq <ahsanfarooq210@gmail.com>
Date: Sat, 11 Apr 2026 16:20:49 +0500
Subject: [PATCH 2/2] Enhance policy update logic to preserve original approval
 mode during AUTO_EDIT transitions in the scheduler. Added a test case to
 verify the correct behavior of mode preservation when a session allow
 triggers AUTO_EDIT.

---
 packages/core/src/scheduler/policy.test.ts | 41 ++++++++++++++++++++++
 packages/core/src/scheduler/policy.ts      |  3 +-
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/packages/core/src/scheduler/policy.test.ts b/packages/core/src/scheduler/policy.test.ts
index c299ab505c2..053e3b38338 100644
--- a/packages/core/src/scheduler/policy.test.ts
+++ b/packages/core/src/scheduler/policy.test.ts
@@ -275,6 +275,47 @@ describe('policy.ts', () => {
       );
     });
 
+    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 () => {
       const mockConfig = {
         getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
diff --git a/packages/core/src/scheduler/policy.ts b/packages/core/src/scheduler/policy.ts
index 2f59fcdaa70..71c5640db9b 100644
--- a/packages/core/src/scheduler/policy.ts
+++ b/packages/core/src/scheduler/policy.ts
@@ -119,6 +119,8 @@ 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);
@@ -127,7 +129,6 @@ export async function updatePolicy(
   // 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.