move pi-mono into companion-cloud as apps/companion-os

- Copy all pi-mono source into apps/companion-os/
- Update Dockerfile to COPY pre-built binary instead of downloading from GitHub Releases
- Update deploy-staging.yml to build pi from source (bun compile) before Docker build
- Add apps/companion-os/** to path triggers
- No more cross-repo dispatch needed

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

packages/agent/src/agent-loop.ts

@@ -0,0 +1,452 @@
+/**
+ * Agent loop that works with AgentMessage throughout.
+ * Transforms to Message[] only at the LLM call boundary.
+ */
+
+import {
+  type AssistantMessage,
+  type Context,
+  EventStream,
+  streamSimple,
+  type ToolResultMessage,
+  validateToolArguments,
+} from "@mariozechner/pi-ai";
+import type {
+  AgentContext,
+  AgentEvent,
+  AgentLoopConfig,
+  AgentMessage,
+  AgentTool,
+  AgentToolResult,
+  StreamFn,
+} from "./types.js";
+
+/**
+ * Start an agent loop with a new prompt message.
+ * The prompt is added to the context and events are emitted for it.
+ */
+export function agentLoop(
+  prompts: AgentMessage[],
+  context: AgentContext,
+  config: AgentLoopConfig,
+  signal?: AbortSignal,
+  streamFn?: StreamFn,
+): EventStream<AgentEvent, AgentMessage[]> {
+  const stream = createAgentStream();
+
+  (async () => {
+    const newMessages: AgentMessage[] = [...prompts];
+    const currentContext: AgentContext = {
+      ...context,
+      messages: [...context.messages, ...prompts],
+    };
+
+    stream.push({ type: "agent_start" });
+    stream.push({ type: "turn_start" });
+    for (const prompt of prompts) {
+      stream.push({ type: "message_start", message: prompt });
+      stream.push({ type: "message_end", message: prompt });
+    }
+
+    await runLoop(
+      currentContext,
+      newMessages,
+      config,
+      signal,
+      stream,
+      streamFn,
+    );
+  })();
+
+  return stream;
+}
+
+/**
+ * Continue an agent loop from the current context without adding a new message.
+ * Used for retries - context already has user message or tool results.
+ *
+ * **Important:** The last message in context must convert to a `user` or `toolResult` message
+ * via `convertToLlm`. If it doesn't, the LLM provider will reject the request.
+ * This cannot be validated here since `convertToLlm` is only called once per turn.
+ */
+export function agentLoopContinue(
+  context: AgentContext,
+  config: AgentLoopConfig,
+  signal?: AbortSignal,
+  streamFn?: StreamFn,
+): EventStream<AgentEvent, AgentMessage[]> {
+  if (context.messages.length === 0) {
+    throw new Error("Cannot continue: no messages in context");
+  }
+
+  if (context.messages[context.messages.length - 1].role === "assistant") {
+    throw new Error("Cannot continue from message role: assistant");
+  }
+
+  const stream = createAgentStream();
+
+  (async () => {
+    const newMessages: AgentMessage[] = [];
+    const currentContext: AgentContext = { ...context };
+
+    stream.push({ type: "agent_start" });
+    stream.push({ type: "turn_start" });
+
+    await runLoop(
+      currentContext,
+      newMessages,
+      config,
+      signal,
+      stream,
+      streamFn,
+    );
+  })();
+
+  return stream;
+}
+
+function createAgentStream(): EventStream<AgentEvent, AgentMessage[]> {
+  return new EventStream<AgentEvent, AgentMessage[]>(
+    (event: AgentEvent) => event.type === "agent_end",
+    (event: AgentEvent) => (event.type === "agent_end" ? event.messages : []),
+  );
+}
+
+/**
+ * Main loop logic shared by agentLoop and agentLoopContinue.
+ */
+async function runLoop(
+  currentContext: AgentContext,
+  newMessages: AgentMessage[],
+  config: AgentLoopConfig,
+  signal: AbortSignal | undefined,
+  stream: EventStream<AgentEvent, AgentMessage[]>,
+  streamFn?: StreamFn,
+): Promise<void> {
+  let firstTurn = true;
+  // Check for steering messages at start (user may have typed while waiting)
+  let pendingMessages: AgentMessage[] =
+    (await config.getSteeringMessages?.()) || [];
+
+  // Outer loop: continues when queued follow-up messages arrive after agent would stop
+  while (true) {
+    let hasMoreToolCalls = true;
+    let steeringAfterTools: AgentMessage[] | null = null;
+
+    // Inner loop: process tool calls and steering messages
+    while (hasMoreToolCalls || pendingMessages.length > 0) {
+      if (!firstTurn) {
+        stream.push({ type: "turn_start" });
+      } else {
+        firstTurn = false;
+      }
+
+      // Process pending messages (inject before next assistant response)
+      if (pendingMessages.length > 0) {
+        for (const message of pendingMessages) {
+          stream.push({ type: "message_start", message });
+          stream.push({ type: "message_end", message });
+          currentContext.messages.push(message);
+          newMessages.push(message);
+        }
+        pendingMessages = [];
+      }
+
+      // Stream assistant response
+      const message = await streamAssistantResponse(
+        currentContext,
+        config,
+        signal,
+        stream,
+        streamFn,
+      );
+      newMessages.push(message);
+
+      if (message.stopReason === "error" || message.stopReason === "aborted") {
+        stream.push({ type: "turn_end", message, toolResults: [] });
+        stream.push({ type: "agent_end", messages: newMessages });
+        stream.end(newMessages);
+        return;
+      }
+
+      // Check for tool calls
+      const toolCalls = message.content.filter((c) => c.type === "toolCall");
+      hasMoreToolCalls = toolCalls.length > 0;
+
+      const toolResults: ToolResultMessage[] = [];
+      if (hasMoreToolCalls) {
+        const toolExecution = await executeToolCalls(
+          currentContext.tools,
+          message,
+          signal,
+          stream,
+          config.getSteeringMessages,
+        );
+        toolResults.push(...toolExecution.toolResults);
+        steeringAfterTools = toolExecution.steeringMessages ?? null;
+
+        for (const result of toolResults) {
+          currentContext.messages.push(result);
+          newMessages.push(result);
+        }
+      }
+
+      stream.push({ type: "turn_end", message, toolResults });
+
+      // Get steering messages after turn completes
+      if (steeringAfterTools && steeringAfterTools.length > 0) {
+        pendingMessages = steeringAfterTools;
+        steeringAfterTools = null;
+      } else {
+        pendingMessages = (await config.getSteeringMessages?.()) || [];
+      }
+    }
+
+    // Agent would stop here. Check for follow-up messages.
+    const followUpMessages = (await config.getFollowUpMessages?.()) || [];
+    if (followUpMessages.length > 0) {
+      // Set as pending so inner loop processes them
+      pendingMessages = followUpMessages;
+      continue;
+    }
+
+    // No more messages, exit
+    break;
+  }
+
+  stream.push({ type: "agent_end", messages: newMessages });
+  stream.end(newMessages);
+}
+
+/**
+ * Stream an assistant response from the LLM.
+ * This is where AgentMessage[] gets transformed to Message[] for the LLM.
+ */
+async function streamAssistantResponse(
+  context: AgentContext,
+  config: AgentLoopConfig,
+  signal: AbortSignal | undefined,
+  stream: EventStream<AgentEvent, AgentMessage[]>,
+  streamFn?: StreamFn,
+): Promise<AssistantMessage> {
+  // Apply context transform if configured (AgentMessage[] → AgentMessage[])
+  let messages = context.messages;
+  if (config.transformContext) {
+    messages = await config.transformContext(messages, signal);
+  }
+
+  // Convert to LLM-compatible messages (AgentMessage[] → Message[])
+  const llmMessages = await config.convertToLlm(messages);
+
+  // Build LLM context
+  const llmContext: Context = {
+    systemPrompt: context.systemPrompt,
+    messages: llmMessages,
+    tools: context.tools,
+  };
+
+  const streamFunction = streamFn || streamSimple;
+
+  // Resolve API key (important for expiring tokens)
+  const resolvedApiKey =
+    (config.getApiKey
+      ? await config.getApiKey(config.model.provider)
+      : undefined) || config.apiKey;
+
+  const response = await streamFunction(config.model, llmContext, {
+    ...config,
+    apiKey: resolvedApiKey,
+    signal,
+  });
+
+  let partialMessage: AssistantMessage | null = null;
+  let addedPartial = false;
+
+  for await (const event of response) {
+    switch (event.type) {
+      case "start":
+        partialMessage = event.partial;
+        context.messages.push(partialMessage);
+        addedPartial = true;
+        stream.push({ type: "message_start", message: { ...partialMessage } });
+        break;
+
+      case "text_start":
+      case "text_delta":
+      case "text_end":
+      case "thinking_start":
+      case "thinking_delta":
+      case "thinking_end":
+      case "toolcall_start":
+      case "toolcall_delta":
+      case "toolcall_end":
+        if (partialMessage) {
+          partialMessage = event.partial;
+          context.messages[context.messages.length - 1] = partialMessage;
+          stream.push({
+            type: "message_update",
+            assistantMessageEvent: event,
+            message: { ...partialMessage },
+          });
+        }
+        break;
+
+      case "done":
+      case "error": {
+        const finalMessage = await response.result();
+        if (addedPartial) {
+          context.messages[context.messages.length - 1] = finalMessage;
+        } else {
+          context.messages.push(finalMessage);
+        }
+        if (!addedPartial) {
+          stream.push({ type: "message_start", message: { ...finalMessage } });
+        }
+        stream.push({ type: "message_end", message: finalMessage });
+        return finalMessage;
+      }
+    }
+  }
+
+  return await response.result();
+}
+
+/**
+ * Execute tool calls from an assistant message.
+ */
+async function executeToolCalls(
+  tools: AgentTool<any>[] | undefined,
+  assistantMessage: AssistantMessage,
+  signal: AbortSignal | undefined,
+  stream: EventStream<AgentEvent, AgentMessage[]>,
+  getSteeringMessages?: AgentLoopConfig["getSteeringMessages"],
+): Promise<{
+  toolResults: ToolResultMessage[];
+  steeringMessages?: AgentMessage[];
+}> {
+  const toolCalls = assistantMessage.content.filter(
+    (c) => c.type === "toolCall",
+  );
+  const results: ToolResultMessage[] = [];
+  let steeringMessages: AgentMessage[] | undefined;
+
+  for (let index = 0; index < toolCalls.length; index++) {
+    const toolCall = toolCalls[index];
+    const tool = tools?.find((t) => t.name === toolCall.name);
+
+    stream.push({
+      type: "tool_execution_start",
+      toolCallId: toolCall.id,
+      toolName: toolCall.name,
+      args: toolCall.arguments,
+    });
+
+    let result: AgentToolResult<any>;
+    let isError = false;
+
+    try {
+      if (!tool) throw new Error(`Tool ${toolCall.name} not found`);
+
+      const validatedArgs = validateToolArguments(tool, toolCall);
+
+      result = await tool.execute(
+        toolCall.id,
+        validatedArgs,
+        signal,
+        (partialResult) => {
+          stream.push({
+            type: "tool_execution_update",
+            toolCallId: toolCall.id,
+            toolName: toolCall.name,
+            args: toolCall.arguments,
+            partialResult,
+          });
+        },
+      );
+    } catch (e) {
+      result = {
+        content: [
+          { type: "text", text: e instanceof Error ? e.message : String(e) },
+        ],
+        details: {},
+      };
+      isError = true;
+    }
+
+    stream.push({
+      type: "tool_execution_end",
+      toolCallId: toolCall.id,
+      toolName: toolCall.name,
+      result,
+      isError,
+    });
+
+    const toolResultMessage: ToolResultMessage = {
+      role: "toolResult",
+      toolCallId: toolCall.id,
+      toolName: toolCall.name,
+      content: result.content,
+      details: result.details,
+      isError,
+      timestamp: Date.now(),
+    };
+
+    results.push(toolResultMessage);
+    stream.push({ type: "message_start", message: toolResultMessage });
+    stream.push({ type: "message_end", message: toolResultMessage });
+
+    // Check for steering messages - skip remaining tools if user interrupted
+    if (getSteeringMessages) {
+      const steering = await getSteeringMessages();
+      if (steering.length > 0) {
+        steeringMessages = steering;
+        const remainingCalls = toolCalls.slice(index + 1);
+        for (const skipped of remainingCalls) {
+          results.push(skipToolCall(skipped, stream));
+        }
+        break;
+      }
+    }
+  }
+
+  return { toolResults: results, steeringMessages };
+}
+
+function skipToolCall(
+  toolCall: Extract<AssistantMessage["content"][number], { type: "toolCall" }>,
+  stream: EventStream<AgentEvent, AgentMessage[]>,
+): ToolResultMessage {
+  const result: AgentToolResult<any> = {
+    content: [{ type: "text", text: "Skipped due to queued user message." }],
+    details: {},
+  };
+
+  stream.push({
+    type: "tool_execution_start",
+    toolCallId: toolCall.id,
+    toolName: toolCall.name,
+    args: toolCall.arguments,
+  });
+  stream.push({
+    type: "tool_execution_end",
+    toolCallId: toolCall.id,
+    toolName: toolCall.name,
+    result,
+    isError: true,
+  });
+
+  const toolResultMessage: ToolResultMessage = {
+    role: "toolResult",
+    toolCallId: toolCall.id,
+    toolName: toolCall.name,
+    content: result.content,
+    details: {},
+    isError: true,
+    timestamp: Date.now(),
+  };
+
+  stream.push({ type: "message_start", message: toolResultMessage });
+  stream.push({ type: "message_end", message: toolResultMessage });
+
+  return toolResultMessage;
+}

