From d1680450d2facb3866bce297b800d95c72688f60 Mon Sep 17 00:00:00 2001 From: roboomp Date: Wed, 12 Aug 2026 08:26:41 +0000 Subject: [PATCH 1/4] fix(vibe): canceled active turn on mode exit Abort an in-flight model turn before unregistering Vibe tools so it cannot emit stale worker calls after cancellation. Add a regression test covering the live streaming transition. Fixes #8326 --- packages/coding-agent/CHANGELOG.md | 4 +++ .../src/modes/interactive-mode.ts | 3 ++ .../test/interactive-mode-vibe-toggle.test.ts | 36 ++++++++++++++++++- 3 files changed, 42 insertions(+), 1 deletion(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index b387e76fd..bf4c6466e 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed `/vibe` cancellation leaving an in-flight model turn unaware that Vibe mode and its tools were removed ([#8326](https://github.com/can1357/oh-my-pi/issues/8326)). + ## [17.2.15] - 2026-08-12 ### Added diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index 66bb5bccb..f623b97b3 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -3483,6 +3483,9 @@ export class InteractiveMode implements InteractiveModeContext { if (!this.vibeModeEnabled) { return; } + if (this.session.isStreaming) { + await this.session.abort(); + } const ownerScope = this.#vibeModeOwnerScope; const killed = await VibeSessionRegistry.global().killAll(this.#vibeParentSession(), ownerScope); await this.session.deactivateVibeTools(this.#vibeModePreviousTools ?? []); diff --git a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts index edc3ce00d..4a6704404 100644 --- a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts +++ b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts @@ -10,7 +10,8 @@ import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test"; import * as path from "node:path"; import { type } from "@oh-my-pi/omptype"; -import { Agent, type AgentTool } from "@oh-my-pi/pi-agent-core"; +import { Agent, type AgentTool, type StreamFn } from "@oh-my-pi/pi-agent-core"; +import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream"; import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { InteractiveMode } from "@oh-my-pi/pi-coding-agent/modes/interactive-mode"; @@ -23,6 +24,7 @@ import { VIBE_TOOL_NAMES } from "@oh-my-pi/pi-coding-agent/tools/vibe"; import { EventBus } from "@oh-my-pi/pi-coding-agent/utils/event-bus"; import { VibeSessionRegistry } from "@oh-my-pi/pi-coding-agent/vibe/runtime"; import { TempDir } from "@oh-my-pi/pi-utils"; +import { createAssistantMessage } from "./helpers/agent-session-setup"; function stubTool(name: string): AgentTool { return { @@ -85,6 +87,7 @@ describe("InteractiveMode vibe mode toggle", () => { let tempDir: TempDir; let authStorage: AuthStorage; let session: AgentSession; + let streamFn: StreamFn | undefined; let mode: InteractiveMode; let modelRegistry: ModelRegistry; let storage: ExitFaultStorage; @@ -113,6 +116,10 @@ describe("InteractiveMode vibe mode toggle", () => { tools: [], messages: [], }, + streamFn: (...args) => { + if (!streamFn) throw new Error("No test stream configured"); + return streamFn(...args); + }, }), sessionManager: SessionManager.create(tempDir.path(), tempDir.path(), storage), settings: Settings.isolated({}), @@ -157,6 +164,33 @@ describe("InteractiveMode vibe mode toggle", () => { expect(session.getAllToolNames().toSorted()).toEqual(["read", "todo"]); }); + it("cancels an in-flight model turn before removing Vibe tools", async () => { + const started = Promise.withResolvers(); + streamFn = (_model, _context, options) => { + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ type: "start", partial: createAssistantMessage("") }); + options?.signal?.addEventListener( + "abort", + () => stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") }), + { once: true }, + ); + started.resolve(); + }); + return stream; + }; + await mode.handleVibeModeCommand(); + const prompt = session.prompt("Delegate this"); + await started.promise; + expect(session.isStreaming).toBe(true); + + await mode.handleVibeModeCommand(); + await prompt; + + expect(session.isStreaming).toBe(false); + expect(session.getToolByName("vibe_spawn")).toBeUndefined(); + }); + it("keeps a same-named non-built-in Todo tool unavailable in Vibe mode", async () => { const model = session.model; if (!model) throw new Error("Expected active model"); From 8edab1af80d2a8cc107140953be467d5c950a85a Mon Sep 17 00:00:00 2001 From: roboomp Date: Wed, 12 Aug 2026 08:37:47 +0000 Subject: [PATCH 2/4] fix(vibe): suppress queued-turn drain during mode exit abort()'s stranded-queue drain could restart a queued user turn on the still-live Vibe tools before teardown removed them. Run the exit teardown with the queued-message drain suppressed, then re-arm it so the queued turn resumes as a clean non-Vibe turn. Fixes #8326 --- .../src/modes/interactive-mode.ts | 19 +++++--- .../coding-agent/src/session/agent-session.ts | 25 +++++++++++ .../test/interactive-mode-vibe-toggle.test.ts | 44 +++++++++++++++++++ 3 files changed, 81 insertions(+), 7 deletions(-) diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index f623b97b3..8b3f3d8ed 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -3483,13 +3483,18 @@ export class InteractiveMode implements InteractiveModeContext { if (!this.vibeModeEnabled) { return; } - if (this.session.isStreaming) { - await this.session.abort(); - } - const ownerScope = this.#vibeModeOwnerScope; - const killed = await VibeSessionRegistry.global().killAll(this.#vibeParentSession(), ownerScope); - await this.session.deactivateVibeTools(this.#vibeModePreviousTools ?? []); - this.session.setVibeModeState(undefined); + // Tear down with the queued-message drain suppressed: aborting the active + // turn would otherwise let a queued user steer/follow-up restart on the + // still-live Vibe tools before this teardown removes them (issue #8326). + let killed = 0; + await this.session.runModeExitTeardown(async () => { + if (this.session.isStreaming) { + await this.session.abort(); + } + killed = await VibeSessionRegistry.global().killAll(this.#vibeParentSession(), this.#vibeModeOwnerScope); + await this.session.deactivateVibeTools(this.#vibeModePreviousTools ?? []); + this.session.setVibeModeState(undefined); + }); this.vibeModeEnabled = false; this.#vibeModePreviousTools = undefined; this.#vibeModeOwnerScope = undefined; diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 5b0475f3c..7cc926a3c 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -5750,6 +5750,31 @@ export class AgentSession { this.#scheduleIdleQueueDrain(); } + /** + * Run a mode-exit `teardown` (abort the active turn, swap the toolset, clear + * mode state) with queued-message auto-resume suppressed, then re-arm the + * drain so a queued user turn resumes cleanly once the previous toolset is + * back. + * + * `abort()`'s stranded-queue drain runs from its own `finally`; without this + * guard a queued steer/follow-up behind the aborted turn would start a fresh + * `agent.continue()` during the teardown's `await`s — while the exiting mode's + * tools/context are still live — and then have those tools removed underneath + * it, reintroducing the mode's stale-tool failure on the restarted turn + * (issue #8326). Suppressing the drain across teardown guarantees the queued + * turn resumes only after teardown, so it runs as a clean non-mode turn. + */ + async runModeExitTeardown(teardown: () => Promise): Promise { + const previouslyBlocked = this.#queuedMessageDrainBlocked; + this.#queuedMessageDrainBlocked = true; + try { + await teardown(); + } finally { + this.#queuedMessageDrainBlocked = previouslyBlocked; + this.#scheduleIdleQueueDrain(); + } + } + async #queueUserMessage( text: string, images: ImageContent[] | undefined, diff --git a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts index 4a6704404..c14c92093 100644 --- a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts +++ b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts @@ -191,6 +191,50 @@ describe("InteractiveMode vibe mode toggle", () => { expect(session.getToolByName("vibe_spawn")).toBeUndefined(); }); + it("resumes a queued user turn without Vibe tools after exiting mid-turn", async () => { + const toolNamesPerCall: string[][] = []; + let firstStarted: PromiseWithResolvers | undefined; + streamFn = (_model, context, options) => { + toolNamesPerCall.push((context.tools ?? []).map(tool => tool.name)); + const isFirst = toolNamesPerCall.length === 1; + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ type: "start", partial: createAssistantMessage("") }); + if (isFirst) { + options?.signal?.addEventListener( + "abort", + () => stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") }), + { once: true }, + ); + firstStarted?.resolve(); + } else { + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Resumed") }); + } + }); + return stream; + }; + firstStarted = Promise.withResolvers(); + + await mode.handleVibeModeCommand(); + const prompt = session.prompt("Delegate this"); + await firstStarted.promise; + // Queue a user steer behind the active Vibe turn. + await session.steer("and then do the other thing"); + + await mode.handleVibeModeCommand(); + await prompt; + await session.waitForIdle(); + + // The aborted Vibe turn plus the resumed queued turn. + expect(toolNamesPerCall.length).toBe(2); + // The resumed turn must not have inherited the torn-down Vibe tools. + for (const name of VIBE_TOOL_NAMES) { + expect(toolNamesPerCall[1]).not.toContain(name); + } + expect(session.getVibeModeState()).toBeUndefined(); + expect(session.getToolByName("vibe_spawn")).toBeUndefined(); + }); + it("keeps a same-named non-built-in Todo tool unavailable in Vibe mode", async () => { const model = session.model; if (!model) throw new Error("Expected active model"); From 7255ffd596efd39c699ed1b8d1136b8c73963b1f Mon Sep 17 00:00:00 2001 From: roboomp Date: Wed, 12 Aug 2026 08:56:05 +0000 Subject: [PATCH 3/4] fix(vibe): isolated mode-exit drain suppression Replace the resettable queued-message block with a nested mode-exit suppression depth that steer and follow-up queueing cannot clear. Cover a steer arriving after abort settles but before Vibe teardown releases. Fixes #8326 --- .../coding-agent/src/session/agent-session.ts | 18 ++++++---- .../test/interactive-mode-vibe-toggle.test.ts | 34 ++++++++++++++----- 2 files changed, 37 insertions(+), 15 deletions(-) diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 7cc926a3c..4564b3eac 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -565,6 +565,7 @@ export class AgentSession { #usageFallbackConfirmer: UsageFallbackConfirmer | undefined; #usagePreflightAbortControllers = new Set(); #queuedMessageDrainBlocked = false; + #modeExitDrainSuppressionDepth = 0; #usagePreflightReadyForNextModelCall = false; #usagePreflightReadyModel: Model | undefined; #detachUsageBeforeQueueDequeue: (() => void) | undefined; @@ -5765,13 +5766,14 @@ export class AgentSession { * turn resumes only after teardown, so it runs as a clean non-mode turn. */ async runModeExitTeardown(teardown: () => Promise): Promise { - const previouslyBlocked = this.#queuedMessageDrainBlocked; - this.#queuedMessageDrainBlocked = true; + this.#modeExitDrainSuppressionDepth++; try { await teardown(); } finally { - this.#queuedMessageDrainBlocked = previouslyBlocked; - this.#scheduleIdleQueueDrain(); + this.#modeExitDrainSuppressionDepth--; + if (this.#modeExitDrainSuppressionDepth === 0) { + this.#scheduleIdleQueueDrain(); + } } } @@ -5823,6 +5825,7 @@ export class AgentSession { #scheduleQueuedMessageDrain(): void { if ( this.#queuedMessageDrainScheduled || + this.#modeExitDrainSuppressionDepth > 0 || this.#queuedMessageDrainBlocked || !this.#canAutoContinueForFollowUp() || !this.agent.hasQueuedMessages() @@ -5832,8 +5835,11 @@ export class AgentSession { this.#queuedMessageDrainScheduled = true; this.#scheduleAgentContinue({ shouldContinue: () => { - this.#queuedMessageDrainScheduled = false; - return this.#canAutoContinueForFollowUp() && this.agent.hasQueuedMessages(); + return ( + this.#modeExitDrainSuppressionDepth === 0 && + this.#canAutoContinueForFollowUp() && + this.agent.hasQueuedMessages() + ); }, onSkip: () => { this.#queuedMessageDrainScheduled = false; diff --git a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts index c14c92093..018b9249d 100644 --- a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts +++ b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts @@ -191,9 +191,9 @@ describe("InteractiveMode vibe mode toggle", () => { expect(session.getToolByName("vibe_spawn")).toBeUndefined(); }); - it("resumes a queued user turn without Vibe tools after exiting mid-turn", async () => { + it("holds a user steer queued during Vibe teardown until the tools are removed", async () => { const toolNamesPerCall: string[][] = []; - let firstStarted: PromiseWithResolvers | undefined; + const firstStarted = Promise.withResolvers(); streamFn = (_model, context, options) => { toolNamesPerCall.push((context.tools ?? []).map(tool => tool.name)); const isFirst = toolNamesPerCall.length === 1; @@ -206,28 +206,44 @@ describe("InteractiveMode vibe mode toggle", () => { () => stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") }), { once: true }, ); - firstStarted?.resolve(); + firstStarted.resolve(); } else { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Resumed") }); } }); return stream; }; - firstStarted = Promise.withResolvers(); await mode.handleVibeModeCommand(); const prompt = session.prompt("Delegate this"); await firstStarted.promise; - // Queue a user steer behind the active Vibe turn. - await session.steer("and then do the other thing"); - await mode.handleVibeModeCommand(); + const abortSettled = Promise.withResolvers(); + const releaseTeardown = Promise.withResolvers(); + const abort = session.abort.bind(session); + vi.spyOn(session, "abort").mockImplementation(async options => { + await abort(options); + abortSettled.resolve(); + await releaseTeardown.promise; + }); + const exit = mode.handleVibeModeCommand(); + await abortSettled.promise; + // Queue while teardown is still guarded. The regular queue path clears its + // retry block, but must not clear the independent mode-exit suppression. + await session.steer("and then do the other thing"); + // Drain the microtasks in which an unguarded schedule calls + // agent.continue(). The queued steer must remain owned by the queue until + // teardown releases. + for (let index = 0; index < 5; index++) await Promise.resolve(); + expect(session.agent.peekSteeringQueue()).toHaveLength(1); + expect(toolNamesPerCall.length).toBe(1); + releaseTeardown.resolve(); + + await exit; await prompt; await session.waitForIdle(); - // The aborted Vibe turn plus the resumed queued turn. expect(toolNamesPerCall.length).toBe(2); - // The resumed turn must not have inherited the torn-down Vibe tools. for (const name of VIBE_TOOL_NAMES) { expect(toolNamesPerCall[1]).not.toContain(name); } From a69ce173986dcd2df37ce48463ff4d2e1592e0c8 Mon Sep 17 00:00:00 2001 From: roboomp Date: Wed, 12 Aug 2026 09:01:07 +0000 Subject: [PATCH 4/4] fix(session): reset queued drain scheduling state Clear the scheduled marker when the queued-message continuation claims its run so later idle steers and follow-ups can schedule another drain. Add coverage for two successive successful idle drains. Fixes #8326 --- .../coding-agent/src/session/agent-session.ts | 1 + .../test/agent-session-steer-idle-drain.test.ts | 17 +++++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 4564b3eac..62cd2ebc2 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -5835,6 +5835,7 @@ export class AgentSession { this.#queuedMessageDrainScheduled = true; this.#scheduleAgentContinue({ shouldContinue: () => { + this.#queuedMessageDrainScheduled = false; return ( this.#modeExitDrainSuppressionDepth === 0 && this.#canAutoContinueForFollowUp() && diff --git a/packages/coding-agent/test/agent-session-steer-idle-drain.test.ts b/packages/coding-agent/test/agent-session-steer-idle-drain.test.ts index 86d68763e..a989e9603 100644 --- a/packages/coding-agent/test/agent-session-steer-idle-drain.test.ts +++ b/packages/coding-agent/test/agent-session-steer-idle-drain.test.ts @@ -106,6 +106,23 @@ describe("AgentSession steer idle drain", () => { expect(continueSpy).toHaveBeenCalledTimes(1); }); + it("delivers successive idle steers after each successful drain", async () => { + await createSession([{ role: "user", content: "hello", timestamp: Date.now() }, createAssistantMessage()]); + const continueSpy = vi.spyOn(session.agent, "continue").mockImplementation(async () => { + session.agent.clearAllQueues(); + }); + + await session.steer("first steer"); + vi.advanceTimersByTime(200); + await session.waitForIdle(); + + await session.steer("second steer"); + vi.advanceTimersByTime(200); + await session.waitForIdle(); + + expect(continueSpy).toHaveBeenCalledTimes(2); + }); + it("delivers a steer queued after an interrupted tool result", async () => { await createSession([ { role: "user", content: "hello", timestamp: Date.now() },