Core Types

interface AgentState {
  systemPrompt: string;
  model: Model<any>;
  thinkingLevel: ThinkingLevel;
  tools: AgentTool<any>[];
  messages: AgentMessage[];
  isStreaming: boolean;
  streamMessage: AgentMessage | null;
  pendingToolCalls: Set<string>;
  error?: string;
}

type ThinkingLevel = 'off' | 'minimal' | 'low' | 'medium' | 'high' | 'xhigh';

type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages];

declare module '@mariozechner/pi-agent-core' {
  interface CustomAgentMessages {
    artifact: ArtifactMessage;
    notification: NotificationMessage;
  }
}

interface CustomAgentMessages {
  // Empty by default - extend via declaration merging
}

// your-types.ts
import '@mariozechner/pi-agent-core';

interface ArtifactMessage {
  role: 'artifact';
  content: string;
  artifactType: 'code' | 'diagram' | 'document';
  timestamp: number;
}

declare module '@mariozechner/pi-agent-core' {
  interface CustomAgentMessages {
    artifact: ArtifactMessage;
  }
}

type AgentEvent =
  | { type: 'agent_start' }
  | { type: 'agent_end'; messages: AgentMessage[] }
  | { type: 'turn_start' }
  | { type: 'turn_end'; message: AgentMessage; toolResults: ToolResultMessage[] }
  | { type: 'message_start'; message: AgentMessage }
  | { type: 'message_update'; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent }
  | { type: 'message_end'; message: AgentMessage }
  | { 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 };

interface AgentContext {
  systemPrompt: string;
  messages: AgentMessage[];
  tools?: AgentTool<any>[];
}

interface AgentLoopConfig extends SimpleStreamOptions {
  model: Model<any>;
  convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
  transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
  getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
  getSteeringMessages?: () => Promise<AgentMessage[]>;
  getFollowUpMessages?: () => Promise<AgentMessage[]>;
}

type StreamFn = (
  ...args: Parameters<typeof streamSimple>
) => ReturnType<typeof streamSimple> | Promise<ReturnType<typeof streamSimple>>;

const customStreamFn: StreamFn = async (model, context, options) => {
  const config = await loadConfig();
  return streamSimple(model, context, { ...options, ...config });
};

type Transport = 'sse' | 'responses';

function streamProxy(
  model: Model<any>,
  context: Context,
  options: ProxyStreamOptions
): ProxyMessageEventStream

import { streamProxy } from '@mariozechner/pi-agent-core';

const agent = new Agent({
  streamFn: (model, context, options) =>
    streamProxy(model, context, {
      ...options,
      authToken: await getAuthToken(),
      proxyUrl: 'https://api.example.com',
    }),
});

interface ProxyStreamOptions extends SimpleStreamOptions {
  authToken: string;
  proxyUrl: string;
}

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: 'stop' | 'length' | 'toolUse'; usage: AssistantMessage['usage'] }
  | { type: 'error'; reason: 'aborted' | 'error'; errorMessage?: string; usage: AssistantMessage['usage'] };

function agentLoop(
  prompts: AgentMessage[],
  context: AgentContext,
  config: AgentLoopConfig,
  signal?: AbortSignal,
  streamFn?: StreamFn
): EventStream<AgentEvent, AgentMessage[]>

import { agentLoop } from '@mariozechner/pi-agent-core';

const stream = agentLoop(
  [{ role: 'user', content: 'Hello!', timestamp: Date.now() }],
  { systemPrompt: 'You are helpful', messages: [], tools: [] },
  config,
  abortSignal
);

for await (const event of stream) {
  console.log('Event:', event);
}

const allMessages = await stream.result();

function agentLoopContinue(
  context: AgentContext,
  config: AgentLoopConfig,
  signal?: AbortSignal,
  streamFn?: StreamFn
): EventStream<AgentEvent, AgentMessage[]>

import { agentLoopContinue } from '@mariozechner/pi-agent-core';

const stream = agentLoopContinue(context, config, abortSignal);

for await (const event of stream) {
  console.log('Event:', event);
}

// 1. Define your custom message type
interface NotificationMessage {
  role: 'notification';
  level: 'info' | 'warning' | 'error';
  message: string;
  timestamp: number;
}

// 2. Extend CustomAgentMessages via declaration merging
declare module '@mariozechner/pi-agent-core' {
  interface CustomAgentMessages {
    notification: NotificationMessage;
  }
}

// 3. Now AgentMessage includes your type
import { Agent } from '@mariozechner/pi-agent-core';

const agent = new Agent({
  convertToLlm: (messages) => {
    return messages.flatMap(m => {
      // Filter out notifications - they're UI-only
      if (m.role === 'notification') {
        return [];
      }
      // Pass through standard messages
      return [m];
    });
  }
});

// 4. Add notification to conversation
const notification: NotificationMessage = {
  role: 'notification',
  level: 'info',
  message: 'Task completed successfully',
  timestamp: Date.now()
};

agent.appendMessage(notification);