@springbrand/agent-runtime 0.1.0

package/src/pi/message/projection.ts

@@ -0,0 +1,604 @@
+import type { AgentEvent } from "@earendil-works/pi-agent-core";
+import type {
+  AssistantMessage,
+  AssistantMessageEvent,
+  ToolCall,
+  ToolResultMessage,
+} from "@earendil-works/pi-ai";
+import type {
+  DynamicToolUIPart,
+  FinishReason,
+  UIMessage,
+  UIMessageChunk,
+} from "ai";
+import { USER_STOP_REASON } from "../../kernel/receipts";
+
+const MAX_OUTPUT_BYTES = 256 * 1024;
+const MAX_OUTPUT_PREVIEW = 16 * 1024;
+
+export interface PiToolApprovalView {
+  readonly id: string;
+  readonly approved?: boolean;
+  readonly reason?: string;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
+}
+
+function safeJson(value: unknown): unknown {
+  const seen = new WeakSet<object>();
+  let body: string | undefined;
+  try {
+    body = JSON.stringify(value, (_key, item: unknown) => {
+      if (typeof item === "bigint") return item.toString();
+      if (
+        item === undefined ||
+        typeof item === "function" ||
+        typeof item === "symbol"
+      ) return null;
+      if (item instanceof Error) {
+        return { name: item.name, message: item.message };
+      }
+      if (item !== null && typeof item === "object") {
+        if (seen.has(item)) return "[Circular]";
+        seen.add(item);
+      }
+      return item;
+    });
+  } catch (error) {
+    return {
+      serializationError:
+        error instanceof Error ? error.message : String(error),
+    };
+  }
+  if (body === undefined) return null;
+  const bytes = new TextEncoder().encode(body).byteLength;
+  if (bytes > MAX_OUTPUT_BYTES) {
+    return {
+      truncated: true,
+      originalBytes: bytes,
+      preview: body.slice(0, MAX_OUTPUT_PREVIEW),
+    };
+  }
+  try {
+    return JSON.parse(body) as unknown;
+  } catch {
+    return null;
+  }
+}
+
+function safeInput(value: unknown): Record<string, unknown> {
+  const safe = safeJson(value);
+  return isRecord(safe) ? safe : {};
+}
+
+function contentText(result: unknown): string | undefined {
+  if (!isRecord(result) || !Array.isArray(result.content)) return undefined;
+  const text = result.content.flatMap((part) =>
+    isRecord(part) && part.type === "text" && typeof part.text === "string"
+      ? [part.text]
+      : []
+  ).join("");
+  return text || undefined;
+}
+
+function toolOutput(result: unknown): unknown {
+  if (!isRecord(result)) return safeJson(result);
+  if (result.details !== undefined) return safeJson(result.details);
+  return safeJson(contentText(result) ?? result.content ?? result);
+}
+
+function errorText(result: unknown): string {
+  const text = contentText(result);
+  if (text) return text;
+  const output = toolOutput(result);
+  if (typeof output === "string") return output;
+  try {
+    return JSON.stringify(output ?? "Tool execution failed");
+  } catch {
+    return "Tool execution failed";
+  }
+}
+
+function toolCallAt(
+  update: AssistantMessageEvent,
+): Partial<ToolCall> | undefined {
+  if (!("contentIndex" in update) || !("partial" in update)) return undefined;
+  const part = update.partial.content[update.contentIndex];
+  return part?.type === "toolCall" ? part : undefined;
+}
+
+function finishReason(reason: string): FinishReason {
+  switch (reason) {
+    case "stop":
+    case "length":
+      return reason;
+    case "toolUse":
+      return "tool-calls";
+    case "error":
+      return "error";
+    default:
+      return "other";
+  }
+}
+
+interface PendingToolInput {
+  toolName?: string;
+  deltas: string[];
+}
+
+interface StreamingToolInput extends PendingToolInput {
+  started: boolean;
+}
+
+/** Stateful Pi AgentEvent to official AI SDK 7 UIMessageChunk encoder. */
+export class PiChunkEncoder {
+  private readonly messageId: string;
+  private readonly startedAt?: number;
+  private started = false;
+  private finished = false;
+  private firstTurnStarted = false;
+  private stepIndex = 0;
+  private readonly activeText = new Set<number>();
+  private readonly activeReasoning = new Set<number>();
+  private readonly pendingToolByIndex = new Map<number, PendingToolInput>();
+  private readonly toolCallIdByIndex = new Map<number, string>();
+  private readonly streamingToolById = new Map<string, StreamingToolInput>();
+  private readonly toolInputs = new Map<
+    string,
+    { toolName: string; input: Record<string, unknown> }
+  >();
+
+  constructor(options: { messageId: string; startedAt?: number }) {
+    this.messageId = options.messageId;
+    this.startedAt = options.startedAt;
+  }
+
+  start(): UIMessageChunk[] {
+    if (this.started) return [];
+    this.started = true;
+    this.firstTurnStarted = true;
+    return [
+      {
+        type: "start",
+        messageId: this.messageId,
+        ...(this.startedAt === undefined
+          ? {}
+          : {
+              messageMetadata: {
+                createdAt: this.startedAt,
+                turnStatus: "running",
+              },
+            }),
+      },
+      { type: "start-step" },
+    ];
+  }
+
+  toolInput(
+    toolCallId: string,
+    toolName: string,
+    input: unknown,
+  ): UIMessageChunk[] {
+    const safe = safeInput(input);
+    this.toolInputs.set(toolCallId, { toolName, input: safe });
+    return [{
+      type: "tool-input-available",
+      toolCallId,
+      toolName,
+      input: safe,
+      dynamic: true,
+    }];
+  }
+
+  approvalRequested(
+    toolCallId: string,
+    approvalId: string,
+  ): UIMessageChunk[] {
+    return [{ type: "tool-approval-request", toolCallId, approvalId }];
+  }
+
+  approvalResponded(
+    toolCallId: string,
+    approvalId: string,
+    approved: boolean,
+    reason?: string,
+  ): UIMessageChunk[] {
+    return [
+      {
+        type: "tool-approval-response",
+        approvalId,
+        approved,
+        ...(reason ? { reason } : {}),
+      },
+      ...(approved
+        ? []
+        : [{ type: "tool-output-denied", toolCallId } as const]),
+    ];
+  }
+
+  encode(event: AgentEvent): UIMessageChunk[] {
+    if (this.finished) return [];
+    switch (event.type) {
+      case "agent_start":
+        return this.start();
+      case "turn_start":
+        if (!this.started) return this.start();
+        if (this.firstTurnStarted) {
+          this.firstTurnStarted = false;
+          return [];
+        }
+        this.stepIndex += 1;
+        return [{ type: "start-step" }];
+      case "message_update":
+        this.firstTurnStarted = false;
+        return this.encodeMessageUpdate(event.assistantMessageEvent);
+      case "tool_execution_start":
+        return this.toolInputs.has(event.toolCallId)
+          ? []
+          : this.toolInput(event.toolCallId, event.toolName, event.args);
+      case "tool_execution_update": {
+        const chunks = this.toolInputs.has(event.toolCallId)
+          ? []
+          : this.toolInput(event.toolCallId, event.toolName, event.args);
+        chunks.push({
+          type: "tool-output-available",
+          toolCallId: event.toolCallId,
+          output: toolOutput(event.partialResult),
+          preliminary: true,
+          dynamic: true,
+        });
+        return chunks;
+      }
+      case "tool_execution_end": {
+        const chunks = this.toolInputs.has(event.toolCallId)
+          ? []
+          : this.toolInput(event.toolCallId, event.toolName, {});
+        chunks.push(event.isError
+          ? {
+              type: "tool-output-error",
+              toolCallId: event.toolCallId,
+              errorText: errorText(event.result),
+              dynamic: true,
+            }
+          : {
+              type: "tool-output-available",
+              toolCallId: event.toolCallId,
+              output: toolOutput(event.result),
+              preliminary: false,
+              dynamic: true,
+            });
+        return chunks;
+      }
+      case "message_end":
+        this.firstTurnStarted = false;
+        return event.message.role === "assistant"
+          ? this.endAssistant(event.message)
+          : [];
+      case "agent_end":
+      case "turn_end":
+      case "message_start":
+        return [];
+      default:
+        throw new Error(
+          `Unsupported Pi AgentEvent: ${String((event as { type?: unknown }).type)}`,
+        );
+    }
+  }
+
+  private encodeMessageUpdate(update: AssistantMessageEvent): UIMessageChunk[] {
+    switch (update.type) {
+      case "text_start":
+        if (this.activeText.has(update.contentIndex)) return [];
+        this.activeText.add(update.contentIndex);
+        return [{ type: "text-start", id: this.partId("text", update.contentIndex) }];
+      case "text_delta":
+        return [
+          ...this.ensureText(update.contentIndex),
+          { type: "text-delta", id: this.partId("text", update.contentIndex), delta: update.delta },
+        ];
+      case "text_end":
+        if (!this.activeText.delete(update.contentIndex)) return [];
+        return [{ type: "text-end", id: this.partId("text", update.contentIndex) }];
+      case "thinking_start":
+        if (this.activeReasoning.has(update.contentIndex)) return [];
+        this.activeReasoning.add(update.contentIndex);
+        return [{ type: "reasoning-start", id: this.partId("reasoning", update.contentIndex) }];
+      case "thinking_delta":
+        return [
+          ...this.ensureReasoning(update.contentIndex),
+          { type: "reasoning-delta", id: this.partId("reasoning", update.contentIndex), delta: update.delta },
+        ];
+      case "thinking_end":
+        if (!this.activeReasoning.delete(update.contentIndex)) return [];
+        return [{ type: "reasoning-end", id: this.partId("reasoning", update.contentIndex) }];
+      case "toolcall_start":
+      case "toolcall_delta":
+      case "toolcall_end":
+        return this.encodeToolInput(update);
+      case "error":
+        return [{ type: "error", errorText: update.error.errorMessage ?? update.reason }];
+      default:
+        return [];
+    }
+  }
+
+  private encodeToolInput(
+    update: Extract<
+      AssistantMessageEvent,
+      { type: "toolcall_start" | "toolcall_delta" | "toolcall_end" }
+    >,
+  ): UIMessageChunk[] {
+    const index = update.contentIndex;
+    if (update.type === "toolcall_start") {
+      this.pendingToolByIndex.set(index, { deltas: [] });
+      this.toolCallIdByIndex.delete(index);
+    }
+
+    const partial = toolCallAt(update);
+    const observedToolCallId = update.type === "toolcall_end"
+      ? update.toolCall.id
+      : partial?.id;
+    const previousToolCallId = this.toolCallIdByIndex.get(index);
+    let pending = this.pendingToolByIndex.get(index) ?? { deltas: [] };
+    if (
+      observedToolCallId &&
+      previousToolCallId &&
+      observedToolCallId !== previousToolCallId
+    ) {
+      pending = { deltas: [] };
+    }
+    const toolName = update.type === "toolcall_end"
+      ? update.toolCall.name
+      : partial?.name;
+    if (toolName) pending.toolName = toolName;
+    if (update.type === "toolcall_delta") pending.deltas.push(update.delta);
+    this.pendingToolByIndex.set(index, pending);
+
+    const toolCallId = observedToolCallId ?? previousToolCallId;
+    if (!toolCallId) return [];
+    this.toolCallIdByIndex.set(index, toolCallId);
+
+    const streaming = this.streamingToolById.get(toolCallId) ?? {
+      deltas: [],
+      started: false,
+    };
+    if (pending.toolName) streaming.toolName = pending.toolName;
+    streaming.deltas.push(...pending.deltas);
+    pending.deltas.length = 0;
+    this.streamingToolById.set(toolCallId, streaming);
+    if (!streaming.toolName) return [];
+
+    const chunks: UIMessageChunk[] = [];
+    const completedOnly = update.type === "toolcall_end" &&
+      !streaming.started &&
+      streaming.deltas.length === 0;
+    if (!streaming.started && !completedOnly) {
+      streaming.started = true;
+      chunks.push({
+        type: "tool-input-start",
+        toolCallId,
+        toolName: streaming.toolName,
+        dynamic: true,
+      });
+    }
+    if (streaming.started) {
+      for (const delta of streaming.deltas) {
+        chunks.push({
+          type: "tool-input-delta",
+          toolCallId,
+          inputTextDelta: delta,
+        });
+      }
+      streaming.deltas.length = 0;
+    }
+    if (update.type === "toolcall_end") {
+      const input = safeInput(update.toolCall.arguments);
+      this.toolInputs.set(toolCallId, {
+        toolName: update.toolCall.name,
+        input,
+      });
+      chunks.push({
+        type: "tool-input-available",
+        toolCallId,
+        toolName: update.toolCall.name,
+        input,
+        dynamic: true,
+      });
+      this.pendingToolByIndex.delete(index);
+      this.toolCallIdByIndex.delete(index);
+    }
+    return chunks;
+  }
+
+  private endAssistant(message: AssistantMessage): UIMessageChunk[] {
+    this.pendingToolByIndex.clear();
+    this.toolCallIdByIndex.clear();
+    const chunks = this.closeOpenParts();
+    if (message.stopReason === "toolUse") {
+      chunks.push({ type: "finish-step" });
+      return chunks;
+    }
+    this.finished = true;
+    if (this.startedAt !== undefined) {
+      const turnStatus = message.stopReason === "aborted"
+        ? "aborted"
+        : message.stopReason === "error"
+          ? "error"
+          : "completed";
+      chunks.push({
+        type: "message-metadata",
+        messageMetadata: {
+          createdAt: this.startedAt,
+          completedAt: message.timestamp,
+          turnDurationMs: Math.max(0, message.timestamp - this.startedAt),
+          turnStatus,
+          ...(message.errorMessage &&
+              (message.stopReason !== "aborted" ||
+                message.errorMessage !== USER_STOP_REASON)
+            ? { error: message.errorMessage }
+            : {}),
+        },
+      });
+    }
+    if (message.stopReason === "aborted") {
+      chunks.push({ type: "abort", reason: message.errorMessage });
+      return chunks;
+    }
+    if (message.stopReason === "error") {
+      chunks.push({
+        type: "error",
+        errorText: message.errorMessage ?? "Pi turn failed",
+      });
+      return chunks;
+    }
+    chunks.push({ type: "finish-step" });
+    chunks.push({
+      type: "finish",
+      finishReason: finishReason(message.stopReason),
+    });
+    return chunks;
+  }
+
+  private closeOpenParts(): UIMessageChunk[] {
+    const chunks: UIMessageChunk[] = [];
+    for (const index of this.activeText) {
+      chunks.push({ type: "text-end", id: this.partId("text", index) });
+    }
+    for (const index of this.activeReasoning) {
+      chunks.push({ type: "reasoning-end", id: this.partId("reasoning", index) });
+    }
+    this.activeText.clear();
+    this.activeReasoning.clear();
+    return chunks;
+  }
+
+  private ensureText(contentIndex: number): UIMessageChunk[] {
+    if (this.activeText.has(contentIndex)) return [];
+    this.activeText.add(contentIndex);
+    return [{ type: "text-start", id: this.partId("text", contentIndex) }];
+  }
+
+  private ensureReasoning(contentIndex: number): UIMessageChunk[] {
+    if (this.activeReasoning.has(contentIndex)) return [];
+    this.activeReasoning.add(contentIndex);
+    return [{ type: "reasoning-start", id: this.partId("reasoning", contentIndex) }];
+  }
+
+  private partId(kind: "text" | "reasoning", index: number): string {
+    return `${this.messageId}:${kind}:${this.stepIndex}:${index}`;
+  }
+}
+
+function projectedTool(
+  call: ToolCall,
+  approval?: PiToolApprovalView,
+): DynamicToolUIPart {
+  const input = safeInput(call.arguments);
+  const base = {
+    type: "dynamic-tool" as const,
+    toolCallId: call.id,
+    toolName: call.name,
+    input,
+  };
+  if (!approval) return { ...base, state: "input-available" };
+  if (approval.approved === false) {
+    return {
+      ...base,
+      state: "output-denied",
+      approval: {
+        id: approval.id,
+        approved: false,
+        ...(approval.reason ? { reason: approval.reason } : {}),
+      },
+    };
+  }
+  if (approval.approved === true) {
+    return {
+      ...base,
+      state: "approval-responded",
+      approval: {
+        id: approval.id,
+        approved: true,
+        ...(approval.reason ? { reason: approval.reason } : {}),
+      },
+    };
+  }
+  return {
+    ...base,
+    state: "approval-requested",
+    approval: { id: approval.id },
+  };
+}
+
+/** Project one canonical assistant message without mutating model context. */
+export function piAssistantToUIMessage(
+  message: AssistantMessage,
+  options: {
+    readonly id: string;
+    readonly metadata?: unknown;
+    readonly approvals?: Readonly<Record<string, PiToolApprovalView | undefined>>;
+  },
+): UIMessage {
+  const parts = message.content.flatMap(
+    (content): Array<UIMessage["parts"][number]> => {
+    if (content.type === "text") {
+      return [{ type: "text", text: content.text, state: "done" }];
+    }
+    if (content.type === "thinking") {
+      return [{ type: "reasoning", text: content.thinking, state: "done" }];
+    }
+    if (content.type === "toolCall") {
+      return [projectedTool(content, options.approvals?.[content.id])];
+    }
+    return [];
+    },
+  );
+  return {
+    id: options.id,
+    role: "assistant",
+    parts,
+    ...(options.metadata === undefined ? {} : { metadata: options.metadata }),
+  };
+}
+
+/** Merge a canonical ToolResult onto its single dynamic Tool UI part. */
+export function applyPiToolResult(
+  message: UIMessage,
+  result: ToolResultMessage,
+): UIMessage {
+  return {
+    ...message,
+    parts: message.parts.map((part): UIMessage["parts"][number] => {
+      if (
+        part.type !== "dynamic-tool" ||
+        part.toolCallId !== result.toolCallId ||
+        part.state === "output-denied"
+      ) return part;
+      const approval = "approval" in part && part.approval?.approved === true
+        ? { ...part.approval, approved: true as const }
+        : undefined;
+      return result.isError
+        ? {
+            type: "dynamic-tool",
+            toolCallId: part.toolCallId,
+            toolName: part.toolName,
+            state: "output-error",
+            input: part.input,
+            errorText: errorText(result),
+            ...(approval ? { approval } : {}),
+          }
+        : {
+            type: "dynamic-tool",
+            toolCallId: part.toolCallId,
+            toolName: part.toolName,
+            state: "output-available",
+            input: part.input,
+            output: toolOutput({ content: result.content, details: result.details }),
+            preliminary: false,
+            ...(approval ? { approval } : {}),
+          };
+    }),
+  };
+}