fix(coding-runtime): close Pi extension lifecycle races

This commit is contained in:
2026-08-23 10:13:10 +08:00
parent 3861c3286a
commit 31de325cd8
9 changed files with 197 additions and 14 deletions

View File

@@ -63,9 +63,11 @@ export class PiManagedExtensionHost {
private readonly leases: PiProjectWriteLeaseCoordinator;
private readonly registrations = new Map<string, WorkerRegistrationRecord>();
private readonly runBindings = new Map<string, string>();
private readonly requestFlights = new Set<Promise<void>>();
private server: Server | null = null;
private bridgeUrl: string | null = null;
private startFlight: Promise<string> | null = null;
private closing = false;
constructor(leases = new PiProjectWriteLeaseCoordinator()) {
this.leases = leases;
@@ -131,16 +133,20 @@ export class PiManagedExtensionHost {
}
async close(): Promise<void> {
for (const record of [...this.registrations.values()]) this.disposeRecord(record);
this.runBindings.clear();
const server = this.server;
this.server = null;
this.bridgeUrl = null;
this.startFlight = null;
if (!server) return;
await new Promise<void>((resolve, reject) => {
this.closing = true;
const closed = server ? new Promise<void>((resolve, reject) => {
server.close((error) => error ? reject(error) : resolve());
});
}) : Promise.resolve();
for (const record of [...this.registrations.values()]) this.disposeRecord(record);
this.runBindings.clear();
await Promise.allSettled([...this.requestFlights]);
server?.closeIdleConnections();
await closed;
this.closing = false;
}
private start(): Promise<string> {
@@ -148,7 +154,10 @@ export class PiManagedExtensionHost {
if (this.startFlight) return this.startFlight;
this.startFlight = new Promise<string>((resolve, reject) => {
const server = createServer((request, response) => {
void this.handle(request, response);
const flight = this.handle(request, response).finally(() => {
this.requestFlights.delete(flight);
});
this.requestFlights.add(flight);
});
server.once('error', reject);
server.listen(0, '127.0.0.1', () => {
@@ -159,6 +168,7 @@ export class PiManagedExtensionHost {
return;
}
this.server = server;
this.closing = false;
this.bridgeUrl = `http://127.0.0.1:${address.port}/v1/worker`;
resolve(this.bridgeUrl);
});
@@ -226,6 +236,7 @@ export class PiManagedExtensionHost {
);
if (record.runId !== value.runId || this.registrations.get(token) !== record) {
lease.release();
this.respond(response, 409, { error: 'Worker run identity is stale' });
return;
}
record.leases.set(value.resourceId, lease);
@@ -266,7 +277,10 @@ export class PiManagedExtensionHost {
private respond(response: ServerResponse, status: number, body: Record<string, unknown>): void {
if (response.writableEnded) return;
response.writeHead(status, { 'content-type': 'application/json; charset=utf-8' });
response.writeHead(status, {
'content-type': 'application/json; charset=utf-8',
...(this.closing ? { connection: 'close' } : {}),
});
response.end(JSON.stringify(body));
}

View File

@@ -5,6 +5,7 @@ const MAX_TITLE_LENGTH = 256;
const MAX_EDITOR_TEXT_LENGTH = 64 * 1024;
const MAX_WIDGET_LINES = 32;
const MAX_WIDGET_LINE_LENGTH = 512;
const MAX_DIAGNOSTICS = 256;
export type PiExtensionUiProjection =
| { kind: 'notify'; conversationId: string; message: string; level: 'info' | 'warning' | 'error' }
@@ -136,7 +137,7 @@ export class PiExtensionUiProjector {
}
if (event.method === 'set_editor_text' && typeof event.text === 'string') {
if (this.getDraftRevision(conversationId) !== run.draftRevision) {
this.diagnostics.push({ method: event.method, reason: 'stale-draft-revision' });
this.recordDiagnostic({ method: event.method, reason: 'stale-draft-revision' });
return null;
}
return {
@@ -151,12 +152,17 @@ export class PiExtensionUiProjector {
}
private unsupported(method: string): null {
this.diagnostics.push({ method: bounded(method, 128), reason: 'unsupported-ui-method' });
this.recordDiagnostic({ method: bounded(method, 128), reason: 'unsupported-ui-method' });
return null;
}
private invalid(method: string): null {
this.diagnostics.push({ method: bounded(method, 128), reason: 'invalid-ui-payload' });
this.recordDiagnostic({ method: bounded(method, 128), reason: 'invalid-ui-payload' });
return null;
}
private recordDiagnostic(diagnostic: PiExtensionUiDiagnostic): void {
this.diagnostics.push(diagnostic);
if (this.diagnostics.length > MAX_DIAGNOSTICS) this.diagnostics.shift();
}
}

View File

@@ -13,6 +13,7 @@ interface StoredInteraction {
interaction: ConversationInteraction;
generation: number;
labels: Map<string, string>;
phase: 'pending' | 'responding';
untrack(): void;
}
@@ -85,6 +86,7 @@ export class PiInteractionStore {
interaction,
generation,
labels,
phase: 'pending',
untrack: () => undefined,
};
stored.untrack = this.transport.trackGenerationResource({
@@ -102,6 +104,7 @@ export class PiInteractionStore {
async respond(conversationId: string, response: PiInteractionResponse): Promise<ConversationInteraction> {
const stored = this.pending.get(this.key(conversationId, response.interactionId));
if (!stored) throw new Error('Pi interaction is not pending');
if (stored.phase === 'responding') throw new Error('Pi interaction response is already in progress');
const state = this.transport.getState(conversationId);
const active = this.transport.getActiveRun(conversationId);
if (state?.generation !== stored.generation
@@ -126,7 +129,18 @@ export class PiInteractionStore {
} else {
throw new Error('Pi interaction response does not match its kind');
}
await this.transport.send(conversationId, command);
stored.phase = 'responding';
try {
await this.transport.send(conversationId, command);
} catch (error) {
if (this.pending.get(this.key(conversationId, stored.interaction.id)) === stored) {
stored.phase = 'pending';
}
throw error;
}
if (this.pending.get(this.key(conversationId, stored.interaction.id)) !== stored) {
throw new Error('Pi interaction belongs to a stale worker run');
}
return this.finish(stored, 'cancelled' in response
? 'cancelled'
: stored.interaction.kind === 'confirm' && 'confirmed' in response && !response.confirmed

View File

@@ -798,6 +798,13 @@ export class PiConversationRuntime implements CodingConversationRuntime {
async recover(conversationId: string): Promise<ConversationRuntimeState> {
await this.waitForProjection(conversationId);
const before = this.snapshot(conversationId);
const generation = this.pool.getState(conversationId)?.generation;
if (before.run.runId && generation) {
await this.interactions.cancelRun(conversationId, before.run.runId, true);
this.extensionUi.endRun(conversationId, before.run.runId);
await this.extensionHost?.clearRun(conversationId, generation, before.run.runId);
}
const state = await this.pool.recover(conversationId);
await this.requestHydration(conversationId, state, false);
this.settleRecoveredRun(conversationId);
@@ -1003,8 +1010,12 @@ export class PiConversationRuntime implements CodingConversationRuntime {
this.replaceWorkerGeneration(event.conversationId, event.state, true);
this.resetProjector(event.conversationId);
const runId = this.states.get(event.conversationId)?.snapshot.run.runId;
if (runId) this.extensionUi.replaceGeneration(event.conversationId, event.generation, runId);
if (this.extensionHost && runId) {
const activeRun = this.pool.getActiveRun(event.conversationId);
const continuesActiveRun = Boolean(runId && activeRun?.runId === runId);
if (runId && continuesActiveRun) {
this.extensionUi.replaceGeneration(event.conversationId, event.generation, runId);
}
if (this.extensionHost && runId && continuesActiveRun) {
void this.extensionHost.bindRun(event.conversationId, event.generation, runId).catch((error) => {
this.recordProjectionFailure(event.conversationId, event.generation, error);
});