getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-15 08:03:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

WIP: Major cleanup - move Attachment to consumers, simplify agent API

Browse source 
- Removed Attachment from agent package (now in web-ui/coding-agent)
- Agent.prompt now takes (text, images?: ImageContent[])
- Removed transports from web-ui (duplicate of agent package)
- Updated coding-agent to use local message types
- Updated mom package for new agent API

Remaining: Fix AgentInterface.ts to compose UserMessageWithAttachments
This commit is contained in:
Mario Zechner 2025-12-28 10:55:12 +01:00
parent f86dea2e4f
commit 6ddc7418da
57 changed files with 167 additions and 1061 deletions
Download patch file Download diff file Expand all files Collapse all files

6
packages/agent/README.md
View file

@ -56,7 +56,6 @@ console.log(agent.state.messages);
The agent internally works with `AgentMessage`, a flexible type that can include:
- Standard LLM messages (`user`, `assistant`, `toolResult`)
- User messages with attachments
- Custom app-specific message types (via declaration merging)
LLMs only understand a subset: `user`, `assistant`, and `toolResult` messages with specific content formats. The `convertToLlm` function bridges this gap.
@ -166,11 +165,14 @@ When you call `prompt(message)`, the agent emits `message_start` and `message_en
```
prompt(userMessage)
→ agent_start
→ turn_start
→ message_start { message: userMessage }
→ message_end { message: userMessage }
→ message_start { message: assistantMessage } // LLM starts responding
→ message_update { message: partialAssistant } // streaming...
→ message_end { message: assistantMessage }
...
```
Queued messages (via `queueMessage()`) emit the same events when injected:
@ -238,7 +240,7 @@ Extend `AgentMessage` for app-specific messages via declaration merging:
```typescript
declare module '@mariozechner/pi-agent-core' {
interface CustomMessages {
interface CustomAgentMessages {
artifact: { role: 'artifact'; code: string; language: string; timestamp: number };
notification: { role: 'notification'; text: string; timestamp: number };
}

2
packages/agent/src/types.ts
View file

@ -151,7 +151,7 @@ export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any
// AgentContext is like Context but uses AgentTool
export interface AgentContext {
systemPrompt: string;
messages: Message[];
messages: AgentMessage[];
tools?: AgentTool<any>[];
}

3
packages/coding-agent/examples/custom-tools/subagent/index.ts
View file

@ -16,7 +16,8 @@ import { spawn } from "node:child_process";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import type { AgentToolResult, Message } from "@mariozechner/pi-ai";
import type { AgentToolResult } from "@mariozechner/pi-agent-core";
import type { Message } from "@mariozechner/pi-ai";
import { StringEnum } from "@mariozechner/pi-ai";
import {
type CustomAgentTool,

4
packages/coding-agent/examples/hooks/custom-compaction.ts
View file

@ -14,7 +14,7 @@
*/
import { complete, getModel } from "@mariozechner/pi-ai";
import { messageTransformer } from "@mariozechner/pi-coding-agent";
import { convertToLlm } from "@mariozechner/pi-coding-agent";
import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";
export default function (pi: HookAPI) {
@ -52,7 +52,7 @@ export default function (pi: HookAPI) {
);
// Transform app messages to pi-ai package format
const transformedMessages = messageTransformer(allMessages);
const transformedMessages = convertToLlm(allMessages);
// Include previous summary context if available
const previousContext = previousSummary ? `\n\nPrevious session summary for context:\n${previousSummary}` : "";

25
packages/coding-agent/src/cli/file-processor.ts
View file

@ -3,21 +3,21 @@
*/
import { access, readFile, stat } from "node:fs/promises";
import type { Attachment } from "@mariozechner/pi-agent-core";
import type { ImageContent } from "@mariozechner/pi-ai";
import chalk from "chalk";
import { resolve } from "path";
import { resolveReadPath } from "../core/tools/path-utils.js";
import { detectSupportedImageMimeTypeFromFile } from "../utils/mime.js";
export interface ProcessedFiles {
textContent: string;
imageAttachments: Attachment[];
text: string;
images: ImageContent[];
}
/** Process @file arguments into text content and image attachments */
export async function processFileArguments(fileArgs: string[]): Promise<ProcessedFiles> {
let textContent = "";
const imageAttachments: Attachment[] = [];
let text = "";
const images: ImageContent[] = [];
for (const fileArg of fileArgs) {
// Expand and resolve path (handles ~ expansion and macOS screenshot Unicode spaces)
@ -45,24 +45,21 @@ export async function processFileArguments(fileArgs: string[]): Promise<Processe
const content = await readFile(absolutePath);
const base64Content = content.toString("base64");
const attachment: Attachment = {
id: `file-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`,
const attachment: ImageContent = {
type: "image",
fileName: absolutePath.split("/").pop() || absolutePath,
mimeType,
size: stats.size,
content: base64Content,
data: base64Content,
};
imageAttachments.push(attachment);
images.push(attachment);
// Add text reference to image
textContent += `<file name="${absolutePath}"></file>\n`;
text += `<file name="${absolutePath}"></file>\n`;
} else {
// Handle text file
try {
const content = await readFile(absolutePath, "utf-8");
textContent += `<file name="${absolutePath}">\n${content}\n</file>\n`;
text += `<file name="${absolutePath}">\n${content}\n</file>\n`;
} catch (error: unknown) {
const message = error instanceof Error ? error.message : String(error);
console.error(chalk.red(`Error: Could not read file ${absolutePath}: ${message}`));
@ -71,5 +68,5 @@ export async function processFileArguments(fileArgs: string[]): Promise<Processe
}
}
return { textContent, imageAttachments };
return { text, images };
}

17
packages/coding-agent/src/core/agent-session.ts
View file

@ -13,15 +13,8 @@
* Modes use this class and add their own I/O layer on top.
*/
import type {
Agent,
AgentEvent,
AgentMessage,
AgentState,
Attachment,
ThinkingLevel,
} from "@mariozechner/pi-agent-core";
import type { AssistantMessage, Message, Model, TextContent } from "@mariozechner/pi-ai";
import type { Agent, AgentEvent, AgentMessage, AgentState, ThinkingLevel } from "@mariozechner/pi-agent-core";
import type { AssistantMessage, ImageContent, Message, Model, TextContent } from "@mariozechner/pi-ai";
import { isContextOverflow, modelsAreEqual, supportsXhigh } from "@mariozechner/pi-ai";
import { getAuthPath } from "../config.js";
import { type BashResult, executeBash as executeBashCommand } from "./bash-executor.js";
@ -83,8 +76,8 @@ export interface AgentSessionConfig {
export interface PromptOptions {
/** Whether to expand file-based slash commands (default: true) */
expandSlashCommands?: boolean;
/** Image/file attachments */
attachments?: Attachment[];
/** Image attachments */
images?: ImageContent[];
}
/** Result from cycleModel() */
@ -492,7 +485,7 @@ export class AgentSession {
// Expand file-based slash commands if requested
const expandedText = expandCommands ? expandSlashCommand(text, [...this._fileCommands]) : text;
await this.agent.prompt(expandedText, options?.attachments);
await this.agent.prompt(expandedText, options?.images);
await this.waitForRetry();
}

6
packages/coding-agent/src/core/compaction.ts
View file

@ -8,7 +8,7 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai";
import { complete } from "@mariozechner/pi-ai";
import { messageTransformer } from "./messages.js";
import { convertToLlm } from "./messages.js";
import { type CompactionEntry, createSummaryMessage, type SessionEntry } from "./session-manager.js";
/**
@ -337,7 +337,7 @@ export async function generateSummary(
: SUMMARIZATION_PROMPT;
// Transform custom messages (like bashExecution) to LLM-compatible messages
const transformedMessages = messageTransformer(currentMessages);
const transformedMessages = convertToLlm(currentMessages);
const summarizationMessages = [
...transformedMessages,
@ -558,7 +558,7 @@ async function generateTurnPrefixSummary(
): Promise<string> {
const maxTokens = Math.floor(0.5 * reserveTokens); // Smaller budget for turn prefix
const transformedMessages = messageTransformer(messages);
const transformedMessages = convertToLlm(messages);
const summarizationMessages = [
...transformedMessages,
{

2
packages/coding-agent/src/core/custom-tools/types.ts
View file

@ -5,7 +5,7 @@
* They can provide custom rendering for tool calls and results in the TUI.
*/
import type { AgentTool, AgentToolResult, AgentToolUpdateCallback } from "@mariozechner/pi-ai";
import type { AgentTool, AgentToolResult, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core";
import type { Component } from "@mariozechner/pi-tui";
import type { Static, TSchema } from "@sinclair/typebox";
import type { Theme } from "../../modes/interactive/theme/theme.js";

4
packages/coding-agent/src/core/hooks/runner.ts
View file

@ -2,7 +2,7 @@
* Hook runner - executes hooks and manages their lifecycle.
*/
import type { Message } from "@mariozechner/pi-ai";
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { ModelRegistry } from "../model-registry.js";
import type { SessionManager } from "../session-manager.js";
import type { AppendEntryHandler, LoadedHook, SendMessageHandler } from "./loader.js";
@ -315,7 +315,7 @@ export class HookRunner {
*
* Note: Messages are already deep-copied by the caller (pi-ai preprocessor).
*/
async emitContext(messages: Message[]): Promise<Message[]> {
async emitContext(messages: AgentMessage[]): Promise<AgentMessage[]> {
const ctx = this.createContext();
let currentMessages = messages;

2
packages/coding-agent/src/core/hooks/tool-wrapper.ts
View file

@ -2,7 +2,7 @@
* Tool wrapper - wraps tools with hook callbacks for interception.
*/
import type { AgentTool, AgentToolUpdateCallback } from "@mariozechner/pi-ai";
import type { AgentTool, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core";
import type { HookRunner } from "./runner.js";
import type { ToolCallEventResult, ToolResultEventResult } from "./types.js";

3
packages/coding-agent/src/core/hooks/types.ts
View file

@ -151,12 +151,11 @@ export type SessionEvent =
* Event data for context event.
* Fired before each LLM call, allowing hooks to modify context non-destructively.
* Original session messages are NOT modified - only the messages sent to the LLM are affected.
* Messages are already in LLM format (Message[], not AgentMessage[]).
*/
export interface ContextEvent {
type: "context";
/** Messages about to be sent to the LLM (deep copy, safe to modify) */
messages: Message[];
messages: AgentMessage[];
}
/**

37
packages/coding-agent/src/core/messages.ts
View file

@ -6,11 +6,7 @@
*/
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { Message } from "@mariozechner/pi-ai";
// ============================================================================
// Custom Message Types
// ============================================================================
import type { ImageContent, Message, TextContent } from "@mariozechner/pi-ai";
/**
* Message type for bash executions via the ! command.
@ -26,8 +22,6 @@ export interface BashExecutionMessage {
timestamp: number;
}
import type { ImageContent, TextContent } from "@mariozechner/pi-ai";
/**
* Message type for hook-injected messages via sendMessage().
* These are custom messages that hooks can inject into the conversation.
@ -41,36 +35,28 @@ export interface HookMessage<T = unknown> {
timestamp: number;
}
// Extend CustomMessages via declaration merging
// Extend CustomAgentMessages via declaration merging
declare module "@mariozechner/pi-agent-core" {
interface CustomMessages {
interface CustomAgentMessages {
bashExecution: BashExecutionMessage;
hookMessage: HookMessage;
}
}
// ============================================================================
// Type Guards
// ============================================================================
/**
* Type guard for BashExecutionMessage.
*/
export function isBashExecutionMessage(msg: AgentMessage | Message): msg is BashExecutionMessage {
return (msg as BashExecutionMessage).role === "bashExecution";
return msg.role === "bashExecution";
}
/**
* Type guard for HookAgentMessage.
* Type guard for HookMessage.
*/
export function isHookMessage(msg: AgentMessage | Message): msg is HookMessage {
return (msg as HookMessage).role === "hookMessage";
return msg.role === "hookMessage";
}
// ============================================================================
// Message Formatting
// ============================================================================
/**
* Convert a BashExecutionMessage to user message text for LLM context.
*/
@ -92,18 +78,15 @@ export function bashExecutionToText(msg: BashExecutionMessage): string {
return text;
}
// ============================================================================
// Message Transformer
// ============================================================================
/**
* Transform AgentMessages (including custom types) to LLM-compatible Messages.
*
* This is used by:
* - Agent's messageTransformer option (for prompt calls)
* - Agent's transormToLlm option (for prompt calls and queued messages)
* - Compaction's generateSummary (for summarization)
* - Custom hooks and tools
*/
export function messageTransformer(messages: AgentMessage[]): Message[] {
export function convertToLlm(messages: AgentMessage[]): Message[] {
return messages
.map((m): Message | null => {
if (isBashExecutionMessage(m)) {
@ -131,5 +114,5 @@ export function messageTransformer(messages: AgentMessage[]): Message[] {
// Filter out unknown message types
return null;
})
.filter((m): m is Message => m !== null);
.filter((m) => m !== null);
}

32
packages/coding-agent/src/core/sdk.ts
View file

@ -29,7 +29,7 @@
* ```
*/
import { Agent, ProviderTransport, type ThinkingLevel } from "@mariozechner/pi-agent-core";
import { Agent, type ThinkingLevel } from "@mariozechner/pi-agent-core";
import type { Model } from "@mariozechner/pi-ai";
import { join } from "path";
import { getAgentDir } from "../config.js";
@ -39,7 +39,7 @@ import { discoverAndLoadCustomTools, type LoadedCustomTool } from "./custom-tool
import type { CustomAgentTool } from "./custom-tools/types.js";
import { discoverAndLoadHooks, HookRunner, type LoadedHook, wrapToolsWithHooks } from "./hooks/index.js";
import type { HookFactory } from "./hooks/types.js";
import { messageTransformer } from "./messages.js";
import { convertToLlm } from "./messages.js";
import { ModelRegistry } from "./model-registry.js";
import { SessionManager } from "./session-manager.js";
import { type Settings, SettingsManager, type SkillsSettings } from "./settings-manager.js";
@ -588,26 +588,24 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
thinkingLevel,
tools: allToolsArray,
},
messageTransformer,
preprocessor: hookRunner
convertToLlm,
transformContext: hookRunner
? async (messages) => {
return hookRunner.emitContext(messages);
}
: undefined,
queueMode: settingsManager.getQueueMode(),
transport: new ProviderTransport({
getApiKey: async () => {
const currentModel = agent.state.model;
if (!currentModel) {
throw new Error("No model selected");
}
const key = await modelRegistry.getApiKey(currentModel);
if (!key) {
throw new Error(`No API key found for provider "${currentModel.provider}"`);
}
return key;
},
}),
getApiKey: async () => {
const currentModel = agent.state.model;
if (!currentModel) {
throw new Error("No model selected");
}
const key = await modelRegistry.getApiKey(currentModel);
if (!key) {
throw new Error(`No API key found for provider "${currentModel.provider}"`);
}
return key;
},
});
time("createAgent");

2
packages/coding-agent/src/core/tools/bash.ts
View file

@ -2,7 +2,7 @@ import { randomBytes } from "node:crypto";
import { createWriteStream } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";
import { spawn } from "child_process";
import { getShellConfig, killProcessTree } from "../../utils/shell.js";

2
packages/coding-agent/src/core/tools/edit.ts
View file

@ -1,4 +1,4 @@
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";
import * as Diff from "diff";
import { constants } from "fs";

2
packages/coding-agent/src/core/tools/find.ts
View file

@ -1,4 +1,4 @@
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";
import { spawnSync } from "child_process";
import { existsSync } from "fs";

2
packages/coding-agent/src/core/tools/grep.ts
View file

@ -1,5 +1,5 @@
import { createInterface } from "node:readline";
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";
import { spawn } from "child_process";
import { readFileSync, type Stats, statSync } from "fs";

3
packages/coding-agent/src/core/tools/index.ts
View file

@ -1,5 +1,3 @@
import type { AgentTool } from "@mariozechner/pi-ai";
export { type BashToolDetails, bashTool, createBashTool } from "./bash.js";
export { createEditTool, editTool } from "./edit.js";
export { createFindTool, type FindToolDetails, findTool } from "./find.js";
@ -9,6 +7,7 @@ export { createReadTool, type ReadToolDetails, readTool } from "./read.js";
export type { TruncationResult } from "./truncate.js";
export { createWriteTool, writeTool } from "./write.js";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { bashTool, createBashTool } from "./bash.js";
import { createEditTool, editTool } from "./edit.js";
import { createFindTool, findTool } from "./find.js";

2
packages/coding-agent/src/core/tools/ls.ts
View file

@ -1,4 +1,4 @@
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";
import { existsSync, readdirSync, statSync } from "fs";
import nodePath from "path";

3
packages/coding-agent/src/core/tools/read.ts
View file

@ -1,4 +1,5 @@
import type { AgentTool, ImageContent, TextContent } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import type { ImageContent, TextContent } from "@mariozechner/pi-ai";
import { Type } from "@sinclair/typebox";
import { constants } from "fs";
import { access, readFile } from "fs/promises";

2
packages/coding-agent/src/core/tools/write.ts
View file

@ -1,4 +1,4 @@
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";
import { mkdir, writeFile } from "fs/promises";
import { dirname } from "path";

2
packages/coding-agent/src/index.ts
View file

@ -73,7 +73,7 @@ export {
isReadToolResult,
isWriteToolResult,
} from "./core/hooks/index.js";
export { messageTransformer } from "./core/messages.js";
export { convertToLlm } from "./core/messages.js";
export { ModelRegistry } from "./core/model-registry.js";
// SDK for programmatic usage
export {

23
packages/coding-agent/src/main.ts
View file

@ -5,8 +5,7 @@
* createAgentSession() options. The SDK does the heavy lifting.
*/
import type { Attachment } from "@mariozechner/pi-agent-core";
import { supportsXhigh } from "@mariozechner/pi-ai";
import { type ImageContent, supportsXhigh } from "@mariozechner/pi-ai";
import chalk from "chalk";
import { existsSync } from "fs";
import { join } from "path";
@ -64,7 +63,7 @@ async function runInteractiveMode(
customTools: LoadedCustomTool[],
setToolUIContext: (uiContext: HookUIContext, hasUI: boolean) => void,
initialMessage?: string,
initialAttachments?: Attachment[],
initialImages?: ImageContent[],
fdPath: string | null = null,
): Promise<void> {
const mode = new InteractiveMode(session, version, changelogMarkdown, customTools, setToolUIContext, fdPath);
@ -93,7 +92,7 @@ async function runInteractiveMode(
if (initialMessage) {
try {
await session.prompt(initialMessage, { attachments: initialAttachments });
await session.prompt(initialMessage, { images: initialImages });
} catch (error: unknown) {
const errorMessage = error instanceof Error ? error.message : "Unknown error occurred";
mode.showError(errorMessage);
@ -122,25 +121,25 @@ async function runInteractiveMode(
async function prepareInitialMessage(parsed: Args): Promise<{
initialMessage?: string;
initialAttachments?: Attachment[];
initialImages?: ImageContent[];
}> {
if (parsed.fileArgs.length === 0) {
return {};
}
const { textContent, imageAttachments } = await processFileArguments(parsed.fileArgs);
const { text, images } = await processFileArguments(parsed.fileArgs);
let initialMessage: string;
if (parsed.messages.length > 0) {
initialMessage = textContent + parsed.messages[0];
initialMessage = text + parsed.messages[0];
parsed.messages.shift();
} else {
initialMessage = textContent;
initialMessage = text;
}
return {
initialMessage,
initialAttachments: imageAttachments.length > 0 ? imageAttachments : undefined,
initialImages: images.length > 0 ? images : undefined,
};
}
@ -330,7 +329,7 @@ export async function main(args: string[]) {
}
const cwd = process.cwd();
const { initialMessage, initialAttachments } = await prepareInitialMessage(parsed);
const { initialMessage, initialImages } = await prepareInitialMessage(parsed);
time("prepareInitialMessage");
const isInteractive = !parsed.print && parsed.mode === undefined;
const mode = parsed.mode || "text";
@ -438,11 +437,11 @@ export async function main(args: string[]) {
customToolsResult.tools,
customToolsResult.setUIContext,
initialMessage,
initialAttachments,
initialImages,
fdPath,
);
} else {
await runPrintMode(session, mode, parsed.messages, initialMessage, initialAttachments);
await runPrintMode(session, mode, parsed.messages, initialMessage, initialImages);
stopThemeWatcher();
if (process.stdout.writableLength > 0) {
await new Promise<void>((resolve) => process.stdout.once("drain", resolve));

9
packages/coding-agent/src/modes/print-mode.ts
View file

@ -6,8 +6,7 @@
* - `pi --mode json "prompt"` - JSON event stream
*/
import type { Attachment } from "@mariozechner/pi-agent-core";
import type { AssistantMessage } from "@mariozechner/pi-ai";
import type { AssistantMessage, ImageContent } from "@mariozechner/pi-ai";
import type { AgentSession } from "../core/agent-session.js";
/**
@ -18,14 +17,14 @@ import type { AgentSession } from "../core/agent-session.js";
* @param mode Output mode: "text" for final response only, "json" for all events
* @param messages Array of prompts to send
* @param initialMessage Optional first message (may contain @file content)
* @param initialAttachments Optional attachments for the initial message
* @param initialImages Optional images for the initial message
*/
export async function runPrintMode(
session: AgentSession,
mode: "text" | "json",
messages: string[],
initialMessage?: string,
initialAttachments?: Attachment[],
initialImages?: ImageContent[],
): Promise<void> {
// Load entries once for session start events
const entries = session.sessionManager.getEntries();
@ -79,7 +78,7 @@ export async function runPrintMode(
// Send initial message with attachments
if (initialMessage) {
await session.prompt(initialMessage, { attachments: initialAttachments });
await session.prompt(initialMessage, { images: initialImages });
}
// Send remaining messages

11
packages/coding-agent/src/modes/rpc/rpc-client.ts
View file

@ -6,7 +6,8 @@
import { type ChildProcess, spawn } from "node:child_process";
import * as readline from "node:readline";
import type { AgentEvent, AgentMessage, Attachment, ThinkingLevel } from "@mariozechner/pi-agent-core";
import type { AgentEvent, AgentMessage, ThinkingLevel } from "@mariozechner/pi-agent-core";
import type { ImageContent } from "@mariozechner/pi-ai";
import type { SessionStats } from "../../core/agent-session.js";
import type { BashResult } from "../../core/bash-executor.js";
import type { CompactionResult } from "../../core/compaction.js";
@ -167,8 +168,8 @@ export class RpcClient {
* Returns immediately after sending; use onEvent() to receive streaming events.
* Use waitForIdle() to wait for completion.
*/
async prompt(message: string, attachments?: Attachment[]): Promise<void> {
await this.send({ type: "prompt", message, attachments });
async prompt(message: string, images?: ImageContent[]): Promise<void> {
await this.send({ type: "prompt", message, images });
}
/**
@ -404,9 +405,9 @@ export class RpcClient {
/**
* Send prompt and wait for completion, returning all events.
*/
async promptAndWait(message: string, attachments?: Attachment[], timeout = 60000): Promise<AgentEvent[]> {
async promptAndWait(message: string, images?: ImageContent[], timeout = 60000): Promise<AgentEvent[]> {
const eventsPromise = this.collectEvents(timeout);
await this.prompt(message, attachments);
await this.prompt(message, images);
return eventsPromise;
}

2
packages/coding-agent/src/modes/rpc/rpc-mode.ts
View file

@ -187,7 +187,7 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
// Hook commands and file slash commands are handled in session.prompt()
session
.prompt(command.message, {
attachments: command.attachments,
images: command.images,
})
.catch((e) => output(error(id, "prompt", e.message)));
return success(id, "prompt");

6
packages/coding-agent/src/modes/rpc/rpc-types.ts
View file

@ -5,8 +5,8 @@
* Responses and events are emitted as JSON lines on stdout.
*/
import type { AgentMessage, Attachment, ThinkingLevel } from "@mariozechner/pi-agent-core";
import type { Model } from "@mariozechner/pi-ai";
import type { AgentMessage, ThinkingLevel } from "@mariozechner/pi-agent-core";
import type { ImageContent, Model } from "@mariozechner/pi-ai";
import type { SessionStats } from "../../core/agent-session.js";
import type { BashResult } from "../../core/bash-executor.js";
import type { CompactionResult } from "../../core/compaction.js";
@ -17,7 +17,7 @@ import type { CompactionResult } from "../../core/compaction.js";
export type RpcCommand =
// Prompting
| { id?: string; type: "prompt"; message: string; attachments?: Attachment[] }
| { id?: string; type: "prompt"; message: string; images?: ImageContent[] }
| { id?: string; type: "queue_message"; message: string }
| { id?: string; type: "abort" }
| { id?: string; type: "reset" }

9
packages/coding-agent/test/agent-session-branching.test.ts
View file

@ -10,7 +10,7 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent, ProviderTransport } from "@mariozechner/pi-agent-core";
import { Agent } from "@mariozechner/pi-agent-core";
import { getModel } from "@mariozechner/pi-ai";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AgentSession } from "../src/core/agent-session.js";
@ -44,13 +44,8 @@ describe.skipIf(!API_KEY)("AgentSession branching", () => {
function createSession(noSession: boolean = false) {
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const transport = new ProviderTransport({
getApiKey: () => API_KEY,
});
const agent = new Agent({
transport,
getApiKey: () => API_KEY,
initialState: {
model,
systemPrompt: "You are a helpful assistant. Be extremely concise, reply with just a few words.",

9
packages/coding-agent/test/agent-session-compaction.test.ts
View file

@ -10,7 +10,7 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent, ProviderTransport } from "@mariozechner/pi-agent-core";
import { Agent } from "@mariozechner/pi-agent-core";
import { getModel } from "@mariozechner/pi-ai";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AgentSession, type AgentSessionEvent } from "../src/core/agent-session.js";
@ -48,13 +48,8 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
function createSession(inMemory = false) {
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const transport = new ProviderTransport({
getApiKey: () => API_KEY,
});
const agent = new Agent({
transport,
getApiKey: () => API_KEY,
initialState: {
model,
systemPrompt: "You are a helpful assistant. Be concise.",

9
packages/coding-agent/test/compaction-hooks.test.ts
View file

@ -5,7 +5,7 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent, ProviderTransport } from "@mariozechner/pi-agent-core";
import { Agent } from "@mariozechner/pi-agent-core";
import { getModel } from "@mariozechner/pi-ai";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AgentSession } from "../src/core/agent-session.js";
@ -72,13 +72,8 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
function createSession(hooks: LoadedHook[]) {
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const transport = new ProviderTransport({
getApiKey: () => API_KEY,
});
const agent = new Agent({
transport,
getApiKey: () => API_KEY,
initialState: {
model,
systemPrompt: "You are a helpful assistant. Be concise.",

18
packages/mom/src/agent.ts
View file

@ -1,15 +1,15 @@
import { Agent, type AgentEvent, type Attachment, ProviderTransport } from "@mariozechner/pi-agent-core";
import { getModel } from "@mariozechner/pi-ai";
import { Agent, type AgentEvent } from "@mariozechner/pi-agent-core";
import { getModel, type ImageContent } from "@mariozechner/pi-ai";
import {
AgentSession,
AuthStorage,
convertToLlm,
formatSkillsForPrompt,
loadSkillsFromDir,
ModelRegistry,
messageTransformer,
type Skill,
} from "@mariozechner/pi-coding-agent";
import { existsSync, readFileSync, statSync } from "fs";
import { existsSync, readFileSync } from "fs";
import { mkdir, writeFile } from "fs/promises";
import { homedir } from "os";
import { join } from "path";
@ -716,7 +716,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
const timestamp = `${now.getFullYear()}-${pad(now.getMonth() + 1)}-${pad(now.getDate())} ${pad(now.getHours())}:${pad(now.getMinutes())}:${pad(now.getSeconds())}${offsetSign}${offsetHours}:${offsetMins}`;
let userMessage = `[${timestamp}] [${ctx.message.userName || "unknown"}]: ${ctx.message.text}`;
const imageAttachments: Attachment[] = [];
const imageAttachments: ImageContent[] = [];
const nonImagePaths: string[] = [];
for (const a of ctx.message.attachments || []) {
@ -725,14 +725,10 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
if (mimeType && existsSync(fullPath)) {
try {
const stats = statSync(fullPath);
imageAttachments.push({
id: a.local,
type: "image",
fileName: a.local.split("/").pop() || a.local,
mimeType,
size: stats.size,
content: readFileSync(fullPath).toString("base64"),
data: readFileSync(fullPath).toString("base64"),
});
} catch {
nonImagePaths.push(fullPath);
@ -755,7 +751,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
};
await writeFile(join(channelDir, "last_prompt.jsonl"), JSON.stringify(debugContext, null, 2));
await session.prompt(userMessage, imageAttachments.length > 0 ? { attachments: imageAttachments } : undefined);
await session.prompt(userMessage, imageAttachments.length > 0 ? { images: imageAttachments } : undefined);
// Wait for queued messages
await queueChain;

2
packages/mom/src/tools/attach.ts
View file

@ -1,4 +1,4 @@
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";
import { basename, resolve as resolvePath } from "path";

2
packages/mom/src/tools/bash.ts
View file

@ -2,7 +2,7 @@ import { randomBytes } from "node:crypto";
import { createWriteStream } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";
import type { Executor } from "../sandbox.js";
import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize, type TruncationResult, truncateTail } from "./truncate.js";

2
packages/mom/src/tools/edit.ts
View file

@ -1,4 +1,4 @@
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";
import * as Diff from "diff";
import type { Executor } from "../sandbox.js";

2
packages/mom/src/tools/index.ts
View file

@ -1,4 +1,4 @@
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import type { Executor } from "../sandbox.js";
import { attachTool } from "./attach.js";
import { createBashTool } from "./bash.js";

3
packages/mom/src/tools/read.ts
View file

@ -1,4 +1,5 @@
import type { AgentTool, ImageContent, TextContent } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import type { ImageContent, TextContent } from "@mariozechner/pi-ai";
import { Type } from "@sinclair/typebox";
import { extname } from "path";
import type { Executor } from "../sandbox.js";

2
packages/mom/src/tools/write.ts
View file

@ -1,4 +1,4 @@
import type { AgentTool } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";
import type { Executor } from "../sandbox.js";

6
packages/web-ui/example/src/custom-messages.ts
View file

@ -18,7 +18,7 @@ export interface SystemNotificationMessage {
// Extend CustomMessages interface via declaration merging
declare module "@mariozechner/pi-web-ui" {
interface CustomMessages {
interface CustomAgentMessages {
"system-notification": SystemNotificationMessage;
}
}
@ -99,8 +99,8 @@ export function customMessageTransformer(messages: AppMessage[]): Message[] {
}
// Strip attachments from user messages
if (m.role === "user") {
const { attachments: _, ...rest } = m as any;
if (m.role === "user-with-attachment") {
const { attachments: _, ...rest } = m;
return rest as Message;
}

5
packages/web-ui/src/ChatPanel.ts
View file

@ -1,9 +1,8 @@
import { Badge } from "@mariozechner/mini-lit/dist/Badge.js";
import { html, LitElement } from "lit";
import { customElement, state } from "lit/decorators.js";
import type { Agent } from "./agent/agent.js";
import "./components/AgentInterface.js";
import type { AgentTool } from "@mariozechner/pi-ai";
import type { Agent, AgentTool } from "@mariozechner/pi-agent-core";
import type { AgentInterface } from "./components/AgentInterface.js";
import { ArtifactsRuntimeProvider } from "./components/sandbox/ArtifactsRuntimeProvider.js";
import { AttachmentsRuntimeProvider } from "./components/sandbox/AttachmentsRuntimeProvider.js";
@ -95,7 +94,7 @@ export class ChatPanel extends LitElement {
const runtimeProvidersFactory = () => {
const attachments: Attachment[] = [];
for (const message of this.agent!.state.messages) {
if (message.role === "user") {
if (message.role === "user-with-attachments") {
message.attachments?.forEach((a) => {
attachments.push(a);
});

341
packages/web-ui/src/agent/agent.ts
View file

@ -1,341 +0,0 @@
import type { Context, QueuedMessage } from "@mariozechner/pi-ai";
import {
type AgentTool,
type AssistantMessage as AssistantMessageType,
getModel,
type ImageContent,
type Message,
type Model,
type TextContent,
} from "@mariozechner/pi-ai";
import type { AppMessage } from "../components/Messages.js";
import type { Attachment } from "../utils/attachment-utils.js";
import type { AgentRunConfig, AgentTransport } from "./transports/types.js";
import type { DebugLogEntry } from "./types.js";
// Default transformer: Keep only LLM-compatible messages, strip app-specific fields
function defaultMessageTransformer(messages: AppMessage[]): Message[] {
return messages
.filter((m) => {
// Only keep standard LLM message roles
return m.role === "user" || m.role === "assistant" || m.role === "toolResult";
})
.map((m) => {
if (m.role === "user") {
// Strip attachments field (app-specific)
// biome-ignore lint/correctness/noUnusedVariables: fine here
const { attachments, ...rest } = m as any;
return rest as Message;
}
return m as Message;
});
}
export type ThinkingLevel = "off" | "minimal" | "low" | "medium" | "high";
export interface AgentState {
systemPrompt: string;
model: Model<any>;
thinkingLevel: ThinkingLevel;
tools: AgentTool<any>[];
messages: AppMessage[];
isStreaming: boolean;
streamMessage: Message | null;
pendingToolCalls: Set<string>;
error?: string;
}
export type AgentEvent =
| { type: "state-update"; state: AgentState }
| { type: "error-no-model" }
| { type: "error-no-api-key"; provider: string }
| { type: "started" }
| { type: "completed" };
export interface AgentOptions {
initialState?: Partial<AgentState>;
debugListener?: (entry: DebugLogEntry) => void;
transport: AgentTransport;
// Transform app messages to LLM-compatible messages before sending to transport
messageTransformer?: (messages: AppMessage[]) => Message[] | Promise<Message[]>;
}
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 transport: AgentTransport;
private debugListener?: (entry: DebugLogEntry) => void;
private messageTransformer: (messages: AppMessage[]) => Message[] | Promise<Message[]>;
private messageQueue: Array<QueuedMessage<AppMessage>> = [];
constructor(opts: AgentOptions) {
this._state = { ...this._state, ...opts.initialState };
this.debugListener = opts.debugListener;
this.transport = opts.transport;
this.messageTransformer = opts.messageTransformer || defaultMessageTransformer;
}
get state(): AgentState {
return this._state;
}
subscribe(fn: (e: AgentEvent) => void): () => void {
this.listeners.add(fn);
fn({ type: "state-update", state: this._state });
return () => this.listeners.delete(fn);
}
// Mutators
setSystemPrompt(v: string) {
this.patch({ systemPrompt: v });
}
setModel(m: Model<any>) {
this.patch({ model: m });
}
setThinkingLevel(l: ThinkingLevel) {
this.patch({ thinkingLevel: l });
}
setTools(t: AgentTool<any>[]) {
this.patch({ tools: t });
}
replaceMessages(ms: AppMessage[]) {
this.patch({ messages: ms.slice() });
}
appendMessage(m: AppMessage) {
this.patch({ messages: [...this._state.messages, m] });
}
async queueMessage(m: AppMessage) {
// Transform message and queue it for injection at next turn
const transformed = await this.messageTransformer([m]);
this.messageQueue.push({
original: m,
llm: transformed[0], // undefined if filtered out
});
}
clearMessages() {
this.patch({ messages: [] });
}
abort() {
this.abortController?.abort();
}
private logState(message: string) {
const { systemPrompt, model, messages } = this._state;
console.log(message, { systemPrompt, model, messages });
}
async prompt(input: string, attachments?: Attachment[]) {
const model = this._state.model;
if (!model) {
this.emit({ type: "error-no-model" });
return;
}
// Build user message with attachments
const content: Array<TextContent | ImageContent> = [{ type: "text", text: input }];
if (attachments?.length) {
for (const a of attachments) {
if (a.type === "image") {
content.push({ type: "image", data: a.content, mimeType: a.mimeType });
} else if (a.type === "document" && a.extractedText) {
content.push({
type: "text",
text: `\n\n[Document: ${a.fileName}]\n${a.extractedText}`,
isDocument: true,
} as TextContent);
}
}
}
const userMessage: AppMessage = {
role: "user",
content,
attachments: attachments?.length ? attachments : undefined,
timestamp: Date.now(),
};
this.abortController = new AbortController();
this.patch({ isStreaming: true, streamMessage: null, error: undefined });
this.emit({ type: "started" });
const reasoning =
this._state.thinkingLevel === "off"
? undefined
: this._state.thinkingLevel === "minimal"
? "low"
: this._state.thinkingLevel;
const cfg: AgentRunConfig = {
systemPrompt: this._state.systemPrompt,
tools: this._state.tools,
model,
reasoning,
getQueuedMessages: async <T>() => {
// Return queued messages (they'll be added to state via message_end event)
const queued = this.messageQueue.slice();
this.messageQueue = [];
return queued as QueuedMessage<T>[];
},
};
try {
let partial: Message | null = null;
let turnDebug: DebugLogEntry | null = null;
let turnStart = 0;
this.logState("prompt started, current state:");
// Transform app messages to LLM-compatible messages (initial set)
const llmMessages = await this.messageTransformer(this._state.messages);
console.log("transformed messages:", llmMessages);
for await (const ev of this.transport.run(
llmMessages,
userMessage as Message,
cfg,
this.abortController.signal,
)) {
switch (ev.type) {
case "turn_start": {
turnStart = performance.now();
// Build request context snapshot (use transformed messages)
const ctx: Context = {
systemPrompt: this._state.systemPrompt,
messages: [...llmMessages],
tools: this._state.tools,
};
turnDebug = {
timestamp: new Date().toISOString(),
request: {
provider: cfg.model.provider,
model: cfg.model.id,
context: { ...ctx },
},
sseEvents: [],
};
break;
}
case "message_start":
case "message_update": {
partial = ev.message;
// Collect SSE-like events for debug (drop heavy partial)
if (ev.type === "message_update" && ev.assistantMessageEvent && turnDebug) {
const copy: any = { ...ev.assistantMessageEvent };
if (copy && "partial" in copy) delete copy.partial;
turnDebug.sseEvents.push(JSON.stringify(copy));
if (!turnDebug.ttft) turnDebug.ttft = performance.now() - turnStart;
}
this.patch({ streamMessage: ev.message });
break;
}
case "message_end": {
partial = null;
this.appendMessage(ev.message as AppMessage);
this.patch({ streamMessage: null });
if (turnDebug) {
if (ev.message.role !== "assistant" && ev.message.role !== "toolResult") {
turnDebug.request.context.messages.push(ev.message);
}
if (ev.message.role === "assistant") turnDebug.response = ev.message as any;
}
break;
}
case "tool_execution_start": {
const s = new Set(this._state.pendingToolCalls);
s.add(ev.toolCallId);
this.patch({ pendingToolCalls: s });
break;
}
case "tool_execution_end": {
const s = new Set(this._state.pendingToolCalls);
s.delete(ev.toolCallId);
this.patch({ pendingToolCalls: s });
break;
}
case "turn_end": {
// finalize current turn
if (turnDebug) {
turnDebug.totalTime = performance.now() - turnStart;
this.debugListener?.(turnDebug);
turnDebug = null;
}
break;
}
case "agent_end": {
this.patch({ streamMessage: null });
break;
}
}
}
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 as AppMessage);
} else {
if (this.abortController?.signal.aborted) {
throw new Error("Request was aborted");
}
}
}
} catch (err: any) {
if (String(err?.message || err) === "no-api-key") {
this.emit({ type: "error-no-api-key", provider: model.provider });
} else {
const msg: AssistantMessageType = {
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(),
};
this.appendMessage(msg as AppMessage);
this.patch({ error: err?.message || String(err) });
}
} finally {
this.patch({ isStreaming: false, streamMessage: null, pendingToolCalls: new Set<string>() });
this.abortController = undefined;
this.emit({ type: "completed" });
}
this.logState("final state:");
}
private patch(p: Partial<AgentState>): void {
this._state = { ...this._state, ...p };
this.emit({ type: "state-update", state: this._state });
}
private emit(e: AgentEvent) {
for (const listener of this.listeners) {
listener(e);
}
}
}

371
packages/web-ui/src/agent/transports/AppTransport.ts
View file

@ -1,371 +0,0 @@
import type {
AgentContext,
AgentLoopConfig,
Api,
AssistantMessage,
AssistantMessageEvent,
Context,
Message,
Model,
SimpleStreamOptions,
ToolCall,
UserMessage,
} from "@mariozechner/pi-ai";
import { agentLoop, agentLoopContinue } from "@mariozechner/pi-ai";
import { AssistantMessageEventStream } from "@mariozechner/pi-ai/dist/utils/event-stream.js";
import { parseStreamingJson } from "@mariozechner/pi-ai/dist/utils/json-parse.js";
import { clearAuthToken, getAuthToken } from "../../utils/auth-token.js";
import { i18n } from "../../utils/i18n.js";
import type { ProxyAssistantMessageEvent } from "./proxy-types.js";
import type { AgentRunConfig, AgentTransport } from "./types.js";
/**
* Stream function that proxies through a server instead of calling providers directly.
* The server strips the partial field from delta events to reduce bandwidth.
* We reconstruct the partial message client-side.
*/
function streamSimpleProxy(
model: Model<any>,
context: Context,
options: SimpleStreamOptions & { authToken: string },
proxyUrl: string,
): AssistantMessageEventStream {
const stream = new AssistantMessageEventStream();
(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;
// Set up abort handler to cancel the reader
const abortHandler = () => {
if (reader) {
reader.cancel("Request aborted by user").catch(() => {});
}
};
if (options.signal) {
options.signal.addEventListener("abort", abortHandler);
}
try {
const response = await fetch(`${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,
// Don't send apiKey or signal - those are added server-side
},
}),
signal: options.signal,
});
if (!response.ok) {
let errorMessage = `Proxy error: ${response.status} ${response.statusText}`;
try {
const errorData = await response.json();
if (errorData.error) {
errorMessage = `Proxy error: ${errorData.error}`;
}
} catch {
// Couldn't parse error response, use default message
}
throw new Error(errorMessage);
}
// Parse SSE stream
reader = response.body!.getReader();
const decoder = new TextDecoder();
let buffer = "";
while (true) {
const { done, value } = await reader.read();
if (done) break;
// Check if aborted after reading
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;
let event: AssistantMessageEvent | undefined;
// Handle different event types
// Server sends events with partial for non-delta events,
// and without partial for delta events
switch (proxyEvent.type) {
case "start":
event = { type: "start", partial };
break;
case "text_start":
partial.content[proxyEvent.contentIndex] = {
type: "text",
text: "",
};
event = { type: "text_start", contentIndex: proxyEvent.contentIndex, partial };
break;
case "text_delta": {
const content = partial.content[proxyEvent.contentIndex];
if (content?.type === "text") {
content.text += proxyEvent.delta;
event = {
type: "text_delta",
contentIndex: proxyEvent.contentIndex,
delta: proxyEvent.delta,
partial,
};
} else {
throw new Error("Received text_delta for non-text content");
}
break;
}
case "text_end": {
const content = partial.content[proxyEvent.contentIndex];
if (content?.type === "text") {
content.textSignature = proxyEvent.contentSignature;
event = {
type: "text_end",
contentIndex: proxyEvent.contentIndex,
content: content.text,
partial,
};
} else {
throw new Error("Received text_end for non-text content");
}
break;
}
case "thinking_start":
partial.content[proxyEvent.contentIndex] = {
type: "thinking",
thinking: "",
};
event = { type: "thinking_start", contentIndex: proxyEvent.contentIndex, partial };
break;
case "thinking_delta": {
const content = partial.content[proxyEvent.contentIndex];
if (content?.type === "thinking") {
content.thinking += proxyEvent.delta;
event = {
type: "thinking_delta",
contentIndex: proxyEvent.contentIndex,
delta: proxyEvent.delta,
partial,
};
} else {
throw new Error("Received thinking_delta for non-thinking content");
}
break;
}
case "thinking_end": {
const content = partial.content[proxyEvent.contentIndex];
if (content?.type === "thinking") {
content.thinkingSignature = proxyEvent.contentSignature;
event = {
type: "thinking_end",
contentIndex: proxyEvent.contentIndex,
content: content.thinking,
partial,
};
} else {
throw new Error("Received thinking_end for non-thinking content");
}
break;
}
case "toolcall_start":
partial.content[proxyEvent.contentIndex] = {
type: "toolCall",
id: proxyEvent.id,
name: proxyEvent.toolName,
arguments: {},
partialJson: "",
} satisfies ToolCall & { partialJson: string } as ToolCall;
event = { type: "toolcall_start", contentIndex: proxyEvent.contentIndex, partial };
break;
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) || {};
event = {
type: "toolcall_delta",
contentIndex: proxyEvent.contentIndex,
delta: proxyEvent.delta,
partial,
};
partial.content[proxyEvent.contentIndex] = { ...content }; // Trigger reactivity
} else {
throw new Error("Received toolcall_delta for non-toolCall content");
}
break;
}
case "toolcall_end": {
const content = partial.content[proxyEvent.contentIndex];
if (content?.type === "toolCall") {
delete (content as any).partialJson;
event = {
type: "toolcall_end",
contentIndex: proxyEvent.contentIndex,
toolCall: content,
partial,
};
}
break;
}
case "done":
partial.stopReason = proxyEvent.reason;
partial.usage = proxyEvent.usage;
event = { type: "done", reason: proxyEvent.reason, message: partial };
break;
case "error":
partial.stopReason = proxyEvent.reason;
partial.errorMessage = proxyEvent.errorMessage;
partial.usage = proxyEvent.usage;
event = { type: "error", reason: proxyEvent.reason, error: partial };
break;
default: {
// Exhaustive check
const _exhaustiveCheck: never = proxyEvent;
console.warn(`Unhandled event type: ${(proxyEvent as any).type}`);
break;
}
}
// Push the event to stream
if (event) {
stream.push(event);
} else {
throw new Error("Failed to create event from proxy event");
}
}
}
}
}
// Check if aborted after reading
if (options.signal?.aborted) {
throw new Error("Request aborted by user");
}
stream.end();
} catch (error) {
const errorMessage = error instanceof Error ? error.message : String(error);
if (errorMessage.toLowerCase().includes("proxy") && errorMessage.includes("Unauthorized")) {
clearAuthToken();
}
partial.stopReason = options.signal?.aborted ? "aborted" : "error";
partial.errorMessage = errorMessage;
stream.push({
type: "error",
reason: partial.stopReason,
error: partial,
} satisfies AssistantMessageEvent);
stream.end();
} finally {
// Clean up abort handler
if (options.signal) {
options.signal.removeEventListener("abort", abortHandler);
}
}
})();
return stream;
}
/**
* Transport that uses an app server with user authentication tokens.
* The server manages user accounts and proxies requests to LLM providers.
*/
export class AppTransport implements AgentTransport {
private readonly proxyUrl = "https://genai.mariozechner.at";
private async getStreamFn() {
const authToken = await getAuthToken();
if (!authToken) {
throw new Error(i18n("Auth token is required for proxy transport"));
}
return <TApi extends Api>(model: Model<TApi>, context: Context, options?: SimpleStreamOptions) => {
return streamSimpleProxy(model, context, { ...options, authToken }, this.proxyUrl);
};
}
private buildContext(messages: Message[], cfg: AgentRunConfig): AgentContext {
return {
systemPrompt: cfg.systemPrompt,
messages,
tools: cfg.tools,
};
}
private buildLoopConfig(cfg: AgentRunConfig): AgentLoopConfig {
return {
model: cfg.model,
reasoning: cfg.reasoning,
getQueuedMessages: cfg.getQueuedMessages,
};
}
async *run(messages: Message[], userMessage: Message, cfg: AgentRunConfig, signal?: AbortSignal) {
const streamFn = await this.getStreamFn();
const context = this.buildContext(messages, cfg);
const pc = this.buildLoopConfig(cfg);
for await (const ev of agentLoop(userMessage as unknown as UserMessage, context, pc, signal, streamFn as any)) {
yield ev;
}
}
async *continue(messages: Message[], cfg: AgentRunConfig, signal?: AbortSignal) {
const streamFn = await this.getStreamFn();
const context = this.buildContext(messages, cfg);
const pc = this.buildLoopConfig(cfg);
for await (const ev of agentLoopContinue(context, pc, signal, streamFn as any)) {
yield ev;
}
}
}