packages/agent/src/agent.ts

@@ -0,0 +1,605 @@
+/**
+ * Agent class that uses the agent-loop directly.
+ * No transport abstraction - calls streamSimple via the loop.
+ */
+
+import {
+  getModel,
+  type ImageContent,
+  type Message,
+  type Model,
+  streamSimple,
+  type TextContent,
+  type ThinkingBudgets,
+  type Transport,
+} from "@mariozechner/pi-ai";
+import { agentLoop, agentLoopContinue } from "./agent-loop.js";
+import type {
+  AgentContext,
+  AgentEvent,
+  AgentLoopConfig,
+  AgentMessage,
+  AgentState,
+  AgentTool,
+  StreamFn,
+  ThinkingLevel,
+} from "./types.js";
+
+/**
+ * Default convertToLlm: Keep only LLM-compatible messages, convert attachments.
+ */
+function defaultConvertToLlm(messages: AgentMessage[]): Message[] {
+  return messages.filter(
+    (m) =>
+      m.role === "user" || m.role === "assistant" || m.role === "toolResult",
+  );
+}
+
+export interface AgentOptions {
+  initialState?: Partial<AgentState>;
+
+  /**
+   * Converts AgentMessage[] to LLM-compatible Message[] before each LLM call.
+   * Default filters to user/assistant/toolResult and converts attachments.
+   */
+  convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
+
+  /**
+   * Optional transform applied to context before convertToLlm.
+   * Use for context pruning, injecting external context, etc.
+   */
+  transformContext?: (
+    messages: AgentMessage[],
+    signal?: AbortSignal,
+  ) => Promise<AgentMessage[]>;
+
+  /**
+   * Steering mode: "all" = send all steering messages at once, "one-at-a-time" = one per turn
+   */
+  steeringMode?: "all" | "one-at-a-time";
+
+  /**
+   * Follow-up mode: "all" = send all follow-up messages at once, "one-at-a-time" = one per turn
+   */
+  followUpMode?: "all" | "one-at-a-time";
+
+  /**
+   * Custom stream function (for proxy backends, etc.). Default uses streamSimple.
+   */
+  streamFn?: StreamFn;
+
+  /**
+   * Optional session identifier forwarded to LLM providers.
+   * Used by providers that support session-based caching (e.g., OpenAI Codex).
+   */
+  sessionId?: string;
+
+  /**
+   * Resolves an API key dynamically for each LLM call.
+   * Useful for expiring tokens (e.g., GitHub Copilot OAuth).
+   */
+  getApiKey?: (
+    provider: string,
+  ) => Promise<string | undefined> | string | undefined;
+
+  /**
+   * Custom token budgets for thinking levels (token-based providers only).
+   */
+  thinkingBudgets?: ThinkingBudgets;
+
+  /**
+   * Preferred transport for providers that support multiple transports.
+   */
+  transport?: Transport;
+
+  /**
+   * Maximum delay in milliseconds to wait for a retry when the server requests a long wait.
+   * If the server's requested delay exceeds this value, the request fails immediately,
+   * allowing higher-level retry logic to handle it with user visibility.
+   * Default: 60000 (60 seconds). Set to 0 to disable the cap.
+   */
+  maxRetryDelayMs?: number;
+}
+
+export class Agent {
+  private _state: AgentState = {
+    systemPrompt: "",
+    model: getModel("google", "gemini-2.5-flash-lite-preview-06-17"),
+    thinkingLevel: "off",
+    tools: [],
+    messages: [],
+    isStreaming: false,
+    streamMessage: null,
+    pendingToolCalls: new Set<string>(),
+    error: undefined,
+  };
+
+  private listeners = new Set<(e: AgentEvent) => void>();
+  private abortController?: AbortController;
+  private convertToLlm: (
+    messages: AgentMessage[],
+  ) => Message[] | Promise<Message[]>;
+  private transformContext?: (
+    messages: AgentMessage[],
+    signal?: AbortSignal,
+  ) => Promise<AgentMessage[]>;
+  private steeringQueue: AgentMessage[] = [];
+  private followUpQueue: AgentMessage[] = [];
+  private steeringMode: "all" | "one-at-a-time";
+  private followUpMode: "all" | "one-at-a-time";
+  public streamFn: StreamFn;
+  private _sessionId?: string;
+  public getApiKey?: (
+    provider: string,
+  ) => Promise<string | undefined> | string | undefined;
+  private runningPrompt?: Promise<void>;
+  private resolveRunningPrompt?: () => void;
+  private _thinkingBudgets?: ThinkingBudgets;
+  private _transport: Transport;
+  private _maxRetryDelayMs?: number;
+
+  constructor(opts: AgentOptions = {}) {
+    this._state = { ...this._state, ...opts.initialState };
+    this.convertToLlm = opts.convertToLlm || defaultConvertToLlm;
+    this.transformContext = opts.transformContext;
+    this.steeringMode = opts.steeringMode || "one-at-a-time";
+    this.followUpMode = opts.followUpMode || "one-at-a-time";
+    this.streamFn = opts.streamFn || streamSimple;
+    this._sessionId = opts.sessionId;
+    this.getApiKey = opts.getApiKey;
+    this._thinkingBudgets = opts.thinkingBudgets;
+    this._transport = opts.transport ?? "sse";
+    this._maxRetryDelayMs = opts.maxRetryDelayMs;
+  }
+
+  /**
+   * Get the current session ID used for provider caching.
+   */
+  get sessionId(): string | undefined {
+    return this._sessionId;
+  }
+
+  /**
+   * Set the session ID for provider caching.
+   * Call this when switching sessions (new session, branch, resume).
+   */
+  set sessionId(value: string | undefined) {
+    this._sessionId = value;
+  }
+
+  /**
+   * Get the current thinking budgets.
+   */
+  get thinkingBudgets(): ThinkingBudgets | undefined {
+    return this._thinkingBudgets;
+  }
+
+  /**
+   * Set custom thinking budgets for token-based providers.
+   */
+  set thinkingBudgets(value: ThinkingBudgets | undefined) {
+    this._thinkingBudgets = value;
+  }
+
+  /**
+   * Get the current preferred transport.
+   */
+  get transport(): Transport {
+    return this._transport;
+  }
+
+  /**
+   * Set the preferred transport.
+   */
+  setTransport(value: Transport) {
+    this._transport = value;
+  }
+
+  /**
+   * Get the current max retry delay in milliseconds.
+   */
+  get maxRetryDelayMs(): number | undefined {
+    return this._maxRetryDelayMs;
+  }
+
+  /**
+   * Set the maximum delay to wait for server-requested retries.
+   * Set to 0 to disable the cap.
+   */
+  set maxRetryDelayMs(value: number | undefined) {
+    this._maxRetryDelayMs = value;
+  }
+
+  get state(): AgentState {
+    return this._state;
+  }
+
+  subscribe(fn: (e: AgentEvent) => void): () => void {
+    this.listeners.add(fn);
+    return () => this.listeners.delete(fn);
+  }
+
+  // State mutators
+  setSystemPrompt(v: string) {
+    this._state.systemPrompt = v;
+  }
+
+  setModel(m: Model<any>) {
+    this._state.model = m;
+  }
+
+  setThinkingLevel(l: ThinkingLevel) {
+    this._state.thinkingLevel = l;
+  }
+
+  setSteeringMode(mode: "all" | "one-at-a-time") {
+    this.steeringMode = mode;
+  }
+
+  getSteeringMode(): "all" | "one-at-a-time" {
+    return this.steeringMode;
+  }
+
+  setFollowUpMode(mode: "all" | "one-at-a-time") {
+    this.followUpMode = mode;
+  }
+
+  getFollowUpMode(): "all" | "one-at-a-time" {
+    return this.followUpMode;
+  }
+
+  setTools(t: AgentTool<any>[]) {
+    this._state.tools = t;
+  }
+
+  replaceMessages(ms: AgentMessage[]) {
+    this._state.messages = ms.slice();
+  }
+
+  appendMessage(m: AgentMessage) {
+    this._state.messages = [...this._state.messages, m];
+  }
+
+  /**
+   * Queue a steering message to interrupt the agent mid-run.
+   * Delivered after current tool execution, skips remaining tools.
+   */
+  steer(m: AgentMessage) {
+    this.steeringQueue.push(m);
+  }
+
+  /**
+   * Queue a follow-up message to be processed after the agent finishes.
+   * Delivered only when agent has no more tool calls or steering messages.
+   */
+  followUp(m: AgentMessage) {
+    this.followUpQueue.push(m);
+  }
+
+  clearSteeringQueue() {
+    this.steeringQueue = [];
+  }
+
+  clearFollowUpQueue() {
+    this.followUpQueue = [];
+  }
+
+  clearAllQueues() {
+    this.steeringQueue = [];
+    this.followUpQueue = [];
+  }
+
+  hasQueuedMessages(): boolean {
+    return this.steeringQueue.length > 0 || this.followUpQueue.length > 0;
+  }
+
+  private dequeueSteeringMessages(): AgentMessage[] {
+    if (this.steeringMode === "one-at-a-time") {
+      if (this.steeringQueue.length > 0) {
+        const first = this.steeringQueue[0];
+        this.steeringQueue = this.steeringQueue.slice(1);
+        return [first];
+      }
+      return [];
+    }
+
+    const steering = this.steeringQueue.slice();
+    this.steeringQueue = [];
+    return steering;
+  }
+
+  private dequeueFollowUpMessages(): AgentMessage[] {
+    if (this.followUpMode === "one-at-a-time") {
+      if (this.followUpQueue.length > 0) {
+        const first = this.followUpQueue[0];
+        this.followUpQueue = this.followUpQueue.slice(1);
+        return [first];
+      }
+      return [];
+    }
+
+    const followUp = this.followUpQueue.slice();
+    this.followUpQueue = [];
+    return followUp;
+  }
+
+  clearMessages() {
+    this._state.messages = [];
+  }
+
+  abort() {
+    this.abortController?.abort();
+  }
+
+  waitForIdle(): Promise<void> {
+    return this.runningPrompt ?? Promise.resolve();
+  }
+
+  reset() {
+    this._state.messages = [];
+    this._state.isStreaming = false;
+    this._state.streamMessage = null;
+    this._state.pendingToolCalls = new Set<string>();
+    this._state.error = undefined;
+    this.steeringQueue = [];
+    this.followUpQueue = [];
+  }
+
+  /** Send a prompt with an AgentMessage */
+  async prompt(message: AgentMessage | AgentMessage[]): Promise<void>;
+  async prompt(input: string, images?: ImageContent[]): Promise<void>;
+  async prompt(
+    input: string | AgentMessage | AgentMessage[],
+    images?: ImageContent[],
+  ) {
+    if (this._state.isStreaming) {
+      throw new Error(
+        "Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.",
+      );
+    }
+
+    const model = this._state.model;
+    if (!model) throw new Error("No model configured");
+
+    let msgs: AgentMessage[];
+
+    if (Array.isArray(input)) {
+      msgs = input;
+    } else if (typeof input === "string") {
+      const content: Array<TextContent | ImageContent> = [
+        { type: "text", text: input },
+      ];
+      if (images && images.length > 0) {
+        content.push(...images);
+      }
+      msgs = [
+        {
+          role: "user",
+          content,
+          timestamp: Date.now(),
+        },
+      ];
+    } else {
+      msgs = [input];
+    }
+
+    await this._runLoop(msgs);
+  }
+
+  /**
+   * Continue from current context (used for retries and resuming queued messages).
+   */
+  async continue() {
+    if (this._state.isStreaming) {
+      throw new Error(
+        "Agent is already processing. Wait for completion before continuing.",
+      );
+    }
+
+    const messages = this._state.messages;
+    if (messages.length === 0) {
+      throw new Error("No messages to continue from");
+    }
+    if (messages[messages.length - 1].role === "assistant") {
+      const queuedSteering = this.dequeueSteeringMessages();
+      if (queuedSteering.length > 0) {
+        await this._runLoop(queuedSteering, { skipInitialSteeringPoll: true });
+        return;
+      }
+
+      const queuedFollowUp = this.dequeueFollowUpMessages();
+      if (queuedFollowUp.length > 0) {
+        await this._runLoop(queuedFollowUp);
+        return;
+      }
+
+      throw new Error("Cannot continue from message role: assistant");
+    }
+
+    await this._runLoop(undefined);
+  }
+
+  /**
+   * Run the agent loop.
+   * If messages are provided, starts a new conversation turn with those messages.
+   * Otherwise, continues from existing context.
+   */
+  private async _runLoop(
+    messages?: AgentMessage[],
+    options?: { skipInitialSteeringPoll?: boolean },
+  ) {
+    const model = this._state.model;
+    if (!model) throw new Error("No model configured");
+
+    this.runningPrompt = new Promise<void>((resolve) => {
+      this.resolveRunningPrompt = resolve;
+    });
+
+    this.abortController = new AbortController();
+    this._state.isStreaming = true;
+    this._state.streamMessage = null;
+    this._state.error = undefined;
+
+    const reasoning =
+      this._state.thinkingLevel === "off"
+        ? undefined
+        : this._state.thinkingLevel;
+
+    const context: AgentContext = {
+      systemPrompt: this._state.systemPrompt,
+      messages: this._state.messages.slice(),
+      tools: this._state.tools,
+    };
+
+    let skipInitialSteeringPoll = options?.skipInitialSteeringPoll === true;
+
+    const config: AgentLoopConfig = {
+      model,
+      reasoning,
+      sessionId: this._sessionId,
+      transport: this._transport,
+      thinkingBudgets: this._thinkingBudgets,
+      maxRetryDelayMs: this._maxRetryDelayMs,
+      convertToLlm: this.convertToLlm,
+      transformContext: this.transformContext,
+      getApiKey: this.getApiKey,
+      getSteeringMessages: async () => {
+        if (skipInitialSteeringPoll) {
+          skipInitialSteeringPoll = false;
+          return [];
+        }
+        return this.dequeueSteeringMessages();
+      },
+      getFollowUpMessages: async () => this.dequeueFollowUpMessages(),
+    };
+
+    let partial: AgentMessage | null = null;
+
+    try {
+      const stream = messages
+        ? agentLoop(
+            messages,
+            context,
+            config,
+            this.abortController.signal,
+            this.streamFn,
+          )
+        : agentLoopContinue(
+            context,
+            config,
+            this.abortController.signal,
+            this.streamFn,
+          );
+
+      for await (const event of stream) {
+        // Update internal state based on events
+        switch (event.type) {
+          case "message_start":
+            partial = event.message;
+            this._state.streamMessage = event.message;
+            break;
+
+          case "message_update":
+            partial = event.message;
+            this._state.streamMessage = event.message;
+            break;
+
+          case "message_end":
+            partial = null;
+            this._state.streamMessage = null;
+            this.appendMessage(event.message);
+            break;
+
+          case "tool_execution_start": {
+            const s = new Set(this._state.pendingToolCalls);
+            s.add(event.toolCallId);
+            this._state.pendingToolCalls = s;
+            break;
+          }
+
+          case "tool_execution_end": {
+            const s = new Set(this._state.pendingToolCalls);
+            s.delete(event.toolCallId);
+            this._state.pendingToolCalls = s;
+            break;
+          }
+
+          case "turn_end":
+            if (
+              event.message.role === "assistant" &&
+              (event.message as any).errorMessage
+            ) {
+              this._state.error = (event.message as any).errorMessage;
+            }
+            break;
+
+          case "agent_end":
+            this._state.isStreaming = false;
+            this._state.streamMessage = null;
+            break;
+        }
+
+        // Emit to listeners
+        this.emit(event);
+      }
+
+      // Handle any remaining partial message
+      if (
+        partial &&
+        partial.role === "assistant" &&
+        partial.content.length > 0
+      ) {
+        const onlyEmpty = !partial.content.some(
+          (c) =>
+            (c.type === "thinking" && c.thinking.trim().length > 0) ||
+            (c.type === "text" && c.text.trim().length > 0) ||
+            (c.type === "toolCall" && c.name.trim().length > 0),
+        );
+        if (!onlyEmpty) {
+          this.appendMessage(partial);
+        } else {
+          if (this.abortController?.signal.aborted) {
+            throw new Error("Request was aborted");
+          }
+        }
+      }
+    } catch (err: any) {
+      const errorMsg: AgentMessage = {
+        role: "assistant",
+        content: [{ type: "text", text: "" }],
+        api: model.api,
+        provider: model.provider,
+        model: model.id,
+        usage: {
+          input: 0,
+          output: 0,
+          cacheRead: 0,
+          cacheWrite: 0,
+          totalTokens: 0,
+          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+        },
+        stopReason: this.abortController?.signal.aborted ? "aborted" : "error",
+        errorMessage: err?.message || String(err),
+        timestamp: Date.now(),
+      } as AgentMessage;
+
+      this.appendMessage(errorMsg);
+      this._state.error = err?.message || String(err);
+      this.emit({ type: "agent_end", messages: [errorMsg] });
+    } finally {
+      this._state.isStreaming = false;
+      this._state.streamMessage = null;
+      this._state.pendingToolCalls = new Set<string>();
+      this.abortController = undefined;
+      this.resolveRunningPrompt?.();
+      this.runningPrompt = undefined;
+      this.resolveRunningPrompt = undefined;
+    }
+  }
+
+  private emit(e: AgentEvent) {
+    for (const listener of this.listeners) {
+      listener(e);
+    }
+  }
+}

