fix(pi): keep active runs alive during background sleep

electron/coding-runtime/pi/runtime.ts

@@ -13,6 +13,7 @@ import {
 import { PiProviderRefreshCoordinator } from './provider-refresh';
 import type {
   CodingConversationRuntime,
+  CodingRuntimeDisposeReason,
   CodingRuntimeCommand,
   CodingRuntimeDiagnostics,
   CodingRuntimePublicError,
@@ -109,6 +110,7 @@ export interface PiConversationRuntimeOptions {
   getDraftRevision?(conversationId: string): number;
   knownExtensionWidgetKeys?: readonly string[];
   onExtensionUiProjection?(projection: PiExtensionUiProjection): void;
+  acquireBackgroundLease?(lease: { id: string; kind: 'coding-run' }): () => void;
 }
 
 export interface PiWorkerProcessAdapter {
@@ -516,6 +518,13 @@ export class PiConversationRuntime implements CodingConversationRuntime {
   private readonly interactions: PiInteractionStore;
   private readonly extensionUi: PiExtensionUiProjector;
   private readonly onExtensionUiProjection: ((projection: PiExtensionUiProjection) => void) | undefined;
+  private readonly acquireBackgroundLease:
+    | ((lease: { id: string; kind: 'coding-run' }) => () => void)
+    | undefined;
+  private readonly runBackgroundLeases = new Map<
+    string,
+    { runId: string; release: () => void }
+  >();
   private readonly states = new Map<string, ConversationReducerState>();
   private readonly inputs = new Map<string, PrepareConversationInput>();
   private readonly listeners = new Set<(patch: ConversationPatchEnvelope) => void>();
@@ -568,6 +577,7 @@ export class PiConversationRuntime implements CodingConversationRuntime {
         : {}),
     });
     this.onExtensionUiProjection = options.onExtensionUiProjection;
+    this.acquireBackgroundLease = options.acquireBackgroundLease;
     if (Boolean(this.isAuthenticationError) !== Boolean(this.refreshCredential)) {
       throw new Error('Provider authentication detection and refresh must be configured together');
     }
@@ -642,6 +652,7 @@ export class PiConversationRuntime implements CodingConversationRuntime {
     this.snapshot(input.conversationId);
     const images = await this.resolveImages(input.attachments);
     const runId = this.id('run');
+    this.acquireRunBackgroundLease(input.conversationId, runId);
     const messageId = `client:${input.clientRequestId}`;
     this.emit(input.conversationId, {
       op: 'message.upsert',
@@ -676,8 +687,9 @@ export class PiConversationRuntime implements CodingConversationRuntime {
     } catch (error) {
       this.extensionUi.endRun(input.conversationId, runId);
       if (this.extensionHost && generation) {
-        await this.extensionHost.clearRun(input.conversationId, generation, runId);
+        await this.extensionHost.clearRun(input.conversationId, generation, runId).catch(() => undefined);
       }
+      this.releaseRunBackgroundLease(input.conversationId, runId);
       throw error;
     }
     this.emit(input.conversationId, {
@@ -734,7 +746,7 @@ export class PiConversationRuntime implements CodingConversationRuntime {
       if (latest.runId === current.runId
         && latest.status === 'aborting'
         && (!worker || worker.state === 'crashed' || worker.generation !== generation)) {
-        this.failRun(conversationId, current.runId!, error, generation);
+        await this.failRun(conversationId, current.runId!, error, generation);
       } else if (latest.runId === current.runId && latest.status === 'aborting') {
         this.emit(conversationId, { op: 'run.state', run: current }, current.runId);
       }
@@ -865,6 +877,7 @@ export class PiConversationRuntime implements CodingConversationRuntime {
   async compact(conversationId: string): Promise<void> {
     await this.waitForProjection(conversationId);
     const runId = this.id('run');
+    this.acquireRunBackgroundLease(conversationId, runId);
     const generation = this.pool.getState(conversationId)?.generation;
     if (generation) this.extensionUi.beginRun(conversationId, generation, runId);
     let ticket;
@@ -880,8 +893,9 @@ export class PiConversationRuntime implements CodingConversationRuntime {
     } catch (error) {
       this.extensionUi.endRun(conversationId, runId);
       if (this.extensionHost && generation) {
-        await this.extensionHost.clearRun(conversationId, generation, runId);
+        await this.extensionHost.clearRun(conversationId, generation, runId).catch(() => undefined);
       }
+      this.releaseRunBackgroundLease(conversationId, runId);
       throw error;
     }
     this.emit(conversationId, {
@@ -891,7 +905,7 @@ export class PiConversationRuntime implements CodingConversationRuntime {
     try {
       await ticket.accepted;
     } catch (error) {
-      this.failRun(conversationId, runId, error);
+      await this.failRun(conversationId, runId, error);
       throw error;
     }
   }
@@ -936,21 +950,69 @@ export class PiConversationRuntime implements CodingConversationRuntime {
     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);
+    try {
+      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);
+      return this.runtimeState(conversationId);
+    } finally {
+      if (before.run.runId) {
+        this.releaseRunBackgroundLease(conversationId, before.run.runId);
+      }
     }
-    const state = await this.pool.recover(conversationId);
-    await this.requestHydration(conversationId, state, false);
-    this.settleRecoveredRun(conversationId);
-    return this.runtimeState(conversationId);
   }
 
-  async dispose(conversationId: string): Promise<void> {
+  async dispose(
+    conversationId: string,
+    reason: CodingRuntimeDisposeReason,
+  ): Promise<void> {
+    await this.waitForProjection(conversationId);
+    const before = this.states.get(conversationId)?.snapshot;
+    if (reason === 'background_sleep'
+      && (this.runBackgroundLeases.has(conversationId)
+        || Boolean(before && !runIsTerminal(before.run.status)))) return;
+
+    const runId = before?.run.runId;
+    if (runId && before && !runIsTerminal(before.run.status)) {
+      await this.enqueueProjection(conversationId, async () => {
+        const current = this.states.get(conversationId)?.snapshot;
+        if (!current || current.run.runId !== runId || runIsTerminal(current.run.status)) return;
+        await this.failRun(
+          conversationId,
+          runId,
+          new CodingRuntimeContractError(
+            'CODING_RUNTIME_START_FAILED',
+            '本地 Agent 已中断，原请求未自动重发。',
+            true,
+          ),
+          current.cursor.workerGeneration,
+          false,
+        );
+      });
+    }
+
     const state = this.pool.getState(conversationId);
-    if (state) await this.interactions.cancelGeneration(conversationId, state.generation);
-    await this.pool.dispose(conversationId);
+    let stopped = reason !== 'background_sleep';
+    try {
+      if (state) await this.interactions.cancelGeneration(conversationId, state.generation);
+      if (reason === 'background_sleep' && this.hasConversationActiveWork(conversationId)) return;
+      const didStop = await this.pool.dispose(
+        conversationId,
+        reason,
+        reason === 'background_sleep'
+          ? () => !this.hasConversationActiveWork(conversationId)
+          : () => true,
+      );
+      stopped = reason === 'background_sleep' ? didStop : true;
+    } finally {
+      if (stopped) this.releaseConversationBackgroundLease(conversationId);
+    }
+    if (reason === 'background_sleep' && !stopped) return;
     this.registry.forget(conversationId);
     this.inputs.delete(conversationId);
     this.states.delete(conversationId);
@@ -1008,14 +1070,27 @@ export class PiConversationRuntime implements CodingConversationRuntime {
     pool: ReturnType<PiWorkerPool['getResilienceProofDiagnostics']>;
     subagents: ReturnType<PiSubagentScheduler['getDiagnostics']> | null;
     extension: ReturnType<PiManagedExtensionHost['getDiagnostics']> | null;
+    backgroundLeases: { active: number };
   } {
     return {
       pool: this.pool.getResilienceProofDiagnostics(),
       subagents: this.subagentScheduler?.getDiagnostics() ?? null,
       extension: this.extensionHost?.getDiagnostics() ?? null,
+      backgroundLeases: { active: this.runBackgroundLeases.size },
     };
   }
 
+  hasActiveWork(): boolean {
+    if (this.runBackgroundLeases.size > 0) return true;
+    return [...this.states.values()].some(({ snapshot }) => !runIsTerminal(snapshot.run.status));
+  }
+
+  private hasConversationActiveWork(conversationId: string): boolean {
+    if (this.runBackgroundLeases.has(conversationId)) return true;
+    const snapshot = this.states.get(conversationId)?.snapshot;
+    return Boolean(snapshot && !runIsTerminal(snapshot.run.status));
+  }
+
   markProviderStale(): void {
     this.pool.markProviderStale();
   }
@@ -1038,11 +1113,23 @@ export class PiConversationRuntime implements CodingConversationRuntime {
 
   async shutdown(): Promise<void> {
     this.unsubscribePool();
-    await this.pool.shutdown();
-    await this.extensionHost?.close();
-    this.projectors.clear();
-    this.projectionChains.clear();
-    this.hydrationFlights.clear();
+    try {
+      const results = await Promise.allSettled([
+        this.pool.shutdown(),
+        this.extensionHost?.close(),
+      ]);
+      const failure = results.find((result): result is PromiseRejectedResult => (
+        result.status === 'rejected'
+      ));
+      if (failure) throw failure.reason;
+    } finally {
+      for (const conversationId of [...this.runBackgroundLeases.keys()]) {
+        this.releaseConversationBackgroundLease(conversationId);
+      }
+      this.projectors.clear();
+      this.projectionChains.clear();
+      this.hydrationFlights.clear();
+    }
   }
 
   private async queue(
@@ -1051,6 +1138,9 @@ export class PiConversationRuntime implements CodingConversationRuntime {
   ): Promise<QueueAcceptance> {
     await this.waitForProjection(input.conversationId);
     const snapshot = this.snapshot(input.conversationId);
+    if (snapshot.run.runId) {
+      this.acquireRunBackgroundLease(input.conversationId, snapshot.run.runId);
+    }
     const images = await this.resolveImages(input.attachments);
     const queuePosition = snapshot.queue.items.length + 1;
     const queueId = this.id('queue');
@@ -1170,7 +1260,7 @@ export class PiConversationRuntime implements CodingConversationRuntime {
         runId,
         failure instanceof Error ? failure : new Error('Prompt acceptance failed'),
       );
-      this.failRun(conversationId, runId, failure);
+      await this.failRun(conversationId, runId, failure);
       throw failure;
     }
   }
@@ -1251,7 +1341,7 @@ export class PiConversationRuntime implements CodingConversationRuntime {
         await this.interactions.cancelGeneration(event.conversationId, event.generation);
         const current = this.snapshot(event.conversationId).run;
         if (current.runId && !runIsTerminal(current.status)) {
-          this.failRun(
+          await this.failRun(
             event.conversationId,
             current.runId,
             event.error,
@@ -1309,25 +1399,32 @@ export class PiConversationRuntime implements CodingConversationRuntime {
         }
       }
       if (event.event.type === 'agent_settled' && snapshot.run.runId) {
-        await this.interactions.cancelRun(event.conversationId, snapshot.run.runId, true);
-        await this.extensionHost?.clearRun(
-          event.conversationId,
-          event.generation,
-          snapshot.run.runId,
-        );
-        this.extensionUi.endRun(event.conversationId, snapshot.run.runId);
+        try {
+          await Promise.allSettled([
+            this.interactions.cancelRun(event.conversationId, snapshot.run.runId, true),
+            this.extensionHost?.clearRun(
+              event.conversationId,
+              event.generation,
+              snapshot.run.runId,
+            ),
+          ]);
+          this.extensionUi.endRun(event.conversationId, snapshot.run.runId);
+        } finally {
+          this.releaseRunBackgroundLease(event.conversationId, snapshot.run.runId);
+        }
       }
     }).catch((error) => {
       this.recordProjectionFailure(event.conversationId, event.generation, error);
     });
   }
 
-  private failRun(
+  private async failRun(
     conversationId: string,
     runId: string,
     error: unknown,
     generation?: number,
-  ): void {
+    releaseBackgroundLease = true,
+  ): Promise<void> {
     const current = this.states.get(conversationId)?.snapshot?.run;
     const snapshotGeneration = this.states.get(conversationId)?.snapshot?.cursor.workerGeneration;
     if (current?.runId !== runId
@@ -1351,10 +1448,18 @@ export class PiConversationRuntime implements CodingConversationRuntime {
       generation: generation ?? snapshotGeneration,
       code: publicError.code,
     });
-    void this.interactions.cancelRun(conversationId, runId, true);
     this.extensionUi.endRun(conversationId, runId);
     const runGeneration = generation ?? this.pool.getState(conversationId)?.generation;
-    if (runGeneration) void this.extensionHost?.clearRun(conversationId, runGeneration, runId);
+    try {
+      await Promise.allSettled([
+        this.interactions.cancelRun(conversationId, runId, true),
+        runGeneration
+          ? this.extensionHost?.clearRun(conversationId, runGeneration, runId)
+          : undefined,
+      ]);
+    } finally {
+      if (releaseBackgroundLease) this.releaseRunBackgroundLease(conversationId, runId);
+    }
   }
 
   private requestHydration(
@@ -1555,6 +1660,35 @@ export class PiConversationRuntime implements CodingConversationRuntime {
     if (input) input.model = clone(model);
   }
 
+  private acquireRunBackgroundLease(conversationId: string, runId: string): void {
+    const current = this.runBackgroundLeases.get(conversationId);
+    if (current?.runId === runId) return;
+    if (current) {
+      throw new CodingRuntimeContractError(
+        'CODING_RUNTIME_START_FAILED',
+        'Conversation already has an active run',
+        true,
+      );
+    }
+    const release = this.acquireBackgroundLease?.({
+      id: `coding-run:${conversationId}:${runId}`,
+      kind: 'coding-run',
+    }) ?? (() => undefined);
+    this.runBackgroundLeases.set(conversationId, { runId, release });
+  }
+
+  private releaseRunBackgroundLease(conversationId: string, runId: string): void {
+    const current = this.runBackgroundLeases.get(conversationId);
+    if (!current || current.runId !== runId) return;
+    this.runBackgroundLeases.delete(conversationId);
+    current.release();
+  }
+
+  private releaseConversationBackgroundLease(conversationId: string): void {
+    const current = this.runBackgroundLeases.get(conversationId);
+    if (current) this.releaseRunBackgroundLease(conversationId, current.runId);
+  }
+
   private id(kind: RuntimeIdKind): string {
     return this.createRuntimeId(kind);
   }