71
packages/web-ui/src/agent/transports/ProviderTransport.ts
View file

@ -1,71 +0,0 @@
import {
type AgentContext,
type AgentLoopConfig,
agentLoop,
agentLoopContinue,
type Message,
type UserMessage,
} from "@mariozechner/pi-ai";
import { getAppStorage } from "../../storage/app-storage.js";
import { applyProxyIfNeeded } from "../../utils/proxy-utils.js";
import type { AgentRunConfig, AgentTransport } from "./types.js";
/**
* Transport that calls LLM providers directly.
* Uses CORS proxy only for providers that require it (Anthropic OAuth, Z-AI).
*/
export class ProviderTransport implements AgentTransport {
private async getModel(cfg: AgentRunConfig) {
const apiKey = await getAppStorage().providerKeys.get(cfg.model.provider);
if (!apiKey) {
throw new Error("no-api-key");
}
const proxyEnabled = await getAppStorage().settings.get<boolean>("proxy.enabled");
const proxyUrl = await getAppStorage().settings.get<string>("proxy.url");
const model = applyProxyIfNeeded(cfg.model, apiKey, proxyEnabled ? proxyUrl || undefined : undefined);
return model;
}
private buildContext(messages: Message[], cfg: AgentRunConfig): AgentContext {
return {
systemPrompt: cfg.systemPrompt,
messages,
tools: cfg.tools,
};
}
private buildLoopConfig(model: AgentRunConfig["model"], cfg: AgentRunConfig): AgentLoopConfig {
return {
model,
reasoning: cfg.reasoning,
// Resolve API key per assistant response (important for expiring OAuth tokens)
getApiKey: async (provider: string) => {
const key = await getAppStorage().providerKeys.get(provider);
return key ?? undefined; // Convert null to undefined for type compatibility
},
getQueuedMessages: cfg.getQueuedMessages,
};
}
async *run(messages: Message[], userMessage: Message, cfg: AgentRunConfig, signal?: AbortSignal) {
const model = await this.getModel(cfg);
const context = this.buildContext(messages, cfg);
const pc = this.buildLoopConfig(model, cfg);
for await (const ev of agentLoop(userMessage as unknown as UserMessage, context, pc, signal)) {
yield ev;
}
}
async *continue(messages: Message[], cfg: AgentRunConfig, signal?: AbortSignal) {
const model = await this.getModel(cfg);
const context = this.buildContext(messages, cfg);
const pc = this.buildLoopConfig(model, cfg);
for await (const ev of agentLoopContinue(context, pc, signal)) {
yield ev;
}
}
}

