@springbrand/agent-runtime 0.1.3-alpha.0 → 0.1.3-alpha.2

package/src/runtime.ts

@@ -1,10 +1,12 @@
 import type { Connection } from "agents";
 import {
+  ChatStreamStalledError,
   MessageType,
   clearChatTerminal,
   parseProtocolMessage,
   recordChatTerminal,
   sendIfOpen,
+  type ChatRecoveryConfig,
 } from "agents/chat";
 import { RecoverableChatAgent } from "./kernel/recoverable-chat-agent";
 import {
@@ -46,6 +48,9 @@ import { connectConfiguredMcpServers } from "./lib/mcp";
 import { installConsoleSink } from "./lib/telemetry-dev";
 import {
   PiRuntimeAdapter,
+  MODEL_STREAM_STALL_TIMEOUT_MS,
+  readModelStreamStallDetails,
+  RetryableModelError,
   type PiCanonicalUserInput,
   type PiCanonicalTranscriptSnapshot,
   type PiChatRecoveryData,
@@ -91,6 +96,9 @@ import {
 
 const SCHEDULED_STABLE_TIMEOUT_MS = 30_000;
 const TURN_EVENT_RETRY_SECONDS = 10;
+const CHAT_RECOVERY_MAX_ATTEMPTS = 5;
+const CHAT_RECOVERY_TERMINAL_MESSAGE =
+  "多次恢复仍未成功，本次生成已停止，当前进度已保留。请发送新消息继续。";
 
 type RuntimeEventOutboxPayload =
   | { readonly type: "model-usage"; readonly event: RuntimeModelUsageEvent }
@@ -110,6 +118,8 @@ interface StoredSubmission extends SubmissionReceipt {
   queuedUiMessageJson?: string | null;
   userMessageId?: string | null;
   regenerateMessageId?: string | null;
+  recoveryErrorCount: number;
+  recoveryReason?: "no_meaningful_model_progress" | "transient_model_error";
 }
 
 interface SubmitMessageOptions {
@@ -210,6 +220,11 @@ interface PendingTemporaryAgentApproval {
 export abstract class AgentRuntimeKernel<
   Env extends Cloudflare.Env = Cloudflare.Env,
 > extends RecoverableChatAgent<Env, RuntimeState, PiChatRecoveryData> {
+  override chatRecovery: ChatRecoveryConfig = {
+    maxAttempts: CHAT_RECOVERY_MAX_ATTEMPTS,
+    terminalMessage: CHAT_RECOVERY_TERMINAL_MESSAGE,
+  };
+
   initialState: RuntimeState = {
     runtimeLoad: { status: "idle", available: false },
   };
@@ -352,15 +367,25 @@ export abstract class AgentRuntimeKernel<
         };
       },
       retryTurn: async (submissionId) => {
-        const receipt = await this.submissions.recover(submissionId);
-        return receipt.status === "completed"
-          ? { status: "completed" as const }
-          : {
-              status: "failed" as const,
-              error:
-                receipt.error ??
-                `Pi recovery ended with status ${receipt.status}`,
-            };
+        try {
+          const receipt = await this.submissions.recover(submissionId);
+          return receipt.status === "completed"
+            ? { status: "completed" as const }
+            : {
+                status: "failed" as const,
+                error:
+                  receipt.error ??
+                  `Pi recovery ended with status ${receipt.status}`,
+              };
+        } catch (error) {
+          if (
+            error instanceof ChatStreamStalledError ||
+            error instanceof RetryableModelError
+          ) {
+            return { status: "scheduled" as const };
+          }
+          throw error;
+        }
       },
       failTurn: async (submissionId, message) => {
         const submission = this.readSubmission(submissionId);
@@ -450,7 +475,6 @@ export abstract class AgentRuntimeKernel<
     const startedAt = output.startedAt ?? Date.now();
     return this.pi.createTurn({
       prepared: this.preparedPi(),
-      baseRevision: this.baseRevision(),
       pinnedDescriptor: submission.assemblyDescriptor,
       submission: {
         id: submission.submissionId,
@@ -833,16 +857,6 @@ export abstract class AgentRuntimeKernel<
       .join("");
   }
 
-  // 作用：返回当前已激活配置的 revision。
-  // 调用：创建 Pi Turn 适配器时调用。
-  // 原因：把初始化不变式收口在一处，避免执行层带着空 revision 继续。
-  private baseRevision(): string {
-    if (!this.runtimeRevision) {
-      throw new Error("Runtime revision is not initialized");
-    }
-    return this.runtimeRevision;
-  }
-
   // 作用：返回当前已准备好的 Pi Runtime。
   // 调用：创建 Turn 适配器或准入固定装配时调用。
   // 原因：明确区分“已有 Snapshot”和“Pi 已完成准备”，防止半初始化状态进入 Turn。
@@ -922,23 +936,24 @@ export abstract class AgentRuntimeKernel<
       await this.hash(submission.assemblyDescriptor)
     ) {
       throw new Error(
-        "Pinned Runtime revision is unavailable for Pi recovery",
+        "Pinned Runtime assembly descriptor is invalid",
       );
     }
   }
 
-  // 作用：按当前内存预算和模型压缩 Pi 上下文。
+  // 作用：按当前模型输入预算压缩 Pi 上下文。
   // 调用：Pi Turn 适配器在发模型请求前通过 `transformContext` 回调。
-  // 原因：压缩必须使用本 Snapshot 的模型和密钥，不能脱离已固定的 Runtime 配置。
+  // 原因：模型窗口必须预留最大输出，压缩也必须使用本 Snapshot 的模型和密钥。
   private async transformPiContext(
     submissionId: string,
     messages: Parameters<PiRuntimeTranscript["compactContext"]>[0],
     signal?: AbortSignal,
   ): ReturnType<PiRuntimeTranscript["compactContext"]> {
     const snapshot = this.assembly();
+    const inputBudget = snapshot.pi.model.contextWindow -
+      snapshot.pi.model.maxTokens;
     const compacted = await this.transcript.compactContext(messages, {
-      compactAfterTokens:
-        snapshot.profile.memory.compactAfterTokens,
+      compactAfterTokens: inputBudget * 0.85,
       model: snapshot.pi.model,
       apiKey: this.pi.resolveApiKey(
         snapshot.bindings.provider,
@@ -1435,6 +1450,19 @@ export abstract class AgentRuntimeKernel<
       return terminal;
     });
     await this.drainRuntimeEvents();
+    try {
+      await this.settleChatRecovery(
+        latest.requestId,
+        effectiveOutcome === "succeeded"
+          ? "completed"
+          : effectiveOutcome === "aborted"
+            ? "skipped"
+            : "failed",
+        effectiveMessage,
+      );
+    } catch (error) {
+      console.error("[chat-recovery] terminal settlement failed", error);
+    }
     let inactiveActivity: RuntimeActivity = "idle";
     if (terminal.status === "completed") {
       try {
@@ -1478,6 +1506,16 @@ export abstract class AgentRuntimeKernel<
   ): Promise<boolean> {
     let decision = await this.materializeRecoveredToolResults(submission);
     for (let step = 0; step < 32; step += 1) {
+      const latest = this.readSubmission(submission.submissionId);
+      if (!latest || isTerminalSubmissionStatus(latest.status)) return false;
+      if (latest.abortReason) {
+        await this.submissions.finish(
+          latest,
+          "aborted",
+          latest.abortReason,
+        );
+        return false;
+      }
       this.db.transaction(() => {
         this.applyPiRecoveryMutations(
           submission,
@@ -1655,13 +1693,24 @@ export abstract class AgentRuntimeKernel<
         submission.abortReason,
       );
     }
-    if (recovery) await this.ensureRuntimeReady();
+    if (recovery) {
+      this.db.submissions.clearRecoveryReason(submissionId);
+      await this.broadcastApprovals();
+      await this.ensureRuntimeReady();
+    }
     try {
       return await this.executeNonTerminalSubmission(
         submission,
         recovery,
       );
     } catch (error) {
+      if (
+        recovery &&
+        (error instanceof ChatStreamStalledError ||
+          error instanceof RetryableModelError)
+      ) {
+        throw error;
+      }
       const failed = await this.submissions.finish(
         submission,
         "failed",
@@ -1711,11 +1760,33 @@ export abstract class AgentRuntimeKernel<
     }
     const recoveryAdapter = this.createSubmissionExecutionAdapter(submission);
     if (recovery) {
-      const ready = await this.prepareRecoveredTurn(
-        submission,
-        recoveryAdapter,
-      );
-      if (!ready) return this.readSubmission(submissionId)!;
+      const deactivatePreparation = this.submissions.activate(submission, {
+        submissionId,
+        requestId: submission.requestId,
+        messageId: submission.assistantMessageId,
+        startedAt: submission.createdAt,
+        continuation: true,
+        agent: recoveryAdapter,
+      });
+      let ready: boolean;
+      try {
+        ready = await this.prepareRecoveredTurn(
+          submission,
+          recoveryAdapter,
+        );
+      } finally {
+        deactivatePreparation();
+      }
+      if (!ready) {
+        const latest = this.readSubmission(submissionId)!;
+        return latest.abortReason
+          ? this.submissions.finish(
+              latest,
+              "aborted",
+              latest.abortReason,
+            )
+          : latest;
+      }
     } else {
       await this.materializeRecoveredToolResults(submission);
     }
@@ -1821,13 +1892,77 @@ export abstract class AgentRuntimeKernel<
       await this.projectTerminal(turn, terminal);
       return terminal;
     } catch (error) {
+      const stalled = error instanceof ChatStreamStalledError;
+      const abortReason = this.readSubmission(submissionId)?.abortReason;
+      if (!abortReason && (stalled || error instanceof RetryableModelError)) {
+        const recoveryErrorCount = this.db.submissions
+          .incrementRecoveryErrorCount(
+            submissionId,
+            stalled
+              ? "no_meaningful_model_progress"
+              : "transient_model_error",
+          );
+        if (stalled) {
+          const stallDetails = readModelStreamStallDetails(errorText(error)) ?? {
+            lastMeaningfulActivityAt:
+              Date.now() - MODEL_STREAM_STALL_TIMEOUT_MS,
+            lastMeaningfulActivityType: "model_stream_started",
+            idleMs: MODEL_STREAM_STALL_TIMEOUT_MS,
+          };
+          this._emit("chat:stream:stalled", {
+            requestId: submission.requestId,
+            submissionId,
+            attempt: recoveryErrorCount,
+            timeoutMs: MODEL_STREAM_STALL_TIMEOUT_MS,
+            ...stallDetails,
+            reason: "no_meaningful_model_progress",
+          });
+        }
+        await this.broadcastApprovals();
+        const stoppedDuringRecovery = this.readSubmission(submissionId)
+          ?.abortReason;
+        if (
+          !stoppedDuringRecovery &&
+          recoveryErrorCount < CHAT_RECOVERY_MAX_ATTEMPTS
+        ) {
+          const recoveryOutcome = await this.scheduleChatRecoveryRetry(
+            {
+              submissionId,
+              requestId: submission.requestId,
+            },
+            async () => {
+              this.broadcast(json({
+                type: MessageType.CF_AGENT_CHAT_MESSAGES,
+                messages: await this.getMessages(),
+              }));
+              if (streamId) this.completeRecoverableStream(streamId);
+              this.sendChatResponse(submission.requestId, "", true);
+            },
+          );
+          if (
+            recoveryOutcome !== "disabled" &&
+            !this.readSubmission(submissionId)?.abortReason
+          ) {
+            if (recovery && recoveryOutcome === "scheduled") throw error;
+            return this.readSubmission(submissionId)!;
+          }
+        } else if (!stoppedDuringRecovery) {
+          terminalIntent = {
+            outcome: "failed",
+            message: CHAT_RECOVERY_TERMINAL_MESSAGE,
+          };
+        }
+      }
       const latest = this.readSubmission(submissionId);
-      const intent = terminalIntent ?? {
-        outcome: latest?.abortReason
-          ? ("aborted" as const)
-          : ("failed" as const),
-        message: errorText(error),
-      };
+      const intent = terminalIntent ?? (latest?.abortReason
+        ? {
+            outcome: "aborted" as const,
+            message: latest.abortReason,
+          }
+        : {
+            outcome: "failed" as const,
+            message: errorText(error),
+          });
       const failed = await this.submissions.finish(
         submission,
         intent.outcome,
@@ -1856,19 +1991,21 @@ export abstract class AgentRuntimeKernel<
     }
   }
 
-  // 作用：先持久化一条 Pi 流记录，再广播给在线客户端。
+  // 作用：持久化一条 Pi 流记录并广播给在线客户端。
   // 调用：活跃 Turn 适配器每产生一条 stream record 时调用。
-  // 原因：耐久写入早于 WebSocket 发送，断线客户端才能从同一序列续传。
+  // 原因：先发 WebSocket 再等持久化完成，让客户端零延迟收到 token，同时
+  // await 保证 Pi 的事件循环不越过尚未落盘的 chunk，续传完整性不受影响。
   private async persistAndBroadcastRecord(
     turn: ActiveTurn,
     chunk: UIMessageChunk,
   ): Promise<void> {
     const streamId = this.streamBySubmission.get(turn.submissionId);
     const body = json(chunk);
-    if (streamId) {
-      await this.appendRecoverableChunk(streamId, body);
-    }
+    const persist = streamId
+      ? this.appendRecoverableChunk(streamId, body)
+      : Promise.resolve();
     this.sendChatResponse(turn.requestId, body, false);
+    await persist;
   }
 
   private sendChatTerminal(requestId: string, error?: string): void {
@@ -2080,7 +2217,7 @@ export abstract class AgentRuntimeKernel<
       if (submission) {
         await this.cancelSubmissionById(
           submission.submissionId,
-          "Client cancelled",
+          USER_STOP_REASON,
         );
       }
       return;
@@ -2821,7 +2958,21 @@ export abstract class AgentRuntimeKernel<
       );
       const turn: RuntimeTurnState = {
         ...(current
-          ? { activeSubmissionId: current.submissionId }
+          ? {
+              activeSubmissionId: current.submissionId,
+              ...(current.requestId
+                ? { activeRequestId: current.requestId }
+                : {}),
+              ...(current.recoveryErrorCount === undefined
+                ? {}
+                : {
+                    recoveryAttempt: current.recoveryErrorCount,
+                    recoveryMax: CHAT_RECOVERY_MAX_ATTEMPTS,
+                    ...(current.recoveryReason
+                      ? { recoveryReason: current.recoveryReason }
+                      : {}),
+                  }),
+            }
           : {}),
         steerable: Boolean(
           current &&