@downcity/agent 1.1.270 → 1.1.276

package/src/session/messages/SessionAssistantMessageWriter.ts

@@ -31,7 +31,10 @@ export class SessionAssistantMessageWriter {
     Pick<SessionAssistantTextPart, "type" | "provider_metadata">
   >();
   private readonly active_text_part_ids = new Map<string, string>();
+  private readonly current_step_part_ids = new Set<string>();
   private write_chain: Promise<void> = Promise.resolve();
+  private step_index = 0;
+  private step_active = false;
   private closed = false;
 
   constructor(recorder: SessionMessages, message_id: string) {
@@ -46,6 +49,70 @@ export class SessionAssistantMessageWriter {
     });
   }
 
+  /** 建立一个独立模型 UI stream 的 canonical step 作用域。 */
+  async begin_step(): Promise<void> {
+    await this.enqueue_write(async () => {
+      if (this.closed) throw new Error("Assistant Message writer is closed");
+      if (this.step_active) {
+        throw new Error("Assistant canonical step is already active");
+      }
+      this.step_index += 1;
+      this.step_active = true;
+      this.current_step_part_ids.clear();
+      this.pending_text_parts.clear();
+      this.active_text_part_ids.clear();
+    });
+  }
+
+  /**
+   * 校验当前 step 的最终快照并原子补充 metadata。
+   *
+   * 最终快照不能创建、删除或重排 Part；任何不一致都表示 canonical chunk
+   * 链路不完整，必须让当前 Turn 失败。
+   */
+  async finish_step(parts: SessionAssistantMessagePart[]): Promise<void> {
+    await this.enqueue_write(async () => {
+      if (!this.step_active) {
+        throw new Error("Assistant canonical step is not active");
+      }
+      const current = this.current_message();
+      const current_step_parts = current.parts.filter(
+        (part) =>
+          this.current_step_part_ids.has(part.part_id) &&
+          part.type !== "step-start",
+      );
+      const final_step_parts = parts.filter((part) => part.type !== "step-start");
+      if (current_step_parts.length !== final_step_parts.length) {
+        throw this.step_snapshot_error(
+          `part count ${current_step_parts.length} != ${final_step_parts.length}`,
+        );
+      }
+
+      const merged_parts = new Map<string, SessionAssistantMessagePart>();
+      for (let index = 0; index < current_step_parts.length; index += 1) {
+        const current_part = current_step_parts[index];
+        const final_part = final_step_parts[index];
+        merged_parts.set(
+          current_part.part_id,
+          this.merge_step_part(current_part, final_part, index),
+        );
+      }
+      await this.recorder.commit_assistant_step(
+        this.message_id,
+        current.parts.map((part) => merged_parts.get(part.part_id) || part),
+      );
+      this.reset_step_state();
+    });
+  }
+
+  /** 释放异常结束的 step 作用域并保留已经写入的 canonical Parts。 */
+  async abort_step(): Promise<void> {
+    await this.enqueue_write(async () => {
+      if (!this.step_active) return;
+      this.reset_step_state();
+    });
+  }
+
   /** 在当前 Assistant writer 的单写者队列中应用原始 chunk。 */
   private async apply_chunk_serialized(chunk: UIMessageChunk): Promise<void> {
     if (this.closed) throw new Error("Assistant Message writer is closed");
@@ -306,7 +373,7 @@ export class SessionAssistantMessageWriter {
         });
         return;
       case "source-url": {
-        const part_id = `source:${chunk.sourceId}`;
+        const part_id = `source:${this.step_index}:${chunk.sourceId}`;
         const current_part = current.parts.find((part) => part.part_id === part_id);
         const provider_metadata = to_session_provider_metadata(chunk.providerMetadata);
         await this.upsert_part({
@@ -324,7 +391,7 @@ export class SessionAssistantMessageWriter {
         return;
       }
       case "source-document": {
-        const part_id = `source:${chunk.sourceId}`;
+        const part_id = `source:${this.step_index}:${chunk.sourceId}`;
         const current_part = current.parts.find((part) => part.part_id === part_id);
         const provider_metadata = to_session_provider_metadata(chunk.providerMetadata);
         await this.upsert_part({
@@ -356,7 +423,9 @@ export class SessionAssistantMessageWriter {
           const data_id = typeof data_chunk.id === "string"
             ? data_chunk.id
             : undefined;
-          const part_id = `data:${data_id ?? generateId()}`;
+          const part_id = data_id
+            ? `data:${this.step_index}:${data_id}`
+            : `data:${generateId()}`;
           const current_part = current.parts.find((part) => part.part_id === part_id);
           await this.upsert_part({
             part_id,
@@ -375,26 +444,7 @@ export class SessionAssistantMessageWriter {
   /** 写入一个完整 Assistant part。 */
   async upsert_part(part: SessionAssistantMessagePart): Promise<void> {
     await this.recorder.update_assistant_part(this.message_id, part);
-  }
-
-  /**
-   * 用 AI SDK 最终 UIMessage 中的 Tool 快照校准流式写入结果。
-   *
-   * 关键点（中文）：最终快照只覆盖实际存在的字段；缺失 metadata 时继续保留
-   * 流式阶段已经写入的 Provider 快照。
-   */
-  async reconcile_final_tool_part(
-    part: SessionAssistantToolPart,
-  ): Promise<void> {
-    await this.enqueue_write(async () => {
-      const current = this.find_tool(part.tool_call_id);
-      await this.upsert_part({
-        ...(current || {}),
-        ...part,
-        part_id: current?.part_id || `tool:${part.tool_call_id}`,
-        sequence: current?.sequence || this.next_part_sequence(),
-      });
-    });
+    if (this.step_active) this.current_step_part_ids.add(part.part_id);
   }
 
   /** Executor 在调用 Tool 实现前写入完整输入。 */
@@ -426,6 +476,7 @@ export class SessionAssistantMessageWriter {
     const existing = current.parts.find(
       (part) =>
         part.type === "file" &&
+        (!this.step_active || this.current_step_part_ids.has(part.part_id)) &&
         part.url === input.url &&
         part.media_type === input.media_type,
     );
@@ -500,6 +551,73 @@ export class SessionAssistantMessageWriter {
     );
   }
 
+  /** 校验并合并同一位置的 canonical Part 与 step 最终快照。 */
+  private merge_step_part(
+    current_part: SessionAssistantMessagePart,
+    final_part: SessionAssistantMessagePart,
+    index: number,
+  ): SessionAssistantMessagePart {
+    if (current_part.type !== final_part.type) {
+      throw this.step_snapshot_error(
+        `part ${index + 1} type ${current_part.type} != ${final_part.type}`,
+      );
+    }
+    if (
+      (current_part.type === "text" || current_part.type === "reasoning") &&
+      (final_part.type === "text" || final_part.type === "reasoning")
+    ) {
+      if (current_part.text !== final_part.text) {
+        throw this.step_snapshot_error(`part ${index + 1} text differs`);
+      }
+    } else if (current_part.type === "tool" && final_part.type === "tool") {
+      if (current_part.tool_call_id !== final_part.tool_call_id) {
+        throw this.step_snapshot_error(`part ${index + 1} tool_call_id differs`);
+      }
+    } else if (current_part.type === "file" && final_part.type === "file") {
+      if (
+        current_part.url !== final_part.url ||
+        current_part.media_type !== final_part.media_type
+      ) {
+        throw this.step_snapshot_error(`part ${index + 1} file identity differs`);
+      }
+    } else if (current_part.type === "source" && final_part.type === "source") {
+      if (
+        current_part.source_type !== final_part.source_type ||
+        current_part.source_id !== final_part.source_id
+      ) {
+        throw this.step_snapshot_error(`part ${index + 1} source identity differs`);
+      }
+    } else if (current_part.type === "data" && final_part.type === "data") {
+      if (
+        current_part.data_type !== final_part.data_type ||
+        current_part.data_id !== final_part.data_id
+      ) {
+        throw this.step_snapshot_error(`part ${index + 1} data identity differs`);
+      }
+    }
+    return {
+      ...current_part,
+      ...final_part,
+      part_id: current_part.part_id,
+      sequence: current_part.sequence,
+    } as SessionAssistantMessagePart;
+  }
+
+  /** 构造不包含正文与工具输出的结构化 step 快照错误。 */
+  private step_snapshot_error(detail: string): Error {
+    return new Error(
+      `Assistant canonical step ${this.step_index} snapshot mismatch: ${detail}`,
+    );
+  }
+
+  /** 清理当前 step 的临时关联状态。 */
+  private reset_step_state(): void {
+    this.step_active = false;
+    this.current_step_part_ids.clear();
+    this.pending_text_parts.clear();
+    this.active_text_part_ids.clear();
+  }
+
   /** 计算下一个不可变 Part 顺序号。 */
   private next_part_sequence(): number {
     return this.current_message().parts.reduce(
@@ -531,7 +649,7 @@ export class SessionAssistantMessageWriter {
     type: "text" | "reasoning",
     chunk_id: string,
   ): string {
-    return `${type}:${chunk_id}`;
+    return `${this.step_index}:${type}:${chunk_id}`;
   }
 
   /** 在首个有效 Delta 到达时才固定文本 Part 的真实顺序。 */
@@ -603,8 +721,7 @@ export class SessionAssistantMessageWriter {
     status: "completed" | "stopped" | "failed",
   ): Promise<void> {
     if (this.closed) return;
-    this.pending_text_parts.clear();
-    this.active_text_part_ids.clear();
+    this.reset_step_state();
     await this.recorder.complete_assistant_message(this.message_id, status);
     this.closed = true;
   }

package/src/types/executor/SessionRunContext.ts

@@ -15,6 +15,9 @@ import type {
 import type {
   SessionAssistantStepCallback,
   SessionUiMessageChunkCallback,
+  SessionUiMessageStepAbortCallback,
+  SessionUiMessageStepFinishCallback,
+  SessionUiMessageStepStartCallback,
 } from "@/executor/types/SessionRun.js";
 import type { SessionUserMessageV1 } from "@/executor/types/SessionRecords.js";
 import type { FileUIPart } from "ai";
@@ -107,6 +110,29 @@ export interface SessionRunContext {
    */
   onUiMessageChunkCallback?: SessionUiMessageChunkCallback;
 
+  /**
+   * 单个模型 UI stream 开始回调。
+   *
+   * 关键点（中文）：Session writer 用它建立独立 step 作用域，确保重复 chunk id
+   * 不会跨模型调用复用同一个 canonical Part。
+   */
+  on_ui_message_step_start?: SessionUiMessageStepStartCallback;
+
+  /**
+   * 单个模型 UI stream 完成快照回调。
+   *
+   * 关键点（中文）：最终快照只校验当前 step 的顺序并补充 metadata，不能创建、
+   * 删除或重排 canonical Part。
+   */
+  on_ui_message_step_finish?: SessionUiMessageStepFinishCallback;
+
+  /**
+   * 单个模型 UI stream 异常结束回调。
+   *
+   * 关键点（中文）：用于释放 step 作用域；已经持久化的流式 Part 继续保留。
+   */
+  on_ui_message_step_abort?: SessionUiMessageStepAbortCallback;
+
   /** Tool 实现开始执行前提交完整输入的顺序屏障。 */
   on_tool_input_ready?: (input: SessionToolInputReady) => Promise<void>;