@springbrand/agent-runtime 0.2.0-alpha.72 → 0.3.0-alpha.10

package/src/conversation/harness-resume.ts

@@ -0,0 +1,288 @@
+import { HarnessClosed, type AgentLane, type LaneSnapshot, type WatchHandle } from "@earendil-works/pi-agent-core";
+import { BACKGROUND_CONTEXT } from "@earendil-works/pi-agent-core/harness/context";
+import type { OperationResultRecord } from "@earendil-works/pi-agent-core/harness/session";
+import type { Connection } from "agents";
+import { MessageType } from "agents/chat";
+import type { UIMessage, UIMessageChunk } from "ai";
+import type { ConversationHistory } from "./history";
+import { HarnessStreamProjection } from "./harness-stream";
+
+interface ActiveConversation {
+  lane: AgentLane;
+  submissionId: string;
+  requestId: string;
+  assistantMessageId: string;
+  userMessageId?: string | null;
+  createdAt: number;
+}
+interface ResumeConnection {
+  connection: Connection;
+  source: ActiveConversation;
+  ready: Promise<void>;
+  watch?: WatchHandle<LaneSnapshot>;
+  projection?: HarnessStreamProjection;
+  acknowledging: boolean;
+  replay?: Promise<void>;
+  terminal?: OperationResultRecord;
+  released?: boolean;
+  finished?: boolean;
+}
+
+
+/** CF's resume handshake, with its replay prefix and subsequent events sourced only from Pi. */
+export class HarnessConversationBridge {
+  private active?: ActiveConversation;
+  private readonly connections = new Map<string, ResumeConnection>();
+
+  constructor(private readonly host: {
+    history: Pick<ConversationHistory, "browserMessages" | "browserSnapshot">;
+    send: (connection: Connection, body: string) => void | Promise<void>;
+    broadcast: (body: string, excludedConnectionIds: string[]) => void;
+  }) {}
+
+  activate(lane: AgentLane, submission: Omit<ActiveConversation, "lane">): void {
+    this.active = { lane, ...submission };
+  }
+
+  private replayMessage(messages: UIMessage[], submissionId: string): UIMessage | undefined {
+    return messages.filter((message) => message.role === "assistant" &&
+      (message.metadata as { turnId?: string } | undefined)?.turnId === submissionId).at(-1);
+  }
+
+  private historyBeforeReplay(messages: UIMessage[], submissionId: string): UIMessage[] {
+    const target = this.replayMessage(messages, submissionId);
+    return target ? messages.filter((message) => message.id !== target.id) : messages;
+  }
+
+  /** Reads never release lifecycle state; hydration leaves the active reply to replay. */
+  async messages(view: "full" | "hydration" = "full"): Promise<UIMessage[]> {
+    for (;;) {
+      const source = this.active;
+      let watch: WatchHandle<LaneSnapshot> | undefined;
+      if (source) {
+        try { watch = await source.lane.watch(BACKGROUND_CONTEXT); }
+        catch (error) {
+          if (!(error instanceof HarnessClosed)) throw error;
+        }
+      }
+      try {
+        const messages = watch
+          ? await this.host.history.browserSnapshot(watch.snapshot)
+          : await this.host.history.browserMessages();
+        if (this.active !== source) continue;
+        return view === "hydration" && source && watch?.snapshot.operation?.id === source.submissionId
+          ? this.historyBeforeReplay(messages, source.submissionId)
+          : messages;
+      } finally { watch?.unsubscribe(); }
+    }
+  }
+
+  private async send(connection: Connection, frame: unknown): Promise<void> {
+    const session = this.connections.get(connection.id);
+    try { await this.host.send(connection, JSON.stringify(frame)); }
+    catch { if (session) this.discard(session); }
+  }
+
+  private async sendHistory(connection: Connection, messages: UIMessage[]): Promise<void> {
+    await this.send(connection, { type: MessageType.CF_AGENT_CHAT_MESSAGES, messages });
+  }
+
+  private async sendIdleHistory(connection: Connection, probeId?: string): Promise<void> {
+    const source = this.active;
+    const session = this.connections.get(connection.id);
+    const messages = await this.host.history.browserMessages();
+    if (this.connections.get(connection.id) !== session) return;
+    if (this.active !== source) return this.resume(connection, probeId);
+    await this.sendHistory(connection, messages);
+    if (this.connections.get(connection.id) !== session) return;
+    if (this.active !== source) return this.resume(connection, probeId);
+    await this.send(connection, { type: MessageType.CF_AGENT_STREAM_RESUME_NONE, ...(probeId ? { probeId } : {}) });
+  }
+
+  private async sendChunks(session: ResumeConnection, chunks: readonly UIMessageChunk[], replay = false): Promise<void> {
+    for (const chunk of chunks) {
+      if (this.connections.get(session.connection.id) !== session) return;
+      await this.send(session.connection, {
+        type: MessageType.CF_AGENT_USE_CHAT_RESPONSE,
+        id: session.source.requestId,
+        body: JSON.stringify(chunk), done: false,
+        ...(replay ? { replay: true, continuation: false } : {}),
+      });
+    }
+  }
+
+  async resume(connection: Connection, probeId?: string): Promise<void> {
+    const source = this.active;
+    if (!source) {
+      return this.sendIdleHistory(connection, probeId);
+    }
+    let session = this.connections.get(connection.id);
+    if (!session || session.source !== source) {
+      this.disconnect(connection);
+      session = { connection, source, ready: Promise.resolve(), acknowledging: false };
+      this.connections.set(connection.id, session);
+      const current = session;
+      current.ready = (async () => {
+        const watch = await source.lane.watch(BACKGROUND_CONTEXT);
+        if (this.connections.get(connection.id) !== current) { watch.unsubscribe(); return; }
+        current.watch = watch;
+      })();
+    }
+    try { await session.ready; }
+    catch (error) {
+      if (this.connections.get(connection.id) !== session) return;
+      this.discard(session);
+      if (this.active !== source) return this.resume(connection, probeId);
+      if (error instanceof HarnessClosed) {
+        await this.sendIdleHistory(connection, probeId);
+      }
+      return;
+    }
+    if (this.connections.get(connection.id) !== session) return;
+    if (this.active !== source) {
+      this.discard(session);
+      return this.resume(connection, probeId);
+    }
+    const snapshot = session.watch?.snapshot;
+    if (!snapshot?.operation || snapshot.operation.id !== source.submissionId) {
+      this.discard(session);
+      return this.sendIdleHistory(connection, probeId);
+    }
+    if (!session.acknowledging) {
+      const messages = await this.host.history.browserSnapshot(snapshot);
+      if (this.connections.get(connection.id) !== session) return;
+      if (this.active !== source) { this.discard(session); return this.resume(connection, probeId); }
+      await this.sendHistory(connection, this.historyBeforeReplay(messages, source.submissionId));
+    }
+    if (this.connections.get(connection.id) !== session) return;
+    if (this.active !== source) { this.discard(session); return this.resume(connection, probeId); }
+    await this.send(connection, {
+      type: MessageType.CF_AGENT_STREAM_RESUMING, id: source.requestId,
+      ...(probeId ? { probeId } : {}),
+    });
+  }
+
+  async acknowledge(connection: Connection, requestId: string): Promise<void> {
+    const session = this.connections.get(connection.id);
+    if (!session || session.source.requestId !== requestId || session.acknowledging) return;
+    if (this.active && this.active.submissionId !== session.source.submissionId) {
+      // A queued Submission may already be live when an older offer is ACKed.
+      // End the old transport, then atomically resume the current lane view.
+      this.disconnect(connection);
+      const current = this.resume(connection);
+      await this.send(connection, { type: MessageType.CF_AGENT_USE_CHAT_RESPONSE, id: requestId, body: "", done: true });
+      await current;
+      return;
+    }
+    session.acknowledging = true;
+    session.replay = this.replay(session);
+    try {
+      await session.replay;
+      if (session.terminal) await this.finish(session, session.terminal);
+      if (session.released) this.discard(session);
+    } catch { this.discard(session); }
+  }
+
+  private async replay(session: ResumeConnection): Promise<void> {
+    const connection = session.connection;
+    const requestId = session.source.requestId;
+    await session.ready;
+    if (this.connections.get(connection.id) !== session) return;
+    if (!session.watch && this.active?.submissionId === session.source.submissionId) {
+      session.source = this.active;
+      session.released = false;
+      const watch = await this.active.lane.watch(BACKGROUND_CONTEXT);
+      if (this.connections.get(connection.id) !== session) { watch.unsubscribe(); return; }
+      session.watch = watch;
+    }
+    const watch = session.watch;
+    const messages = watch ? await this.host.history.browserSnapshot(watch.snapshot) : await this.host.history.browserMessages();
+    if (this.connections.get(connection.id) !== session) return;
+    const source = session.source;
+    const message = this.replayMessage(messages, source.submissionId) ?? {
+      id: source.assistantMessageId, role: "assistant" as const, parts: [],
+      metadata: { turnId: source.submissionId, userMessageId: source.userMessageId,
+        createdAt: source.createdAt, turnStatus: "running" },
+    };
+    const projection = new HarnessStreamProjection({
+      messageId: message.id, startedAt: source.createdAt, turnId: source.submissionId,
+      userMessageId: source.userMessageId,
+    });
+    session.projection = projection;
+    const prefix = projection.restore(message, Boolean(watch?.snapshot.operation?.streamingMessage));
+    await this.sendChunks(session, prefix, true);
+    if (this.connections.get(connection.id) !== session) return;
+    await this.send(connection, { type: MessageType.CF_AGENT_USE_CHAT_RESPONSE,
+      id: requestId, body: "", done: false, replayComplete: true });
+    if (this.connections.get(connection.id) !== session) return;
+    watch?.start(async (event) => {
+      if ("runId" in event && event.runId !== source.submissionId) return;
+      await this.sendChunks(session, projection.event(event));
+    });
+  }
+
+  private async finish(session: ResumeConnection, record: OperationResultRecord): Promise<void> {
+    if (session.finished || !session.projection) return;
+    session.finished = true;
+    await session.watch?.drain();
+    await this.sendChunks(session, session.projection.finish(record));
+    if (this.connections.get(session.connection.id) === session) {
+      await this.send(session.connection, { type: MessageType.CF_AGENT_USE_CHAT_RESPONSE,
+        id: session.source.requestId, body: "", done: true });
+    }
+  }
+
+  /** Called only after the Host has persisted the Submission's terminal result. */
+  async settle(record: OperationResultRecord): Promise<void> {
+    await Promise.all([...this.connections.values()].filter((session) =>
+      session.source.submissionId === record.operationId).map(async (session) => {
+      session.terminal = record;
+      try {
+        if (session.replay) await session.replay;
+        await this.finish(session, record);
+      } catch { this.discard(session); }
+    }));
+    if (this.active?.submissionId === record.operationId) {
+      this.broadcastResponse({ id: this.active.requestId, body: "", done: true });
+    }
+  }
+
+  /** Called after a drive returns, outside Pi event listeners. */
+  async release(lane: AgentLane): Promise<void> {
+    if (this.active?.lane === lane) this.active = undefined;
+    await Promise.all([...this.connections.values()].filter((session) => session.source.lane === lane)
+      .map(async (session) => {
+        try {
+          await session.ready;
+          if (session.replay) await session.replay;
+          if (session.projection) {
+            await session.watch?.drain();
+            this.discard(session);
+          } else {
+            session.watch?.unsubscribe();
+            session.watch = undefined;
+            session.released = true;
+          }
+        } catch { this.discard(session); }
+      }));
+  }
+
+  disconnect(connection: Pick<Connection, "id">): void {
+    const session = this.connections.get(connection.id);
+    if (session) this.discard(session);
+  }
+
+  private discard(session: ResumeConnection): void {
+    if (this.connections.get(session.connection.id) !== session) return;
+    session.watch?.unsubscribe();
+    this.connections.delete(session.connection.id);
+  }
+
+  broadcastResponse(frame: { id: string; body: string; done: boolean; error?: boolean; continuation?: boolean }): void {
+    // Approval is a Host fact, so it has no duplicate Pi event for the watch to project.
+    const isApproval = frame.body && JSON.parse(frame.body).type === "tool-approval-request";
+    const excluded = isApproval ? [] : [...this.connections.values()]
+      .filter((session) => session.source.requestId === frame.id).map(({ connection }) => connection.id);
+    this.host.broadcast(JSON.stringify({ type: MessageType.CF_AGENT_USE_CHAT_RESPONSE, ...frame }), excluded);
+  }
+}