3
packages/web-ui/src/agent/transports/index.ts
View file

@ -1,3 +0,0 @@
export * from "./AppTransport.js";
export * from "./ProviderTransport.js";
export * from "./types.js";

15
packages/web-ui/src/agent/transports/proxy-types.ts
View file

@ -1,15 +0,0 @@
import type { StopReason, Usage } from "@mariozechner/pi-ai";
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: Usage }
| { type: "error"; reason: Extract<StopReason, "aborted" | "error">; errorMessage: string; usage: Usage };

26
packages/web-ui/src/agent/transports/types.ts
View file

@ -1,26 +0,0 @@
import type { AgentEvent, AgentTool, Message, Model, QueuedMessage } from "@mariozechner/pi-ai";
// The minimal configuration needed to run a turn.
export interface AgentRunConfig {
systemPrompt: string;
tools: AgentTool<any>[];
model: Model<any>;
reasoning?: "low" | "medium" | "high";
getQueuedMessages?: <T>() => Promise<QueuedMessage<T>[]>;
}
// Events yielded by transports must match the @mariozechner/pi-ai prompt() events.
// We re-export the Message type above; consumers should use the upstream AgentEvent type.
export interface AgentTransport {
/** Run with a new user message */
run(
messages: Message[],
userMessage: Message,
config: AgentRunConfig,
signal?: AbortSignal,
): AsyncIterable<AgentEvent>;
/** Continue from current context (no new user message) */
continue(messages: Message[], config: AgentRunConfig, signal?: AbortSignal): AsyncIterable<AgentEvent>;
}

