diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 26cce8295..ba792977d 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -1,6 +1,14 @@ # Changelog ## [Unreleased] +### Added + +- Added `close()` method to `AuthStorage` for properly closing the underlying credential store +- Added support for `initiatorOverride` option in OpenAI and Anthropic providers to customize message attribution + +### Fixed + +- Fixed resource cleanup in `AuthCredentialStore.close()` to properly finalize all prepared statements before closing the database ## [13.11.1] - 2026-03-13 diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 3756c9d14..cb933604e 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -292,6 +292,7 @@ export class AuthStorage { #fallbackResolver?: (provider: string) => string | undefined; #store: AuthCredentialStore; #configValueResolver: (config: string) => Promise; + #closed = false; constructor(store: AuthCredentialStore, options: AuthStorageOptions = {}) { this.#store = store; @@ -319,6 +320,17 @@ export class AuthStorage { return new AuthStorage(store, options); } + /** + * Close the underlying credential store. + * + * After calling this, the instance must not be reused. + */ + close(): void { + if (this.#closed) return; + this.#closed = true; + this.#store.close(); + } + /** * Set a runtime API key override (not persisted to disk). * Used for CLI --api-key flag. @@ -2132,6 +2144,7 @@ export class AuthCredentialStore { #getCacheStmt: Statement; #upsertCacheStmt: Statement; #deleteExpiredCacheStmt: Statement; + #closed = false; constructor(db: Database) { this.#db = db; @@ -2578,6 +2591,19 @@ export class AuthCredentialStore { } close(): void { + if (this.#closed) return; + this.#closed = true; + this.#listActiveStmt.finalize(); + this.#listActiveByProviderStmt.finalize(); + this.#listDisabledByProviderStmt.finalize(); + this.#insertStmt.finalize(); + this.#updateStmt.finalize(); + this.#deleteStmt.finalize(); + this.#deleteByProviderStmt.finalize(); + this.#hardDeleteStmt.finalize(); + this.#getCacheStmt.finalize(); + this.#upsertCacheStmt.finalize(); + this.#deleteExpiredCacheStmt.finalize(); this.#db.close(); } } diff --git a/packages/ai/src/providers/anthropic.ts b/packages/ai/src/providers/anthropic.ts index 499f2b0d8..779a8145e 100644 --- a/packages/ai/src/providers/anthropic.ts +++ b/packages/ai/src/providers/anthropic.ts @@ -570,6 +570,7 @@ export const streamAnthropic: StreamFunction<"anthropic-messages"> = ( hasImages: hasCopilotVisionInput(context.messages), premiumMultiplier: model.premiumMultiplier, headers: { ...(model.headers ?? {}), ...(options?.headers ?? {}) }, + initiatorOverride: options?.initiatorOverride, }) : undefined; const output: AssistantMessage = { diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 4f1b5bcaa..c1ebe5456 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -16,6 +16,7 @@ import { type Context, isSpecialServiceTier, type Message, + type MessageAttribution, type Model, type OpenAICompat, type ServiceTier, @@ -201,6 +202,7 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions"> = ( context, apiKey, options?.headers, + options?.initiatorOverride, ); const params = buildParams(model, context, options); options?.onPayload?.(params); @@ -515,7 +517,12 @@ async function createClient( context: Context, apiKey?: string, extraHeaders?: Record, -) { + initiatorOverride?: MessageAttribution, +): Promise<{ + client: OpenAI; + copilotPremiumRequests: number | undefined; + baseUrl: string | undefined; +}> { if (!apiKey) { if (!$env.OPENAI_API_KEY) { throw new Error( @@ -539,6 +546,7 @@ async function createClient( hasImages, premiumMultiplier: model.premiumMultiplier, headers, + initiatorOverride, }); Object.assign(headers, copilot.headers); copilotPremiumRequests = copilot.premiumRequests; diff --git a/packages/ai/src/providers/openai-responses.ts b/packages/ai/src/providers/openai-responses.ts index dcf94980e..aa8e67bf1 100644 --- a/packages/ai/src/providers/openai-responses.ts +++ b/packages/ai/src/providers/openai-responses.ts @@ -12,6 +12,7 @@ import { type CacheRetention, type Context, isSpecialServiceTier, + type MessageAttribution, type Model, type ServiceTier, type StreamFunction, @@ -117,7 +118,13 @@ export const streamOpenAIResponses: StreamFunction<"openai-responses"> = ( try { // Create OpenAI client const apiKey = options?.apiKey || getEnvApiKey(model.provider) || ""; - const { client, copilotPremiumRequests, baseUrl } = createClient(model, context, apiKey, options?.headers); + const { client, copilotPremiumRequests, baseUrl } = createClient( + model, + context, + apiKey, + options?.headers, + options?.initiatorOverride, + ); const { params } = buildParams(model, context, options); const requestAbortController = new AbortController(); const requestSignal = options?.signal @@ -190,7 +197,12 @@ function createClient( context: Context, apiKey?: string, extraHeaders?: Record, -) { + initiatorOverride?: MessageAttribution, +): { + client: OpenAI; + copilotPremiumRequests: number | undefined; + baseUrl: string | undefined; +} { if (!apiKey) { if (!$env.OPENAI_API_KEY) { throw new Error( @@ -211,6 +223,7 @@ function createClient( hasImages, premiumMultiplier: model.premiumMultiplier, headers, + initiatorOverride, }); Object.assign(headers, copilot.headers); copilotPremiumRequests = copilot.premiumRequests; diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 4be39de0f..c664bb84e 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,6 +1,14 @@ # Changelog ## [Unreleased] +### Added + +- Added `close()` method to SessionManager for properly closing persistent writers after flushing pending data + +### Fixed + +- Fixed timeout handling in RpcClient to properly clear timeouts and prevent resource leaks +- Fixed AgentSession disposal to call SessionManager's `close()` method when available, ensuring proper cleanup of persistent writers ## [13.11.1] - 2026-03-13 diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index 3fa485aa7..2e8056958 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -152,16 +152,19 @@ export class RpcClient { }); // Timeout to prevent hanging forever - void Bun.sleep(30000).then(() => { - if (!readySettled) { - readySettled = true; - readyReject( - new Error(`Timeout waiting for agent to become ready. Stderr: ${this.#process?.peekStderr() ?? ""}`), - ); - } + const readyTimeout = this.#startTimeout(30000, () => { + if (readySettled) return; + readySettled = true; + readyReject( + new Error(`Timeout waiting for agent to become ready. Stderr: ${this.#process?.peekStderr() ?? ""}`), + ); }); - await readyPromise; + try { + await readyPromise; + } finally { + clearTimeout(readyTimeout); + } } /** @@ -207,6 +210,12 @@ export class RpcClient { return this.#process?.peekStderr() ?? ""; } + #startTimeout(timeoutMs: number, onTimeout: () => void): NodeJS.Timeout { + const timer = setTimeout(onTimeout, timeoutMs); + timer.unref(); + return timer; + } + // ========================================================================= // Command Methods // ========================================================================= @@ -440,17 +449,17 @@ export class RpcClient { if (event.type === "agent_end") { settled = true; unsubscribe(); + clearTimeout(timeoutId); resolve(); } }); - void (async () => { - await Bun.sleep(timeout); + const timeoutId = this.#startTimeout(timeout, () => { if (settled) return; settled = true; unsubscribe(); reject(new Error(`Timeout waiting for agent to become idle. Stderr: ${this.#process?.peekStderr() ?? ""}`)); - })(); + }); return promise; } @@ -466,17 +475,17 @@ export class RpcClient { if (event.type === "agent_end") { settled = true; unsubscribe(); + clearTimeout(timeoutId); resolve(events); } }); - void (async () => { - await Bun.sleep(timeout); + const timeoutId = this.#startTimeout(timeout, () => { if (settled) return; settled = true; unsubscribe(); reject(new Error(`Timeout collecting events. Stderr: ${this.#process?.peekStderr() ?? ""}`)); - })(); + }); return promise; } @@ -522,25 +531,26 @@ export class RpcClient { const fullCommand = { ...command, id } as RpcCommand; const { promise, resolve, reject } = Promise.withResolvers(); let settled = false; - void (async () => { - await Bun.sleep(30000); + const timeoutId = this.#startTimeout(30000, () => { if (settled) return; this.#pendingRequests.delete(id); settled = true; reject( new Error(`Timeout waiting for response to ${command.type}. Stderr: ${this.#process?.peekStderr() ?? ""}`), ); - })(); + }); this.#pendingRequests.set(id, { resolve: response => { if (settled) return; settled = true; + clearTimeout(timeoutId); resolve(response); }, reject: error => { if (settled) return; settled = true; + clearTimeout(timeoutId); reject(error); }, }); @@ -555,6 +565,7 @@ export class RpcClient { this.#pendingRequests.delete(id); if (settled) return; settled = true; + clearTimeout(timeoutId); reject(err); }); } diff --git a/packages/coding-agent/src/prompts/system/system-prompt.md b/packages/coding-agent/src/prompts/system/system-prompt.md index 0e44b0787..6b65c84f8 100644 --- a/packages/coding-agent/src/prompts/system/system-prompt.md +++ b/packages/coding-agent/src/prompts/system/system-prompt.md @@ -297,6 +297,7 @@ When a tool call fails, read the full error before doing anything else. When a f - You **MUST** exhaust tools/context/files first — explore. ## 7. Verification - Test everything rigorously → Future contributor cannot break behavior without failure. Prefer unit/e2e. +- You **MUST NOT** rely on mocks — they invent behaviors that never happen in production and hide real bugs. - You **SHOULD** run only tests you added/modified unless asked otherwise. - You **MUST NOT** yield without proof when non-trivial work, self-assessment is deceptive: tests, linters, type checks, repro steps… exhaust all external verification. diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 8a04d9584..c91a9b9db 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -1539,7 +1539,7 @@ export class AgentSession { if (drained === false && deliveryState) { logger.warn("Async job completion deliveries still pending during dispose", { ...deliveryState }); } - await this.sessionManager.flush(); + await this.sessionManager.close(); for (const state of this.#providerSessionState.values()) { state.close(); } diff --git a/packages/coding-agent/src/session/session-manager.ts b/packages/coding-agent/src/session/session-manager.ts index aa33dff3e..1029b35a9 100644 --- a/packages/coding-agent/src/session/session-manager.ts +++ b/packages/coding-agent/src/session/session-manager.ts @@ -1706,6 +1706,16 @@ export class SessionManager { if (this.#persistError) throw this.#persistError; } + /** Close the persistent writer after flushing all pending data. */ + async close(): Promise { + if (!this.#persistWriter) return; + await this.#queuePersistTask(async () => { + await this.#closePersistWriterInternal(); + this.#flushed = true; + }); + if (this.#persistError) throw this.#persistError; + } + getCwd(): string { return this.cwd; } diff --git a/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts b/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts index 93bceabe6..0520e3455 100644 --- a/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts +++ b/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts @@ -126,6 +126,7 @@ describe("AgentSession auto-compaction queue resume", () => { afterEach(async () => { await session.dispose(); + authStorage.close(); tempDir.removeSync(); vi.useRealTimers(); getRuntimeSignals().length = 0; diff --git a/packages/coding-agent/test/agent-session-auto-compaction-x-initiator.test.ts b/packages/coding-agent/test/agent-session-auto-compaction-x-initiator.test.ts index 61fbb28a4..7b92e3ee1 100644 --- a/packages/coding-agent/test/agent-session-auto-compaction-x-initiator.test.ts +++ b/packages/coding-agent/test/agent-session-auto-compaction-x-initiator.test.ts @@ -1,109 +1,122 @@ -import { afterEach, beforeEach, describe, expect, it, mock, vi } from "bun:test"; +import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import * as path from "node:path"; +import type { AssistantMessage, Context, SimpleStreamOptions } from "@oh-my-pi/pi-ai"; +import * as ai from "@oh-my-pi/pi-ai"; import { TempDir } from "@oh-my-pi/pi-utils"; +import { Settings } from "../src/config/settings"; +import { createAgentSession } from "../src/sdk"; import type { AgentSession } from "../src/session/agent-session"; +import { AuthStorage } from "../src/session/auth-storage"; +import { SessionManager } from "../src/session/session-manager"; -const compactMock = vi.fn(); - -mock.module("../src/session/compaction", () => ({ - calculateContextTokens: (usage: { - totalTokens?: number; - input: number; - output: number; - cacheRead: number; - cacheWrite: number; - }) => usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite, - calculatePromptTokens: () => 0, - collectEntriesForBranchSummary: () => [], - compact: compactMock, - estimateTokens: () => 0, - generateBranchSummary: async () => "Branch summary", - prepareCompaction: (entries: Array<{ id?: string }>) => ({ - firstKeptEntryId: entries[entries.length - 1]?.id ?? "missing-entry", - messagesToSummarize: [], - turnPrefixMessages: [], - recentMessages: [], - isSplitTurn: false, - tokensBefore: 321, - previousSummary: undefined, - previousPreserveData: undefined, - fileOps: { read: new Set(), edited: new Set() }, - settings: { - enabled: true, - strategy: "context-full", - thresholdPercent: 80, - reserveTokens: 16384, - keepRecentTokens: 20000, - autoContinue: false, - remoteEnabled: false, +function createAssistantMessage(text: string): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "openai-completions", + provider: "github-copilot", + model: "gpt-4o", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, - }), - shouldCompact: () => true, -})); + stopReason: "stop", + timestamp: Date.now(), + }; +} -describe("AgentSession compaction Copilot initiator attribution", async () => { - const { getBundledModel } = await import("@oh-my-pi/pi-ai"); - const { Settings } = await import("../src/config/settings"); - const { createAgentSession } = await import("../src/sdk"); - const { AuthStorage } = await import("../src/session/auth-storage"); - const { SessionManager } = await import("../src/session/session-manager"); +function contextContainsMarker(context: Context, marker: string): boolean { + return context.messages.some(message => { + if (typeof message.content === "string") { + return message.content.includes(marker); + } + return message.content.some(block => { + if (block.type === "text") { + return block.text.includes(marker); + } + if (block.type === "thinking") { + return block.thinking.includes(marker); + } + return false; + }); + }); +} +function captureCompactionCalls(marker: string) { + const capturedOptions: Array = []; + const originalCompleteSimple = ai.completeSimple; + vi.spyOn(ai, "completeSimple").mockImplementation(async (...args) => { + const [model, context, options] = args; + if (model.provider === "github-copilot" && contextContainsMarker(context, marker)) { + capturedOptions.push(options); + return createAssistantMessage("Compacted summary") as never; + } + return originalCompleteSimple(...args); + }); + return capturedOptions; +} + +describe("AgentSession compaction Copilot initiator attribution", () => { let tempDir: TempDir; const sessions: Array<{ dispose: () => Promise }> = []; + const authStorages: AuthStorage[] = []; beforeEach(() => { tempDir = TempDir.createSync("@pi-auto-compaction-x-initiator-"); - compactMock.mockReset(); - compactMock.mockImplementation((preparation: { firstKeptEntryId: string }) => ({ - summary: "Compacted summary", - firstKeptEntryId: preparation.firstKeptEntryId, - tokensBefore: 321, - })); }); afterEach(async () => { for (const session of sessions.splice(0)) { await session.dispose(); } + for (const authStorage of authStorages.splice(0)) { + authStorage.close(); + } vi.restoreAllMocks(); tempDir.removeSync(); }); - async function createSession(taskDepth: number) { - const model = getBundledModel("github-copilot", "gpt-4o"); + async function createSession(taskDepth: number, marker: string) { + const model = ai.getBundledModel("github-copilot", "gpt-4o"); if (!model) { throw new Error("Expected github-copilot/gpt-4o model to exist"); } const authStorage = await AuthStorage.create(path.join(tempDir.path(), `testauth-${taskDepth}.db`)); + authStorages.push(authStorage); authStorage.setRuntimeApiKey("github-copilot", "test-key"); const sessionManager = SessionManager.inMemory(); sessionManager.appendMessage({ role: "user", - content: "Initial request with enough text to summarize later.", + content: `Initial request with enough text to summarize later. ${marker}`, timestamp: Date.now() - 3, }); sessionManager.appendMessage({ role: "assistant", - content: [{ type: "text", text: "Initial response with extra context for compaction." }], + content: [{ type: "text", text: `Initial response with extra context for compaction. ${marker}` }], api: model.api, provider: model.provider, model: model.id, stopReason: "stop", usage: { - input: 100, - output: 50, + // Keep this large so manual compaction remains eligible even if defaults are used. + input: 120_000, + output: 2_000, cacheRead: 0, cacheWrite: 0, - totalTokens: 150, + totalTokens: 122_000, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, timestamp: Date.now() - 2, }); sessionManager.appendMessage({ role: "user", - content: "Latest request before the oversized assistant turn.", + content: `Latest request before the oversized assistant turn. ${marker}`, timestamp: Date.now(), }); @@ -115,6 +128,7 @@ describe("AgentSession compaction Copilot initiator attribution", async () => { sessionManager, settings: Settings.isolated({ "compaction.autoContinue": false, + "compaction.keepRecentTokens": 1, }), disableExtensionDiscovery: true, skills: [], @@ -133,6 +147,16 @@ describe("AgentSession compaction Copilot initiator attribution", async () => { expect(model.headers?.["X-Initiator"]).toBeUndefined(); } + function expectInitiatorOverride( + capturedOptions: Array, + expected: "agent" | undefined, + ) { + expect(capturedOptions.length).toBeGreaterThan(0); + for (const options of capturedOptions) { + expect(options?.initiatorOverride).toBe(expected); + } + } + async function triggerAutoCompaction( session: Pick, model: { api: string; provider: string; id: string; contextWindow: number }, @@ -170,74 +194,54 @@ describe("AgentSession compaction Copilot initiator attribution", async () => { } it("keeps main-session manual compaction user-attributed", async () => { - const { model, session } = await createSession(0); + const marker = `main-manual-${Date.now()}`; + const capturedOptions = captureCompactionCalls(marker); + const { model, session } = await createSession(0, marker); await session.compact(); - expect(compactMock).toHaveBeenCalledTimes(1); - const compactModel = compactMock.mock.calls[0]?.[1] as { - provider: string; - id: string; - headers?: Record; - }; - const compactOptions = compactMock.mock.calls[0]?.[5] as { initiatorOverride?: string } | undefined; - expect(compactModel.provider).toBe("github-copilot"); - expect(compactModel.id).toBe(model.id); - expectNoForcedCopilotHeader(compactModel); - expect(compactOptions?.initiatorOverride).toBeUndefined(); + expect(model.provider).toBe("github-copilot"); + expect(model.id).toBe("gpt-4o"); + expectNoForcedCopilotHeader(model); + expectInitiatorOverride(capturedOptions, undefined); }); it("uses agent attribution for main-session auto-compaction", async () => { - const { model, session } = await createSession(0); + const marker = `main-auto-${Date.now()}`; + const capturedOptions = captureCompactionCalls(marker); + const { model, session } = await createSession(0, marker); await triggerAutoCompaction(session, model); - expect(compactMock).toHaveBeenCalledTimes(1); - const compactModel = compactMock.mock.calls[0]?.[1] as { - provider: string; - id: string; - headers?: Record; - }; - const compactOptions = compactMock.mock.calls[0]?.[5] as { initiatorOverride?: string } | undefined; - expect(compactModel.provider).toBe("github-copilot"); - expect(compactModel.id).toBe(model.id); - expectNoForcedCopilotHeader(compactModel); - expect(compactOptions?.initiatorOverride).toBe("agent"); + expect(model.provider).toBe("github-copilot"); + expect(model.id).toBe("gpt-4o"); + expectNoForcedCopilotHeader(model); + expectInitiatorOverride(capturedOptions, "agent"); }); it("keeps subagent manual compaction user-attributed", async () => { - const { model, session } = await createSession(1); + const marker = `subagent-manual-${Date.now()}`; + const capturedOptions = captureCompactionCalls(marker); + const { model, session } = await createSession(1, marker); await session.compact(); - expect(compactMock).toHaveBeenCalledTimes(1); - const compactModel = compactMock.mock.calls[0]?.[1] as { - provider: string; - id: string; - headers?: Record; - }; - const compactOptions = compactMock.mock.calls[0]?.[5] as { initiatorOverride?: string } | undefined; - expect(compactModel.provider).toBe("github-copilot"); - expect(compactModel.id).toBe(model.id); - expectNoForcedCopilotHeader(compactModel); - expect(compactOptions?.initiatorOverride).toBeUndefined(); + expect(model.provider).toBe("github-copilot"); + expect(model.id).toBe("gpt-4o"); + expectNoForcedCopilotHeader(model); + expectInitiatorOverride(capturedOptions, undefined); }); it("uses agent attribution for subagent auto-compaction", async () => { - const { model, session } = await createSession(1); + const marker = `subagent-auto-${Date.now()}`; + const capturedOptions = captureCompactionCalls(marker); + const { model, session } = await createSession(1, marker); await triggerAutoCompaction(session, model); - expect(compactMock).toHaveBeenCalledTimes(1); - const compactModel = compactMock.mock.calls[0]?.[1] as { - provider: string; - id: string; - headers?: Record; - }; - const compactOptions = compactMock.mock.calls[0]?.[5] as { initiatorOverride?: string } | undefined; - expect(compactModel.provider).toBe("github-copilot"); - expect(compactModel.id).toBe(model.id); - expectNoForcedCopilotHeader(compactModel); - expect(compactOptions?.initiatorOverride).toBe("agent"); + expect(model.provider).toBe("github-copilot"); + expect(model.id).toBe("gpt-4o"); + expectNoForcedCopilotHeader(model); + expectInitiatorOverride(capturedOptions, "agent"); }); }); diff --git a/packages/coding-agent/test/agent-session-before-agent-start-attribution.test.ts b/packages/coding-agent/test/agent-session-before-agent-start-attribution.test.ts index ef8c4c226..6b8decc11 100644 --- a/packages/coding-agent/test/agent-session-before-agent-start-attribution.test.ts +++ b/packages/coding-agent/test/agent-session-before-agent-start-attribution.test.ts @@ -39,12 +39,13 @@ describe("AgentSession before_agent_start attribution fallback", () => { let tempDir: TempDir; let session: AgentSession; let modelRegistry: ModelRegistry; + let authStorage: AuthStorage | undefined; const injectedText = "before-agent-start injected message"; beforeEach(async () => { tempDir = TempDir.createSync("@pi-before-agent-start-attribution-"); - const authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); + authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); authStorage.setRuntimeApiKey("anthropic", "test-key"); modelRegistry = new ModelRegistry(authStorage); }); @@ -54,6 +55,8 @@ describe("AgentSession before_agent_start attribution fallback", () => { if (session) { await session.dispose(); } + authStorage?.close(); + authStorage = undefined; tempDir.removeSync(); }); diff --git a/packages/coding-agent/test/agent-session-branching.test.ts b/packages/coding-agent/test/agent-session-branching.test.ts index 43c72546f..59f3354f9 100644 --- a/packages/coding-agent/test/agent-session-branching.test.ts +++ b/packages/coding-agent/test/agent-session-branching.test.ts @@ -26,6 +26,7 @@ describe.skipIf(!e2eApiKey("ANTHROPIC_API_KEY"))("AgentSession branching", () => let session: AgentSession; let tempDir: string; let sessionManager: SessionManager; + let authStorage: AuthStorage | undefined; beforeEach(() => { // Create temp directory for session files @@ -37,6 +38,8 @@ describe.skipIf(!e2eApiKey("ANTHROPIC_API_KEY"))("AgentSession branching", () => if (session) { await session.dispose(); } + authStorage?.close(); + authStorage = undefined; if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } @@ -64,7 +67,7 @@ describe.skipIf(!e2eApiKey("ANTHROPIC_API_KEY"))("AgentSession branching", () => sessionManager = noSession ? SessionManager.inMemory() : SessionManager.create(tempDir); const settings = Settings.isolated(); - const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); + authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); session = new AgentSession({ diff --git a/packages/coding-agent/test/agent-session-compaction.test.ts b/packages/coding-agent/test/agent-session-compaction.test.ts index 8f3ffe62a..242b90b8d 100644 --- a/packages/coding-agent/test/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/agent-session-compaction.test.ts @@ -27,6 +27,7 @@ describe.skipIf(!e2eApiKey("ANTHROPIC_API_KEY"))("AgentSession compaction e2e", let tempDir: string; let sessionManager: SessionManager; let events: AgentSessionEvent[]; + let authStorage: AuthStorage | undefined; beforeEach(() => { // Create temp directory for session files @@ -41,6 +42,8 @@ describe.skipIf(!e2eApiKey("ANTHROPIC_API_KEY"))("AgentSession compaction e2e", if (session) { await session.dispose(); } + authStorage?.close(); + authStorage = undefined; if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } @@ -68,7 +71,7 @@ describe.skipIf(!e2eApiKey("ANTHROPIC_API_KEY"))("AgentSession compaction e2e", sessionManager = inMemory ? SessionManager.inMemory() : SessionManager.create(tempDir); const settings = Settings.isolated({ "compaction.keepRecentTokens": 1 }); - const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); + authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); const modelRegistry = new ModelRegistry(authStorage); session = new AgentSession({ diff --git a/packages/coding-agent/test/agent-session-concurrent.test.ts b/packages/coding-agent/test/agent-session-concurrent.test.ts index ec6629008..88408f1d6 100644 --- a/packages/coding-agent/test/agent-session-concurrent.test.ts +++ b/packages/coding-agent/test/agent-session-concurrent.test.ts @@ -45,6 +45,7 @@ function createAssistantMessage(text: string): AssistantMessage { describe("AgentSession concurrent prompt guard", () => { let session: AgentSession; let tempDir: string; + const authStorages: AuthStorage[] = []; beforeEach(() => { tempDir = path.join(os.tmpdir(), `pi-concurrent-test-${Snowflake.next()}`); @@ -55,6 +56,9 @@ describe("AgentSession concurrent prompt guard", () => { if (session) { await session.dispose(); } + for (const authStorage of authStorages.splice(0)) { + authStorage.close(); + } if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } @@ -93,6 +97,7 @@ describe("AgentSession concurrent prompt guard", () => { const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); + authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); @@ -181,6 +186,7 @@ describe("AgentSession concurrent prompt guard", () => { const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); + authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); @@ -205,6 +211,7 @@ describe("AgentSession concurrent prompt guard", () => { describe("AgentSession TTSR resume gate", () => { let session: AgentSession; let tempDir: string; + const authStorages: AuthStorage[] = []; beforeEach(() => { tempDir = path.join(os.tmpdir(), `pi-ttsr-gate-test-${Snowflake.next()}`); @@ -215,6 +222,9 @@ describe("AgentSession TTSR resume gate", () => { if (session) { await session.dispose(); } + for (const authStorage of authStorages.splice(0)) { + authStorage.close(); + } if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } @@ -318,6 +328,7 @@ describe("AgentSession TTSR resume gate", () => { const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-int.db")); + authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); @@ -401,6 +412,7 @@ describe("AgentSession TTSR resume gate", () => { const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-def.db")); + authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); @@ -470,6 +482,7 @@ describe("AgentSession TTSR resume gate", () => { const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-abt.db")); + authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); @@ -604,6 +617,7 @@ describe("AgentSession TTSR resume gate", () => { const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-tool.db")); + authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); @@ -628,6 +642,7 @@ describe("AgentSession TTSR resume gate", () => { }); it("prompt() waits for context-promotion continuation to finish", async () => { const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-promo.db")); + authStorages.push(authStorage); authStorage.setRuntimeApiKey("openai-codex", "test-key"); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); diff --git a/packages/coding-agent/test/agent-session-context-promotion.test.ts b/packages/coding-agent/test/agent-session-context-promotion.test.ts index fb254d79e..f00d1f6fb 100644 --- a/packages/coding-agent/test/agent-session-context-promotion.test.ts +++ b/packages/coding-agent/test/agent-session-context-promotion.test.ts @@ -26,6 +26,7 @@ describe("AgentSession context promotion", () => { if (session) { await session.dispose(); } + authStorage.close(); tempDir.removeSync(); }); diff --git a/packages/coding-agent/test/agent-session-eager-todo.test.ts b/packages/coding-agent/test/agent-session-eager-todo.test.ts index e9d7341ad..155b0d630 100644 --- a/packages/coding-agent/test/agent-session-eager-todo.test.ts +++ b/packages/coding-agent/test/agent-session-eager-todo.test.ts @@ -107,6 +107,7 @@ describe("AgentSession eager todo enforcement", () => { let session: AgentSession; let streamCallCount = 0; let scriptedResponses: AssistantMessage[] = []; + let authStorage: AuthStorage | undefined; const observedCalls: ObservedPromptCall[] = []; beforeEach(async () => { @@ -118,7 +119,7 @@ describe("AgentSession eager todo enforcement", () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5"); if (!model) throw new Error("Expected claude-sonnet-4-5 model to exist"); - const authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); + authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); authStorage.setRuntimeApiKey("anthropic", "test-key"); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir.path(), "models.yml")); const settings = Settings.isolated({ @@ -199,6 +200,8 @@ describe("AgentSession eager todo enforcement", () => { if (session) { await session.dispose(); } + authStorage?.close(); + authStorage = undefined; tempDir.removeSync(); }); diff --git a/packages/coding-agent/test/agent-session-handoff.test.ts b/packages/coding-agent/test/agent-session-handoff.test.ts index f98b6886f..29a75102e 100644 --- a/packages/coding-agent/test/agent-session-handoff.test.ts +++ b/packages/coding-agent/test/agent-session-handoff.test.ts @@ -86,6 +86,7 @@ describe("AgentSession handoff", () => { if (session) { await session.dispose(); } + authStorage.close(); try { await tempDir.remove(); } catch {} diff --git a/packages/coding-agent/test/agent-session-new-session-todos.test.ts b/packages/coding-agent/test/agent-session-new-session-todos.test.ts index e914866f1..7f9f42d27 100644 --- a/packages/coding-agent/test/agent-session-new-session-todos.test.ts +++ b/packages/coding-agent/test/agent-session-new-session-todos.test.ts @@ -24,6 +24,7 @@ describe("AgentSession newSession clears todo artifacts", () => { let tempDir: string; let session: AgentSession; let sessionManager: SessionManager; + let authStorage: AuthStorage | undefined; beforeEach(async () => { tempDir = path.join(os.tmpdir(), `pi-new-session-todos-test-${Snowflake.next()}`); @@ -31,7 +32,7 @@ describe("AgentSession newSession clears todo artifacts", () => { sessionManager = SessionManager.create(tempDir); const settings = Settings.isolated(); - const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); + authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); const model = getBundledModel("anthropic", "claude-sonnet-4-5"); @@ -71,6 +72,8 @@ describe("AgentSession newSession clears todo artifacts", () => { if (session) { await session.dispose(); } + authStorage?.close(); + authStorage = undefined; if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } diff --git a/packages/coding-agent/test/agent-session-resolve-reminder.test.ts b/packages/coding-agent/test/agent-session-resolve-reminder.test.ts index 8f635c044..fbf729367 100644 --- a/packages/coding-agent/test/agent-session-resolve-reminder.test.ts +++ b/packages/coding-agent/test/agent-session-resolve-reminder.test.ts @@ -40,6 +40,7 @@ describe("AgentSession resolve reminder", () => { let tempDir: string; let pendingActionStore: PendingActionStore; let streamCallCount = 0; + let authStorage: AuthStorage | undefined; beforeEach(async () => { tempDir = path.join(os.tmpdir(), `pi-resolve-reminder-test-${Snowflake.next()}`); @@ -52,7 +53,7 @@ describe("AgentSession resolve reminder", () => { throw new Error("Test model not found in registry"); } - const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); + authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorage.setRuntimeApiKey("anthropic", "test-key"); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); @@ -92,6 +93,8 @@ describe("AgentSession resolve reminder", () => { afterEach(async () => { await session.dispose(); + authStorage?.close(); + authStorage = undefined; if (fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true, force: true }); } diff --git a/packages/coding-agent/test/agent-session-role-thinking.test.ts b/packages/coding-agent/test/agent-session-role-thinking.test.ts index 81a7221a3..7aee5cdf6 100644 --- a/packages/coding-agent/test/agent-session-role-thinking.test.ts +++ b/packages/coding-agent/test/agent-session-role-thinking.test.ts @@ -13,6 +13,7 @@ describe("AgentSession role model thinking behavior", () => { let tempDir: TempDir; let session: AgentSession; let sessionSettings: Settings; + const authStorages: AuthStorage[] = []; beforeEach(() => { tempDir = TempDir.createSync("@pi-role-thinking-"); @@ -22,6 +23,9 @@ describe("AgentSession role model thinking behavior", () => { if (session) { await session.dispose(); } + for (const authStorage of authStorages.splice(0)) { + authStorage.close(); + } tempDir.removeSync(); }); @@ -47,6 +51,7 @@ describe("AgentSession role model thinking behavior", () => { }, }); const authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); + authStorages.push(authStorage); authStorage.setRuntimeApiKey("anthropic", "test-key"); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir.path(), "models.yml")); @@ -181,6 +186,7 @@ describe("AgentSession role model thinking behavior", () => { }, }); const authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth-non-xhigh.db")); + authStorages.push(authStorage); authStorage.setRuntimeApiKey("anthropic", "test-key"); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir.path(), "models-non-xhigh.yml")); @@ -210,6 +216,7 @@ describe("AgentSession role model thinking behavior", () => { }, }); const authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth-cycle-thinking.db")); + authStorages.push(authStorage); authStorage.setRuntimeApiKey("anthropic", "test-key"); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir.path(), "models-cycle-thinking.yml")); diff --git a/packages/coding-agent/test/agent-session-user-shortcut-hooks.test.ts b/packages/coding-agent/test/agent-session-user-shortcut-hooks.test.ts index f20544555..025c844c5 100644 --- a/packages/coding-agent/test/agent-session-user-shortcut-hooks.test.ts +++ b/packages/coding-agent/test/agent-session-user-shortcut-hooks.test.ts @@ -16,10 +16,11 @@ describe("AgentSession user shortcut hooks", () => { let tempDir: TempDir; let session: AgentSession; let modelRegistry: ModelRegistry; + let authStorage: AuthStorage | undefined; beforeEach(async () => { tempDir = TempDir.createSync("@pi-user-shortcut-hooks-"); - const authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); + authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); modelRegistry = new ModelRegistry(authStorage); }); @@ -28,6 +29,8 @@ describe("AgentSession user shortcut hooks", () => { if (session) { await session.dispose(); } + authStorage?.close(); + authStorage = undefined; tempDir.removeSync(); }); diff --git a/packages/coding-agent/test/auth-storage-minimax-login.test.ts b/packages/coding-agent/test/auth-storage-minimax-login.test.ts index aeb2dc38b..0e40fd257 100644 --- a/packages/coding-agent/test/auth-storage-minimax-login.test.ts +++ b/packages/coding-agent/test/auth-storage-minimax-login.test.ts @@ -18,6 +18,7 @@ describe("AuthStorage MiniMax login", () => { afterEach(() => { vi.restoreAllMocks(); + authStorage.close(); if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } diff --git a/packages/coding-agent/test/auth-storage-rotation.test.ts b/packages/coding-agent/test/auth-storage-rotation.test.ts index 1185cc93b..19d1fdd89 100644 --- a/packages/coding-agent/test/auth-storage-rotation.test.ts +++ b/packages/coding-agent/test/auth-storage-rotation.test.ts @@ -53,6 +53,7 @@ describe("AuthStorage account rotation", () => { afterEach(() => { vi.restoreAllMocks(); + authStorage.close(); if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } diff --git a/packages/coding-agent/test/compaction-thinking-model.test.ts b/packages/coding-agent/test/compaction-thinking-model.test.ts index 990bce3e1..95dcd7929 100644 --- a/packages/coding-agent/test/compaction-thinking-model.test.ts +++ b/packages/coding-agent/test/compaction-thinking-model.test.ts @@ -30,6 +30,7 @@ const HAS_ANTHROPIC_AUTH = !!e2eApiKey("ANTHROPIC_API_KEY"); describe.skipIf(!HAS_ANTIGRAVITY_AUTH)("Compaction with thinking models (Antigravity)", () => { let session: AgentSession; let tempDir: string; + let authStorage: AuthStorage | undefined; beforeEach(() => { tempDir = path.join(os.tmpdir(), `pi-thinking-compaction-test-${Snowflake.next()}`); @@ -40,6 +41,8 @@ describe.skipIf(!HAS_ANTIGRAVITY_AUTH)("Compaction with thinking models (Antigra if (session) { await session.dispose(); } + authStorage?.close(); + authStorage = undefined; if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } @@ -76,7 +79,7 @@ describe.skipIf(!HAS_ANTIGRAVITY_AUTH)("Compaction with thinking models (Antigra const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); - const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); + authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); const modelRegistry = new ModelRegistry(authStorage); session = new AgentSession({ @@ -141,6 +144,7 @@ describe.skipIf(!HAS_ANTIGRAVITY_AUTH)("Compaction with thinking models (Antigra describe.skipIf(!HAS_ANTHROPIC_AUTH)("Compaction with thinking models (Anthropic)", () => { let session: AgentSession; let tempDir: string; + let authStorage: AuthStorage | undefined; beforeEach(() => { tempDir = path.join(os.tmpdir(), `pi-thinking-compaction-anthropic-test-${Snowflake.next()}`); @@ -151,6 +155,8 @@ describe.skipIf(!HAS_ANTHROPIC_AUTH)("Compaction with thinking models (Anthropic if (session) { await session.dispose(); } + authStorage?.close(); + authStorage = undefined; if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } @@ -179,7 +185,7 @@ describe.skipIf(!HAS_ANTHROPIC_AUTH)("Compaction with thinking models (Anthropic const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); - const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); + authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); const modelRegistry = new ModelRegistry(authStorage); session = new AgentSession({ diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index a3861f7ab..dd0b6d72d 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -1,4 +1,4 @@ -import { afterEach, beforeEach, describe, expect, it, mock, vi } from "bun:test"; +import { afterAll, afterEach, beforeEach, describe, expect, it, mock, vi } from "bun:test"; import * as path from "node:path"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; import { getBundledModel } from "@oh-my-pi/pi-ai/models"; @@ -118,6 +118,10 @@ beforeEach(() => { resetEntryCounter(); }); +afterAll(() => { + mock.restore(); +}); + afterEach(() => { completeSimpleMock.mockReset(); vi.restoreAllMocks(); diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index b90cd9523..5a5880afa 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -18,17 +18,19 @@ describe("ExtensionRunner", () => { let extensionsDir: string; let sessionManager: SessionManager; let modelRegistry: ModelRegistry; + let authStorage: AuthStorage; beforeEach(async () => { tempDir = TempDir.createSync("@pi-runner-test-"); extensionsDir = path.join(getProjectAgentDir(tempDir.path()), "extensions"); fs.mkdirSync(extensionsDir, { recursive: true }); sessionManager = SessionManager.inMemory(); - const authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); + authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); modelRegistry = new ModelRegistry(authStorage); }); afterEach(() => { + authStorage.close(); tempDir.removeSync(); }); diff --git a/packages/coding-agent/test/lm-studio-fix.test.ts b/packages/coding-agent/test/lm-studio-fix.test.ts index fd420c68a..e16c40432 100644 --- a/packages/coding-agent/test/lm-studio-fix.test.ts +++ b/packages/coding-agent/test/lm-studio-fix.test.ts @@ -19,6 +19,7 @@ describe("ModelRegistry LM Studio Fixes", () => { }); afterEach(() => { + authStorage.close(); if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } diff --git a/packages/coding-agent/test/model-registry-runtime-provider.test.ts b/packages/coding-agent/test/model-registry-runtime-provider.test.ts index dd9dcf09a..88e6d3c70 100644 --- a/packages/coding-agent/test/model-registry-runtime-provider.test.ts +++ b/packages/coding-agent/test/model-registry-runtime-provider.test.ts @@ -34,6 +34,7 @@ describe("ModelRegistry runtime provider registration", () => { for (const sourceId of sourceIds) { unregisterOAuthProviders(sourceId); } + authStorage.close(); if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true, force: true }); } diff --git a/packages/coding-agent/test/model-registry.test.ts b/packages/coding-agent/test/model-registry.test.ts index a71237aed..df4e80e4c 100644 --- a/packages/coding-agent/test/model-registry.test.ts +++ b/packages/coding-agent/test/model-registry.test.ts @@ -25,6 +25,7 @@ describe("ModelRegistry", () => { }); afterEach(() => { + authStorage.close(); if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); } diff --git a/packages/coding-agent/test/streaming-edit-abort.test.ts b/packages/coding-agent/test/streaming-edit-abort.test.ts index 55a09ad26..5c9d5bffa 100644 --- a/packages/coding-agent/test/streaming-edit-abort.test.ts +++ b/packages/coding-agent/test/streaming-edit-abort.test.ts @@ -79,7 +79,11 @@ function chunkStringRandomly(text: string, seed: number): string[] { return chunks; } -async function createSession(tempDir: string, streamFn: Agent["streamFn"], tool: AgentTool): Promise { +async function createSession( + tempDir: string, + streamFn: Agent["streamFn"], + tool: AgentTool, +): Promise<{ session: AgentSession; authStorage: AuthStorage }> { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const agent = new Agent({ getApiKey: () => "test-key", @@ -97,12 +101,15 @@ async function createSession(tempDir: string, streamFn: Agent["streamFn"], tool: authStorage.setRuntimeApiKey("anthropic", "test-key"); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); - return new AgentSession({ - agent, - sessionManager, - settings, - modelRegistry, - }); + return { + session: new AgentSession({ + agent, + sessionManager, + settings, + modelRegistry, + }), + authStorage, + }; } function buildEditTool(): AgentTool { @@ -215,14 +222,21 @@ describe("streaming edit abort", () => { const chunks = chunkStringRandomly(diff, seed); const abortSignalRef: { current?: AbortSignal } = {}; const streamFn = createStreamForDiff("sample.txt", chunks, abortSignalRef); - const session = await createSession(tempDir, streamFn, editTool); + const { session, authStorage } = await createSession(tempDir, streamFn, editTool); - await session.prompt("apply patch"); + try { + await session.prompt("apply patch"); - const lastAssistant = lastAssistantMessage(session.state.messages); - expect(lastAssistant?.stopReason).not.toBe("aborted"); - expect(abortSignalRef.current?.aborted ?? false).toBe(false); - await session.dispose(); + const lastAssistant = lastAssistantMessage(session.state.messages); + expect(lastAssistant?.stopReason).not.toBe("aborted"); + expect(abortSignalRef.current?.aborted ?? false).toBe(false); + } finally { + try { + await session.dispose(); + } finally { + authStorage.close(); + } + } } }); @@ -234,14 +248,21 @@ describe("streaming edit abort", () => { const chunks = chunkStringRandomly(diff, seed); const abortSignalRef: { current?: AbortSignal } = {}; const streamFn = createStreamForDiff("sample.txt", chunks, abortSignalRef); - const session = await createSession(tempDir, streamFn, editTool); + const { session, authStorage } = await createSession(tempDir, streamFn, editTool); - await session.prompt("apply patch"); + try { + await session.prompt("apply patch"); - const lastAssistant = lastAssistantMessage(session.state.messages); - expect(lastAssistant?.stopReason).toBe("aborted"); - expect(abortSignalRef.current?.aborted ?? false).toBe(true); - await session.dispose(); + const lastAssistant = lastAssistantMessage(session.state.messages); + expect(lastAssistant?.stopReason).toBe("aborted"); + expect(abortSignalRef.current?.aborted ?? false).toBe(true); + } finally { + try { + await session.dispose(); + } finally { + authStorage.close(); + } + } } }); }); diff --git a/packages/coding-agent/test/utilities.ts b/packages/coding-agent/test/utilities.ts index 0d11979bc..446731bc6 100644 --- a/packages/coding-agent/test/utilities.ts +++ b/packages/coding-agent/test/utilities.ts @@ -113,6 +113,7 @@ export async function createTestSession(options: TestSessionOptions = {}): Promi const cleanup = async () => { await session.dispose(); + authStorage.close(); if (tempDir && fs.existsSync(tempDir)) { fs.rmSync(tempDir, { recursive: true }); }