packages/agent/src/index.ts

@@ -0,0 +1,8 @@
+// Core Agent
+export * from "./agent.js";
+// Loop functions
+export * from "./agent-loop.js";
+// Proxy utilities
+export * from "./proxy.js";
+// Types
+export * from "./types.js";

packages/agent/src/proxy.ts

@@ -0,0 +1,369 @@
+/**
+ * Proxy stream function for apps that route LLM calls through a server.
+ * The server manages auth and proxies requests to LLM providers.
+ */
+
+// Internal import for JSON parsing utility
+import {
+  type AssistantMessage,
+  type AssistantMessageEvent,
+  type Context,
+  EventStream,
+  type Model,
+  parseStreamingJson,
+  type SimpleStreamOptions,
+  type StopReason,
+  type ToolCall,
+} from "@mariozechner/pi-ai";
+
+// Create stream class matching ProxyMessageEventStream
+class ProxyMessageEventStream extends EventStream<
+  AssistantMessageEvent,
+  AssistantMessage
+> {
+  constructor() {
+    super(
+      (event) => event.type === "done" || event.type === "error",
+      (event) => {
+        if (event.type === "done") return event.message;
+        if (event.type === "error") return event.error;
+        throw new Error("Unexpected event type");
+      },
+    );
+  }
+}
+
+/**
+ * Proxy event types - server sends these with partial field stripped to reduce bandwidth.
+ */
+export type ProxyAssistantMessageEvent =
+  | { type: "start" }
+  | { type: "text_start"; contentIndex: number }
+  | { type: "text_delta"; contentIndex: number; delta: string }
+  | { type: "text_end"; contentIndex: number; contentSignature?: string }
+  | { type: "thinking_start"; contentIndex: number }
+  | { type: "thinking_delta"; contentIndex: number; delta: string }
+  | { type: "thinking_end"; contentIndex: number; contentSignature?: string }
+  | {
+      type: "toolcall_start";
+      contentIndex: number;
+      id: string;
+      toolName: string;
+    }
+  | { type: "toolcall_delta"; contentIndex: number; delta: string }
+  | { type: "toolcall_end"; contentIndex: number }
+  | {
+      type: "done";
+      reason: Extract<StopReason, "stop" | "length" | "toolUse">;
+      usage: AssistantMessage["usage"];
+    }
+  | {
+      type: "error";
+      reason: Extract<StopReason, "aborted" | "error">;
+      errorMessage?: string;
+      usage: AssistantMessage["usage"];
+    };
+
+export interface ProxyStreamOptions extends SimpleStreamOptions {
+  /** Auth token for the proxy server */
+  authToken: string;
+  /** Proxy server URL (e.g., "https://genai.example.com") */
+  proxyUrl: string;
+}
+
+/**
+ * Stream function that proxies through a server instead of calling LLM providers directly.
+ * The server strips the partial field from delta events to reduce bandwidth.
+ * We reconstruct the partial message client-side.
+ *
+ * Use this as the `streamFn` option when creating an Agent that needs to go through a proxy.
+ *
+ * @example
+ * ```typescript
+ * const agent = new Agent({
+ *   streamFn: (model, context, options) =>
+ *     streamProxy(model, context, {
+ *       ...options,
+ *       authToken: await getAuthToken(),
+ *       proxyUrl: "https://genai.example.com",
+ *     }),
+ * });
+ * ```
+ */
+export function streamProxy(
+  model: Model<any>,
+  context: Context,
+  options: ProxyStreamOptions,
+): ProxyMessageEventStream {
+  const stream = new ProxyMessageEventStream();
+
+  (async () => {
+    // Initialize the partial message that we'll build up from events
+    const partial: AssistantMessage = {
+      role: "assistant",
+      stopReason: "stop",
+      content: [],
+      api: model.api,
+      provider: model.provider,
+      model: model.id,
+      usage: {
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 0,
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+      },
+      timestamp: Date.now(),
+    };
+
+    let reader: ReadableStreamDefaultReader<Uint8Array> | undefined;
+
+    const abortHandler = () => {
+      if (reader) {
+        reader.cancel("Request aborted by user").catch(() => {});
+      }
+    };
+
+    if (options.signal) {
+      options.signal.addEventListener("abort", abortHandler);
+    }
+
+    try {
+      const response = await fetch(`${options.proxyUrl}/api/stream`, {
+        method: "POST",
+        headers: {
+          Authorization: `Bearer ${options.authToken}`,
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+          model,
+          context,
+          options: {
+            temperature: options.temperature,
+            maxTokens: options.maxTokens,
+            reasoning: options.reasoning,
+          },
+        }),
+        signal: options.signal,
+      });
+
+      if (!response.ok) {
+        let errorMessage = `Proxy error: ${response.status} ${response.statusText}`;
+        try {
+          const errorData = (await response.json()) as { error?: string };
+          if (errorData.error) {
+            errorMessage = `Proxy error: ${errorData.error}`;
+          }
+        } catch {
+          // Couldn't parse error response
+        }
+        throw new Error(errorMessage);
+      }
+
+      reader = response.body!.getReader();
+      const decoder = new TextDecoder();
+      let buffer = "";
+
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+
+        if (options.signal?.aborted) {
+          throw new Error("Request aborted by user");
+        }
+
+        buffer += decoder.decode(value, { stream: true });
+        const lines = buffer.split("\n");
+        buffer = lines.pop() || "";
+
+        for (const line of lines) {
+          if (line.startsWith("data: ")) {
+            const data = line.slice(6).trim();
+            if (data) {
+              const proxyEvent = JSON.parse(data) as ProxyAssistantMessageEvent;
+              const event = processProxyEvent(proxyEvent, partial);
+              if (event) {
+                stream.push(event);
+              }
+            }
+          }
+        }
+      }
+
+      if (options.signal?.aborted) {
+        throw new Error("Request aborted by user");
+      }
+
+      stream.end();
+    } catch (error) {
+      const errorMessage =
+        error instanceof Error ? error.message : String(error);
+      const reason = options.signal?.aborted ? "aborted" : "error";
+      partial.stopReason = reason;
+      partial.errorMessage = errorMessage;
+      stream.push({
+        type: "error",
+        reason,
+        error: partial,
+      });
+      stream.end();
+    } finally {
+      if (options.signal) {
+        options.signal.removeEventListener("abort", abortHandler);
+      }
+    }
+  })();
+
+  return stream;
+}
+
+/**
+ * Process a proxy event and update the partial message.
+ */
+function processProxyEvent(
+  proxyEvent: ProxyAssistantMessageEvent,
+  partial: AssistantMessage,
+): AssistantMessageEvent | undefined {
+  switch (proxyEvent.type) {
+    case "start":
+      return { type: "start", partial };
+
+    case "text_start":
+      partial.content[proxyEvent.contentIndex] = { type: "text", text: "" };
+      return {
+        type: "text_start",
+        contentIndex: proxyEvent.contentIndex,
+        partial,
+      };
+
+    case "text_delta": {
+      const content = partial.content[proxyEvent.contentIndex];
+      if (content?.type === "text") {
+        content.text += proxyEvent.delta;
+        return {
+          type: "text_delta",
+          contentIndex: proxyEvent.contentIndex,
+          delta: proxyEvent.delta,
+          partial,
+        };
+      }
+      throw new Error("Received text_delta for non-text content");
+    }
+
+    case "text_end": {
+      const content = partial.content[proxyEvent.contentIndex];
+      if (content?.type === "text") {
+        content.textSignature = proxyEvent.contentSignature;
+        return {
+          type: "text_end",
+          contentIndex: proxyEvent.contentIndex,
+          content: content.text,
+          partial,
+        };
+      }
+      throw new Error("Received text_end for non-text content");
+    }
+
+    case "thinking_start":
+      partial.content[proxyEvent.contentIndex] = {
+        type: "thinking",
+        thinking: "",
+      };
+      return {
+        type: "thinking_start",
+        contentIndex: proxyEvent.contentIndex,
+        partial,
+      };
+
+    case "thinking_delta": {
+      const content = partial.content[proxyEvent.contentIndex];
+      if (content?.type === "thinking") {
+        content.thinking += proxyEvent.delta;
+        return {
+          type: "thinking_delta",
+          contentIndex: proxyEvent.contentIndex,
+          delta: proxyEvent.delta,
+          partial,
+        };
+      }
+      throw new Error("Received thinking_delta for non-thinking content");
+    }
+
+    case "thinking_end": {
+      const content = partial.content[proxyEvent.contentIndex];
+      if (content?.type === "thinking") {
+        content.thinkingSignature = proxyEvent.contentSignature;
+        return {
+          type: "thinking_end",
+          contentIndex: proxyEvent.contentIndex,
+          content: content.thinking,
+          partial,
+        };
+      }
+      throw new Error("Received thinking_end for non-thinking content");
+    }
+
+    case "toolcall_start":
+      partial.content[proxyEvent.contentIndex] = {
+        type: "toolCall",
+        id: proxyEvent.id,
+        name: proxyEvent.toolName,
+        arguments: {},
+        partialJson: "",
+      } satisfies ToolCall & { partialJson: string } as ToolCall;
+      return {
+        type: "toolcall_start",
+        contentIndex: proxyEvent.contentIndex,
+        partial,
+      };
+
+    case "toolcall_delta": {
+      const content = partial.content[proxyEvent.contentIndex];
+      if (content?.type === "toolCall") {
+        (content as any).partialJson += proxyEvent.delta;
+        content.arguments =
+          parseStreamingJson((content as any).partialJson) || {};
+        partial.content[proxyEvent.contentIndex] = { ...content }; // Trigger reactivity
+        return {
+          type: "toolcall_delta",
+          contentIndex: proxyEvent.contentIndex,
+          delta: proxyEvent.delta,
+          partial,
+        };
+      }
+      throw new Error("Received toolcall_delta for non-toolCall content");
+    }
+
+    case "toolcall_end": {
+      const content = partial.content[proxyEvent.contentIndex];
+      if (content?.type === "toolCall") {
+        delete (content as any).partialJson;
+        return {
+          type: "toolcall_end",
+          contentIndex: proxyEvent.contentIndex,
+          toolCall: content,
+          partial,
+        };
+      }
+      return undefined;
+    }
+
+    case "done":
+      partial.stopReason = proxyEvent.reason;
+      partial.usage = proxyEvent.usage;
+      return { type: "done", reason: proxyEvent.reason, message: partial };
+
+    case "error":
+      partial.stopReason = proxyEvent.reason;
+      partial.errorMessage = proxyEvent.errorMessage;
+      partial.usage = proxyEvent.usage;
+      return { type: "error", reason: proxyEvent.reason, error: partial };
+
+    default: {
+      const _exhaustiveCheck: never = proxyEvent;
+      console.warn(`Unhandled proxy event type: ${(proxyEvent as any).type}`);
+      return undefined;
+    }
+  }
+}