11
packages/web-ui/src/agent/types.ts
View file

@ -1,11 +0,0 @@
import type { AssistantMessage, Context } from "@mariozechner/pi-ai";
export interface DebugLogEntry {
timestamp: string;
request: { provider: string; model: string; context: Context };
response?: AssistantMessage;
error?: unknown;
sseEvents: string[];
ttft?: number;
totalTime?: number;
}

13
packages/web-ui/src/components/AgentInterface.ts
View file

@ -6,9 +6,9 @@ import type { MessageEditor } from "./MessageEditor.js";
import "./MessageEditor.js";
import "./MessageList.js";
import "./Messages.js"; // Import for side effects to register the custom elements
import type { Agent, AgentEvent } from "../agent/agent.js";
import { getAppStorage } from "../storage/app-storage.js";
import "./StreamingMessageContainer.js";
import type { Agent, AgentEvent } from "@mariozechner/pi-agent-core";
import type { Attachment } from "../utils/attachment-utils.js";
import { formatUsage } from "../utils/format.js";
import { i18n } from "../utils/i18n.js";
@ -130,16 +130,13 @@ export class AgentInterface extends LitElement {
}
if (!this.session) return;
this._unsubscribeSession = this.session.subscribe(async (ev: AgentEvent) => {
if (ev.type === "state-update") {
if (ev.type === "message_update") {
if (this._streamingContainer) {
this._streamingContainer.isStreaming = ev.state.isStreaming;
this._streamingContainer.setMessage(ev.state.streamMessage, !ev.state.isStreaming);
const isStreaming = this.session?.state.isStreaming || false;
this._streamingContainer.isStreaming = isStreaming;
this._streamingContainer.setMessage(ev.message, !isStreaming);
}
this.requestUpdate();
} else if (ev.type === "error-no-model") {
// TODO show some UI feedback
} else if (ev.type === "error-no-api-key") {
// Handled by onApiKeyRequired callback
}
});
}

