diff --git a/CHANGELOG.md b/CHANGELOG.md index 16befe2..0059e2a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Add optional platform-driver disposal and forward mobile disposal to lifecycle-capable device-mcp backends. + +### Fixed + +- Dispose session-owned drivers after consumer cleanup, including cleanup errors, forced relaunch, failed launch, daemon shutdown, and server startup rollback. Retain failed disposals in one server-owned lifecycle, so the next cleanup or daemon stop retries them. + ## [0.12.0] ### Added diff --git a/README.md b/README.md index bcde817..1228a1c 100644 --- a/README.md +++ b/README.md @@ -236,6 +236,16 @@ type ISessionManager = { }; ``` +### Session-scoped device teardown + +Core captures `sessionManager.getPlatformDriver?.()` before calling the consumer's `cleanup()`, then calls the driver's optional `dispose()` hook. Disposal runs even when consumer cleanup fails. This applies to `mm cleanup`, forced relaunch, failed launch, daemon shutdown, and server startup rollback. A failed disposal fails that cleanup call and is retained for a later attempt. The daemon owns that retention (one per `createServer` instance), so a later `mm cleanup`, forced relaunch, or daemon stop retries it, whatever wrapper of the session manager a tool receives. Only the current session's cleanup and disposal can fail a call. Drivers retained from earlier sessions are retried best-effort and in parallel. When they fail again, the daemon writes a `[WARN]` line to its log, and the current cleanup or launch continues. Tools reach this retention through `ToolContext.cleanupSession()`; without it (a context built outside the daemon), the tools fall back to retention per session-manager object. + +`MobilePlatformDriver.dispose()` delegates to its backend's optional `dispose()` hook. With a device-mcp version that implements private IDB companions, this stops only that backend's companion and recording. Older device-mcp versions remain compatible, but do not gain companion teardown. Adopt the device-mcp lifecycle release in the consumer after it is published; this change does not update the core dependency range. + +Consumers must expose a session-owned driver through `getPlatformDriver()` and create a fresh driver/backend for each launch. If launch fails before the driver is exposed, the consumer must dispose that locally held driver in its own failure path. Direct calls to a consumer's `cleanup()` also remain the consumer's responsibility. + +`mm cleanup --shutdown` additionally stops the worktree's daemon. Driver disposal does not shut down a simulator, stop Metro, or terminate shared IDB companions. Use separate worktrees/daemons and explicit device IDs for independent agent sessions. Separate companions isolate connection teardown even for the same simulator; app state on that simulator remains shared. + ### Workflow Context & Capabilities The `WorkflowContext` aggregates optional capabilities that consumers inject through the `contextFactory`. The tool system checks for capabilities at runtime — tools that depend on missing capabilities return clear errors. diff --git a/SKILL.md b/SKILL.md index f6be665..94b08ea 100644 --- a/SKILL.md +++ b/SKILL.md @@ -18,6 +18,8 @@ mm describe-screen # 4. Re-describe after every action to get fresh refs mm cleanup --shutdown # 5. Clean up when done ``` +Cleanup also disposes the session driver's owned resources after consumer app/session cleanup, including when cleanup fails. With a lifecycle-capable device-mcp backend, this releases its private companion and recording. `--shutdown` stops the worktree's daemon; it does not shut down a simulator or stop Metro. Independent agents should use separate worktrees/daemons and explicit device IDs. + **Critical rules:** - **Always `describe-screen` before interacting.** Refs like `e1`, `e2` are ephemeral — use the refs from the latest response (`describe-screen` or the `observations` of your last action). diff --git a/src/platform/mobile-platform-driver.dispose.test.ts b/src/platform/mobile-platform-driver.dispose.test.ts new file mode 100644 index 0000000..d7875ee --- /dev/null +++ b/src/platform/mobile-platform-driver.dispose.test.ts @@ -0,0 +1,46 @@ +import type { DeviceBackend } from '@metamask/device-mcp'; +import { describe, expect, it, vi } from 'vitest'; + +import { MobilePlatformDriver } from './mobile-platform-driver.js'; + +describe('MobilePlatformDriver disposal', () => { + it('disposes only its own backend once across concurrent and repeated calls', async () => { + const disposeA = vi.fn().mockResolvedValue(undefined); + const disposeB = vi.fn().mockResolvedValue(undefined); + const a = new MobilePlatformDriver({ + platform: 'ios', + dispose: disposeA, + } as unknown as DeviceBackend); + const b = new MobilePlatformDriver({ + platform: 'ios', + dispose: disposeB, + } as unknown as DeviceBackend); + await Promise.all([a.dispose(), a.dispose()]); + await a.dispose(); + expect(disposeA).toHaveBeenCalledTimes(1); + expect(disposeB).not.toHaveBeenCalled(); + await b.dispose(); + expect(disposeB).toHaveBeenCalledTimes(1); + }); + + it('supports device-mcp versions without disposal', async () => { + const driver = new MobilePlatformDriver({ + platform: 'ios', + } as DeviceBackend); + expect(await driver.dispose()).toBeUndefined(); + }); + + it('allows a retry after failed disposal', async () => { + const dispose = vi + .fn() + .mockRejectedValueOnce(new Error('failed')) + .mockResolvedValue(undefined); + const driver = new MobilePlatformDriver({ + platform: 'ios', + dispose, + } as unknown as DeviceBackend); + await expect(driver.dispose()).rejects.toThrowError('failed'); + await driver.dispose(); + expect(dispose).toHaveBeenCalledTimes(2); + }); +}); diff --git a/src/platform/mobile-platform-driver.ts b/src/platform/mobile-platform-driver.ts index e1180cf..7cf96a2 100644 --- a/src/platform/mobile-platform-driver.ts +++ b/src/platform/mobile-platform-driver.ts @@ -92,6 +92,8 @@ export class MobilePlatformDriver implements IPlatformDriver { #windowSize: Viewport | undefined; + #disposing: Promise | undefined; + /** * @param backend - The device-mcp backend to delegate to. * @param bundleId - The app bundle ID for getAppState calls. @@ -112,6 +114,20 @@ export class MobilePlatformDriver implements IPlatformDriver { this.#excludeDevOverlay = options.excludeDevOverlay ?? true; } + /** Release only this backend's resources; compatible with older device-mcp releases. */ + async dispose(): Promise { + const backend = this.#backend as DeviceBackend & { + dispose?(): Promise; + }; + this.#disposing ??= Promise.resolve() + .then(async () => backend.dispose?.()) + .catch((error: unknown) => { + this.#disposing = undefined; + throw error; + }); + await this.#disposing; + } + /** * @param targetType - Target identifier type. * @param targetValue - Target value for element lookup. diff --git a/src/platform/types.ts b/src/platform/types.ts index 0a4d837..02c808d 100644 --- a/src/platform/types.ts +++ b/src/platform/types.ts @@ -64,6 +64,9 @@ export type ScreenCapture = { }; export type IPlatformDriver = { + /** Release resources owned by this driver. Idempotent; does not shut down the device. */ + dispose?(): Promise; + /** * Waits for the target, then clicks it. When `signal` is aborted while the * driver waits for the target, the driver throws the abort reason and does diff --git a/src/server/cleanup-session.test.ts b/src/server/cleanup-session.test.ts new file mode 100644 index 0000000..a5904c5 --- /dev/null +++ b/src/server/cleanup-session.test.ts @@ -0,0 +1,287 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { cleanupSession, DriverLifecycle } from './cleanup-session.js'; +import type { IPlatformDriver } from '../platform/types.js'; +import { createMockSessionManager } from '../tools/test-utils/mock-factories.js'; + +describe('cleanupSession', () => { + it('captures the driver before consumer cleanup clears it and disposes afterward', async () => { + const dispose = vi.fn().mockResolvedValue(undefined); + let driver = { dispose } as unknown as IPlatformDriver | undefined; + const manager = createMockSessionManager(); + manager.getPlatformDriver = () => driver; + vi.spyOn(manager, 'cleanup').mockImplementation(async () => { + expect(dispose).not.toHaveBeenCalled(); + driver = undefined; + return true; + }); + + expect(await cleanupSession(manager)).toBe(true); + expect(dispose).toHaveBeenCalledTimes(1); + }); + + it('disposes the driver even when consumer cleanup fails', async () => { + const dispose = vi.fn().mockResolvedValue(undefined); + const manager = createMockSessionManager(); + manager.getPlatformDriver = () => + ({ dispose }) as unknown as IPlatformDriver; + vi.spyOn(manager, 'cleanup').mockRejectedValue( + new Error('app close failed'), + ); + + await expect(cleanupSession(manager)).rejects.toThrowError( + 'app close failed', + ); + expect(dispose).toHaveBeenCalledTimes(1); + }); + + it('retries failed disposal even after the consumer has cleared the driver', async () => { + const dispose = vi + .fn() + .mockRejectedValueOnce(new Error('dispose failed')) + .mockResolvedValue(undefined); + let driver = { dispose } as unknown as IPlatformDriver | undefined; + const manager = createMockSessionManager(); + manager.getPlatformDriver = () => driver; + vi.spyOn(manager, 'cleanup').mockImplementation(async () => { + driver = undefined; + return true; + }); + + await expect(cleanupSession(manager)).rejects.toThrowError( + 'dispose failed', + ); + expect(await cleanupSession(manager)).toBe(true); + expect(dispose).toHaveBeenCalledTimes(2); + }); + + it('does not dispose another session manager’s live driver through a shared lifecycle', async () => { + const lifecycle = new DriverLifecycle(); + const a = createMockSessionManager(); + const b = createMockSessionManager(); + const disposeA = vi.fn().mockResolvedValue(undefined); + const disposeB = vi.fn().mockResolvedValue(undefined); + a.getPlatformDriver = () => + ({ dispose: disposeA }) as unknown as IPlatformDriver; + b.getPlatformDriver = () => + ({ dispose: disposeB }) as unknown as IPlatformDriver; + + expect(await lifecycle.cleanupSession(a)).toBe(true); + expect(disposeA).toHaveBeenCalledTimes(1); + expect(disposeB).not.toHaveBeenCalled(); + expect(b.cleanup).not.toHaveBeenCalled(); + + expect(await lifecycle.cleanupSession(b)).toBe(true); + expect(disposeA).toHaveBeenCalledTimes(1); + expect(disposeB).toHaveBeenCalledTimes(1); + expect(lifecycle.pendingCount).toBe(0); + }); + + it('supports managers and drivers without disposal hooks', async () => { + const manager = createMockSessionManager(); + manager.getPlatformDriver = undefined; + expect(await cleanupSession(manager)).toBe(true); + }); + + it('reports both failures while still attempting every pending disposal', async () => { + const dispose = vi.fn().mockRejectedValue(new Error('dispose failed')); + const manager = createMockSessionManager(); + manager.getPlatformDriver = () => + ({ dispose }) as unknown as IPlatformDriver; + vi.spyOn(manager, 'cleanup').mockRejectedValue(new Error('session failed')); + await expect(cleanupSession(manager)).rejects.toMatchObject({ + message: + 'Session cleanup and driver disposal failed: session failed; dispose failed', + errors: [ + expect.objectContaining({ message: 'session failed' }), + expect.objectContaining({ message: 'dispose failed' }), + ], + }); + }); + + describe('drivers retained from earlier sessions', () => { + /** + * Build a manager whose driver is replaced by each test, and whose + * cleanup clears the driver like a real consumer does. + * + * @returns The manager and a setter for its current driver. + */ + function sessionWithSwappableDriver(): { + manager: ReturnType; + setDriver: (dispose: () => Promise) => void; + } { + let driver: IPlatformDriver | undefined; + const manager = createMockSessionManager(); + manager.getPlatformDriver = () => driver; + vi.spyOn(manager, 'cleanup').mockImplementation(async () => { + driver = undefined; + return true; + }); + return { + manager, + setDriver: (dispose) => { + driver = { dispose } as unknown as IPlatformDriver; + }, + }; + } + + it('retries a stuck older driver without failing later cleanups', async () => { + const onStaleDisposalError = vi.fn(); + const lifecycle = new DriverLifecycle({ onStaleDisposalError }); + const { manager, setDriver } = sessionWithSwappableDriver(); + const stuck = vi.fn().mockRejectedValue(new Error('stuck companion')); + setDriver(stuck); + + // The session that owns the stuck driver still reports its failure. + await expect(lifecycle.cleanupSession(manager)).rejects.toThrowError( + 'stuck companion', + ); + expect(lifecycle.pendingCount).toBe(1); + expect(onStaleDisposalError).not.toHaveBeenCalled(); + + for (let attempt = 1; attempt <= 2; attempt++) { + const healthy = vi.fn().mockResolvedValue(undefined); + setDriver(healthy); + expect(await lifecycle.cleanupSession(manager)).toBe(true); + expect(healthy).toHaveBeenCalledTimes(1); + expect(stuck).toHaveBeenCalledTimes(1 + attempt); + expect(onStaleDisposalError).toHaveBeenCalledTimes(attempt); + expect(lifecycle.pendingCount).toBe(1); + } + expect(onStaleDisposalError).toHaveBeenLastCalledWith( + expect.objectContaining({ message: 'stuck companion' }), + ); + }); + + it('drops an older driver from pending once its retry succeeds', async () => { + const lifecycle = new DriverLifecycle(); + const { manager, setDriver } = sessionWithSwappableDriver(); + const flaky = vi + .fn() + .mockRejectedValueOnce(new Error('stuck companion')) + .mockResolvedValue(undefined); + setDriver(flaky); + await expect(lifecycle.cleanupSession(manager)).rejects.toThrowError( + 'stuck companion', + ); + + setDriver(vi.fn().mockResolvedValue(undefined)); + expect(await lifecycle.cleanupSession(manager)).toBe(true); + expect(flaky).toHaveBeenCalledTimes(2); + expect(lifecycle.pendingCount).toBe(0); + + expect(await lifecycle.cleanupSession(manager)).toBe(true); + expect(flaky).toHaveBeenCalledTimes(2); + }); + + it('still fails the call for the current driver when an older driver is stuck', async () => { + const lifecycle = new DriverLifecycle(); + const { manager, setDriver } = sessionWithSwappableDriver(); + setDriver(vi.fn().mockRejectedValue(new Error('stuck companion'))); + await expect(lifecycle.cleanupSession(manager)).rejects.toThrowError( + 'stuck companion', + ); + + setDriver(vi.fn().mockRejectedValue(new Error('current dispose failed'))); + await expect(lifecycle.cleanupSession(manager)).rejects.toThrowError( + new Error('current dispose failed'), + ); + expect(lifecycle.pendingCount).toBe(2); + }); + + it('disposes a current driver that is also pending only once per call', async () => { + const lifecycle = new DriverLifecycle(); + const dispose = vi + .fn() + .mockRejectedValueOnce(new Error('dispose failed')) + .mockResolvedValue(undefined); + const driver = { dispose } as unknown as IPlatformDriver; + const manager = createMockSessionManager(); + // The consumer keeps exposing the same driver after a failed attempt. + manager.getPlatformDriver = () => driver; + + await expect(lifecycle.cleanupSession(manager)).rejects.toThrowError( + 'dispose failed', + ); + expect(await lifecycle.cleanupSession(manager)).toBe(true); + expect(dispose).toHaveBeenCalledTimes(2); + expect(lifecycle.pendingCount).toBe(0); + }); + + it('retries older drivers in parallel', async () => { + const lifecycle = new DriverLifecycle(); + const { manager, setDriver } = sessionWithSwappableDriver(); + const resolvers: (() => void)[] = []; + let deferRetries = false; + const olderDisposals = [1, 2].map(() => + vi.fn(async () => { + if (!deferRetries) { + throw new Error('stuck companion'); + } + await new Promise((resolve) => { + resolvers.push(resolve); + }); + }), + ); + for (const dispose of olderDisposals) { + setDriver(dispose); + await expect(lifecycle.cleanupSession(manager)).rejects.toThrowError( + 'stuck companion', + ); + } + expect(lifecycle.pendingCount).toBe(2); + olderDisposals.forEach((dispose) => dispose.mockClear()); + + deferRetries = true; + setDriver(vi.fn().mockResolvedValue(undefined)); + const pending = lifecycle.cleanupSession(manager); + await vi.waitFor(() => expect(resolvers).toHaveLength(2)); + // Both retries started before either one settled. + for (const dispose of olderDisposals) { + expect(dispose).toHaveBeenCalledTimes(1); + } + resolvers.forEach((resolve) => resolve()); + expect(await pending).toBe(true); + expect(lifecycle.pendingCount).toBe(0); + }); + }); + + it('retains failed disposals across per-request wrappers of one manager', async () => { + // The daemon hands each request a new wrapper (a Proxy on the stack) of + // the same session manager. Retention must not depend on that identity. + const lifecycle = new DriverLifecycle(); + const dispose = vi + .fn() + .mockRejectedValueOnce(new Error('dispose failed')) + .mockResolvedValue(undefined); + let driver = { dispose } as unknown as IPlatformDriver | undefined; + const manager = createMockSessionManager(); + manager.getPlatformDriver = () => driver; + vi.spyOn(manager, 'cleanup').mockImplementation(async () => { + driver = undefined; + return true; + }); + const wrap = (): typeof manager => new Proxy(manager, {}); + + await expect(lifecycle.cleanupSession(wrap())).rejects.toThrowError( + 'dispose failed', + ); + expect(lifecycle.pendingCount).toBe(1); + expect(await lifecycle.cleanupSession(wrap())).toBe(true); + expect(dispose).toHaveBeenCalledTimes(2); + expect(lifecycle.pendingCount).toBe(0); + }); + + it('uses the given lifecycle instead of the per-manager fallback', async () => { + const lifecycle = new DriverLifecycle(); + const dispose = vi.fn().mockRejectedValueOnce(new Error('dispose failed')); + const manager = createMockSessionManager(); + manager.getPlatformDriver = () => + ({ dispose }) as unknown as IPlatformDriver; + + await expect(cleanupSession(manager, lifecycle)).rejects.toThrowError( + 'dispose failed', + ); + expect(lifecycle.pendingCount).toBe(1); + }); +}); diff --git a/src/server/cleanup-session.ts b/src/server/cleanup-session.ts new file mode 100644 index 0000000..87b51b5 --- /dev/null +++ b/src/server/cleanup-session.ts @@ -0,0 +1,165 @@ +import type { ISessionManager } from './session-manager.js'; +import type { IPlatformDriver } from '../platform/types.js'; +import { debugWarn, extractErrorMessage } from '../utils'; + +/** + * Options for a {@link DriverLifecycle}. + */ +export type DriverLifecycleOptions = { + /** + * Called for each driver from an earlier session whose retried disposal + * fails again. These failures never fail the current call, so use this + * hook to make them visible (for example in the daemon log). + */ + onStaleDisposalError?: (error: unknown) => void; +}; + +/** + * Owns the drivers whose disposal failed, so a later cleanup can retry them. + * + * A consumer can clear its driver during cleanup before disposal fails; the + * driver is then reachable only from here. Retention is keyed by this + * instance, not by the session manager object, because callers can see the + * same manager through different wrappers (for example a per-request proxy). + * `createServer` owns one instance and uses it for the cleanup tool, forced + * and failed launches, daemon stop and startup rollback. + * + * Only the current session decides whether a call fails. Drivers retained + * from earlier sessions are retried best-effort on every call: their failures + * are reported to `onStaleDisposalError` and never fail the current cleanup + * or launch. + */ +export class DriverLifecycle { + readonly #pending = new Set(); + + readonly #onStaleDisposalError: ((error: unknown) => void) | undefined; + + /** + * @param options - Hooks for retries of drivers from earlier sessions. + */ + constructor(options: DriverLifecycleOptions = {}) { + this.#onStaleDisposalError = options.onStaleDisposalError; + } + + /** + * @returns How many drivers still wait for a successful disposal. + */ + get pendingCount(): number { + return this.#pending.size; + } + + /** + * Close the app/session before releasing the driver resources it may still + * use. Disposal also runs when consumer cleanup fails, including partial + * launches. Drivers retained from earlier sessions are retried in parallel, + * best-effort. + * + * @param manager - The session manager whose resources should be released. + * @returns Whether the consumer cleaned up a session. + * @throws The consumer cleanup error or the current driver's disposal + * error; an `AggregateError` when both fail. + */ + async cleanupSession(manager: ISessionManager): Promise { + // Read the driver first: consumers may clear it during cleanup. + const driver = manager.getPlatformDriver?.(); + const current = driver?.dispose ? driver : undefined; + if (current) { + // If it is already pending from an earlier failed attempt, it is + // treated as current and is not retried again as a stale driver. + this.#pending.add(current); + } + + const errors: unknown[] = []; + let cleanedUp = false; + try { + cleanedUp = await manager.cleanup(); + } catch (error) { + errors.push(error); + } + + const stale = [...this.#pending].filter((pending) => pending !== current); + const [currentFailure] = await Promise.all([ + current ? this.#dispose(current) : undefined, + this.#retryStale(stale), + ]); + if (currentFailure) { + errors.push(currentFailure.error); + } + + if (errors.length === 1) { + throw errors[0]; + } + if (errors.length > 1) { + throw new AggregateError( + errors, + `Session cleanup and driver disposal failed: ${errors + .map(extractErrorMessage) + .join('; ')}`, + ); + } + return cleanedUp; + } + + /** + * Dispose one driver and drop it from the pending set on success. + * + * @param driver - The driver to dispose. + * @returns The failure, or undefined when disposal succeeded. + */ + async #dispose( + driver: IPlatformDriver, + ): Promise<{ error: unknown } | undefined> { + try { + await driver.dispose?.(); + this.#pending.delete(driver); + return undefined; + } catch (error) { + return { error }; + } + } + + /** + * Retry drivers from earlier sessions in parallel. Failures are logged and + * reported, but never thrown. + * + * @param drivers - The retained drivers to retry. + */ + async #retryStale(drivers: IPlatformDriver[]): Promise { + const results = await Promise.allSettled( + drivers.map(async (driver) => driver.dispose?.()), + ); + results.forEach((result, index) => { + if (result.status === 'fulfilled') { + this.#pending.delete(drivers[index]); + return; + } + debugWarn('driver-lifecycle.retry', result.reason); + this.#onStaleDisposalError?.(result.reason); + }); + } +} + +// Fallback for callers without a server-owned lifecycle (for example a +// ToolContext built outside createServer). +const fallbackLifecycles = new WeakMap(); + +/** + * Clean up a session and dispose its driver. Prefer passing the server-owned + * {@link DriverLifecycle}; without one, failed disposals are retained per + * session manager object. + * + * @param manager - The session manager whose resources should be released. + * @param lifecycle - The lifecycle that retains failed disposals. + * @returns Whether the consumer cleaned up a session. + */ +export async function cleanupSession( + manager: ISessionManager, + lifecycle?: DriverLifecycle, +): Promise { + let owner = lifecycle ?? fallbackLifecycles.get(manager); + if (!owner) { + owner = new DriverLifecycle(); + fallbackLifecycles.set(manager, owner); + } + return owner.cleanupSession(manager); +} diff --git a/src/server/create-server.test.ts b/src/server/create-server.test.ts index 45aed5b..8a51306 100644 --- a/src/server/create-server.test.ts +++ b/src/server/create-server.test.ts @@ -511,6 +511,8 @@ describe('createServer integration', () => { await server.stop(); const mockSM = createMockSessionManager(); + const dispose = vi.fn().mockResolvedValue(undefined); + mockSM.getPlatformDriver.mockReturnValue({ dispose }); const customServer = createServer( buildConfig({ sessionManager: mockSM as unknown as ServerConfig['sessionManager'], @@ -524,6 +526,7 @@ describe('createServer integration', () => { try { await expect(customServer.start()).rejects.toThrowError(/EACCES/u); expect(mockSM.cleanup).toHaveBeenCalled(); + expect(dispose).toHaveBeenCalledTimes(1); } finally { await fs.chmod(tmpDir, 0o755).catch(() => {}); } @@ -622,6 +625,126 @@ describe('createServer integration', () => { expect(await server.stop()).toBeUndefined(); }); + it('disposes the session driver on daemon stop even when consumer cleanup fails', async () => { + await server.stop(); + const mockSM = createMockSessionManager(); + const dispose = vi.fn().mockResolvedValue(undefined); + mockSM.getPlatformDriver.mockReturnValue({ dispose }); + mockSM.cleanup.mockRejectedValue(new Error('app close failed')); + const customServer = createServer( + buildConfig({ + sessionManager: mockSM as unknown as ServerConfig['sessionManager'], + }), + ); + await customServer.start(); + await customServer.stop(); + expect(dispose).toHaveBeenCalledTimes(1); + expect(mockSM.cleanup.mock.invocationCallOrder[0]).toBeLessThan( + dispose.mock.invocationCallOrder[0], + ); + }); + + describe('driver disposal retries (server-owned)', () => { + function flakyDriverSession(): { + mockSM: ReturnType; + dispose: ReturnType; + } { + const mockSM = createMockSessionManager(); + const dispose = vi + .fn() + .mockRejectedValueOnce(new Error('companion stop timed out')) + .mockResolvedValue(undefined); + let driver: + | { dispose: typeof dispose; getPlatform: () => 'ios' } + | undefined = { dispose, getPlatform: () => 'ios' }; + mockSM.hasActiveSession.mockImplementation(() => driver !== undefined); + mockSM.getPlatformDriver.mockImplementation(() => driver); + mockSM.cleanup.mockImplementation(async () => { + // The consumer resets its session state, dropping the driver. + driver = undefined; + return true; + }); + return { mockSM, dispose }; + } + + async function postCleanup(port: number): Promise { + const res = await httpRequest(`http://127.0.0.1:${port}/tool/cleanup`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({}), + }); + await res.json(); + return res.status; + } + + it('retries a failed disposal on the next /tool/cleanup', async () => { + await server.stop(); + const { mockSM, dispose } = flakyDriverSession(); + const customServer = createServer( + buildConfig({ + sessionManager: mockSM as unknown as ServerConfig['sessionManager'], + }), + ); + const started = await customServer.start(); + try { + expect(await postCleanup(started.port)).toBe(500); + expect(dispose).toHaveBeenCalledTimes(1); + expect(await postCleanup(started.port)).toBe(200); + expect(dispose).toHaveBeenCalledTimes(2); + } finally { + await customServer.stop(); + } + expect(dispose).toHaveBeenCalledTimes(2); + }); + + it('retries a failed /tool/cleanup disposal on daemon stop', async () => { + await server.stop(); + const { mockSM, dispose } = flakyDriverSession(); + const customServer = createServer( + buildConfig({ + sessionManager: mockSM as unknown as ServerConfig['sessionManager'], + }), + ); + const started = await customServer.start(); + expect(await postCleanup(started.port)).toBe(500); + expect(dispose).toHaveBeenCalledTimes(1); + await customServer.stop(); + expect(dispose).toHaveBeenCalledTimes(2); + }); + + it('keeps /tool/cleanup healthy and logs on stop while an older driver stays stuck', async () => { + await server.stop(); + const { mockSM, dispose } = flakyDriverSession(); + dispose.mockReset(); + dispose.mockRejectedValue(new Error('stuck companion')); + const logPath = path.join(tmpDir, 'stuck-driver.log'); + const customServer = createServer( + buildConfig({ + sessionManager: mockSM as unknown as ServerConfig['sessionManager'], + logFilePath: logPath, + }), + ); + const started = await customServer.start(); + expect(await postCleanup(started.port)).toBe(500); + expect(await postCleanup(started.port)).toBe(200); + expect(dispose).toHaveBeenCalledTimes(2); + + await customServer.stop(); + expect(dispose).toHaveBeenCalledTimes(3); + await vi.waitFor(async () => { + const log = await fs.readFile(logPath, 'utf-8'); + expect(log).toContain( + '[WARN] Retained driver disposal failed again: stuck companion', + ); + expect(log).toContain('[INFO] Daemon stopped'); + }); + const log = await fs.readFile(logPath, 'utf-8'); + // One warning per retry: the second /tool/cleanup and stop(). + expect(log.match(/\[WARN\] Retained driver disposal/gu)).toHaveLength(2); + expect(log).not.toContain('[ERROR] Cleanup failed'); + }); + }); + describe('POST /tool/:name input validation', () => { it('returns 400 for missing required field', async () => { const res = await httpRequest( diff --git a/src/server/create-server.ts b/src/server/create-server.ts index 84ef8e8..71f412b 100644 --- a/src/server/create-server.ts +++ b/src/server/create-server.ts @@ -5,6 +5,7 @@ import * as fs from 'node:fs/promises'; import * as http from 'node:http'; import { setTimeout as delay } from 'node:timers/promises'; +import { DriverLifecycle } from './cleanup-session.js'; import { writeDaemonState, removeDaemonState } from './daemon-state.js'; import { compactObservation } from './observation-compaction.js'; import { RefStore } from './ref-store.js'; @@ -263,6 +264,17 @@ export function createServer(config: ServerConfig): ServerInstance { const queue = new RequestQueue(config.requestTimeoutMs); const refStore = new RefStore(config.sessionManager); const nonce = randomUUID(); + // One owner for driver disposal retries across tools, stop() and rollback. + // Keyed by this server, not by the session manager object a caller sees. + // Retries of drivers from earlier sessions never fail the current call, so + // log them here to keep repeated failures visible. + const driverLifecycle = new DriverLifecycle({ + onStaleDisposalError: (error): void => + appendLog( + config.logFilePath, + `[WARN] Retained driver disposal failed again: ${extractErrorMessage(error)}`, + ), + }); let httpServer: http.Server | null = null; let worktreeRoot = ''; @@ -359,6 +371,8 @@ export function createServer(config: ServerConfig): ServerInstance { toolRegistry, driver, signal, + cleanupSession: async () => + driverLifecycle.cleanupSession(config.sessionManager), }; } @@ -942,7 +956,7 @@ export function createServer(config: ServerConfig): ServerInstance { }); } try { - await config.sessionManager.cleanup(); + await driverLifecycle.cleanupSession(config.sessionManager); } catch { // Swallow — we're already propagating startupError. } @@ -994,7 +1008,7 @@ export function createServer(config: ServerConfig): ServerInstance { // 4. Clean up session try { - await config.sessionManager.cleanup(); + await driverLifecycle.cleanupSession(config.sessionManager); } catch (error) { appendLog( config.logFilePath, diff --git a/src/server/session-manager.ts b/src/server/session-manager.ts index b7fa4d4..ea04402 100644 --- a/src/server/session-manager.ts +++ b/src/server/session-manager.ts @@ -328,6 +328,11 @@ export type ISessionManager = { // Platform Driver (Optional — consumers that support mobile implement these) // ----------------------------------------------------------------------------- + /** + * Return the session-owned driver, including during partial launch when possible. + * Core captures it before cleanup and disposes its resources afterward. Consumers + * must dispose drivers that fail before they can be exposed here. + */ getPlatformDriver?(): IPlatformDriver | undefined; setPlatformDriver?(driver: IPlatformDriver): void; diff --git a/src/tools/cleanup.test.ts b/src/tools/cleanup.test.ts index f96f6da..985cb44 100644 --- a/src/tools/cleanup.test.ts +++ b/src/tools/cleanup.test.ts @@ -8,6 +8,7 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { cleanupTool } from './cleanup.js'; import { createMockSessionManager } from './test-utils/mock-factories.js'; +import { DriverLifecycle } from '../server/cleanup-session.js'; import type { ToolContext } from '../types/http.js'; function createMockContext(hasActive = false): ToolContext { @@ -70,4 +71,51 @@ describe('cleanupTool', () => { expect(context.sessionManager.cleanup).toHaveBeenCalledTimes(2); }); + + it('disposes the active platform driver through the cleanup tool', async () => { + const context = createMockContext(true); + const dispose = vi.fn().mockResolvedValue(undefined); + context.sessionManager.getPlatformDriver = () => + ({ dispose }) as unknown as NonNullable; + await cleanupTool({}, context); + expect(dispose).toHaveBeenCalledTimes(1); + }); + + it('uses the server-owned cleanup when the context provides one', async () => { + const context = createMockContext(true); + const serverCleanup = vi.fn().mockResolvedValue(true); + const result = await cleanupTool( + {}, + { ...context, cleanupSession: serverCleanup }, + ); + expect(serverCleanup).toHaveBeenCalledTimes(1); + expect(context.sessionManager.cleanup).not.toHaveBeenCalled(); + expect(result).toStrictEqual({ ok: true, result: { cleanedUp: true } }); + }); + + it('succeeds while a stuck driver from an earlier session keeps failing', async () => { + const context = createMockContext(true); + const lifecycle = new DriverLifecycle(); + const { sessionManager } = context; + const stuck = vi.fn().mockRejectedValue(new Error('stuck companion')); + sessionManager.getPlatformDriver = () => + ({ dispose: stuck }) as unknown as NonNullable; + const serverContext: ToolContext = { + ...context, + cleanupSession: async () => lifecycle.cleanupSession(sessionManager), + }; + await expect(cleanupTool({}, serverContext)).rejects.toThrowError( + 'stuck companion', + ); + + const healthy = vi.fn().mockResolvedValue(undefined); + sessionManager.getPlatformDriver = () => + ({ dispose: healthy }) as unknown as NonNullable; + const result = await cleanupTool({}, serverContext); + + expect(result).toStrictEqual({ ok: true, result: { cleanedUp: true } }); + expect(healthy).toHaveBeenCalledTimes(1); + expect(stuck).toHaveBeenCalledTimes(2); + expect(lifecycle.pendingCount).toBe(1); + }); }); diff --git a/src/tools/cleanup.ts b/src/tools/cleanup.ts index cae4b36..8a43b4e 100644 --- a/src/tools/cleanup.ts +++ b/src/tools/cleanup.ts @@ -1,5 +1,6 @@ import type { CleanupInput, CleanupResult } from './types'; import { createToolSuccess } from './utils.js'; +import { cleanupSession } from '../server/cleanup-session.js'; import type { ToolContext, ToolResponse } from '../types/http.js'; /** @@ -13,7 +14,8 @@ export async function cleanupTool( _input: CleanupInput, context: ToolContext, ): Promise> { - const cleanedUp = await context.sessionManager.cleanup(); + const cleanedUp = await (context.cleanupSession?.() ?? + cleanupSession(context.sessionManager)); return createToolSuccess({ cleanedUp }); } diff --git a/src/tools/launch.test.ts b/src/tools/launch.test.ts index 4774a05..b56e84b 100644 --- a/src/tools/launch.test.ts +++ b/src/tools/launch.test.ts @@ -9,6 +9,7 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { launchTool } from './launch.js'; import type { LaunchInput } from './types'; import type { ExtensionState } from '../capabilities/types.js'; +import { DriverLifecycle } from '../server/cleanup-session.js'; import type { SessionLaunchResult } from '../server/session-manager.js'; import { createMockSessionManager } from './test-utils/mock-factories.js'; import { ErrorCodes } from './types/errors.js'; @@ -211,6 +212,9 @@ describe('launchTool', () => { it('cleans up and relaunches when force is true', async () => { const context = createMockContext({ hasActive: true }); vi.spyOn(context.sessionManager, 'cleanup').mockResolvedValue(true); + const dispose = vi.fn().mockResolvedValue(undefined); + context.sessionManager.getPlatformDriver = () => + ({ dispose }) as unknown as NonNullable; const result = await launchTool( { stateMode: 'default', force: true }, @@ -219,11 +223,91 @@ describe('launchTool', () => { expect(result.ok).toBe(true); expect(context.sessionManager.cleanup).toHaveBeenCalled(); + expect(dispose).toHaveBeenCalledTimes(1); + expect(dispose.mock.invocationCallOrder[0]).toBeLessThan( + vi.mocked(context.sessionManager.launch).mock.invocationCallOrder[0], + ); expect(context.sessionManager.launch).toHaveBeenCalled(); }); + + it('relaunches with force while a stuck driver from an earlier session keeps failing', async () => { + const context = createMockContext({ hasActive: true }); + const lifecycle = new DriverLifecycle(); + const { sessionManager } = context; + const stuck = vi.fn().mockRejectedValue(new Error('stuck companion')); + sessionManager.getPlatformDriver = () => + ({ dispose: stuck }) as unknown as NonNullable; + const serverContext: ToolContext = { + ...context, + cleanupSession: async () => lifecycle.cleanupSession(sessionManager), + }; + await expect( + lifecycle.cleanupSession(sessionManager), + ).rejects.toThrowError('stuck companion'); + + const healthy = vi.fn().mockResolvedValue(undefined); + sessionManager.getPlatformDriver = () => + ({ dispose: healthy }) as unknown as NonNullable; + const result = await launchTool( + { stateMode: 'default', force: true }, + serverContext, + ); + + expect(result.ok).toBe(true); + expect(sessionManager.launch).toHaveBeenCalledTimes(1); + expect(healthy).toHaveBeenCalledTimes(1); + expect(stuck).toHaveBeenCalledTimes(2); + }); }); describe('launch failures', () => { + it('disposes an exposed partial-launch driver even if consumer cleanup fails', async () => { + const context = createMockContext({ hasActive: false }); + const dispose = vi.fn().mockResolvedValue(undefined); + context.sessionManager.getPlatformDriver = () => + ({ dispose }) as unknown as NonNullable; + vi.spyOn(context.sessionManager, 'launch').mockRejectedValue( + Object.assign(new Error('device failed'), { + code: ErrorCodes.MM_DEPENDENCIES_MISSING, + }), + ); + vi.spyOn(context.sessionManager, 'cleanup').mockRejectedValue( + new Error('cleanup failed'), + ); + + const result = await launchTool({ stateMode: 'default' }, context); + expect(dispose).toHaveBeenCalledTimes(1); + expect(result).toMatchObject({ + ok: false, + error: { + code: ErrorCodes.MM_DEPENDENCIES_MISSING, + message: 'device failed', + }, + }); + }); + + it('uses the server-owned cleanup for forced and failed launches', async () => { + const serverCleanup = vi.fn().mockResolvedValue(true); + const forced = createMockContext({ hasActive: true }); + await launchTool( + { stateMode: 'default', force: true }, + { ...forced, cleanupSession: serverCleanup }, + ); + expect(serverCleanup).toHaveBeenCalledTimes(1); + expect(forced.sessionManager.cleanup).not.toHaveBeenCalled(); + + const failed = createMockContext({ hasActive: false }); + vi.spyOn(failed.sessionManager, 'launch').mockRejectedValue( + new Error('device failed'), + ); + const result = await launchTool( + { stateMode: 'default' }, + { ...failed, cleanupSession: serverCleanup }, + ); + expect(result.ok).toBe(false); + expect(serverCleanup).toHaveBeenCalledTimes(2); + }); + it('returns port conflict error for EADDRINUSE', async () => { const context = createMockContext({ hasActive: false }); vi.spyOn(context.sessionManager, 'launch').mockRejectedValue( diff --git a/src/tools/launch.ts b/src/tools/launch.ts index b8b8e04..7c0e2ed 100644 --- a/src/tools/launch.ts +++ b/src/tools/launch.ts @@ -2,8 +2,9 @@ import type { LaunchInput, LaunchPrerequisite, LaunchResult } from './types'; import { ErrorCodes } from './types'; import type { ErrorCode } from './types/errors.js'; import { createToolError, createToolSuccess } from './utils.js'; +import { cleanupSession } from '../server/cleanup-session.js'; import type { ToolContext, ToolResponse } from '../types/http.js'; -import { extractErrorMessage } from '../utils'; +import { debugWarn, extractErrorMessage } from '../utils'; const PROD_MODE_PREREQUISITES: LaunchPrerequisite[] = [ { @@ -35,11 +36,13 @@ export async function launchTool( context: ToolContext, ): Promise> { const { sessionManager } = context; + const cleanup = async (): Promise => + context.cleanupSession?.() ?? cleanupSession(sessionManager); try { if (sessionManager.hasActiveSession()) { if (input.force) { - await sessionManager.cleanup(); + await cleanup(); } else { return createToolError( ErrorCodes.MM_SESSION_ALREADY_RUNNING, @@ -52,7 +55,18 @@ export async function launchTool( sessionManager.setContext(input.context); } - const result = await sessionManager.launch(input); + const result = await sessionManager + .launch(input) + .catch(async (error: unknown) => { + // Preserve the launch error after best-effort cleanup of a partial session. + try { + await cleanup(); + } catch (rollbackError) { + // The original launch error is the actionable failure. + debugWarn('launch.rollback', rollbackError); + } + throw error; + }); const isProdMode = sessionManager.getEnvironmentMode() === 'prod'; return createToolSuccess({ diff --git a/src/tools/test-utils/mock-factories.ts b/src/tools/test-utils/mock-factories.ts index 5585172..1ea1520 100644 --- a/src/tools/test-utils/mock-factories.ts +++ b/src/tools/test-utils/mock-factories.ts @@ -11,6 +11,7 @@ import { vi } from 'vitest'; import type { ExtensionState } from '../../capabilities/types.js'; +import type { IPlatformDriver } from '../../platform/types.js'; import type { TrackedPage, SessionLaunchResult, @@ -120,6 +121,13 @@ export function createMockSessionManager( getContractSeedingCapability: vi.fn().mockReturnValue(undefined), getStateSnapshotCapability: vi.fn().mockReturnValue(undefined), + // Platform Driver + getPlatformDriver: vi + .fn<() => IPlatformDriver | undefined>() + .mockReturnValue(undefined) as + | (() => IPlatformDriver | undefined) + | undefined, + // Environment setWorkflowContext: vi.fn(), getEnvironmentMode: vi diff --git a/src/types/http.ts b/src/types/http.ts index 5322fd4..7c36564 100644 --- a/src/types/http.ts +++ b/src/types/http.ts @@ -34,6 +34,12 @@ export type ToolContext = { * stop. Absent when the context was not built by the daemon. */ signal?: AbortSignal; + /** + * Cleans up the session and disposes its driver through the server-owned + * lifecycle, so failed disposals are retried by later cleanups and by + * daemon stop. Absent when the context was not built by the daemon. + */ + cleanupSession?: () => Promise; }; /** diff --git a/vitest.config.mts b/vitest.config.mts index b7cf922..01f8640 100644 --- a/vitest.config.mts +++ b/vitest.config.mts @@ -35,10 +35,10 @@ export default defineConfig({ // Auto-update the coverage thresholds when running locally. // Disabled in CI to prevent non-deterministic config changes. autoUpdate: !process.env.CI, - branches: 91.16, - functions: 93.64, - lines: 96.35, - statements: 96.11, + branches: 94.91, + functions: 96.73, + lines: 97.68, + statements: 97.53, }, },