packages/agent/src/types.ts

@@ -0,0 +1,237 @@
+import type {
+  AssistantMessageEvent,
+  ImageContent,
+  Message,
+  Model,
+  SimpleStreamOptions,
+  streamSimple,
+  TextContent,
+  Tool,
+  ToolResultMessage,
+} from "@mariozechner/pi-ai";
+import type { Static, TSchema } from "@sinclair/typebox";
+
+/** Stream function - can return sync or Promise for async config lookup */
+export type StreamFn = (
+  ...args: Parameters<typeof streamSimple>
+) => ReturnType<typeof streamSimple> | Promise<ReturnType<typeof streamSimple>>;
+
+/**
+ * Configuration for the agent loop.
+ */
+export interface AgentLoopConfig extends SimpleStreamOptions {
+  model: Model<any>;
+
+  /**
+   * Converts AgentMessage[] to LLM-compatible Message[] before each LLM call.
+   *
+   * Each AgentMessage must be converted to a UserMessage, AssistantMessage, or ToolResultMessage
+   * that the LLM can understand. AgentMessages that cannot be converted (e.g., UI-only notifications,
+   * status messages) should be filtered out.
+   *
+   * @example
+   * ```typescript
+   * convertToLlm: (messages) => messages.flatMap(m => {
+   *   if (m.role === "custom") {
+   *     // Convert custom message to user message
+   *     return [{ role: "user", content: m.content, timestamp: m.timestamp }];
+   *   }
+   *   if (m.role === "notification") {
+   *     // Filter out UI-only messages
+   *     return [];
+   *   }
+   *   // Pass through standard LLM messages
+   *   return [m];
+   * })
+   * ```
+   */
+  convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
+
+  /**
+   * Optional transform applied to the context before `convertToLlm`.
+   *
+   * Use this for operations that work at the AgentMessage level:
+   * - Context window management (pruning old messages)
+   * - Injecting context from external sources
+   *
+   * @example
+   * ```typescript
+   * transformContext: async (messages) => {
+   *   if (estimateTokens(messages) > MAX_TOKENS) {
+   *     return pruneOldMessages(messages);
+   *   }
+   *   return messages;
+   * }
+   * ```
+   */
+  transformContext?: (
+    messages: AgentMessage[],
+    signal?: AbortSignal,
+  ) => Promise<AgentMessage[]>;
+
+  /**
+   * Resolves an API key dynamically for each LLM call.
+   *
+   * Useful for short-lived OAuth tokens (e.g., GitHub Copilot) that may expire
+   * during long-running tool execution phases.
+   */
+  getApiKey?: (
+    provider: string,
+  ) => Promise<string | undefined> | string | undefined;
+
+  /**
+   * Returns steering messages to inject into the conversation mid-run.
+   *
+   * Called after each tool execution to check for user interruptions.
+   * If messages are returned, remaining tool calls are skipped and
+   * these messages are added to the context before the next LLM call.
+   *
+   * Use this for "steering" the agent while it's working.
+   */
+  getSteeringMessages?: () => Promise<AgentMessage[]>;
+
+  /**
+   * Returns follow-up messages to process after the agent would otherwise stop.
+   *
+   * Called when the agent has no more tool calls and no steering messages.
+   * If messages are returned, they're added to the context and the agent
+   * continues with another turn.
+   *
+   * Use this for follow-up messages that should wait until the agent finishes.
+   */
+  getFollowUpMessages?: () => Promise<AgentMessage[]>;
+}
+
+/**
+ * Thinking/reasoning level for models that support it.
+ * Note: "xhigh" is only supported by OpenAI gpt-5.1-codex-max, gpt-5.2, gpt-5.2-codex, gpt-5.3, and gpt-5.3-codex models.
+ */
+export type ThinkingLevel =
+  | "off"
+  | "minimal"
+  | "low"
+  | "medium"
+  | "high"
+  | "xhigh";
+
+/**
+ * Extensible interface for custom app messages.
+ * Apps can extend via declaration merging:
+ *
+ * @example
+ * ```typescript
+ * declare module "@mariozechner/agent" {
+ *   interface CustomAgentMessages {
+ *     artifact: ArtifactMessage;
+ *     notification: NotificationMessage;
+ *   }
+ * }
+ * ```
+ */
+export interface CustomAgentMessages {
+  // Empty by default - apps extend via declaration merging
+}
+
+/**
+ * AgentMessage: Union of LLM messages + custom messages.
+ * This abstraction allows apps to add custom message types while maintaining
+ * type safety and compatibility with the base LLM messages.
+ */
+export type AgentMessage =
+  | Message
+  | CustomAgentMessages[keyof CustomAgentMessages];
+
+/**
+ * Agent state containing all configuration and conversation data.
+ */
+export interface AgentState {
+  systemPrompt: string;
+  model: Model<any>;
+  thinkingLevel: ThinkingLevel;
+  tools: AgentTool<any>[];
+  messages: AgentMessage[]; // Can include attachments + custom message types
+  isStreaming: boolean;
+  streamMessage: AgentMessage | null;
+  pendingToolCalls: Set<string>;
+  error?: string;
+}
+
+export interface AgentToolResult<T> {
+  // Content blocks supporting text and images
+  content: (TextContent | ImageContent)[];
+  // Details to be displayed in a UI or logged
+  details: T;
+}
+
+// Callback for streaming tool execution updates
+export type AgentToolUpdateCallback<T = any> = (
+  partialResult: AgentToolResult<T>,
+) => void;
+
+// AgentTool extends Tool but adds the execute function
+export interface AgentTool<
+  TParameters extends TSchema = TSchema,
+  TDetails = any,
+> extends Tool<TParameters> {
+  // A human-readable label for the tool to be displayed in UI
+  label: string;
+  execute: (
+    toolCallId: string,
+    params: Static<TParameters>,
+    signal?: AbortSignal,
+    onUpdate?: AgentToolUpdateCallback<TDetails>,
+  ) => Promise<AgentToolResult<TDetails>>;
+}
+
+// AgentContext is like Context but uses AgentTool
+export interface AgentContext {
+  systemPrompt: string;
+  messages: AgentMessage[];
+  tools?: AgentTool<any>[];
+}
+
+/**
+ * Events emitted by the Agent for UI updates.
+ * These events provide fine-grained lifecycle information for messages, turns, and tool executions.
+ */
+export type AgentEvent =
+  // Agent lifecycle
+  | { type: "agent_start" }
+  | { type: "agent_end"; messages: AgentMessage[] }
+  // Turn lifecycle - a turn is one assistant response + any tool calls/results
+  | { type: "turn_start" }
+  | {
+      type: "turn_end";
+      message: AgentMessage;
+      toolResults: ToolResultMessage[];
+    }
+  // Message lifecycle - emitted for user, assistant, and toolResult messages
+  | { type: "message_start"; message: AgentMessage }
+  // Only emitted for assistant messages during streaming
+  | {
+      type: "message_update";
+      message: AgentMessage;
+      assistantMessageEvent: AssistantMessageEvent;
+    }
+  | { type: "message_end"; message: AgentMessage }
+  // Tool execution lifecycle
+  | {
+      type: "tool_execution_start";
+      toolCallId: string;
+      toolName: string;
+      args: any;
+    }
+  | {
+      type: "tool_execution_update";
+      toolCallId: string;
+      toolName: string;
+      args: any;
+      partialResult: any;
+    }
+  | {
+      type: "tool_execution_end";
+      toolCallId: string;
+      toolName: string;
+      result: any;
+      isError: boolean;
+    };