5
packages/web-ui/src/components/MessageList.ts
View file

@ -1,16 +1,15 @@
import type { AgentMessage, AgentTool } from "@mariozechner/pi-agent-core";
import type {
AgentTool,
AssistantMessage as AssistantMessageType,
ToolResultMessage as ToolResultMessageType,
} from "@mariozechner/pi-ai";
import { html, LitElement, type TemplateResult } from "lit";
import { property } from "lit/decorators.js";
import { repeat } from "lit/directives/repeat.js";
import type { AppMessage } from "./Messages.js";
import { renderMessage } from "./message-renderer-registry.js";
export class MessageList extends LitElement {
@property({ type: Array }) messages: AppMessage[] = [];
@property({ type: Array }) messages: AgentMessage[] = [];
@property({ type: Array }) tools: AgentTool[] = [];
@property({ type: Object }) pendingToolCalls?: Set<string>;
@property({ type: Boolean }) isStreaming: boolean = false;

37
packages/web-ui/src/components/Messages.ts
View file

@ -1,6 +1,7 @@
import type {
AgentTool,
AssistantMessage as AssistantMessageType,
ImageContent,
TextContent,
ToolCall,
ToolResultMessage as ToolResultMessageType,
UserMessage as UserMessageType,
@ -12,8 +13,14 @@ import type { Attachment } from "../utils/attachment-utils.js";
import { formatUsage } from "../utils/format.js";
import { i18n } from "../utils/i18n.js";
import "./ThinkingBlock.js";
import type { AgentTool } from "@mariozechner/pi-agent-core";
export type UserMessageWithAttachments = UserMessageType & { attachments?: Attachment[] };
export type UserMessageWithAttachments = {
role: "user-with-attachments";
content: string | (TextContent | ImageContent)[];
timestamp: number;
attachments?: Attachment[];
};
// Artifact message type for session persistence
export interface ArtifactMessage {
@ -25,26 +32,16 @@ export interface ArtifactMessage {
timestamp: string;
}
// Base message union
type BaseMessage = AssistantMessageType | UserMessageWithAttachments | ToolResultMessageType | ArtifactMessage;
// Extensible interface - apps can extend via declaration merging
// Example:
// declare module "@mariozechner/pi-web-ui" {
// interface CustomMessages {
// "system-notification": SystemNotificationMessage;
// }
// }
export interface CustomMessages {
// Empty by default - apps extend via declaration merging
declare module "@mariozechner/pi-agent-core" {
interface CustomAgentMessages {
"user-with-attachment": UserMessageWithAttachments;
artifact: ArtifactMessage;
}
}
// AppMessage is union of base messages + custom messages
export type AppMessage = BaseMessage | CustomMessages[keyof CustomMessages];
@customElement("user-message")
export class UserMessage extends LitElement {
@property({ type: Object }) message!: UserMessageWithAttachments;
@property({ type: Object }) message!: UserMessageWithAttachments | UserMessageType;
protected override createRenderRoot(): HTMLElement | DocumentFragment {
return this;
@ -66,7 +63,9 @@ export class UserMessage extends LitElement {
<div class="user-message-container py-2 px-4 rounded-xl">
<markdown-block .content=${content}></markdown-block>
${
this.message.attachments && this.message.attachments.length > 0
this.message.role === "user-with-attachments" &&
this.message.attachments &&
this.message.attachments.length > 0
? html`
<div class="mt-3 flex flex-wrap gap-2">
${this.message.attachments.map(

9
packages/web-ui/src/components/StreamingMessageContainer.ts
View file

@ -1,4 +1,5 @@
import type { AgentTool, Message, ToolResultMessage } from "@mariozechner/pi-ai";
import type { AgentMessage, AgentTool } from "@mariozechner/pi-agent-core";
import type { ToolResultMessage } from "@mariozechner/pi-ai";
import { html, LitElement } from "lit";
import { property, state } from "lit/decorators.js";
@ -9,8 +10,8 @@ export class StreamingMessageContainer extends LitElement {
@property({ type: Object }) toolResultsById?: Map<string, ToolResultMessage>;
@property({ attribute: false }) onCostClick?: () => void;
@state() private _message: Message | null = null;
private _pendingMessage: Message | null = null;
@state() private _message: AgentMessage | null = null;
private _pendingMessage: AgentMessage | null = null;
private _updateScheduled = false;
private _immediateUpdate = false;
@ -24,7 +25,7 @@ export class StreamingMessageContainer extends LitElement {
}
// Public method to update the message with batching for performance
public setMessage(message: Message | null, immediate = false) {
public setMessage(message: AgentMessage | null, immediate = false) {
// Store the latest message
this._pendingMessage = message;

10
packages/web-ui/src/components/message-renderer-registry.ts
View file

@ -1,11 +1,11 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { TemplateResult } from "lit";
import type { AppMessage } from "./Messages.js";
// Extract role type from AppMessage union
export type MessageRole = AppMessage["role"];
export type MessageRole = AgentMessage["role"];
// Generic message renderer typed to specific message type
export interface MessageRenderer<TMessage extends AppMessage = AppMessage> {
export interface MessageRenderer<TMessage extends AgentMessage = AgentMessage> {
render(message: TMessage): TemplateResult;
}
@ -14,7 +14,7 @@ const messageRenderers = new Map<MessageRole, MessageRenderer<any>>();
export function registerMessageRenderer<TRole extends MessageRole>(
role: TRole,
renderer: MessageRenderer<Extract<AppMessage, { role: TRole }>>,
renderer: MessageRenderer<Extract<AgentMessage, { role: TRole }>>,
): void {
messageRenderers.set(role, renderer);
}
@ -23,6 +23,6 @@ export function getMessageRenderer(role: MessageRole): MessageRenderer | undefin
return messageRenderers.get(role);
}
export function renderMessage(message: AppMessage): TemplateResult | undefined {
export function renderMessage(message: AgentMessage): TemplateResult | undefined {
return messageRenderers.get(message.role)?.render(message);
}

12
packages/web-ui/src/index.ts
View file

@ -1,13 +1,7 @@
// Main chat interface
export type { AgentState, ThinkingLevel } from "./agent/agent.js";
// State management
export { Agent } from "./agent/agent.js";
// Transports
export { AppTransport } from "./agent/transports/AppTransport.js";
export { ProviderTransport } from "./agent/transports/ProviderTransport.js";
export type { ProxyAssistantMessageEvent } from "./agent/transports/proxy-types.js";
export type { AgentRunConfig, AgentTransport } from "./agent/transports/types.js";
export type { Agent, AgentMessage, AgentState, ThinkingLevel } from "@mariozechner/pi-agent-core";
export type { Model } from "@mariozechner/pi-ai";
export { ChatPanel } from "./ChatPanel.js";
// Components
export { AgentInterface } from "./components/AgentInterface.js";
@ -18,7 +12,7 @@ export { Input } from "./components/Input.js";
export { MessageEditor } from "./components/MessageEditor.js";
export { MessageList } from "./components/MessageList.js";
// Message components
export type { AppMessage, CustomMessages, UserMessageWithAttachments } from "./components/Messages.js";
export type { UserMessageWithAttachments } from "./components/Messages.js";
export { AssistantMessage, ToolMessage, UserMessage } from "./components/Messages.js";
// Message renderer registry
export {

2
packages/web-ui/src/storage/stores/sessions-store.ts
View file

@ -1,4 +1,4 @@
import type { AgentState } from "../../agent/agent.js";
import type { AgentState } from "@mariozechner/pi-agent-core";
import { Store } from "../store.js";
import type { SessionData, SessionMetadata, StoreConfig } from "../types.js";

5
packages/web-ui/src/storage/types.ts
View file

@ -1,6 +1,5 @@
import type { AgentMessage, ThinkingLevel } from "@mariozechner/pi-agent-core";
import type { Model } from "@mariozechner/pi-ai";
import type { ThinkingLevel } from "../agent/agent.js";
import type { AppMessage } from "../components/Messages.js";
/**
* Transaction interface for atomic operations across stores.
@ -159,7 +158,7 @@ export interface SessionData {
thinkingLevel: ThinkingLevel;
/** Full conversation history (with attachments inline) */
messages: AppMessage[];
messages: AgentMessage[];
/** ISO 8601 UTC timestamp of creation */
createdAt: string;

8
packages/web-ui/src/tools/artifacts/artifacts.ts
View file

@ -1,13 +1,13 @@
import { icon } from "@mariozechner/mini-lit";
import "@mariozechner/mini-lit/dist/MarkdownBlock.js";
import { Button } from "@mariozechner/mini-lit/dist/Button.js";
import { type AgentTool, type Message, StringEnum, type ToolCall } from "@mariozechner/pi-ai";
import type { Agent, AgentMessage, AgentTool } from "@mariozechner/pi-agent-core";
import { StringEnum, type ToolCall } from "@mariozechner/pi-ai";
import { type Static, Type } from "@sinclair/typebox";
import { html, LitElement, type TemplateResult } from "lit";
import { customElement, property, state } from "lit/decorators.js";
import { createRef, type Ref, ref } from "lit/directives/ref.js";
import { X } from "lucide";
import type { Agent } from "../../agent/agent.js";
import type { ArtifactMessage } from "../../components/Messages.js";
import { ArtifactsRuntimeProvider } from "../../components/sandbox/ArtifactsRuntimeProvider.js";
import { AttachmentsRuntimeProvider } from "../../components/sandbox/AttachmentsRuntimeProvider.js";
@ -85,7 +85,7 @@ export class ArtifactsPanel extends LitElement {
if (this.agent) {
const attachments: Attachment[] = [];
for (const message of this.agent.state.messages) {
if (message.role === "user" && message.attachments) {
if (message.role === "user-with-attachments" && message.attachments) {
attachments.push(...message.attachments);
}
}
@ -292,7 +292,7 @@ export class ArtifactsPanel extends LitElement {
// Re-apply artifacts by scanning a message list (optional utility)
public async reconstructFromMessages(
messages: Array<Message | { role: "aborted" } | { role: "artifact" }>,
messages: Array<AgentMessage | { role: "aborted" } | { role: "artifact" }>,
): Promise<void> {
const toolCalls = new Map<string, ToolCall>();
const artifactToolName = "artifacts";

3
packages/web-ui/src/tools/extract-document.ts
View file

@ -1,4 +1,5 @@
import type { AgentTool, ToolResultMessage } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import type { ToolResultMessage } from "@mariozechner/pi-ai";
import { type Static, Type } from "@sinclair/typebox";
import { html } from "lit";
import { createRef, ref } from "lit/directives/ref.js";

3
packages/web-ui/src/tools/javascript-repl.ts
View file

@ -1,5 +1,6 @@
import { i18n } from "@mariozechner/mini-lit";
import type { AgentTool, ToolResultMessage } from "@mariozechner/pi-ai";
import type { AgentTool } from "@mariozechner/pi-agent-core";
import type { ToolResultMessage } from "@mariozechner/pi-ai";
import { type Static, Type } from "@sinclair/typebox";
import { html } from "lit";
import { createRef, ref } from "lit/directives/ref.js";