package/src/conversation/harness-stream.ts

@@ -0,0 +1,99 @@
+import type { AgentEvent, Events, HarnessEvent } from "@earendil-works/pi-agent-core";
+import type { OperationResultRecord } from "@earendil-works/pi-agent-core/harness/session";
+import type { AssistantMessage } from "@earendil-works/pi-ai";
+import type { UIMessage, UIMessageChunk } from "ai";
+import { PiChunkEncoder, type ToolInputPreparer } from "./message";
+
+export interface HarnessStreamOptions {
+  readonly messageId: string;
+  readonly startedAt: number;
+  readonly turnId: string;
+  readonly userMessageId?: string | null;
+  readonly prepareToolInput?: ToolInputPreparer;
+}
+
+/** Project the final Host decision without rewriting Pi's canonical operation result. */
+export function projectSubmissionOutcome(
+  record: OperationResultRecord,
+  submission: { status: string; completedAt?: number | null; error?: string | null },
+): OperationResultRecord {
+  const status = submission.status === "aborted" ? "aborted"
+    : submission.status === "error" ? "failed"
+    : submission.status === "skipped" ? "declined" : "completed";
+  const { error: _previousError, ...result } = record;
+  return {
+    ...result, status, endedAt: submission.completedAt ?? record.endedAt,
+    ...((status === "failed" || status === "aborted") && submission.error
+      ? { error: { code: status, message: submission.error } } : {}),
+  };
+}
+
+/** One projection for both broadcast events and a resumed Pi watch. */
+export class HarnessStreamProjection {
+  private readonly encoder: PiChunkEncoder;
+  private terminal: AssistantMessage | undefined;
+
+  constructor(options: HarnessStreamOptions) {
+    this.encoder = new PiChunkEncoder(options);
+  }
+
+  restore(message: UIMessage, streaming: boolean): UIMessageChunk[] {
+    return this.encoder.restore(message, streaming);
+  }
+
+  event(event: HarnessEvent): UIMessageChunk[] {
+    switch (event.type) {
+      case "run_start":
+      case "run_resume": return this.encoder.start();
+      case "turn_start":
+        this.terminal = undefined;
+        return this.encoder.encode({ type: "turn_start" } as AgentEvent);
+      case "message_update":
+        return event.event.type === "error" ? [] : this.encoder.encode({
+          type: "message_update", message: event.message, assistantMessageEvent: event.event,
+        } as AgentEvent);
+      case "tool_start": return this.encoder.encode({
+        type: "tool_execution_start", toolCallId: event.toolCallId, toolName: event.toolName, args: event.args,
+      } as AgentEvent);
+      case "tool_update": return this.encoder.encode({
+        type: "tool_execution_update", toolCallId: event.toolCallId, toolName: event.toolName, args: {}, partialResult: event.partialResult,
+      } as AgentEvent);
+      case "tool_end": return this.encoder.encode({
+        type: "tool_execution_end", toolCallId: event.toolCallId, toolName: event.toolName, result: event.result, isError: event.isError,
+      } as AgentEvent);
+      case "message_end":
+        if (event.message.role !== "assistant") return [];
+        this.terminal = event.message;
+        return this.encoder.encode({ type: "message_end", message: event.message } as AgentEvent, { terminal: false });
+      default: return [];
+    }
+  }
+
+  approvalRequested(toolCallId: string, approvalId: string): UIMessageChunk[] {
+    return this.encoder.approvalRequested(toolCallId, approvalId);
+  }
+
+  finish(result: Pick<OperationResultRecord, "status" | "endedAt" | "error">): UIMessageChunk[] {
+    return this.encoder.finish({
+      stopReason: result.status === "aborted" ? "aborted" : result.status === "failed" ? "error" :
+        this.terminal?.stopReason === "length" ? "length" : "stop",
+      timestamp: result.endedAt,
+      ...(result.error?.message ? { errorMessage: result.error.message } : {}),
+    });
+  }
+}
+
+const STREAM_EVENTS = ["run_start", "run_resume", "turn_start", "message_update", "tool_start", "tool_update", "tool_end", "message_end"] as const;
+
+export function bindHarnessStream(
+  events: Events,
+  options: HarnessStreamOptions & { readonly publish: (chunks: readonly UIMessageChunk[]) => Promise<void> },
+) {
+  const projection = new HarnessStreamProjection(options);
+  const unsubscribers = STREAM_EVENTS.map((type) => events.on(type, (event) => options.publish(projection.event(event))));
+  return {
+    publishApprovalRequest: (toolCallId: string, approvalId: string) => options.publish(projection.approvalRequested(toolCallId, approvalId)),
+    flushTerminal: (result: OperationResultRecord) => options.publish(projection.finish(result)),
+    unsubscribe: () => { for (const unsubscribe of unsubscribers) unsubscribe(); },
+  };
+}