fix(coding-runtime): close Pi extension lifecycle races
This commit is contained in:
@@ -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));
|
||||
}
|
||||
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user