diff --git a/apps/vscode/test/kimi-harness.integration.test.ts b/apps/vscode/test/kimi-harness.integration.test.ts index fe8e7596c..6614606f4 100644 --- a/apps/vscode/test/kimi-harness.integration.test.ts +++ b/apps/vscode/test/kimi-harness.integration.test.ts @@ -127,7 +127,7 @@ async function createRuntimeRig(extraAliases: readonly string[] = []): Promise { async function createMcpHandlerRig(): Promise { const homeDir = await mkdtemp(join(tmpdir(), "kimi-vscode-mcp-handler-")); - cleanups.push(() => rm(homeDir, { recursive: true, force: true })); + cleanups.push(() => rm(homeDir, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 })); const harness = await createPlainHarness(homeDir); const broadcasts: BroadcastRecord[] = []; const logs: LogRecord[] = []; @@ -506,7 +506,7 @@ describe("VS Code Kimi harness integration (shares one in-process SDK home)", () it("keeps project-layer servers in the list refreshed after every mutation", async () => { const rig = await createMcpHandlerRig(); const project = await mkdtemp(join(tmpdir(), "kimi-vscode-mcp-project-")); - cleanups.push(() => rm(project, { recursive: true, force: true })); + cleanups.push(() => rm(project, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 })); await mkdir(join(project, ".git"), { recursive: true }); await writeFile( join(project, ".mcp.json"), diff --git a/packages/acp-server/src/start.ts b/packages/acp-server/src/start.ts index be3ecf93b..f7f66e117 100644 --- a/packages/acp-server/src/start.ts +++ b/packages/acp-server/src/start.ts @@ -16,6 +16,7 @@ import { Readable, Writable } from 'node:stream'; import { ndJsonStream, type AgentConnection, type Stream } from '@agentclientprotocol/sdk'; import { bootstrap, + drainLogCloses, drainQueryStoreDisposals, drainSessionIndexMirror, drainSessionMetadataWrites, @@ -186,8 +187,9 @@ export async function runAcpServerWithStream( // Flush the append-log write-behind before disposing, so a clean shutdown // never races a pending drain against teardown (and doesn't drop the last // persisted ops). Best-effort: a flush failure must not block disposal. + const appendLogStore = core.accessor.get(IAppendLogStore); try { - await core.accessor.get(IAppendLogStore).flush(); + await appendLogStore.flush(); } catch { // ignore — disposal proceeds regardless } @@ -201,10 +203,13 @@ export async function runAcpServerWithStream( // `core.dispose()` runs the mirror's and the query store's synchronous // `dispose()`, whose drains/closes are asynchronous — await them so an // embedding host that removes homeDir right after close() never races - // an in-flight shard close (ENOTEMPTY on teardown). + // an in-flight shard close (ENOTEMPTY on teardown). The same window + // exists for the append-log retirement flushes released by disposal. + await appendLogStore.drainRetirements(); await drainSessionIndexMirror(); await drainQueryStoreDisposals(); await drainSessionMetadataWrites(); + await drainLogCloses(); })(); return closePromise; }; diff --git a/packages/agent-core-v2/src/_base/log/logService.ts b/packages/agent-core-v2/src/_base/log/logService.ts index c5a0c5dc2..232a22075 100644 --- a/packages/agent-core-v2/src/_base/log/logService.ts +++ b/packages/agent-core-v2/src/_base/log/logService.ts @@ -16,6 +16,23 @@ import { import { createFileLogWriter, type FileLogWriter } from './fileLog'; import { ILogOptions } from './logConfig'; +const pendingLogCloses = new Set>(); + +export function trackLogClose(close: Promise): void { + const tracked = close.then( + () => undefined, + () => undefined, + ); + pendingLogCloses.add(tracked); + void tracked.finally(() => pendingLogCloses.delete(tracked)); +} + +export async function drainLogCloses(): Promise { + while (pendingLogCloses.size > 0) { + await Promise.all(pendingLogCloses); + } +} + interface ExtractedPayload { readonly ctx?: LogContext; readonly error?: LogEntryError; @@ -150,7 +167,7 @@ export class AppLogService extends BoundLogger implements ILogService { override dispose(): void { this.sink.flushSync(); - void this.sink.close(); + trackLogClose(this.sink.close()); super.dispose(); } } diff --git a/packages/agent-core-v2/src/persistence/backends/node-fs/appendLogStore.ts b/packages/agent-core-v2/src/persistence/backends/node-fs/appendLogStore.ts index 49bdb73e6..31554bcf6 100644 --- a/packages/agent-core-v2/src/persistence/backends/node-fs/appendLogStore.ts +++ b/packages/agent-core-v2/src/persistence/backends/node-fs/appendLogStore.ts @@ -11,6 +11,14 @@ import { const textEncoder = new TextEncoder(); +const pendingRetirements = new Set>(); + +export async function drainAppendLogRetirements(): Promise { + while (pendingRetirements.size > 0) { + await Promise.all(pendingRetirements); + } +} + interface LogState { pending: unknown[]; flushPromise: Promise | undefined; @@ -118,6 +126,10 @@ export class AppendLogStore implements IAppendLogStore { await this.flush(); } + drainRetirements(): Promise { + return drainAppendLogRetirements(); + } + acquire(scope: string, key: string): IDisposable { const state = this.state(scope, key); state.refCount++; @@ -171,7 +183,10 @@ export class AppendLogStore implements IAppendLogStore { state.refCount--; if (state.refCount > 0) return; state.retired = true; - state.retirement = this.settleRetiredState(scope, key, state).catch(() => undefined); + const retirement = this.settleRetiredState(scope, key, state).catch(() => undefined); + state.retirement = retirement; + pendingRetirements.add(retirement); + void retirement.finally(() => pendingRetirements.delete(retirement)); } private async settleRetiredState(scope: string, key: string, state: LogState): Promise { diff --git a/packages/agent-core-v2/src/persistence/interface/appendLogStore.ts b/packages/agent-core-v2/src/persistence/interface/appendLogStore.ts index 016b462df..c8f6b3fca 100644 --- a/packages/agent-core-v2/src/persistence/interface/appendLogStore.ts +++ b/packages/agent-core-v2/src/persistence/interface/appendLogStore.ts @@ -30,6 +30,7 @@ export interface IAppendLogStore { flush(): Promise; close(): Promise; acquire(scope: string, key: string): IDisposable; + drainRetirements(): Promise; } export const IAppendLogStore: ServiceIdentifier = diff --git a/packages/agent-core-v2/src/session/sessionLog/sessionLogService.ts b/packages/agent-core-v2/src/session/sessionLog/sessionLogService.ts index 279e9a3a3..ca3d97eb6 100644 --- a/packages/agent-core-v2/src/session/sessionLog/sessionLogService.ts +++ b/packages/agent-core-v2/src/session/sessionLog/sessionLogService.ts @@ -8,7 +8,7 @@ import { ISessionStateService } from '#/session/state/sessionState'; import { ILogService, type LogLevel } from '#/_base/log/log'; import { createFileLogWriter, type FileLogWriter } from '#/_base/log/fileLog'; import { ILogOptions, resolveSessionLogPath } from '#/_base/log/logConfig'; -import { BoundLogger, type LogLevelState } from '#/_base/log/logService'; +import { BoundLogger, trackLogClose, type LogLevelState } from '#/_base/log/logService'; export const sessionLogRootLevelKey = defineState('sessionLog.rootLevel', () => ({ level: 'info', @@ -61,7 +61,7 @@ export class SessionLogService extends BoundLogger implements ILogService { override dispose(): void { this.sink.flushSync(); - void this.sink.close(); + trackLogClose(this.sink.close()); super.dispose(); } } diff --git a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts index 1873ecce8..0e7f1bfa1 100644 --- a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts +++ b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts @@ -10,6 +10,7 @@ import { } from '#/_base/di/scope'; import { unwrapErrorCause } from '#/_base/errors/errors'; import { AsyncEmitter, Emitter, type Event, type IWaitUntil } from '#/_base/event'; +import { drainLogCloses } from '#/_base/log/logService'; import { DEFAULT_PLAN_MODE_SECTION } from '#/features/plan/configSection'; import { IAgentPlanService } from '#/features/plan/plan'; import { LifecycleScope } from '#/app/scopes'; @@ -379,9 +380,11 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec await this.announceWillClose({ sessionId, handle, reason: 'exit' }); this.sessions.delete(sessionId); await this.drainAgents(handle); + await this.appendLogStore.drainRetirements(); await drainSessionMetadataWrites(); await this.indexMirror.drain(); handle.dispose(); + await drainLogCloses(); this._onDidCloseSession.fire({ sessionId }); } @@ -391,6 +394,7 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec const meta = handle.accessor.get(ISessionMetadata); await meta.setArchived(true); await this.drainAgents(handle); + await this.appendLogStore.drainRetirements(); this.event.publish( new SessionArchived({ payload: { sessionId, workspaceId: this.workspaceContext.workspaceId }, @@ -401,6 +405,7 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec await drainSessionMetadataWrites(); await this.indexMirror.drain(); handle.dispose(); + await drainLogCloses(); this._onDidArchiveSession.fire({ sessionId }); } diff --git a/packages/agent-core-v2/test/app/git/gitService.test.ts b/packages/agent-core-v2/test/app/git/gitService.test.ts index 687ca8d65..bb7fc2182 100644 --- a/packages/agent-core-v2/test/app/git/gitService.test.ts +++ b/packages/agent-core-v2/test/app/git/gitService.test.ts @@ -85,7 +85,7 @@ describe('GitService', () => { expect(result.additions).toBe(0); expect(result.deletions).toBe(0); expect(result.pullRequest).toBeNull(); - }); + }, 15000); it('reports a modified file with numstat', async () => { writeFileSync(join(repo, 'a.txt'), 'line1\n'); diff --git a/packages/agent-core-v2/test/harness/agent.ts b/packages/agent-core-v2/test/harness/agent.ts index 3c795dc29..edfa3f662 100644 --- a/packages/agent-core-v2/test/harness/agent.ts +++ b/packages/agent-core-v2/test/harness/agent.ts @@ -1004,6 +1004,10 @@ class PersistenceAppendLogStore implements IAppendLogStore { return toDisposable(() => { }); } + drainRetirements(): Promise { + return Promise.resolve(); + } + snapshot(): WireRecord[] { return this.persistence.records.map(cloneRecord); } diff --git a/packages/agent-core-v2/test/persistence/backends/node-fs/appendLogStore.test.ts b/packages/agent-core-v2/test/persistence/backends/node-fs/appendLogStore.test.ts index 24e1e48ba..49eedc331 100644 --- a/packages/agent-core-v2/test/persistence/backends/node-fs/appendLogStore.test.ts +++ b/packages/agent-core-v2/test/persistence/backends/node-fs/appendLogStore.test.ts @@ -192,6 +192,41 @@ describe('AppendLogStore', () => { replacementOwner.dispose(); }); + it('final release retirement is awaited by drainRetirements', async () => { + let markAppendStarted!: () => void; + const appendStarted = new Promise((resolve) => { + markAppendStarted = resolve; + }); + let releaseAppend!: () => void; + const appendGate = new Promise((resolve) => { + releaseAppend = resolve; + }); + const originalAppend = storage.append.bind(storage); + storage.append = async (...args) => { + markAppendStarted(); + await appendGate; + return originalAppend(...args); + }; + + const owner = record.acquire(SCOPE, KEY); + record.append(SCOPE, KEY, { n: 1 }); + await appendStarted; + owner.dispose(); + + let drained = false; + const draining = record.drainRetirements().then(() => { + drained = true; + }); + await Promise.resolve(); + await Promise.resolve(); + expect(drained).toBe(false); + + releaseAppend(); + await draining; + expect(drained).toBe(true); + expect(await collect(SCOPE, KEY)).toEqual([{ n: 1 }]); + }); + it('keeps a sticky failure until every acquired owner releases it', async () => { const failure = new Error('shared append failed'); let reportFailure!: (error: unknown) => void; diff --git a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts index 84464e108..a70bffd52 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -155,6 +155,7 @@ function recordingAppendLog(initial: readonly WireRecord[] = []): { flush: () => Promise.resolve(), close: () => Promise.resolve(), acquire: () => ({ dispose: () => {} }), + drainRetirements: () => Promise.resolve(), }; return { appended, diff --git a/packages/agent-core-v2/test/session/sessionLog/sessionLogService.test.ts b/packages/agent-core-v2/test/session/sessionLog/sessionLogService.test.ts index 93ff7603b..55e90d0fd 100644 --- a/packages/agent-core-v2/test/session/sessionLog/sessionLogService.test.ts +++ b/packages/agent-core-v2/test/session/sessionLog/sessionLogService.test.ts @@ -10,13 +10,14 @@ import { registerScopedService, } from '#/_base/di/scope'; import { createScopedTestHost } from '#/_base/di/test'; +import type { FileLogWriter } from '#/_base/log/fileLog'; import { ILogService } from '#/_base/log/log'; import { logSeed, resolveLoggingConfig, resolveSessionLogPath, } from '#/_base/log/logConfig'; -import { AppLogService } from '#/_base/log/logService'; +import { AppLogService, drainLogCloses } from '#/_base/log/logService'; import { SessionLogService } from '#/session/sessionLog/sessionLogService'; import { makeSessionContext, sessionContextSeed } from '#/session/sessionContext/sessionContext'; import { ISessionStateService } from '#/session/state/sessionState'; @@ -140,6 +141,32 @@ describe('SessionLogService', () => { expect(text).toContain('on-dispose'); }); }); + + it('dispose tracks the sink close so drainLogCloses waits for it', async () => { + const host = buildHost(); + const session = host.child(LifecycleScope.Session, 's1', testSessionSeed()); + const log = session.accessor.get(ILogService) as SessionLogService; + const sink = (log as unknown as { sink: FileLogWriter }).sink; + const originalClose = sink.close.bind(sink); + let releaseClose!: () => void; + const closeGate = new Promise((resolve) => { + releaseClose = resolve; + }); + sink.close = () => originalClose().then(() => closeGate); + log.info('drain-me'); + host.dispose(); + + let drained = false; + const draining = drainLogCloses().then(() => { + drained = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(drained).toBe(false); + + releaseClose(); + await draining; + expect(drained).toBe(true); + }); }); describe('ILogService cross-scope resolution', () => { diff --git a/packages/agent-core-v2/test/wire/stubs.ts b/packages/agent-core-v2/test/wire/stubs.ts index c6fb70ece..5bc06cf9f 100644 --- a/packages/agent-core-v2/test/wire/stubs.ts +++ b/packages/agent-core-v2/test/wire/stubs.ts @@ -30,6 +30,7 @@ const noopLog: IAppendLogStore = { flush: async () => {}, close: async () => {}, acquire: () => toDisposable(() => {}), + drainRetirements: () => Promise.resolve(), }; const noopBlob: IAgentBlobService = { @@ -146,5 +147,6 @@ export function recordingWireLog( flush: async () => {}, close: async () => {}, acquire: () => toDisposable(() => {}), + drainRetirements: () => Promise.resolve(), }; } diff --git a/packages/agent-core/AGENTS.md b/packages/agent-core/AGENTS.md index 411873145..da6745d5d 100644 --- a/packages/agent-core/AGENTS.md +++ b/packages/agent-core/AGENTS.md @@ -9,4 +9,4 @@ - `src/mcp/registry.ts` (`McpServerRegistry`) is the single config view for MCP servers: `global` (layered mcp.json files) / `plugin` (manifests, read-only, final effective config via `PluginManager.mcpServerEntries`) / `caller` (SDK session injection). All management lookups in `src/rpc/core-impl.ts` go through it; mutations only accept mutable (user-level) entries and push changes into live sessions. - Live-session sync is one path: `KimiCore.reconcileMcpServerInSession` recomputes the registry runtime target (`resolveRuntimeTarget`: enabled plugin > project layer > user file; caller injection shadows everything and is never touched) per (session, name) and drives the session to it. Never add mutation-path-specific connect/remove logic — extend the reconciliation. - Wire-facing config DTOs are redacted: session `McpServerEntry`/`McpServerInfo.config` and read-only `McpManagedServerInfo` entries carry the `src/mcp/config-view.ts` projection (`envKeys`/`headerKeys` instead of literal `env`/`headers` values, which may hold credentials). Mutable user-level management entries keep full values for edit UIs. Core-internal code compares full configs via `McpConnectionManager.getRawEntry`. -- One process-wide `McpOAuthService` lives on `KimiCore` and is shared with every `Session`; each `Session` subscribes to its credential events (save/invalidate/refresh-failed) in its constructor and unsubscribes on close, so even sessions still initializing see every event. Never construct a per-scope OAuth service in new code. Token writes go through the process-local `OAuthTokenTransaction` (`@moonshot-ai/kimi-code-oauth`), which serializes refresh grants per credential identity and stamps `obtained_at` on every durable write. Interactive authorization flows are serialized per credential too: a second `beginAuthorization` for the same identity joins the in-flight flow instead of resetting the shared provider's PKCE/state. Proactive refresh timers and in-flight flows live and die with `McpOAuthService.shutdown()`, which `KimiCore.shutdown()` awaits. +- One process-wide `McpOAuthService` lives on `KimiCore` and is shared with every `Session`; each `Session` subscribes to its credential events (save/invalidate/refresh-failed) in its constructor and unsubscribes on close, so even sessions still initializing see every event. Never construct a per-scope OAuth service in new code. Token writes go through the process-local `OAuthTokenTransaction` (`@moonshot-ai/kimi-code-oauth`), which serializes refresh grants per credential identity and stamps `obtained_at` on every durable write. Interactive authorization flows are serialized per credential too: a second `beginAuthorization` for the same identity joins the in-flight flow instead of resetting the shared provider's PKCE/state. Proactive refresh timers, their in-flight refreshes, and interactive flows live and die with `McpOAuthService.shutdown()`, which `KimiCore.shutdown()` awaits. diff --git a/packages/agent-core/src/agent/background/index.ts b/packages/agent-core/src/agent/background/index.ts index bb36e504c..a4efb9b72 100644 --- a/packages/agent-core/src/agent/background/index.ts +++ b/packages/agent-core/src/agent/background/index.ts @@ -548,6 +548,19 @@ export class BackgroundManager { return results.filter((info): info is BackgroundTaskInfo => info !== undefined); } + /** + * Await every queued `output.log` append and persisted task-record write. + * Tasks that reached a terminal state are already drained by `finalizeTask`; + * this covers writes still queued on live (kept-alive) tasks when the owning + * session closes. + */ + async drainWrites(): Promise { + const entries = Array.from(this.tasks.values()); + await Promise.all( + entries.flatMap((entry) => [entry.outputWriteQueue, entry.persistWriteQueue]), + ); + } + /** * Wait for a task to reach a terminal state. * Returns immediately if already terminal. Times out after `timeoutMs`. @@ -1009,6 +1022,7 @@ export class BackgroundManager { entry.pendingOutput = []; entry.pendingOutputBytes = 0; } + await entry.outputWriteQueue; this.fireTerminalEffects(entry); entry.foregroundRelease?.resolve('terminal'); entry.terminal.resolve(); diff --git a/packages/agent-core/src/mcp/oauth/service.ts b/packages/agent-core/src/mcp/oauth/service.ts index 130ff19b0..92cb1033a 100644 --- a/packages/agent-core/src/mcp/oauth/service.ts +++ b/packages/agent-core/src/mcp/oauth/service.ts @@ -152,6 +152,9 @@ export class McpOAuthService { private readonly listeners = new Set(); private readonly refreshes = new Map>(); private readonly refreshTimers = new Map(); + /** In-flight timer-triggered proactive refreshes, awaited by {@link shutdown}. */ + private readonly pendingProactiveRefreshes = new Set>(); + private shutdownStarted = false; /** In-flight interactive flows by credential store key; values resolve to the shared flow. */ private readonly activeAuthorizations = new Map>(); @@ -261,11 +264,15 @@ export class McpOAuthService { /** * Release everything the service owns: pending proactive-refresh timers, - * in-flight interactive flows (closing their callback listeners), event - * listeners, and cached providers. Idempotent. + * in-flight proactive refreshes (awaited so their token writes and events + * land before listeners are dropped), in-flight interactive flows (closing + * their callback listeners), event listeners, and cached providers. + * Idempotent. */ async shutdown(): Promise { + this.shutdownStarted = true; this.stopProactiveRefresh(); + await Promise.all(this.pendingProactiveRefreshes); const inFlight = [...this.activeAuthorizations.values()]; this.activeAuthorizations.clear(); await Promise.all( @@ -531,6 +538,7 @@ export class McpOAuthService { } private scheduleRefresh(serverName: string, serverUrl: string | URL, expiresAt: number): void { + if (this.shutdownStarted) return; const canonicalUrl = canonicalMcpOAuthResource(serverUrl); const storeKey = mcpOAuthStoreKey(serverName, canonicalUrl); this.cancelScheduledRefresh(serverName, canonicalUrl); @@ -557,7 +565,7 @@ export class McpOAuthService { timer = setTimeout( () => { this.refreshTimers.delete(storeKey); - void this.refresh(serverName, canonicalUrl).catch((error: unknown) => { + const pending = this.refresh(serverName, canonicalUrl).catch((error: unknown) => { this.emit({ type: 'refresh-failed', serverName, @@ -565,6 +573,10 @@ export class McpOAuthService { error: error instanceof Error ? error.message : String(error), }); }); + this.pendingProactiveRefreshes.add(pending); + void pending.finally(() => { + this.pendingProactiveRefreshes.delete(pending); + }); }, Math.max(delay, 0), ); diff --git a/packages/agent-core/src/session/index.ts b/packages/agent-core/src/session/index.ts index 103b7e4f1..08a6667d4 100644 --- a/packages/agent-core/src/session/index.ts +++ b/packages/agent-core/src/session/index.ts @@ -564,6 +564,7 @@ export class Session { ); await this.cancelActiveTurnsOnClose(); await this.stopBackgroundTasksOnExit(); + await this.drainBackgroundTaskWrites(); await this.flushMetadata(); await this.triggerSessionEnd('exit'); } finally { @@ -661,6 +662,12 @@ export class Session { ); } + private async drainBackgroundTaskWrites(): Promise { + await Promise.all( + Array.from(this.readyAgents(), (agent) => agent.background.drainWrites()), + ); + } + /** * Wait for all still-running background tasks (across every agent) to reach a * terminal state before a `kimi -p` (print) run exits. diff --git a/packages/agent-core/test/agent/background/output-access.test.ts b/packages/agent-core/test/agent/background/output-access.test.ts index 16bbc0cde..379095cd7 100644 --- a/packages/agent-core/test/agent/background/output-access.test.ts +++ b/packages/agent-core/test/agent/background/output-access.test.ts @@ -2,7 +2,8 @@ * BackgroundManager output retrieval surface. */ -import { mkdtempSync, rmSync } from 'node:fs'; +import { mkdtempSync } from 'node:fs'; +import { rm } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { Readable } from 'node:stream'; import type { Writable } from 'node:stream'; @@ -43,8 +44,8 @@ describe('BackgroundManager — readOutput / getOutputSnapshot', () => { persistence = fixture.persistence!; }); - afterEach(() => { - rmSync(sessionDir, { recursive: true, force: true }); + afterEach(async () => { + await rm(sessionDir, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); }); it('getOutputSnapshot returns output.log path when persisted output exists', async () => { diff --git a/packages/agent-core/test/harness/skill-session.test.ts b/packages/agent-core/test/harness/skill-session.test.ts index 987380961..91e537f89 100644 --- a/packages/agent-core/test/harness/skill-session.test.ts +++ b/packages/agent-core/test/harness/skill-session.test.ts @@ -15,6 +15,7 @@ import { type SDKAPI, type TelemetryClient, } from '../../src'; +import { __resetRootLoggerForTest } from '../../src/logging/logger'; import { recordingContextTelemetry, type TelemetryContextRecord, @@ -38,7 +39,8 @@ describe('HarnessAPI session skills', () => { }); afterEach(async () => { - await rm(tmp, { recursive: true, force: true }); + await __resetRootLoggerForTest(); + await rm(tmp, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); vi.unstubAllEnvs(); }); diff --git a/packages/agent-core/test/mcp/oauth-service.test.ts b/packages/agent-core/test/mcp/oauth-service.test.ts index 810abdaa6..d09191aeb 100644 --- a/packages/agent-core/test/mcp/oauth-service.test.ts +++ b/packages/agent-core/test/mcp/oauth-service.test.ts @@ -316,6 +316,10 @@ describe('McpOAuthService single-flight refresh', () => { hasTokens: true, expired: false, }); + await waitFor( + () => fixture.events.filter((event) => event.type === 'tokens-saved').length === 2, + 'the refreshed tokens to be saved', + ); expect(fixture.events.filter((event) => event.type === 'tokens-saved')).toHaveLength(2); }, 15000); @@ -642,9 +646,49 @@ describe('McpOAuthService proactive refresh scheduling', () => { }); await waitFor(() => authServer.counts.refresh === 1, 'an immediate proactive refresh'); + await waitFor( + () => fixture.events.filter((event) => event.type === 'tokens-saved').length === 2, + 'the refreshed tokens to be saved', + ); expect(fixture.events.filter((event) => event.type === 'tokens-saved')).toHaveLength(2); }, 15000); + it('does not arm new proactive timers after shutdown', async () => { + const fixture = makeFixture(); + cleanups.push(() => rm(fixture.storeDir, { recursive: true, force: true })); + const authServer = await startFakeAuthServer(); + + const provider = fixture.service.getProvider(SERVER_NAME, SERVER_URL); + const state = authServerState(authServer.url); + provider.saveDiscoveryState(state.discovery); + provider.saveClientInformation(state.client); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + expires_in: 60, + }); + await waitFor(() => authServer.counts.refresh === 1, 'an immediate proactive refresh'); + await waitFor( + () => fixture.events.filter((event) => event.type === 'tokens-saved').length === 2, + 'the refreshed tokens to be saved', + ); + + await fixture.service.shutdown(); + const timers = ( + fixture.service as unknown as { refreshTimers: ReadonlyMap } + ).refreshTimers; + expect(timers.size).toBe(0); + + await provider.saveTokens({ + access_token: 'post-shutdown-token', + refresh_token: 'post-shutdown-refresh-token', + token_type: 'Bearer', + expires_in: 60, + }); + expect(timers.size).toBe(0); + }, 15000); + it('re-arms scheduling for expiries beyond the setTimeout limit', async () => { const fixture = makeFixture(); cleanups.push(() => rm(fixture.storeDir, { recursive: true, force: true })); diff --git a/packages/kap-server/src/start.ts b/packages/kap-server/src/start.ts index 79f956331..b5bb30e4a 100644 --- a/packages/kap-server/src/start.ts +++ b/packages/kap-server/src/start.ts @@ -3,8 +3,10 @@ import { drainQueryStoreDisposals, drainSessionMetadataWrites, drainSessionIndexMirror, + drainLogCloses, ConfigWarning, CapabilityChanged, + IAppendLogStore, IConfigService, IEventService, IOAuthService, @@ -357,11 +359,14 @@ export async function startServer(opts: ServerStartOptions): Promise {}); + void entry.handle + .close() + .catch(() => {}) + .finally(() => this.writerOps.leave()); }, this.opts.lockHoldMs); timer.unref(); } diff --git a/packages/minidb/test/cluster/lock.test.ts b/packages/minidb/test/cluster/lock.test.ts index ffd114702..6e54a9247 100644 --- a/packages/minidb/test/cluster/lock.test.ts +++ b/packages/minidb/test/cluster/lock.test.ts @@ -10,6 +10,7 @@ import fs from 'node:fs/promises'; import path from 'node:path'; import { ClusterDb } from '../../src/cluster/index.js'; import { ShardLockPool } from '../../src/cluster/lock-pool.js'; +import { ShardHandle } from '../../src/cluster/shard.js'; import { shardDirName } from '../../src/cluster/utils.js'; import { tmpDir, rmrf } from '../e2e/helpers/tmp.js'; import { keyOnShard, sleep } from './helpers.js'; @@ -244,3 +245,66 @@ test('closeAll() drains in-flight callbacks before closing handles — no MiniDb await rmrf(dir); } }); + +test('closeAll() waits for the lockHold timer’s in-flight writer close', async () => { + const dir = await tmpDir('minidb-cluster-'); + try { + const pool = new ShardLockPool({ + writerOpts: { valueCodec: 'json' }, + readerOpts: { valueCodec: 'json' }, + lockRenewMs: 0, + lockAcquireTimeoutMs: 1_000, + lockHoldMs: 80, + maxWriters: 4, + maxReaders: 4, + readOnly: false, + applyDefs: async () => {}, + }); + const shardDir = path.join(dir, shardDirName(1, 4)); + await pool.withWriter(1, shardDir, (db) => db.set('k', { v: 1 })); + assert.equal(pool.writersCached, 1); + + const closeStarted = deferred(); + const closeGate = deferred(); + let closeFinished = false; + const origClose = ShardHandle.prototype.close; + ShardHandle.prototype.close = async function (this: ShardHandle) { + closeStarted.resolve(); + await closeGate.promise; + await origClose.call(this); + closeFinished = true; + }; + try { + await closeStarted.promise; + assert.equal(pool.writersCached, 0, 'the hold timer already dropped the writer entry'); + + const closing = pool.closeAll(); + let closeReturned = false; + void closing.then(() => (closeReturned = true)); + for (let i = 0; i < 5; i++) await new Promise((r) => setImmediate(r)); + assert.equal(closeReturned, false, 'closeAll waits for the timer-fired close to drain'); + + closeGate.resolve(); + await closing; + assert.equal(closeFinished, true, 'the timer-fired close completed before closeAll returned'); + } finally { + ShardHandle.prototype.close = origClose; + } + + const pool2 = new ShardLockPool({ + writerOpts: { valueCodec: 'json' }, + readerOpts: { valueCodec: 'json' }, + lockRenewMs: 0, + lockAcquireTimeoutMs: 300, + lockHoldMs: 0, + maxWriters: 4, + maxReaders: 4, + readOnly: false, + applyDefs: async () => {}, + }); + assert.deepEqual(await pool2.withWriter(1, shardDir, (db) => db.get('k')), { v: 1 }, 'the shard lock was released'); + await pool2.closeAll(); + } finally { + await rmrf(dir); + } +}); diff --git a/packages/node-sdk/src/sdk-rpc-client-v2.ts b/packages/node-sdk/src/sdk-rpc-client-v2.ts index 773d6f4cc..39c410d8f 100644 --- a/packages/node-sdk/src/sdk-rpc-client-v2.ts +++ b/packages/node-sdk/src/sdk-rpc-client-v2.ts @@ -163,12 +163,14 @@ import { } from '@moonshot-ai/agent-core-v2/mcpCore/oauth/service'; import { createMcpOAuthStore } from '@moonshot-ai/agent-core-v2/app/mcpConfig/oauthStore'; import { canonicalMcpOAuthResource } from '@moonshot-ai/agent-core-v2/mcpCore/oauth/store'; +import { IAppendLogStore } from '@moonshot-ai/agent-core-v2/persistence/interface/appendLogStore'; import { IAtomicDocumentStore } from '@moonshot-ai/agent-core-v2/persistence/interface/atomicDocumentStore'; import { loadMcpServers } from '@moonshot-ai/agent-core-v2/workspace/workspaceMcpConfig/internal/config-loader'; import type { McpServerConfig as WorkspaceMcpServerConfig } from '@moonshot-ai/agent-core-v2/mcpCore/config-schema'; import { bootstrap, DEFAULT_AGENT_PROFILE_NAME, + drainLogCloses, drainQueryStoreDisposals, drainSessionIndexMirror, ensureKimiHome, @@ -549,9 +551,12 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { // disposal fires — a host that removes homeDir right after close() must // not race an in-flight shard close (ENOTEMPTY on teardown). await this.app.accessor.get(ISessionIndexMirror).drain(); + const appendLogStore = this.app.accessor.get(IAppendLogStore); this.app.dispose(); + await appendLogStore.drainRetirements(); await drainSessionIndexMirror(); await drainQueryStoreDisposals(); + await drainLogCloses(); } /** diff --git a/packages/node-sdk/test/sdk-rpc-client-v2.test.ts b/packages/node-sdk/test/sdk-rpc-client-v2.test.ts index 02e7271ab..433572b90 100644 --- a/packages/node-sdk/test/sdk-rpc-client-v2.test.ts +++ b/packages/node-sdk/test/sdk-rpc-client-v2.test.ts @@ -538,8 +538,14 @@ key = "${titleOAuthRef.key}" // The resumed session is a fresh, fully usable scope — not the handle // the temporary path just tore down. await client.renameSession({ id: 'ses_title_race', title: 'Resumed title' }); - const sessions = await client.listSessions({ workDir }); - expect(sessions.find((item) => item.id === 'ses_title_race')?.title).toBe('Resumed title'); + await expect + .poll( + async () => + (await client.listSessions({ workDir })).find((item) => item.id === 'ses_title_race') + ?.title, + { interval: 50, timeout: 4000 }, + ) + .toBe('Resumed title'); } finally { await client.close(); fetchSpy.mockRestore();