pum-agent 0.1.0-beta.3

package/src/transcript.tsx

@@ -0,0 +1,393 @@
+import { StyledText, fg, type SyntaxStyle } from "@opentui/core";
+import {
+  useBlinkingText,
+  useMarkdownCaret,
+  useShimmerText,
+  useSpinner,
+} from "./animation";
+import type { Theme } from "./theme";
+import type { ToolCall } from "./tool-line";
+
+export type Role = "user" | "assistant" | "thinking" | "system" | "error";
+
+export type Line =
+  | { kind: "text"; role: Role; text: string }
+  | { kind: "tool"; call: ToolCall }
+  | { kind: "agent-message"; sender: string; recipient: string; text: string };
+
+export type PendingLine = {
+  id: string;
+  line: Extract<Line, { kind: "text" | "agent-message" }>;
+  /** Text used to match pi's message_start event. */
+  deliveryText?: string;
+  /** Pi inserted the message, but the active streamed message must finish first. */
+  delivered?: boolean;
+};
+
+export type PendingTranscriptState = {
+  lines: Line[];
+  stream: { kind: "assistant" | "thinking"; text: string } | null;
+  pending: PendingLine[];
+};
+
+/** Resolve a delivered message without splitting the active streamed output. */
+export function resolvePendingDelivery<T extends PendingTranscriptState>(value: T, id: string): T {
+  const pending = value.pending.find((item) => item.id === id);
+  if (!pending) return value;
+  if (value.stream) {
+    return {
+      ...value,
+      pending: value.pending.map((item) => item.id === id ? { ...item, delivered: true } : item),
+    };
+  }
+  return {
+    ...value,
+    lines: [...value.lines, pending.line],
+    pending: value.pending.filter((item) => item.id !== id),
+  };
+}
+
+/** Finish the stream, then insert messages that arrived while it was active. */
+export function settleTranscriptMessage<T extends PendingTranscriptState>(value: T): T {
+  const lines = [...value.lines];
+  if (value.stream?.text.trim()) {
+    lines.push({ kind: "text", role: value.stream.kind, text: value.stream.text.trim() });
+  }
+  const delivered = value.pending.filter((item) => item.delivered);
+  return {
+    ...value,
+    lines: [...lines, ...delivered.map((item) => item.line)],
+    stream: null,
+    pending: value.pending.filter((item) => !item.delivered),
+  };
+}
+
+type LineGroup = "tool" | "thinking" | "other";
+
+const lineGroup = (line: Line): LineGroup => {
+  if (line.kind === "tool") return "tool";
+  return line.kind === "text" && line.role === "thinking" ? "thinking" : "other";
+};
+
+/** Exactly one gap around tool/thinking groups, but none inside either group. */
+export function needsTranscriptGap(prev: Line | undefined, line: Line): boolean {
+  if (!prev) return false;
+  const prevGroup = lineGroup(prev);
+  const group = lineGroup(line);
+
+  if (prevGroup !== group && (prevGroup !== "other" || group !== "other")) return true;
+  if (group !== "other") return false;
+
+  // Preserve the normal turn layout for rows outside tool/thinking groups.
+  const isUser = line.kind === "text" && line.role === "user";
+  const prevIsUser = prev.kind === "text" && prev.role === "user";
+  const isAnswer = line.kind === "text" && line.role === "assistant";
+  const isAgentMessage = line.kind === "agent-message" || prev.kind === "agent-message";
+  return isUser || prevIsUser || isAnswer || isAgentMessage;
+}
+
+const GUTTER = "  ";
+const PROMPT = "❯ ";
+
+/** Remove provider trace wrappers and compact adjacent thinking entries. */
+export function normalizeThinkingText(text: string): string {
+  return text
+    .replaceAll("**", "")
+    .replace(/\n[ \t]*\n+/g, "\n");
+}
+
+export function roleColor(theme: Theme, role: Role): string {
+  switch (role) {
+    case "user":
+      return theme.user;
+    case "assistant":
+      return theme.assistant;
+    case "thinking":
+      return theme.thinking;
+    case "system":
+      return theme.dim;
+    case "error":
+      return theme.error;
+  }
+}
+
+/**
+ * A one-row glyph cell beside a growing text cell. The text wraps inside its
+ * own narrower column, so every wrapped row lines up under the first and the
+ * glyph appears only once. The background comes from the box, because a text's
+ * own `bg` paints glyph cells only.
+ */
+function Row({
+  glyph,
+  glyphColor,
+  background,
+  children,
+}: {
+  glyph: string;
+  glyphColor: string;
+  background?: string;
+  children: React.ReactNode;
+}) {
+  return (
+    <box
+      style={{
+        flexDirection: "row",
+        width: "100%",
+        backgroundColor: background ?? "transparent",
+      }}
+    >
+      {/* A numeric width pins the gutter: a whitespace-only <text> measures
+          inconsistently once the message column wraps, losing a column. */}
+      <box style={{ width: 2, flexShrink: 0 }}>
+        {glyph.trim() ? <text content={glyph} fg={glyphColor} /> : null}
+      </box>
+      {/* The nested flex item gives every transcript type the same measured
+          remaining-width column as the tool-row body. */}
+      <box style={{ flexDirection: "row", flexGrow: 1, flexShrink: 1, minWidth: 0 }}>
+        {children}
+      </box>
+    </box>
+  );
+}
+
+export function TextLine({
+  theme,
+  syntaxStyle,
+  role,
+  text,
+  workingCaret = false,
+}: {
+  theme: Theme;
+  syntaxStyle: SyntaxStyle;
+  role: Role;
+  text: string;
+  workingCaret?: boolean;
+}) {
+  const color = roleColor(theme, role);
+  const isUser = role === "user";
+  const isAssistant = role === "assistant";
+  const displayText = role === "thinking" ? normalizeThinkingText(text) : text;
+  const textCaret = useBlinkingText({
+    chunks: [fg(color)(displayText)],
+    contentKey: `${role}:${displayText}`,
+    caretColor: color,
+    active: workingCaret && !isAssistant,
+  });
+  const markdownCaret = useMarkdownCaret(text, workingCaret && isAssistant);
+
+  // Assistant text is already styled while streaming, so settlement only
+  // finalizes Markdown parsing and does not swap through a plain-text phase.
+  // User messages are static Markdown. They never get a streaming caret.
+  if (isAssistant || isUser) {
+    return (
+      <Row
+        glyph={isUser ? PROMPT : GUTTER}
+        glyphColor={color}
+        background={isUser ? theme.userBg : undefined}
+      >
+        <markdown
+          ref={isAssistant && workingCaret ? markdownCaret : undefined}
+          content={isAssistant && workingCaret ? undefined : text}
+          streaming={false}
+          syntaxStyle={syntaxStyle}
+          fg={color}
+          style={{ flexGrow: 1, flexShrink: 1, minWidth: 0, width: "100%" }}
+        />
+      </Row>
+    );
+  }
+
+  return (
+    <Row glyph={GUTTER} glyphColor={color}>
+      <text
+        ref={workingCaret ? textCaret : undefined}
+        content={workingCaret ? undefined : displayText}
+        fg={color}
+        wrapMode="word"
+        style={{ flexGrow: 1, flexShrink: 1, minWidth: 0, width: "100%" }}
+      />
+    </Row>
+  );
+}
+
+/** The line currently streaming in: shimmered, with a caret riding the end. */
+export function StreamLine({
+  theme,
+  syntaxStyle,
+  role,
+  text,
+}: {
+  theme: Theme;
+  syntaxStyle: SyntaxStyle;
+  role: "assistant" | "thinking";
+  text: string;
+}) {
+  const color = roleColor(theme, role);
+  const displayText = role === "thinking" ? normalizeThinkingText(text) : text;
+  const shimmer = useShimmerText({
+    text: displayText,
+    color,
+    highlight: theme.highlight,
+    active: role === "thinking",
+    caret: true,
+  });
+  const markdown = useMarkdownCaret(text, role === "assistant");
+
+  return (
+    <Row glyph={GUTTER} glyphColor={color}>
+      {role === "assistant" ? (
+        <markdown
+          ref={markdown}
+          streaming
+          syntaxStyle={syntaxStyle}
+          fg={color}
+          style={{ flexGrow: 1, flexShrink: 1, minWidth: 0 }}
+        />
+      ) : (
+        <text
+          ref={shimmer}
+          wrapMode="word"
+          style={{ flexGrow: 1, flexShrink: 1, minWidth: 0, width: "100%" }}
+        />
+      )}
+    </Row>
+  );
+}
+
+export function PendingMessageLine({
+  theme,
+  syntaxStyle,
+  pending,
+}: {
+  theme: Theme;
+  syntaxStyle: SyntaxStyle;
+  pending: PendingLine;
+}) {
+  const line = pending.line;
+  if (line.kind === "agent-message") {
+    return (
+      <Row glyph="◇ " glyphColor={theme.dim} background={theme.agentMessageBg}>
+        <box style={{ flexDirection: "column", flexGrow: 1, flexShrink: 1, minWidth: 0 }}>
+          <text
+            content={`${line.sender} → ${line.recipient} · queued`}
+            fg={theme.dim}
+            wrapMode="word"
+            style={{ width: "100%", flexShrink: 1, minWidth: 0 }}
+          />
+          <markdown
+            content={line.text}
+            streaming={false}
+            syntaxStyle={syntaxStyle}
+            fg={theme.dim}
+            style={{ width: "100%", flexGrow: 1, flexShrink: 1, minWidth: 0 }}
+          />
+        </box>
+      </Row>
+    );
+  }
+
+  return (
+    <Row glyph="○ " glyphColor={theme.dim} background={theme.userBg}>
+      <markdown
+        content={line.text}
+        streaming={false}
+        syntaxStyle={syntaxStyle}
+        fg={theme.dim}
+        style={{ flexGrow: 1, flexShrink: 1, minWidth: 0, width: "100%" }}
+      />
+    </Row>
+  );
+}
+
+export function AgentMessageLine({
+  theme,
+  syntaxStyle,
+  line,
+}: {
+  theme: Theme;
+  syntaxStyle: SyntaxStyle;
+  line: Extract<Line, { kind: "agent-message" }>;
+}) {
+  return (
+    <Row glyph="◇ " glyphColor={theme.agentMessage} background={theme.agentMessageBg}>
+      <box style={{ flexDirection: "column", flexGrow: 1, flexShrink: 1, minWidth: 0 }}>
+        <text
+          content={`${line.sender} → ${line.recipient}`}
+          fg={theme.agentMessage}
+          wrapMode="word"
+          style={{ width: "100%", flexShrink: 1, minWidth: 0 }}
+        />
+        <markdown
+          content={line.text}
+          streaming={false}
+          syntaxStyle={syntaxStyle}
+          fg={theme.agentMessage}
+          style={{ width: "100%", flexGrow: 1, flexShrink: 1, minWidth: 0 }}
+        />
+      </box>
+    </Row>
+  );
+}
+
+export function toolStateGlyph(state: ToolCall["state"]): string {
+  if (state === "ok") return "✓";
+  if (state === "rejected") return "!";
+  return "✗";
+}
+
+export function ToolLine({
+  theme,
+  call,
+  workingCaret = false,
+}: {
+  theme: Theme;
+  call: ToolCall;
+  workingCaret?: boolean;
+}) {
+  const spinner = useSpinner(call.state === "running");
+  const failed = call.state === "error";
+  const rejected = call.state === "rejected";
+  const toolColor = failed ? theme.error : rejected ? theme.warn : theme.tool;
+  const argColor = failed ? theme.error : rejected ? theme.warn : theme.toolArg;
+  const detailColor = failed ? theme.error : rejected ? theme.warn : theme.dim;
+
+  const prefix = call.arg
+    ? new StyledText([fg(toolColor)(call.name), fg(detailColor)(" · ")])
+    : null;
+  const bodyChunks = call.arg
+    ? [fg(argColor)(call.arg)]
+    : [fg(toolColor)(call.name)];
+  if (call.detail) bodyChunks.push(fg(detailColor)(`  ${call.detail}`));
+
+  const caret = useBlinkingText({
+    chunks: bodyChunks,
+    contentKey: `${call.name}:${call.arg}:${call.state}:${call.detail ?? ""}`,
+    caretColor: failed ? theme.error : rejected ? theme.warn : theme.accent,
+    active: workingCaret,
+  });
+
+  return (
+    <Row glyph={GUTTER} glyphColor={toolColor}>
+      <box style={{ flexDirection: "row", flexGrow: 1, flexShrink: 1, minWidth: 0 }}>
+        {prefix ? <text content={prefix} style={{ flexShrink: 0 }} /> : null}
+        <text
+          ref={workingCaret ? caret : undefined}
+          content={workingCaret ? undefined : new StyledText(bodyChunks)}
+          wrapMode="word"
+          style={{ flexGrow: 1, flexShrink: 1, minWidth: 0 }}
+        />
+      </box>
+      <box style={{ width: 1, flexShrink: 0 }} />
+      <box style={{ width: 1, flexShrink: 0 }}>
+        {call.state === "running" ? (
+          <text ref={spinner} fg={theme.accent} />
+        ) : (
+          <text
+            content={toolStateGlyph(call.state)}
+            fg={call.state === "ok" ? theme.success : call.state === "rejected" ? theme.warn : theme.error}
+          />
+        )}
+      </box>
+    </Row>
+  );
+}

