Skip to content
15 changes: 15 additions & 0 deletions apps/desktop/src/main/__tests__/main-startup-lifetime.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,21 @@ test('does not release renderer IPC before persistent handlers are registered',
assert.doesNotMatch(readyBody, /runtimeHostManager\?\.start\(\)/u);
});

test('wires Artifact deletion and purge events to preview invalidation for the active scope', () => {
const subscribe = bootSource.indexOf('client.subscribeArtifactChanges(');
const unsubscribe = bootSource.indexOf('unsubscribeArtifactChanges();', subscribe);
assert.ok(subscribe >= 0 && unsubscribe > subscribe);
const handler = bootSource.slice(subscribe, unsubscribe);
assert.match(handler, /managedArtifactPreview\.revoke\(scope\.targetEpoch,\s*frame\.sessionId,\s*frame\.artifactId\)/u);
assert.match(handler, /managedArtifactPreview\.releaseSession\(scope\.targetEpoch,\s*frame\.sessionId\)/u);
assert.match(bootSource, /managedArtifactPreview\.openScope\(scope\.targetEpoch\)/u);
assert.match(bootSource, /await managedArtifactPreview\.closeScope\(scope\.targetEpoch\)/u);
assert.ok(
bootSource.indexOf('await managedArtifactPreview.closeScope(scope.targetEpoch)') <
bootSource.indexOf('unsubscribeArtifactChanges();'),
);
});

