Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
19 kB · 617 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618/** * Custom message types and transformers for the coding agent. * * Extends the base AgentMessage type with coding-agent specific message types, * and provides a transformer to convert them to LLM-compatible messages. */
import type { AgentMessage } from "@earendil-works/pi-agent-core";import type { ImageContent, Message, TextContent } from "@earendil-works/pi-ai";import type { AgentCronJob } from "./cron-jobs.js";import type { AppliedRefinementEdit, HarnessScope, RefinementResult } from "./refinement/refinement.js";import { isSessionSlashCommandName, parseSessionSlashCommand, type SessionSlashCommand } from "./slash-commands.js";import { formatToolDuration } from "./tools/render-utils.js";
export const COMPACTION_SUMMARY_PREFIX = `The conversation history before this point was compacted into the following summary:
<summary>`;
export const COMPACTION_SUMMARY_SUFFIX = `</summary>`;
export const BRANCH_SUMMARY_PREFIX = `The following is a summary of a branch that this conversation came back from:
<summary>`;
export const BRANCH_SUMMARY_SUFFIX = `</summary>`;
export const HEARTBEAT_PROMPT_CUSTOM_TYPE = "heartbeat_prompt";export const HEARTBEAT_PROMPT_PREVIEW_LABEL = "Heartbeat prompt";export const IPYTHON_STATE_RESTORED_CUSTOM_TYPE = "ipython_state_restored";export const SESSION_SLASH_COMMAND_CUSTOM_TYPE = "session_slash_command";export const SESSION_SLASH_COMMAND_RESULT_CUSTOM_TYPE = "session_slash_command_result";export const COMPACTION_OUTCOME_CUSTOM_TYPE = "compaction_outcome";export const REFINEMENT_OUTCOME_CUSTOM_TYPE = "refinement_outcome";export const RLM_CHILD_FAILURE_CUSTOM_TYPE = "rlm_child_failure";export const RLM_CHILD_TERMINAL_NOTICE_CUSTOM_TYPE = "rlm_child_terminal_notice";export const ASYNC_BASH_COMPLETION_CUSTOM_TYPE = "async_bash_completion";export const ASYNC_BASH_COMPLETION_PREVIEW_LABEL = "Shell message received";export const MANUAL_CONTINUATION_CUSTOM_TYPE = "manual_continuation";export const MANUAL_CONTINUATION_PROMPT = `<system-notice>Continue the current task from where you stopped. Resume the most recent unfinished step. Do not pause to summarize progress, reconfirm the plan, or ask whether to proceed.</system-notice>`;
export interface SessionSlashCommandDetails { command: SessionSlashCommand; commandEntryId?: string;}
export interface SessionSlashCommandResultDetails { command: SessionSlashCommand; success: boolean; severity: "info" | "warning" | "error"; error?: string; commandEntryId?: string;}
export interface SessionSlashCommandMessage extends CustomMessage<SessionSlashCommandDetails> { customType: typeof SESSION_SLASH_COMMAND_CUSTOM_TYPE; content: string; details: SessionSlashCommandDetails;}
export interface SessionSlashCommandResultMessage extends CustomMessage<SessionSlashCommandResultDetails> { customType: typeof SESSION_SLASH_COMMAND_RESULT_CUSTOM_TYPE; content: string; details: SessionSlashCommandResultDetails;}
export type CompactionOutcomeReason = "threshold" | "overflow" | "requested";export type CompactionOutcome = "skipped" | "cancelled" | "failed";
export interface CompactionOutcomeDetails { reason: CompactionOutcomeReason; outcome: CompactionOutcome;}
export interface CompactionOutcomeMessage extends CustomMessage<CompactionOutcomeDetails> { customType: typeof COMPACTION_OUTCOME_CUSTOM_TYPE; content: string; details: CompactionOutcomeDetails;}
export interface RefinementOutcomeDetails { refinementId: string; summary: string; scope: HarnessScope; rollbackOf?: string; edits: AppliedRefinementEdit[];}
export interface RefinementOutcomeMessage extends CustomMessage<RefinementOutcomeDetails> { customType: typeof REFINEMENT_OUTCOME_CUSTOM_TYPE; content: string; details: RefinementOutcomeDetails;}
export function createManualContinuationMessage(timestamp = Date.now()): CustomMessage { return { role: "custom", customType: MANUAL_CONTINUATION_CUSTOM_TYPE, content: MANUAL_CONTINUATION_PROMPT, display: false, timestamp, };}
export interface RlmChildFailureDetails { childId: string; sessionName: string; error: string;}
export type RlmChildTerminalNoticeDetails = | { kind: "cancelled"; childId: string; sessionName: string; reason?: string; } | { kind: "completed_without_reply"; childId: string; sessionName: string; lastAssistantTextPreview?: string; };
export interface AsyncBashCompletionDetails { pid: number; command: string; exitCode: number;}
interface AsyncBashCompletionMessage extends CustomMessage<AsyncBashCompletionDetails> { customType: typeof ASYNC_BASH_COMPLETION_CUSTOM_TYPE; content: string;}
export function createAsyncBashCompletionMessage( details: AsyncBashCompletionDetails, timestamp = Date.now(),): AsyncBashCompletionMessage { return { role: "custom", customType: ASYNC_BASH_COMPLETION_CUSTOM_TYPE, content: `${ASYNC_BASH_COMPLETION_PREVIEW_LABEL}.Source: bashCommand completed (pid ${details.pid}, exit code ${details.exitCode}).Command: ${JSON.stringify(details.command)}
Inspect the saved BashHandle with .poll(), .output(), or .tail(), then continue the task.`, display: true, details, timestamp, };}
export function createRlmChildFailureMessage( details: RlmChildFailureDetails, timestamp = Date.now(),): CustomMessage<RlmChildFailureDetails> { return { role: "custom", customType: RLM_CHILD_FAILURE_CUSTOM_TYPE, content: `RLM child ${details.sessionName} (${details.childId}) failed: ${details.error}`, display: true, details, timestamp, };}
export function createRlmChildTerminalNoticeMessage( details: RlmChildTerminalNoticeDetails, timestamp = Date.now(),): CustomMessage<RlmChildTerminalNoticeDetails> { const content = details.kind === "cancelled" ? `RLM child ${details.sessionName} (${details.childId}) was cancelled${details.reason ? `: ${details.reason}` : ""}` : `RLM child ${details.sessionName} (${details.childId}) completed without sending a reply${details.lastAssistantTextPreview ? `. Last assistant text: ${details.lastAssistantTextPreview}` : ""}`; return { role: "custom", customType: RLM_CHILD_TERMINAL_NOTICE_CUSTOM_TYPE, content, display: true, details, timestamp, };}
/** * Message type for bash executions via the ! command. */export interface BashExecutionMessage { role: "bashExecution"; command: string; output: string; exitCode: number | undefined; cancelled: boolean; truncated: boolean; fullOutputPath?: string; timestamp: number; /** If true, this message is excluded from LLM context (!! prefix) */ excludeFromContext?: boolean;}
/** * Message type for extension-injected messages via sendMessage(). * These are custom messages that extensions can inject into the conversation. */export interface CustomMessage<T = unknown> { role: "custom"; customType: string; content: string | (TextContent | ImageContent)[]; display: boolean; details?: T; timestamp: number;}
export interface HeartbeatPromptDetails { jobId: string; schedule: string; status: AgentCronJob["status"]; runCount: number; nextRunAt?: string; lastRunAt?: string;}
export interface IpythonStateRestoredDetails { restored: boolean;}
export interface BranchSummaryMessage { role: "branchSummary"; summary: string; fromId: string; timestamp: number;}
export interface CompactionSummaryMessage { role: "compactionSummary"; summary: string; tokensBefore: number; /** Number of retained messages that precede this summary in transcript presentation. */ retainedMessageCount?: number; /** User instructions that guided the summary (from `/compact <instructions>`) */ customInstructions?: string; timestamp: number;}
declare module "@earendil-works/pi-agent-core" { interface CustomAgentMessages { bashExecution: BashExecutionMessage; custom: CustomMessage; branchSummary: BranchSummaryMessage; compactionSummary: CompactionSummaryMessage; }}
/** * Format bash output for LLM context. The fence must be longer than any * backtick run in the output so command output cannot terminate it early. */export function bashOutputToText( msg: Pick<BashExecutionMessage, "output" | "exitCode" | "cancelled" | "truncated" | "fullOutputPath">,): string { let text = ""; if (msg.output) { let longestBacktickRun = 0; for (const match of msg.output.matchAll(/`+/g)) { longestBacktickRun = Math.max(longestBacktickRun, match[0].length); } const fence = "`".repeat(Math.max(3, longestBacktickRun + 1)); text += `${fence}\n${msg.output}\n${fence}`; } else { text += "(no output)"; } if (msg.cancelled) { text += "\n\n(command cancelled)"; } else if (msg.exitCode !== null && msg.exitCode !== undefined && msg.exitCode !== 0) { text += `\n\nCommand exited with code ${msg.exitCode}`; } if (msg.truncated) { text += msg.fullOutputPath ? `\n\n[Output truncated. Full output: ${msg.fullOutputPath}]` : "\n\n[Output truncated.]"; } return text;}
/** * Convert a BashExecutionMessage to user message text for LLM context. */export function bashExecutionToText(msg: BashExecutionMessage): string { return `Ran \`${msg.command}\`\n${bashOutputToText(msg)}`;}
export function createBranchSummaryMessage(summary: string, fromId: string, timestamp: string): BranchSummaryMessage { return { role: "branchSummary", summary, fromId, timestamp: new Date(timestamp).getTime(), };}
export function createCompactionSummaryMessage( summary: string, tokensBefore: number, timestamp: string, customInstructions?: string, retainedMessageCount?: number,): CompactionSummaryMessage { return { role: "compactionSummary", summary, tokensBefore, retainedMessageCount, customInstructions, timestamp: new Date(timestamp).getTime(), };}
/** Convert CustomMessageEntry to AgentMessage format */export function createCustomMessage( customType: string, content: string | (TextContent | ImageContent)[], display: boolean, details: unknown | undefined, timestamp: string,): CustomMessage { return { role: "custom", customType, content, display, details, timestamp: new Date(timestamp).getTime(), };}
export function createSessionSlashCommandMessage( command: SessionSlashCommand, details: Omit<SessionSlashCommandDetails, "command"> = {}, display = true, timestamp = Date.now(),): SessionSlashCommandMessage { return { role: "custom", customType: SESSION_SLASH_COMMAND_CUSTOM_TYPE, content: command.text, display, details: { ...details, command: { ...command } }, timestamp, };}
export function createSessionSlashCommandResultMessage( content: string, details: SessionSlashCommandResultDetails, display = true, timestamp = Date.now(),): SessionSlashCommandResultMessage { return { role: "custom", customType: SESSION_SLASH_COMMAND_RESULT_CUSTOM_TYPE, content, display, details: { ...details, command: { ...details.command } }, timestamp, };}
export function createCompactionOutcomeMessage( content: string, details: CompactionOutcomeDetails, display = true, timestamp = Date.now(),): CompactionOutcomeMessage { return { role: "custom", customType: COMPACTION_OUTCOME_CUSTOM_TYPE, content, display, details: { ...details }, timestamp, };}
export function createRefinementOutcomeMessage( result: RefinementResult, display = true, timestamp = Date.now(),): RefinementOutcomeMessage { return { role: "custom", customType: REFINEMENT_OUTCOME_CUSTOM_TYPE, content: `Refinement complete: ${result.summary}`, display, details: { refinementId: result.id, summary: result.summary, scope: result.scope ?? "local", ...(result.rollbackOf ? { rollbackOf: result.rollbackOf } : {}), edits: result.appliedEdits, }, timestamp, };}
function isRecord(value: unknown): value is Record<string, unknown> { return typeof value === "object" && value !== null;}
function hasValidCustomMessageEnvelope(message: Record<string, unknown>, customType: string): boolean { return ( message.role === "custom" && message.customType === customType && typeof message.content === "string" && typeof message.display === "boolean" && typeof message.timestamp === "number" && Number.isFinite(message.timestamp) );}
export function isSessionSlashCommand(value: unknown): value is SessionSlashCommand { if ( !isRecord(value) || !isSessionSlashCommandName(value.name) || typeof value.args !== "string" || typeof value.text !== "string" ) { return false; } const parsed = parseSessionSlashCommand(value.text); return ( parsed !== undefined && parsed.name === value.name && parsed.args === value.args && parsed.text === value.text );}
function isValidCommandEntryId(value: unknown): value is string | undefined { return value === undefined || (typeof value === "string" && value.length > 0);}
export function isSessionSlashCommandMessage(message: unknown): message is SessionSlashCommandMessage { if ( !isRecord(message) || !hasValidCustomMessageEnvelope(message, SESSION_SLASH_COMMAND_CUSTOM_TYPE) || typeof message.content !== "string" ) { return false; } if (!isRecord(message.details) || !isSessionSlashCommand(message.details.command)) return false; return message.content === message.details.command.text && isValidCommandEntryId(message.details.commandEntryId);}
export function isSessionSlashCommandResultMessage(message: unknown): message is SessionSlashCommandResultMessage { if (!isRecord(message) || !hasValidCustomMessageEnvelope(message, SESSION_SLASH_COMMAND_RESULT_CUSTOM_TYPE)) return false; if (!isRecord(message.details) || !isSessionSlashCommand(message.details.command)) return false; return ( typeof message.details.success === "boolean" && (message.details.severity === "info" || message.details.severity === "warning" || message.details.severity === "error") && (message.details.error === undefined || typeof message.details.error === "string") && isValidCommandEntryId(message.details.commandEntryId) );}
export function isCompactionOutcomeMessage(message: unknown): message is CompactionOutcomeMessage { if (!isRecord(message) || !hasValidCustomMessageEnvelope(message, COMPACTION_OUTCOME_CUSTOM_TYPE)) return false; if (!isRecord(message.details)) return false; return ( (message.details.reason === "threshold" || message.details.reason === "overflow" || message.details.reason === "requested") && (message.details.outcome === "skipped" || message.details.outcome === "cancelled" || message.details.outcome === "failed") );}
function isAppliedRefinementEdit(value: unknown): value is AppliedRefinementEdit { return ( isRecord(value) && (value.action === "create" || value.action === "update" || value.action === "delete") && typeof value.kind === "string" && typeof value.id === "string" && typeof value.applied === "boolean" );}
export function isRefinementOutcomeMessage(message: unknown): message is RefinementOutcomeMessage { if (!isRecord(message) || !hasValidCustomMessageEnvelope(message, REFINEMENT_OUTCOME_CUSTOM_TYPE)) return false; if (!isRecord(message.details)) return false; return ( typeof message.details.summary === "string" && (message.details.scope === "local" || message.details.scope === "global") && Array.isArray(message.details.edits) && message.details.edits.every(isAppliedRefinementEdit) );}
export function createHeartbeatPromptMessage( job: AgentCronJob, timestamp = Date.now(),): CustomMessage<HeartbeatPromptDetails> { return { role: "custom", customType: HEARTBEAT_PROMPT_CUSTOM_TYPE, content: job.prompt, display: true, details: { jobId: job.id, schedule: job.schedule.expression, status: job.status, runCount: job.runCount, nextRunAt: job.nextRunAt, lastRunAt: job.lastRunAt, }, timestamp, };}
/** * Transform AgentMessages (including custom types) to LLM-compatible Messages. * * This is used by: * - Agent's transormToLlm option (for prompt calls and queued messages) * - Compaction's generateSummary (for summarization) * - Custom extensions and tools */function withToolDurationAnnotation( content: (TextContent | ImageContent)[], durationMs: number,): (TextContent | ImageContent)[] { const annotation = `[Took ${formatToolDuration(durationMs)}]`; let lastTextIndex = -1; for (let index = content.length - 1; index >= 0; index--) { if (content[index].type === "text") { lastTextIndex = index; break; } } if (lastTextIndex === -1) { return [...content, { type: "text", text: annotation }]; } return content.map((block, index) => { if (index !== lastTextIndex || block.type !== "text") { return block; } return { type: "text", text: block.text ? `${block.text}\n\n${annotation}` : annotation }; });}
export function convertToLlm(messages: AgentMessage[]): Message[] { return messages .map((m): Message | undefined => { switch (m.role) { case "bashExecution": if (m.excludeFromContext) { return undefined; } return { role: "user", content: [{ type: "text", text: bashExecutionToText(m) }], timestamp: m.timestamp, }; case "custom": { if ( m.customType === SESSION_SLASH_COMMAND_CUSTOM_TYPE || m.customType === SESSION_SLASH_COMMAND_RESULT_CUSTOM_TYPE || m.customType === COMPACTION_OUTCOME_CUSTOM_TYPE || m.customType === REFINEMENT_OUTCOME_CUSTOM_TYPE ) { return undefined; } const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content; return { role: "user", content, timestamp: m.timestamp, }; } case "branchSummary": return { role: "user", content: [{ type: "text" as const, text: BRANCH_SUMMARY_PREFIX + m.summary + BRANCH_SUMMARY_SUFFIX }], timestamp: m.timestamp, }; case "compactionSummary": return { role: "user", content: [ { type: "text" as const, text: COMPACTION_SUMMARY_PREFIX + m.summary + COMPACTION_SUMMARY_SUFFIX }, ], timestamp: m.timestamp, }; case "user": case "assistant": return m; case "toolResult": { if (m.durationMs === undefined) { return m; } // The model sees the duration; the stored message stays untouched. return { ...m, content: withToolDurationAnnotation(m.content, m.durationMs) }; } default: // biome-ignore lint/correctness/noSwitchDeclarations: fine const _exhaustiveCheck: never = m; return undefined; } }) .filter((m) => m !== undefined);}