From ac21ae0794f016ab443f0df5e59697ecf80edb5e Mon Sep 17 00:00:00 2001
From: kingwl <kingwenlu@gmail.com>
Date: Wed, 22 Jul 2026 10:12:09 +0800
Subject: [PATCH] fix(mode): dedupe the boundary flush wrapper and cover its
 disposed branch

CI caught two gaps in the review-fix commit: the two post-next() flush
listeners were literal clones (jscpd), and the disposed-skip branch inside
them had no covering test. Extract one shared flushAfter wrapper and pin the
skip with a captured-continuation disposal test (a downstream listener
disposes the fiber mid-waterfall; the resumed wrapper appends nothing).
---
 packages/mode/mode/src/index.ts       | 19 ++++++-------------
 packages/mode/mode/tests/mode.spec.ts | 23 +++++++++++++++++++++++
 2 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/packages/mode/mode/src/index.ts b/packages/mode/mode/src/index.ts
index ff7c2edaa6..c83b29d181 100644
--- a/packages/mode/mode/src/index.ts
+++ b/packages/mode/mode/src/index.ts
@@ -241,7 +241,7 @@ export class ModesService extends Service {
     // the same): downstream listeners may await, and a `session/set_mode`
     // arriving during that window must still shape the request this boundary
     // precedes — a pre-next() flush would apply it one request late.
-    ctx.on('agent/prompt-submit', async (agent, _content, _source, _signal, next) => {
+    const flushAfter = async <T>(agent: Agent, next: () => Promise<T>): Promise<T> => {
       const decision = await next()
       if (!disposed) {
         try {
@@ -251,18 +251,11 @@ export class ModesService extends Service {
         }
       }
       return decision
-    }, { prepend: true })
-    ctx.on('agent/turn-continuation', async (agent, _turn, _decision, _signal, next) => {
-      const decision = await next()
-      if (!disposed) {
-        try {
-          this.onBoundary(agent)
-        } catch (error) {
-          ctx.logger.warn('dsh-mode: boundary flush failed: %o', error)
-        }
-      }
-      return decision
-    }, { prepend: true })
+    }
+    ctx.on('agent/prompt-submit', (agent, _content, _source, _signal, next) =>
+      flushAfter(agent, next), { prepend: true })
+    ctx.on('agent/turn-continuation', (agent, _turn, _decision, _signal, next) =>
+      flushAfter(agent, next), { prepend: true })
     ctx.on('agent/request-error', async (
       agent,
       _turn,
diff --git a/packages/mode/mode/tests/mode.spec.ts b/packages/mode/mode/tests/mode.spec.ts
index 0af2a4b9b1..1245fc0df6 100644
--- a/packages/mode/mode/tests/mode.spec.ts
+++ b/packages/mode/mode/tests/mode.spec.ts
@@ -258,6 +258,29 @@ describe('the boundary flush', () => {
     expect(ctx.modes.get(agent)).toEqual({ current: PLAN_MODE })
   })
 
+  it('skips the flush after the plugin fiber is disposed (a captured wrapper must not write into a dead service)', async () => {
+    const ctx = new Context()
+    await ctx.plugin(SystemPrompt)
+    await ctx.plugin(ToolRegistry)
+    const fiber = await ctx.plugin(ModesService, PLAN_CONFIG)
+    const agent = await agentWithSession(ctx)
+    ctx.modes.set(agent, PLAN_MODE)
+    // A downstream listener captured before disposal keeps the waterfall
+    // continuation alive across the unload; the resumed wrapper must not
+    // append through the disposed service.
+    ctx.on('agent/turn-continuation', async (_agent, _turn, decision, _signal, next) => {
+      await fiber.dispose()
+      await next()
+      return decision
+    })
+    agent.session.append('step/end', { turn: 1, step: 1 })
+    await agentEvents(ctx, agent).waterfall(
+      'agent/turn-continuation', 1, { action: 'stop' }, new AbortController().signal,
+      () => Promise.resolve({ action: 'stop' }),
+    )
+    expect(agent.session.events.some(event => event.type === 'mode/set')).toBe(false)
+  })
+
   it('flushes at step/end too (a mid-turn flip lands on the following step)', async () => {
     const ctx = await setup()
     const agent = await agentWithSession(ctx)