test('resolves persisted locale before first post-settings recovery prompt', () => {
const rendererRecoveryStart = earlyWindowSource.indexOf('onRendererProcessGone: async');
const rendererRecovery = earlyWindowSource.slice(
Expand Down
109 changes: 105 additions & 4 deletions apps/desktop/src/main/__tests__/managed-artifact-preview.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,8 @@ test('isolates leases by origin and rejects credentials for another preview', as
await assert.rejects(fetch(first.url));
assert.equal(await (await fetch(second.url)).text(), 'second');
await assert.rejects(service.prepare('host1', client(), 's1', 'a1'), /closed/);
service.openScope('host1');
assert.equal(await (await fetch((await service.prepare('host1', client(), 's1', 'a1')).url)).text(), html);
await service.revoke('host2', 's1', 'a1');
await assert.rejects(fetch(second.url));
} finally { await service.close(); }
Expand Down Expand Up @@ -136,18 +138,117 @@ test('close and cancellation during a stream cannot publish a live endpoint', as
}
});

test('bounds concurrent preparations before allocating buffers or ports', async () => {
test('bounds previews per session instead of starving another session', async () => {
const service = new ManagedArtifactPreview();
let resume!: () => void;
const gate = new Promise<void>((resolve) => { resume = resolve; });
const source = client();
const slow = { ...source, getArtifact: async (s: string, a: string) => { await gate; return source.getArtifact(s, a); } };
const slow = { ...source, getArtifact: async () => { await gate; return source.getArtifact('s1', 'a1'); } };
const pending = Array.from({ length: 16 }, () => service.prepare('h', slow, 's1', 'a1'));
try {
await assert.rejects(service.prepare('h', slow, 's1', 'a1'), /Too many/);
const other = service.prepare('h', client('other'), 's2', 'a1');
resume();
assert.equal((await Promise.all(pending)).length, 16);
} finally { resume(); await Promise.allSettled(pending); await service.close(); }
assert.equal((await other).reachable, true);
await Promise.all(pending);
} finally {
resume();
await Promise.allSettled(pending);
await service.close();
}
});

test('rejects previews above the global limit without invalidating existing sessions', async () => {
const service = new ManagedArtifactPreview();
try {
const endpoints = [];
for (let index = 0; index < 64; index += 1) {
endpoints.push(await service.prepare('h', client(`preview-${index}`), `s${index}`, 'a1'));
}
await assert.rejects(service.prepare('h', client('replacement'), 's64', 'a1'), /across the Desktop/);
assert.equal(await (await fetch(endpoints[0]!.url)).text(), 'preview-0');

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Keep an aggregate bound on retained previews. This new assertion explicitly keeps the first lease alive after creating a 65th session-scoped preview, while ManagedArtifactPreview.prepare() now checks only 16 leases per session (managed-artifact-preview.ts:75-84). Each lease retains up to 8 MiB in its HTTP handler (:25,104-124) for 30 minutes (:27,160-163), and creates a separate listening server. Five sessions can now retain 80 maximum-sized previews (640 MiB); more sessions have no process-wide bound. The prior 64-lease eviction was imperfect for UX, but removing it without a replacement lets repeated preview creation exhaust Desktop memory/sockets. Please keep a global resource budget (with an admission/eviction policy that does not silently starve other sessions) and test that budget.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for catching this. You’re right that the per-Session limit alone left aggregate preview resources unbounded.

I’ve added a global limit of 64 active previews and a 128 MiB aggregate content budget. In-flight preparations reserve capacity too; requests that exceed either limit are rejected without evicting or disrupting existing previews. The regression tests cover both limits and verify that failed preparations release their reservations.

The fix is pushed in f4bb8bc71. Thanks for the careful review!

assert.equal(await (await fetch(endpoints[63]!.url)).text(), 'preview-63');
await service.releaseUrl(endpoints[0]!.url);
const replacement = await service.prepare('h', client('replacement'), 's64', 'a1');
assert.equal(await (await fetch(replacement.url)).text(), 'replacement');
} finally {
await service.close();
}
});

test('reserves aggregate preview bytes before streaming and releases reservations on failure', async () => {
const service = new ManagedArtifactPreview();
let resume!: () => void;
const gate = new Promise<void>((resolve) => { resume = resolve; });
const large = {
getArtifact: async () => ({ id: 'a1', sessionId: 's1', turnId: 't1', createdAt: 0, name: 'large.html', kind: 'html' as const, sizeBytes: PREVIEW_MAX_BYTES, source: 'tool_result' as const }),
streamArtifact: async () => { await gate; throw new Error('test stream failure'); },
};
const successfulLarge = {
getArtifact: large.getArtifact,
streamArtifact: async (_sessionId: string, _artifactId: string, write: (chunk: Uint8Array) => Promise<void>) => {
const bytes = Buffer.alloc(PREVIEW_MAX_BYTES);
await write(bytes);
return bytes.length;
},
};
const pending = Array.from({ length: 16 }, (_, index) => service.prepare('h', large, `s${index + 1}`, 'a1'));
try {
// Each in-flight Artifact reserves 8 MiB, reaching the 128 MiB aggregate budget.
await new Promise((resolve) => setImmediate(resolve));
await assert.rejects(service.prepare('h', large, 's17', 'a1'), /memory limit/);
resume();
await Promise.all(pending.map((preparation) => assert.rejects(preparation, /test stream failure/)));
const recovered = await service.prepare('h', successfulLarge, 's1', 'a1');
assert.equal(recovered.reachable, true);
await service.releaseUrl(recovered.url);
} finally {
resume();
await Promise.allSettled(pending);
await service.close();
}
});

test('releases every preview for a purged session', async () => {
const service = new ManagedArtifactPreview();
try {
const first = await service.prepare('h', client(), 's1', 'a1');
const second = await service.prepare('h', client('second'), 's1', 'a2');
const otherScope = await service.prepare('other-host', client('other scope'), 's1', 'a1');
await service.releaseSession('h', 's1');
await assert.rejects(fetch(first.url));
await assert.rejects(fetch(second.url));
assert.equal(await (await fetch(otherScope.url)).text(), 'other scope');
assert.equal((await service.prepare('h', client(), 's1', 'a1')).reachable, true);
} finally { await service.close(); }
});

test('delete and Session purge cancel previews that are still preparing', async () => {
for (const release of [
(service: ManagedArtifactPreview) => service.revoke('h', 's1', 'a1'),
(service: ManagedArtifactPreview) => service.releaseSession('h', 's1'),
]) {
const service = new ManagedArtifactPreview();
let resume!: () => void;
const gate = new Promise<void>((resolve) => { resume = resolve; });
const source = client();
const preparing = service.prepare('h', {
...source,
getArtifact: async (...args) => {
await gate;
return source.getArtifact(...args);
},
}, 's1', 'a1');
try {
await release(service);
resume();
await assert.rejects(preparing, /closed/);
} finally {
resume();
await Promise.allSettled([preparing]);
await service.close();
}
}
});

test('tool binds to the admitted session and returns endpoint evidence only', async () => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ import { RuntimeHostReconnectingIpcMain } from '../runtime-host-reconnecting-ipc
import { desktopSessionResourceKey } from '../../shared/runtime-host-identity.js';
import { waitFor as pollFor } from '@maka/core/test-only/async-primitives';
import { canRepairManagedRuntimeHostStartup } from '../runtime-host-startup-recovery.js';
import { ManagedArtifactPreview } from '../managed-artifact-preview.js';

const TEST_HOST_ID = 'a'.repeat(64);
const TEST_TARGET_EPOCH = 'test-target-epoch';
Expand Down Expand Up @@ -544,6 +545,67 @@ test('tears down the whole candidate when the Host connection closes', async ()
assert.equal(host.closeCalls, 1);
});

test('closes old managed Artifact previews and reopens the scope after reconnect', async () => {
const ipc = ipcHarness();
const firstHost = connectionHarness('preview-first');
const preview = new ManagedArtifactPreview();
const bytes = Buffer.from('<!doctype html><title>Preview</title>');
const candidateDeps = {
...deps(ipc),
registerClientIpc: (_client, _ipc, _controls, _target, scope) => {
preview.openScope(scope.targetEpoch);
return () => preview.closeScope(scope.targetEpoch);
},
} satisfies DesktopRuntimeHostCandidateDeps;
const firstCandidate = await createDesktopRuntimeHostCandidate(firstHost.connection, candidateDeps);
let secondCandidate: Awaited<ReturnType<typeof createDesktopRuntimeHostCandidate>> | undefined;
const source = {
getArtifact: async () => ({
id: 'artifact-1',
sessionId: 'session-1',
turnId: 'turn-1',
createdAt: 0,
name: 'preview.html',
kind: 'html' as const,
sizeBytes: bytes.length,
source: 'tool_result' as const,
}),
streamArtifact: async (_sessionId: string, _artifactId: string, write: (chunk: Uint8Array) => Promise<void>) => {
await write(bytes);
return bytes.length;
},
};

try {
const endpoint = await preview.prepare(
TEST_TARGET_EPOCH,
source,
'session-1',
'artifact-1',
);
assert.equal(await (await fetch(endpoint.url)).text(), bytes.toString());

firstHost.disconnect();
await firstCandidate.closed;

await assert.rejects(fetch(endpoint.url));

const secondHost = connectionHarness('preview-second');
secondCandidate = await createDesktopRuntimeHostCandidate(secondHost.connection, candidateDeps);
const replacement = await preview.prepare(
TEST_TARGET_EPOCH,
source,
'session-1',
'artifact-1',
);
assert.equal(await (await fetch(replacement.url)).text(), bytes.toString());
} finally {
await firstCandidate.close();
await secondCandidate?.close();
await preview.close();
}
});

test('preserves supported IPC when the connection closes before candidate startup returns', { timeout: 5_000 }, async (t) => {
const ipc = ipcHarness();
const router = new RuntimeHostReconnectingIpcMain(ipc);
Expand Down
38 changes: 35 additions & 3 deletions apps/desktop/src/main/managed-artifact-preview.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ import type { DesktopRuntimeHostClient } from './runtime-host-client.js';

export const PREVIEW_MAX_BYTES = 8 * 1024 * 1024;
const MAX_PREVIEWS = 16;
const MAX_GLOBAL_PREVIEWS = 64;
const MAX_GLOBAL_PREVIEW_BYTES = 128 * 1024 * 1024;
const PREVIEW_TTL_MS = 30 * 60 * 1000;
const READ_DEADLINE_MS = 30_000;

Expand All @@ -42,16 +44,22 @@ interface Lease {
server: Server;
timer?: ReturnType<typeof setTimeout>;
url?: string;
reservedBytes: number;
}

/** Desktop-owned, bounded, ephemeral HTML snapshots. No workspace directory is served. */
export class ManagedArtifactPreview {
private readonly leases = new Set<Lease>();
private readonly retiredScopes = new Set<string>();
private reservedBytes = 0;
private closed = false;

constructor(private readonly ttlMs = PREVIEW_TTL_MS) {}

openScope(scope: string): void {
this.retiredScopes.delete(scope);
}

async releaseUrl(url: string): Promise<void> {
const lease = [...this.leases].find((entry) => entry.url === url);
if (lease) await this.release(lease);
Expand All @@ -68,10 +76,18 @@ export class ManagedArtifactPreview {
throw new Error('Invalid Artifact identity');
}
if (this.closed || this.retiredScopes.has(scope)) throw new Error('Preview owner is closed');
if (this.leases.size >= MAX_PREVIEWS) throw new Error('Too many active previews; wait for expiry');
const sessionLeases = [...this.leases].filter(
(lease) => lease.scope === scope && lease.sessionId === sessionId,
);
if (sessionLeases.length >= MAX_PREVIEWS) {
throw new Error('Too many active previews; wait for expiry');
}
if (this.leases.size >= MAX_GLOBAL_PREVIEWS) {
throw new Error('Too many active previews across the Desktop; wait for expiry');
}
signal?.throwIfAborted();
// Reserve before asynchronous reads, so concurrent preparations cannot exceed the bound.
const lease: Lease = { scope, sessionId, artifactId, server: createServer() };
const lease: Lease = { scope, sessionId, artifactId, server: createServer(), reservedBytes: 0 };
this.leases.add(lease);
const assertActive = () => {
signal?.throwIfAborted();
Expand All @@ -84,6 +100,11 @@ export class ManagedArtifactPreview {
if (!Number.isSafeInteger(artifact.sizeBytes) || artifact.sizeBytes < 0 || artifact.sizeBytes > PREVIEW_MAX_BYTES) {
throw new Error('HTML preview exceeds the 8 MiB limit; use Save As instead');
}
if (this.reservedBytes + artifact.sizeBytes > MAX_GLOBAL_PREVIEW_BYTES) {
throw new Error('The Desktop preview memory limit has been reached; wait for previews to expire');
}
lease.reservedBytes = artifact.sizeBytes;
this.reservedBytes += lease.reservedBytes;
const chunks: Buffer[] = [];
let size = 0;
const total = await withDeadline(client.streamArtifact(sessionId, artifactId, async (chunk) => {
Expand Down Expand Up @@ -156,6 +177,14 @@ export class ManagedArtifactPreview {
await Promise.all([...this.leases].filter((lease) => lease.scope === scope && lease.sessionId === sessionId && lease.artifactId === artifactId).map((lease) => this.release(lease)));
}

async releaseSession(scope: string, sessionId: string): Promise<void> {
await Promise.all(
[...this.leases]
.filter((lease) => lease.scope === scope && lease.sessionId === sessionId)
.map((lease) => this.release(lease)),
);
}

async closeScope(scope: string): Promise<void> {
this.retiredScopes.add(scope);
await Promise.all([...this.leases].filter((lease) => lease.scope === scope).map((lease) => this.release(lease)));
Expand All @@ -167,7 +196,10 @@ export class ManagedArtifactPreview {
}

private async release(lease: Lease): Promise<void> {
this.leases.delete(lease);
if (this.leases.delete(lease)) {
this.reservedBytes -= lease.reservedBytes;
lease.reservedBytes = 0;
}
clearTimeout(lease.timer);
await new Promise<void>((resolve) => {
lease.server.close(() => resolve());
Expand Down
2 changes: 2 additions & 0 deletions apps/desktop/src/main/runtime-host-artifacts-ipc-main.ts
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,8 @@ export function registerRuntimeHostArtifactsIpc(
"artifacts:delete",
async (_event, sessionId: string, artifactId: string) => {
const result = await deps.client.deleteArtifact(sessionId, artifactId);
// Keep the direct revoke: the Host also publishes artifact.changed, but
// stopping the bytes here must not depend on feed delivery to this Client.
await deps.preview?.service.revoke(deps.preview.scope, sessionId, artifactId);
return result;
},
Expand Down
9 changes: 9 additions & 0 deletions apps/desktop/src/main/runtime-host-boot.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1575,6 +1575,13 @@ function registerHostClientIpc(
const unsubscribeSessionCatalogChanges = client.subscribeSessionCatalogChanges(
({ sessionId }) => emitTargetSessionsChanged("updated", sessionId),
);
const unsubscribeArtifactChanges = client.subscribeArtifactChanges((frame) => {
Comment thread
me2seeks marked this conversation as resolved.
if (frame.reason === 'deleted') {
void managedArtifactPreview.revoke(scope.targetEpoch, frame.sessionId, frame.artifactId);
} else {
void managedArtifactPreview.releaseSession(scope.targetEpoch, frame.sessionId);
}
});
const unsubscribeProjectCatalogChanges = client.subscribeProjectCatalogChanges(() => {
sendToRenderer("projects:changed");
});
Expand Down Expand Up @@ -1847,12 +1854,14 @@ function registerHostClientIpc(
});
registerOnboardingIpc({ onboardingService, ipcMain: scopedIpc });
registerTaskSubmissionReadinessIpc(taskSubmissionReadinessService, scopedIpc);
managedArtifactPreview.openScope(scope.targetEpoch);
return async () => {
clientPluginTransport.release(client);
unsubscribeConfigurationChanges();
await managedArtifactPreview.closeScope(scope.targetEpoch);
unsubscribeConnectionCatalogChanges();
unsubscribeSessionCatalogChanges();
unsubscribeArtifactChanges();
unsubscribeProjectCatalogChanges();
unsubscribeScheduledTaskChanges();
runtimePolicyTargets.delete(target);
Expand Down
6 changes: 6 additions & 0 deletions apps/desktop/src/main/runtime-host-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ import {
type ProjectDirectoryEntry,
type ProjectDirectoryRoot,
SESSION_TRANSCRIPT_BOOTSTRAP_MAX_BYTES,
type ArtifactChangedFrame,
type SessionCatalogChangedFrame,
type ScheduledTaskChangedFrame,
type SessionCatalogItem,
Expand Down Expand Up @@ -395,6 +396,11 @@ export class DesktopRuntimeHostClient {
return this.connection.subscribeConfigurationChanges(listener);
}

subscribeArtifactChanges(listener: (frame: ArtifactChangedFrame) => void): () => void {
this.#assertOpen();
return this.connection.subscribeArtifactChanges(listener);
}

subscribeConnectionCatalogChanges(listener: (revision: number) => void): () => void {
this.#assertOpen();
return this.connection.subscribeConnectionCatalogChanges(listener);
Expand Down
1 change: 1 addition & 0 deletions packages/cli/src/__tests__/acp-stdio-server.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -708,6 +708,7 @@ describe('Maka ACP stdio server', () => {
return subscription;
},
subscribeConfigurationChanges: () => () => undefined,
subscribeArtifactChanges: () => () => undefined,
subscribeConnectionCatalogChanges: () => () => undefined,
subscribeProjectCatalogChanges: () => () => undefined,
subscribeSessionCatalogChanges: () => () => undefined,
Expand Down
Loading