From 3cb9258875daa6da6582a8274c7f381aa12c78ff Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 16 Jul 2026 07:17:27 +0000 Subject: [PATCH] fix(session): aborted title generation during dispose - Routed automatic first-input and replan title requests through AgentSession lifecycle cancellation. - Propagated disposal aborts to online provider and local tiny-model title generation. - Added a regression test proving an in-flight title request settles when disposal begins. Fixes #5666 --- packages/coding-agent/CHANGELOG.md | 4 ++ .../src/modes/controllers/input-controller.ts | 13 +--- .../coding-agent/src/session/agent-session.ts | 20 ++++-- .../coding-agent/src/utils/title-generator.ts | 11 +-- ...t-session-title-generation-dispose.test.ts | 70 +++++++++++++++++++ 5 files changed, 99 insertions(+), 19 deletions(-) create mode 100644 packages/coding-agent/test/agent-session-title-generation-dispose.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 0e753ced7..58e2dec97 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed `/quit` and `/exit` leaving failed or stalled automatic title-generation requests alive during session teardown; disposal now aborts both online provider and local tiny-model title requests ([#5666](https://github.com/can1357/oh-my-pi/issues/5666)). + ## [17.0.1] - 2026-07-16 ### Changed diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index af4485175..120df86e3 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -35,7 +35,6 @@ import { EnhancedPasteController } from "../../utils/enhanced-paste"; import { getEditorCommand, openInEditor } from "../../utils/external-editor"; import { ensureSupportedImageInput, ImageInputTooLargeError, loadImageInput } from "../../utils/image-loading"; import { resizeImage } from "../../utils/image-resize"; -import { generateSessionTitle } from "../../utils/title-generator"; /** * Slash commands that may carry secrets in their arguments should never be @@ -821,16 +820,8 @@ export class InputController { // chance, so titling defers past "hi" instead of latching onto it. if (!this.ctx.sessionManager.getSessionName() && !$env.PI_NO_TITLE && !isLowSignalTitleInput(text)) { this.#showTinyTitleDownloadProgress(this.ctx.settings.get("providers.tinyModel")); - const registry = this.ctx.session.modelRegistry; - generateSessionTitle( - text, - registry, - this.ctx.settings, - this.ctx.session.sessionId, - this.ctx.session.model, - provider => this.ctx.session.agent.metadataForProvider(provider), - this.ctx.session.titleSystemPrompt, - ) + this.ctx.session + .generateTitle(text) .then(async title => { // Re-check: a concurrent attempt for an earlier message may have // already named the session. Don't clobber it. Terminal title and diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index d2ab574dc..dea857ffa 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -1860,6 +1860,7 @@ export class AgentSession { * generation path. Refresh via {@link AgentSession.setTitleSystemPrompt} when * the session cwd changes. */ #titleSystemPrompt: string | undefined; + #titleGenerationAbortController = new AbortController(); #toolChoiceQueue = new ToolChoiceQueue(); // Bash execution state @@ -6226,6 +6227,7 @@ export class AgentSession { */ beginDispose(): void { this.#isDisposed = true; + this.#titleGenerationAbortController.abort(); this.#flushPendingIrcAsides(); this.yieldQueue.clear(); this.agent.setAsideMessageProvider(undefined); @@ -8982,16 +8984,26 @@ export class AgentSession { this.#replanTitleRefreshInFlight = refresh; } - async #refreshTitleAfterReplan(context: string, sessionId: string): Promise { - const title = await generateSessionTitle( - context, + /** + * Generate an automatic session title tied to this session's lifecycle. + * Input and replan callers share the signal so disposal cancels provider and + * local-worker requests instead of leaving background inference alive. + */ + generateTitle(firstMessage: string): Promise { + return generateSessionTitle( + firstMessage, this.#modelRegistry, this.settings, - sessionId, + this.sessionManager.getSessionId(), this.model, provider => this.agent.metadataForProvider(provider), this.#titleSystemPrompt, + this.#titleGenerationAbortController.signal, ); + } + + async #refreshTitleAfterReplan(context: string, sessionId: string): Promise { + const title = await this.generateTitle(context); if (!title) return; if (this.sessionManager.getSessionId() !== sessionId) return; if (!this.settings.get("title.refreshOnReplan")) return; diff --git a/packages/coding-agent/src/utils/title-generator.ts b/packages/coding-agent/src/utils/title-generator.ts index 6e7d19e24..bdce9ebc1 100644 --- a/packages/coding-agent/src/utils/title-generator.ts +++ b/packages/coding-agent/src/utils/title-generator.ts @@ -68,6 +68,7 @@ function getTitleModel(registry: ModelRegistry, settings: Settings, currentModel * resolver instead of a pre-evaluated value ensures the metadata's account_uuid * reflects the credential actually selected for this request. * @param customSystemPrompt Optional title-specific system prompt override + * @param signal Session-lifecycle cancellation for background title requests */ export async function generateSessionTitle( firstMessage: string, @@ -77,6 +78,7 @@ export async function generateSessionTitle( currentModel?: Model, metadataResolver?: (provider: string) => Record | undefined, customSystemPrompt?: string, + signal?: AbortSignal, ): Promise { // Defer titling for greetings / acknowledgements / empty input. The default // tiny title model can't reliably decline trivial input, so this happens @@ -97,7 +99,7 @@ export async function generateSessionTitle( sessionId, currentModel, metadataResolver, - undefined, + signal, titleSystemPrompt, ); } @@ -117,9 +119,10 @@ export async function generateSessionTitle( return null; } try { - const localTitle = titleSystemPrompt - ? await tinyTitleClient.generate(tinyModel, firstMessage, { systemPrompt: titleSystemPrompt }) - : await tinyTitleClient.generate(tinyModel, firstMessage); + const localTitle = await tinyTitleClient.generate(tinyModel, firstMessage, { + signal, + systemPrompt: titleSystemPrompt, + }); if (!localTitle) { logger.warn("title-generator: local tiny model produced no title; skipping (no online fallback)", { sessionId, diff --git a/packages/coding-agent/test/agent-session-title-generation-dispose.test.ts b/packages/coding-agent/test/agent-session-title-generation-dispose.test.ts new file mode 100644 index 000000000..9a0f0e858 --- /dev/null +++ b/packages/coding-agent/test/agent-session-title-generation-dispose.test.ts @@ -0,0 +1,70 @@ +import { afterEach, describe, expect, it, vi } from "bun:test"; +import * as path from "node:path"; +import { Agent } from "@oh-my-pi/pi-agent-core"; +import * as ai from "@oh-my-pi/pi-ai"; +import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock"; +import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { TempDir } from "@oh-my-pi/pi-utils"; +import { createAssistantMessage } from "./helpers/agent-session-setup"; + +let session: AgentSession | undefined; +let authStorage: AuthStorage | undefined; +let tempDir: TempDir | undefined; + +afterEach(async () => { + vi.restoreAllMocks(); + await session?.dispose(); + authStorage?.close(); + tempDir?.removeSync(); + session = undefined; + authStorage = undefined; + tempDir = undefined; +}); + +describe("AgentSession title generation disposal", () => { + it("aborts an in-flight automatic title request when disposal begins", async () => { + tempDir = TempDir.createSync("@pi-title-dispose-"); + authStorage = await AuthStorage.create(path.join(tempDir.path(), "auth.db")); + authStorage.setRuntimeApiKey("anthropic", "test-key"); + const model = getBundledModel("anthropic", "claude-sonnet-4-5"); + if (!model) throw new Error("Expected claude-sonnet-4-5 model to exist"); + + const settings = Settings.isolated({ + "compaction.enabled": false, + "providers.tinyModel": "online", + }); + settings.overrideModelRoles({ smol: `${model.provider}/${model.id}` }); + const agent = new Agent({ + getApiKey: () => "test-key", + initialState: { model, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: createMockModel({ responses: [{ content: ["Done"] }] }).stream, + }); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry: new ModelRegistry(authStorage), + }); + const started = Promise.withResolvers(); + const response = Promise.withResolvers(); + let requestSignal: AbortSignal | undefined; + vi.spyOn(ai, "completeSimple").mockImplementation((_model, _context, options) => { + requestSignal = options?.signal; + requestSignal?.addEventListener("abort", () => response.resolve(createAssistantMessage("")), { once: true }); + started.resolve(); + return response.promise; + }); + + const generation = session.generateTitle("Investigate shutdown"); + await started.promise; + session.beginDispose(); + + expect(requestSignal?.aborted).toBe(true); + expect(await generation).toBeNull(); + }); +});