package/src/web-search.ts

@@ -0,0 +1,157 @@
+import type { Provider } from "@earendil-works/pi-ai";
+import type { ModelRuntime, SessionManager } from "@earendil-works/pi-coding-agent";
+import { AsyncLocalStorage } from "node:async_hooks";
+
+/**
+ * Web search through the Codex subscription.
+ *
+ * OpenAI's `web_search` is a hosted tool. The model calls it, OpenAI runs it,
+ * and the answer comes back already informed by the results.
+ */
+const HOSTED_SEARCH_PROVIDERS = ["openai-codex"];
+
+/** Session entry type used for searches, which pi does not represent as tools. */
+export const WEB_SEARCH_CUSTOM_TYPE = "pum.web_search";
+
+/** Mutable so the Ctrl+P toggle takes effect without rebuilding the provider. */
+export const webSearch = { enabled: false };
+
+const searchRoute = new AsyncLocalStorage<string>();
+let socketObserverInstalled = false;
+
+function addSearchTool(payload: unknown): unknown | undefined {
+  if (!webSearch.enabled || !payload || typeof payload !== "object") return undefined;
+  const body = payload as { tools?: unknown[] };
+  const tools = Array.isArray(body.tools) ? body.tools : [];
+  if (tools.some((t) => (t as { type?: string })?.type === "web_search")) return undefined;
+  return { ...body, tools: [...tools, { type: "web_search" }] };
+}
+
+/** Delegates through the prototype so the provider's other members survive. */
+function wrapProvider(base: Provider): Provider {
+  const wrapped: Provider = Object.create(base);
+  wrapped.stream = ((model: any, context: any, options: any) =>
+    base.stream(model, context, { ...options, onPayload: addSearchTool })) as Provider["stream"];
+  wrapped.streamSimple = ((model: any, context: any, options: any) =>
+    base.streamSimple(model, context, {
+      ...options,
+      onPayload: addSearchTool,
+    })) as Provider["streamSimple"];
+  return wrapped;
+}
+
+export type SearchCall =
+  | { phase: "start"; id: string; query: string }
+  | { phase: "end"; id: string; query: string; ok: boolean };
+
+export type SearchCallRecord = {
+  id: string;
+  query: string;
+  state: "running" | "ok" | "error";
+};
+
+export class SearchCallRouter {
+  private readonly listeners = new Map<string, Set<(call: SearchCall) => void>>();
+
+  subscribe(route: string, listener: (call: SearchCall) => void): () => void {
+    const listeners = this.listeners.get(route) ?? new Set();
+    listeners.add(listener);
+    this.listeners.set(route, listeners);
+    return () => {
+      listeners.delete(listener);
+      if (listeners.size === 0) this.listeners.delete(route);
+    };
+  }
+
+  emit(route: string, call: SearchCall): void {
+    for (const listener of this.listeners.get(route) ?? []) listener(call);
+  }
+}
+
+const searchCalls = new SearchCallRouter();
+
+/** Keep the route active through the asynchronous agent and provider call chain. */
+export function withSearchRoute<T>(route: string, operation: () => T): T {
+  return searchRoute.run(route, operation);
+}
+
+/** Subscribe one transcript to searches from only its own agent session. */
+export function observeSearchCalls(
+  route: string,
+  onCall: (call: SearchCall) => void,
+): () => void {
+  installSocketObserver();
+  return searchCalls.subscribe(route, onCall);
+}
+
+/** Persist an out-of-band search as session metadata, not LLM context. */
+export function persistSearchCall(
+  sessionManager: Pick<SessionManager, "appendCustomEntry">,
+  call: SearchCall,
+): void {
+  const record: SearchCallRecord = {
+    id: call.id,
+    query: call.query,
+    state: call.phase === "start" ? "running" : call.ok ? "ok" : "error",
+  };
+  sessionManager.appendCustomEntry(WEB_SEARCH_CUSTOM_TYPE, record);
+}
+
+/**
+ * pi drops `web_search_call` items. Observe WebSocket frames without changing
+ * the cached WebSocket transport. Each socket captures its current session route.
+ */
+function installSocketObserver(): void {
+  if (socketObserverInstalled) return;
+  const Original = globalThis.WebSocket;
+  if (!Original || (Original as { __pumPatched?: boolean }).__pumPatched) return;
+  socketObserverInstalled = true;
+
+  const Patched = new Proxy(Original, {
+    construct(target, args: any[]) {
+      const route = searchRoute.getStore();
+      const socket = new (target as any)(...args);
+      const seen = new Map<string, string>();
+      socket.addEventListener?.("message", (ev: any) => {
+        try {
+          const raw = ev?.data;
+          if (!route || typeof raw !== "string" || !raw.includes("web_search_call")) return;
+          const event = JSON.parse(raw);
+          const item = event?.item;
+          if (item?.type !== "web_search_call") return;
+          const id = String(item.id ?? "");
+          const query = String(item.action?.query ?? seen.get(id) ?? "");
+          if (query) seen.set(id, query);
+          if (event.type === "response.output_item.added") {
+            searchCalls.emit(route, { phase: "start", id, query });
+          } else if (event.type === "response.output_item.done") {
+            searchCalls.emit(route, {
+              phase: "end",
+              id,
+              query,
+              ok: item.status !== "failed",
+            });
+            seen.delete(id);
+          }
+        } catch {
+          // Search observation must never break an agent turn.
+        }
+      });
+      return socket;
+    },
+  });
+  (Patched as unknown as { __pumPatched: boolean }).__pumPatched = true;
+  globalThis.WebSocket = Patched as typeof WebSocket;
+}
+
+/** Returns the provider ids that now carry the hosted search tool. */
+export function installWebSearch(runtime: ModelRuntime): string[] {
+  const installed: string[] = [];
+  for (const id of HOSTED_SEARCH_PROVIDERS) {
+    const base = runtime.getProvider(id);
+    if (!base) continue;
+    runtime.registerNativeProvider(wrapProvider(base));
+    installed.push(id);
+  }
+  return installed;
+}

package/src/worktree-command.ts

@@ -0,0 +1,39 @@
+import type { ToolCall } from "./tool-line";
+import type { SubagentManager } from "./subagents/manager";
+
+export function runWorktreeCommand({
+  name,
+  manager,
+  append,
+  patch,
+  settled,
+}: {
+  name?: string;
+  manager: SubagentManager;
+  append: (call: ToolCall) => void;
+  patch: (id: string, patch: Partial<ToolCall>) => void;
+  settled: () => void;
+}): void {
+  const id = `worktree-command-${Date.now()}`;
+  const call: ToolCall = {
+    id,
+    name: "worktree",
+    arg: name ? `create ${name}` : "create",
+    state: "running",
+  };
+  append(call);
+  manager.persistToolEvent(call);
+  manager
+    .createStandaloneWorktree(name)
+    .then((record) => {
+      const update: Partial<ToolCall> = { state: "ok", detail: record.branch };
+      patch(id, update);
+      manager.persistToolEvent({ ...call, ...update });
+    })
+    .catch((error) => {
+      const update: Partial<ToolCall> = { state: "error", detail: String(error) };
+      patch(id, update);
+      manager.persistToolEvent({ ...call, ...update });
+    })
+    .finally(settled);
+}