omnirush 0.9.1 → 0.10.1

package/assets/extensions/omnirush/pi-engine.ts

@@ -49,13 +49,35 @@
 // Ids are stable across processes (derived from the pi entry ids), so a
 // resumed session or a repeated flush merges instead of duplicating.
 //
+// The 2.x engine's own fields ride along, the way the desktop's 2.x engine
+// adds them to the same 1.x records (omnirush-gui engine2/native.ts, whose
+// merge rules `mergeNative` repeats): a 2.x field the 1.x record does not
+// have is added under its 2.x name (info.type, info.text, info.model,
+// info.time.streamed, info.rawFinish, info.providerState, info.snapshot,
+// text/reasoning `state`, reasoning `time`, tool `name` / `executed` /
+// `providerState` / `time` {created, ran, completed}, tool state `content`,
+// the shell and compaction records' fields); objects both have (time,
+// tokens, error, tool state) are merged key by key; a 2.x value under a name
+// the 1.x record already uses with another value goes under `v2.<name>` on
+// the same record (tokens.v2.output: 2.x counts visible output without the
+// reasoning tokens pi's output includes; a tool state's v2.error); and 2.x
+// records 1.x never listed become parts of the nearest message, typed by
+// their 2.x type: "model-switched" (pi's model and thinking-level changes,
+// and a prompt sent to another model than the last one recorded) and "idle"
+// (the run's outcome). What pi does not record is left out: the
+// step snapshots' hashes, provider-executed tools and their result state,
+// the service tier, encrypted reasoning and signatures (kept out of the
+// trace, as before), retries, agent and location switches. Stream timing
+// (block and tool-call starts and ends, a tool's run start, the run's idle
+// outcome) comes from stream-timing.ts's `omnirush-timing` entries.
+//
 // Also here, pure and shared with the tests: the helpers the desktop's turn
 // observer (apps/server/src/collector-observer.ts) applies to the engine's
 // messages when a turn settles: the turn's model, the message budget of a
 // flush and its split into `turn.messages` parts.
 
 import { createHash } from "node:crypto";
-import { isAbsolute, resolve } from "node:path";
+import { isAbsolute, relative, resolve } from "node:path";
 
 /** customType of bgshell.ts's finished-background-command message. */
 export const BACKGROUND_BASH_TYPE = "omnirush-bash-result";
@@ -67,6 +89,8 @@ export const BACKGROUND_BASH_TYPE = "omnirush-bash-result";
 export const SUBAGENT_ENTRY_TYPE = "omnirush-subagent";
 /** A sub-agent's prompt: the task the parent agent wrote, after this prefix (agents-lib.ts buildChildArgs). */
 export const SUBAGENT_TASK_PREFIX = "Task: ";
+/** customType of stream-timing.ts's entries (a model step's stream timing, a run's idle outcome). */
+export const TIMING_ENTRY_TYPE = "omnirush-timing";
 
 export type EngineMessage = { info: Record<string, unknown>; parts: Array<Record<string, unknown>> };
 
@@ -262,6 +286,172 @@ export function spawnedSessionIds(details: unknown): string[] {
 
 type ToolResult = { message: Record<string, any>; at: number | null };
 
+// --- the 2.x engine's own fields (omnirush-gui apps/server/src/engine2/native.ts) ---
+
+/** Where a 2.x value goes when its name is taken on the 1.x record by another value. */
+export const NATIVE_COLLISIONS = "v2";
+const INFO_NESTED = { time: {}, tokens: { cache: {} }, error: { data: {} } };
+const STATE_NESTED = { metadata: {}, time: {} };
+
+function deepEqual(a: unknown, b: unknown): boolean {
+  if (a === b) return true;
+  if (typeof a !== typeof b || a === null || b === null || typeof a !== "object") return false;
+  if (Array.isArray(a) !== Array.isArray(b)) return false;
+  if (Array.isArray(a)) {
+    const other = b as unknown[];
+    return a.length === other.length && a.every((item, index) => deepEqual(item, other[index]));
+  }
+  const ka = Object.keys(a as object);
+  const kb = Object.keys(b as object);
+  return ka.length === kb.length && ka.every((key) => deepEqual((a as Record<string, unknown>)[key], (b as Record<string, unknown>)[key]));
+}
+
+/**
+ * Adds 2.x fields (`native`) to a 1.x record, as the desktop's 2.x engine
+ * does: a new name is added as-is, an equal value is not repeated, the
+ * objects `nested` names are merged key by key (its value: the keys merged
+ * one level further down), and any other value under a taken name goes under
+ * `v2.<name>`. Objects of `target` are copied before they change.
+ */
+export function mergeNative(
+  target: Record<string, unknown>,
+  native: Record<string, unknown>,
+  nested: Record<string, Record<string, unknown>> = {},
+): Record<string, unknown> {
+  for (const [key, value] of Object.entries(native)) {
+    if (value === undefined || key === NATIVE_COLLISIONS) continue;
+    if (!(key in target)) {
+      target[key] = value;
+      continue;
+    }
+    const current = target[key];
+    if (deepEqual(current, value)) continue;
+    if (key in nested && isRecord(current) && isRecord(value)) {
+      target[key] = mergeNative({ ...current }, value, nested[key] as Record<string, Record<string, unknown>>);
+      continue;
+    }
+    const bucket: Record<string, unknown> = isRecord(target[NATIVE_COLLISIONS]) ? { ...(target[NATIVE_COLLISIONS] as Record<string, unknown>) } : {};
+    bucket[key] = value;
+    target[NATIVE_COLLISIONS] = bucket;
+  }
+  return target;
+}
+
+/** `record` without its undefined values. */
+function defined<T extends Record<string, unknown>>(record: T): T {
+  for (const key of Object.keys(record)) if (record[key] === undefined) delete record[key];
+  return record;
+}
+
+function parseJson(text: string): any {
+  try {
+    return JSON.parse(text);
+  } catch {
+    return undefined;
+  }
+}
+
+/** OpenAI Responses APIs: pi keeps their item and response ids, the 2.x engine's provider state. */
+function responsesApi(api: unknown): boolean {
+  return typeof api === "string" && api.includes("responses");
+}
+
+/** A text block's 2.x provider state: its Responses message item id and phase (pi's textSignature). */
+function textState(block: Record<string, any>, api: unknown): Record<string, unknown> | undefined {
+  const signature = typeof block.textSignature === "string" ? block.textSignature : "";
+  if (!responsesApi(api) || !signature) return undefined;
+  if (!signature.startsWith("{")) return { itemId: signature };
+  const parsed = parseJson(signature);
+  if (!isRecord(parsed) || typeof parsed.id !== "string") return undefined;
+  return { itemId: parsed.id, ...(typeof parsed.phase === "string" ? { phase: parsed.phase } : {}) };
+}
+
+/**
+ * A thinking block's 2.x provider state: its Responses reasoning item id.
+ * The encrypted reasoning the item also holds stays out of the trace.
+ */
+function reasoningState(block: Record<string, any>, api: unknown): Record<string, unknown> | undefined {
+  if (!responsesApi(api) || typeof block.thinkingSignature !== "string") return undefined;
+  const item = parseJson(block.thinkingSignature);
+  return isRecord(item) && typeof item.id === "string" ? { itemId: item.id } : undefined;
+}
+
+/** A pi tool call id as the 2.x call id and item id (Responses calls are `<call_id>|<item id>`). */
+function nativeCallId(id: string, api: unknown): { id: string; itemId?: string } {
+  if (!responsesApi(api)) return { id };
+  const [callId, itemId] = id.split("|");
+  return callId && itemId ? { id: callId, itemId } : { id };
+}
+
+/** The provider's own finish reason, as the 2.x engine keeps it (Responses: the incomplete reason only). */
+function nativeRawFinish(message: Record<string, any>): string | undefined {
+  const raw = typeof message.rawStopReason === "string" ? message.rawStopReason : "";
+  if (!raw) return undefined;
+  if (!responsesApi(message.api)) return raw;
+  const incomplete = /^incomplete\.(.+)$/.exec(raw);
+  return incomplete ? incomplete[1] : undefined;
+}
+
+/** A failed or aborted step's 2.x error ({type, message}). */
+function nativeError(message: Record<string, any>): Record<string, unknown> | undefined {
+  const text = typeof message.errorMessage === "string" && message.errorMessage ? message.errorMessage : undefined;
+  if (message.stopReason === "aborted") return { type: "aborted", message: text ?? "aborted" };
+  if (message.stopReason === "error") return { type: "provider.unknown", message: text ?? "error" };
+  return undefined;
+}
+
+/** pi usage as the 2.x step tokens: output without the reasoning tokens pi counts in it. */
+function nativeTokens(usage: unknown): Record<string, unknown> | undefined {
+  if (!isRecord(usage)) return undefined;
+  const tokens = stepTokens(usage) as { input: number; output: number; reasoning: number; cache: { read: number; write: number } };
+  return { input: tokens.input, output: Math.max(0, tokens.output - tokens.reasoning), reasoning: tokens.reasoning, cache: tokens.cache };
+}
+
+/** A tool result's content as the 2.x tool content (images: the 1.x attachment holds their description). */
+function nativeToolContent(content: unknown): Array<Record<string, unknown>> {
+  const out: Array<Record<string, unknown>> = [];
+  let images = 0;
+  for (const block of Array.isArray(content) ? content : typeof content === "string" ? [content] : []) {
+    if (typeof block === "string") {
+      if (block) out.push({ type: "text", text: block });
+    } else if (isRecord(block) && block.type === "text" && typeof block.text === "string") {
+      out.push({ type: "text", text: block.text });
+    } else if (isRecord(block) && block.type === "image") {
+      const file = imagePart(block, images);
+      out.push({ type: "file", uri: `sameAs:state.attachments[${images}].url`, mime: file.mime, name: file.filename });
+      images += 1;
+    }
+  }
+  return out;
+}
+
+/** Stream timing a step's `omnirush-timing` entry recorded (stream-timing.ts). */
+type StepTiming = { streamed?: number; blocks: Map<number, { created?: number; completed?: number }>; tools: Record<string, number> };
+
+function stepTimings(entries: readonly PiEntry[]): Map<number, StepTiming> {
+  const timings = new Map<number, StepTiming>();
+  for (const entry of entries) {
+    if (entry.type !== "custom" || entry.customType !== TIMING_ENTRY_TYPE || !isRecord(entry.data)) continue;
+    const data = entry.data;
+    if (typeof data.step !== "number") continue;
+    const blocks = new Map<number, { created?: number; completed?: number }>();
+    for (const block of Array.isArray(data.blocks) ? data.blocks : []) {
+      if (!isRecord(block) || typeof block.index !== "number") continue;
+      blocks.set(block.index, defined({ created: epochMs(block.created) ?? undefined, completed: epochMs(block.completed) ?? undefined }));
+    }
+    const tools: Record<string, number> = {};
+    if (isRecord(data.tools)) for (const [id, at] of Object.entries(data.tools)) if (epochMs(at) !== null) tools[id] = epochMs(at)!;
+    timings.set(data.step, { ...(epochMs(data.streamed) !== null ? { streamed: epochMs(data.streamed)! } : {}), blocks, tools });
+  }
+  return timings;
+}
+
+/** A user message the viewer never takes for a prompt (only synthetic text or a compaction): context records do not join it. */
+function engineWritten(message: EngineMessage): boolean {
+  return message.info.role === "user"
+    && message.parts.every((part) => (part.type === "text" && part.synthetic === true) || part.type === "compaction");
+}
+
 /**
  * The engine messages of a pi session branch (entries root to leaf, as
  * SessionManager.getBranch() returns them or a session file lists them).
@@ -291,15 +481,69 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
     if (message?.role === "user") upcoming = null;
   }
 
+  const timings = stepTimings(entries);
+
   const out: EngineMessage[] = [];
   let model: { providerID: string; modelID: string } | null = options.initialModel ?? null;
   let thinkingLevel: string | null = options.initialThinkingLevel ?? null;
   let lastUserId: string | null = null;
   const variant = () => (thinkingLevel && thinkingLevel !== "off" ? thinkingLevel : null);
 
+  // 2.x context records (model switches) wait for the next message that takes them.
+  const pendingRecords: Array<Record<string, unknown>> = [];
+  const contextPart = (record: Record<string, unknown>, host: EngineMessage) => ({ ...record, sessionID: sessionId, messageID: host.info.id });
+  /** Every message goes out through here: waiting context records join the front of the first one the viewer does not take for engine-written. */
+  const emit = (message: EngineMessage): void => {
+    if (pendingRecords.length > 0 && !engineWritten(message)) {
+      message.parts = [...pendingRecords.splice(0).map((record) => contextPart(record, message)), ...message.parts];
+    }
+    out.push(message);
+  };
+  /** A 2.x record at the end of the latest message that takes one (the run's idle marker). */
+  const appendRecord = (record: Record<string, unknown>): void => {
+    for (let index = out.length - 1; index >= 0; index -= 1) {
+      const host = out[index]!;
+      if (engineWritten(host)) continue;
+      host.parts = [...host.parts, contextPart(record, host)];
+      return;
+    }
+  };
+  // 2.x model-switched records: a model or thinking-level change pi recorded
+  // (the latest of a run of changes, once a message follows), or a prompt
+  // sent to another model than the last one recorded (pi records no change
+  // for `--model` on a resumed session). Tracked apart from the 1.x state.
+  type ModelState = { model: { providerID: string; modelID: string } | null; thinkingLevel: string | null };
+  let recorded: ModelState = { model, thinkingLevel };
+  const modelRef = (state: ModelState) => {
+    if (!state.model) return undefined;
+    const effort = state.thinkingLevel && state.thinkingLevel !== "off" ? state.thinkingLevel : undefined;
+    return { id: state.model.modelID, providerID: state.model.providerID, ...(effort ? { variant: effort } : {}) };
+  };
+  const recordSwitch = (id: string, at: number | null, next: ModelState) => {
+    const now = modelRef(next);
+    const previous = modelRef(recorded);
+    recorded = next;
+    if (!now || deepEqual(now, previous)) return;
+    pendingRecords.push(defined({ id, time: { created: at ?? 0 }, type: "model-switched", model: now, previous }));
+  };
+  let pendingSwitch: { entry: PiEntry; at: number | null } | null = null;
+  const settleSwitch = () => {
+    if (!pendingSwitch) return;
+    const { entry, at } = pendingSwitch;
+    pendingSwitch = null;
+    recordSwitch(engineMessageId(String(entry.id)), at, { model, thinkingLevel });
+  };
+
   // The effort of the prompt the current turn answers (a live setting when known).
   let turnVariant: string | null = null;
-  const userMessage = (id: string, created: number | null, parts: Array<Record<string, unknown>>, index: number, extraInfo: Record<string, unknown> = {}): string | null => {
+  const userMessage = (
+    id: string,
+    created: number | null,
+    parts: Array<Record<string, unknown>>,
+    index: number,
+    extraInfo: Record<string, unknown> = {},
+    native?: Record<string, unknown>,
+  ): string | null => {
     const live = typeof entries[index]?.id === "string" ? options.promptSettings?.get(entries[index]!.id as string) : undefined;
     const messageModel = nextAssistantModel.get(index) ?? live?.model ?? model ?? null;
     const messageVariant = live && live.variant !== undefined ? live.variant : variant();
@@ -313,10 +557,23 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
       ...(messageVariant ? { variant: messageVariant } : {}),
       ...extraInfo,
     };
-    out.push({
+    // A prompt sent to another model than the last one recorded: that switch first.
+    if (native?.type === "user") recordSwitch(engineMessageId(String(entries[index]?.id ?? id), "_model"), created, { model: messageModel, thinkingLevel: messageVariant ?? null });
+    const message: EngineMessage = {
       info,
       parts: parts.map((part, partIndex) => ({ id: partId(id, partIndex), sessionID: sessionId, messageID: id, ...part })),
-    });
+    };
+    if (native) {
+      // 2.x prompt fields (engine2/native.ts enrichPrompt): an attached image's
+      // entry points at its 1.x file part (the bytes are not embedded), and its
+      // other fields join that part.
+      const files = message.parts.flatMap((part, partIndex) => part.type === "file"
+        ? [defined({ uri: `sameAs:parts[${partIndex}].url`, mime: part.mime as string | undefined, name: part.filename as string | undefined })]
+        : []);
+      message.info = mergeNative({ ...message.info }, { ...native, ...(files.length > 0 ? { files } : {}) }, INFO_NESTED);
+      for (const part of message.parts) if (part.type === "file") mergeNative(part, defined({ mime: part.mime, name: part.filename }));
+    }
+    emit(message);
     lastUserId = id;
     return messageVariant ?? null;
   };
@@ -356,7 +613,7 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
         },
       };
     });
-    out.push({
+    emit({
       info: {
         id: shellId,
         sessionID: sessionId,
@@ -380,13 +637,28 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
     if (!entryId) return;
     const at = epochMs(entry.timestamp);
     if (entry.type === "model_change") {
-      if (typeof entry.provider === "string" && typeof entry.modelId === "string") model = { providerID: entry.provider, modelID: entry.modelId };
+      if (typeof entry.provider === "string" && typeof entry.modelId === "string") {
+        model = { providerID: entry.provider, modelID: entry.modelId };
+        pendingSwitch = { entry, at };
+      }
       return;
     }
     if (entry.type === "thinking_level_change") {
-      if (typeof entry.thinkingLevel === "string") thinkingLevel = entry.thinkingLevel;
+      if (typeof entry.thinkingLevel === "string") {
+        thinkingLevel = entry.thinkingLevel;
+        pendingSwitch = { entry, at };
+      }
+      return;
+    }
+    if (entry.type === "custom" && entry.customType === TIMING_ENTRY_TYPE) {
+      // The run went idle here: its 2.x idle marker closes the run's last message.
+      const idle = isRecord(entry.data) && isRecord(entry.data.idle) ? entry.data.idle : null;
+      const outcome = idle && typeof idle.outcome === "string" ? idle.outcome : null;
+      if (idle && outcome) appendRecord({ id: engineMessageId(entryId), time: { created: epochMs(idle.at) ?? at ?? 0 }, type: "idle", outcome });
       return;
     }
+    // A message follows: the model switch before it is recorded.
+    settleSwitch();
     if (entry.type === "compaction") {
       const id = engineMessageId(entryId);
       userMessage(id, at, [{ type: "compaction" }], index);
@@ -394,8 +666,9 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
       const summaryParts: Array<Record<string, unknown>> = [];
       if (typeof entry.summary === "string" && entry.summary) summaryParts.push({ type: "text", text: entry.summary });
       if (isRecord(entry.usage)) summaryParts.push({ type: "step-finish", reason: "stop", tokens: stepTokens(entry.usage), cost: stepCost(entry.usage) });
-      out.push({
-        info: {
+      emit({
+        // The 2.x compaction record's fields on its summary reply, as the desktop keeps them.
+        info: mergeNative({
           id: summaryId,
           sessionID: sessionId,
           role: "assistant",
@@ -407,21 +680,26 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
           path: { cwd, root: cwd },
           ...(isRecord(entry.usage) ? { tokens: stepTokens(entry.usage), cost: stepCost(entry.usage) } : {}),
           finish: "stop",
-        },
+        }, defined({
+          type: "compaction",
+          status: "completed",
+          summary: typeof entry.summary === "string" ? entry.summary : undefined,
+          tokens: nativeTokens(entry.usage),
+        }), INFO_NESTED),
         parts: summaryParts.map((part, partIndex) => ({ id: partId(summaryId, partIndex), sessionID: sessionId, messageID: summaryId, ...part })),
       });
       return;
     }
     if (entry.type === "branch_summary") {
       if (typeof entry.summary === "string" && entry.summary) {
-        userMessage(engineMessageId(entryId), at, [{ type: "text", text: entry.summary, synthetic: true, metadata: { branchSummary: true } }], index);
+        userMessage(engineMessageId(entryId), at, [{ type: "text", text: entry.summary, synthetic: true, metadata: { branchSummary: true } }], index, {}, { type: "synthetic", text: entry.summary });
       }
       return;
     }
     if (entry.type === "custom_message") {
       const parts = contentParts(entry.content, { synthetic: true, metadata: { customType: entry.customType ?? null } });
       const id = engineMessageId(entryId);
-      if (parts.length > 0) userMessage(id, at, parts, index);
+      if (parts.length > 0) userMessage(id, at, parts, index, {}, { type: "synthetic", text: contentText(entry.content) });
       if (entry.customType === BACKGROUND_BASH_TYPE) backgroundBash(id, entryId, (entry as any).details, at, index);
       return;
     }
@@ -433,19 +711,20 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
       case "user": {
         const injected = options.injectedPrompts?.get(entryId);
         const parts = contentParts(message.content, injected ? { synthetic: true, metadata: { ...injected } } : {});
-        if (parts.length > 0) turnVariant = userMessage(id, created, parts, index);
+        const native = { type: injected ? "synthetic" : "user", text: contentText(message.content) };
+        if (parts.length > 0) turnVariant = userMessage(id, created, parts, index, {}, native);
         return;
       }
       case "custom": {
         const parts = contentParts(message.content, { synthetic: true, metadata: { customType: message.customType ?? null } });
-        if (parts.length > 0) userMessage(id, created, parts, index);
+        if (parts.length > 0) userMessage(id, created, parts, index, {}, { type: "synthetic", text: contentText(message.content) });
         if (message.customType === BACKGROUND_BASH_TYPE) backgroundBash(id, entryId, message.details, created, index);
         return;
       }
       case "branchSummary":
       case "compactionSummary": {
         if (typeof message.summary === "string" && message.summary) {
-          userMessage(id, created, [{ type: "text", text: message.summary, synthetic: true, metadata: { [message.role]: true } }], index);
+          userMessage(id, created, [{ type: "text", text: message.summary, synthetic: true, metadata: { [message.role]: true } }], index, {}, { type: "synthetic", text: message.summary });
         }
         return;
       }
@@ -478,8 +757,10 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
             },
           },
         ];
-        out.push({
-          info: {
+        const outputBytes = Buffer.byteLength(output);
+        emit({
+          // The 2.x shell record's fields on the reply, as the desktop keeps them (pi keeps no shell id).
+          info: mergeNative({
             id: shellId,
             sessionID: sessionId,
             role: "assistant",
@@ -492,32 +773,54 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
             cost: 0,
             tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
             finish: "stop",
-          },
+          }, defined({
+            type: "shell",
+            command,
+            status: message.cancelled === true ? "killed" : "exited",
+            exit: typeof message.exitCode === "number" ? message.exitCode : undefined,
+            output: { output, cursor: outputBytes, size: outputBytes, truncated: message.truncated === true },
+          }), INFO_NESTED),
           parts: parts.map((part, partIndex) => ({ id: partId(shellId, partIndex), sessionID: sessionId, messageID: shellId, ...part })),
         });
         return;
       }
       case "assistant": {
         const completed = at ?? created;
+        const api = message.api;
+        const timing = typeof message.timestamp === "number" ? timings.get(message.timestamp) : undefined;
         const parts: Array<Record<string, unknown>> = [{ type: "step-start" }];
+        // The 2.x content entry of each part (by part index), added once the parts have their ids.
+        const natives = new Map<number, Record<string, unknown>>();
         const changed: string[] = [];
+        // When the step's tools settled (2.x time.completed): the last result, once every call has one.
+        let settled: number | null = completed;
+        let pending = false;
         const content = Array.isArray(message.content) ? message.content : [];
-        for (const block of content) {
-          if (!isRecord(block)) continue;
+        content.forEach((block: unknown, blockIndex: number) => {
+          if (!isRecord(block)) return;
+          const blockTime = timing?.blocks.get(blockIndex);
           if (block.type === "thinking") {
             if (typeof block.thinking === "string" && block.thinking.trim()) {
               parts.push({ type: "reasoning", text: block.thinking, ...(block.redacted === true ? { metadata: { redacted: true } } : {}) });
+              const time = blockTime && (blockTime.created !== undefined || blockTime.completed !== undefined) ? { ...blockTime } : undefined;
+              natives.set(parts.length - 1, defined({ type: "reasoning", text: block.thinking, state: reasoningState(block, api), time }));
             }
           } else if (block.type === "text") {
-            if (typeof block.text === "string" && block.text) parts.push({ type: "text", text: block.text });
+            if (typeof block.text === "string" && block.text) {
+              parts.push({ type: "text", text: block.text });
+              natives.set(parts.length - 1, defined({ type: "text", text: block.text, state: textState(block, api) }));
+            }
           } else if (block.type === "toolCall") {
             const callId = typeof block.id === "string" ? block.id : "";
             const input = isRecord(block.arguments) ? block.arguments : {};
             const result = callId ? results.get(callId) : undefined;
             const start = (callId ? options.toolStarts?.get(callId) : undefined) ?? completed ?? 0;
             let state: Record<string, unknown>;
+            let nativeState: Record<string, unknown>;
             if (!result) {
               state = { status: "running", input, time: { start } };
+              nativeState = { status: "running", input };
+              pending = true;
             } else {
               const text = contentText(result.message.content);
               const attachments = contentParts(result.message.content).filter((part) => part.type === "file");
@@ -540,6 +843,11 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
                     ...(attachments.length > 0 ? { attachments } : {}),
                     time,
                   };
+              const toolContent = nativeToolContent(result.message.content);
+              nativeState = result.message.isError === true
+                ? { status: "error", input, error: { type: "tool.execution", message: text } }
+                : { status: "completed", input, ...(toolContent.length > 0 ? { content: toolContent } : {}) };
+              if (result.at !== null && (settled === null || result.at > settled)) settled = result.at;
               const target = typeof input.path === "string" ? input.path : null;
               if (result.message.isError !== true && FILE_TOOLS.has(String(block.name)) && target) {
                 const absolute = isAbsolute(target) ? target : resolve(cwd, target);
@@ -547,8 +855,21 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
               }
             }
             parts.push({ type: "tool", callID: callId, tool: typeof block.name === "string" ? block.name : "tool", state });
+            const call = nativeCallId(callId, api);
+            const ran = (callId ? timing?.tools[callId] ?? options.toolStarts?.get(callId) : undefined) ?? undefined;
+            const time = defined({ created: blockTime?.created, ran, completed: result?.at ?? undefined });
+            natives.set(parts.length - 1, defined({
+              type: "tool",
+              id: call.id,
+              name: typeof block.name === "string" ? block.name : "tool",
+              // pi runs every tool itself, none provider-side.
+              executed: false,
+              providerState: call.itemId ? { itemId: call.itemId } : undefined,
+              state: nativeState,
+              time: Object.keys(time).length > 0 ? time : undefined,
+            }));
           }
-        }
+        });
         if (changed.length > 0) parts.push({ type: "patch", files: changed });
         const finish = finishReason(message.stopReason);
         const tokens = stepTokens(message.usage);
@@ -559,29 +880,59 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
         }
         const error = assistantError(message);
         const effort = typeof message.providerThinkingLevel === "string" && message.providerThinkingLevel ? message.providerThinkingLevel : turnVariant ?? variant();
-        out.push({
-          info: {
-            id,
-            sessionID: sessionId,
-            role: "assistant",
-            ...(lastUserId ? { parentID: lastUserId } : {}),
-            time: { created: created ?? 0, completed: completed ?? 0 },
-            ...(typeof message.model === "string" ? { modelID: message.model } : {}),
-            ...(typeof message.provider === "string" ? { providerID: message.provider } : {}),
-            ...(typeof message.responseModel === "string" && message.responseModel && message.responseModel !== message.model
-              ? { responseModelID: message.responseModel }
-              : {}),
-            agent,
-            mode: agent,
-            ...(effort ? { variant: effort } : {}),
-            path: { cwd, root: cwd },
-            cost,
-            tokens,
-            ...(finish ? { finish } : {}),
-            ...(error ? { error } : {}),
-          },
-          parts: parts.map((part, partIndex) => ({ id: partId(id, partIndex), sessionID: sessionId, messageID: id, ...part })),
-        });
+        const stepParts = parts.map((part, partIndex) => ({ id: partId(id, partIndex), sessionID: sessionId, messageID: id, ...part }));
+        // The 2.x content entries' own fields on their parts (engine2/native.ts enrichPart).
+        for (const [partIndex, native] of natives) {
+          const part = stepParts[partIndex]!;
+          if (part.type !== "tool") {
+            mergeNative(part, native, { time: {} });
+            continue;
+          }
+          const { id: nativeId, state: nativeState, ...rest } = native;
+          mergeNative(part, rest);
+          if (nativeId !== part.callID) mergeNative(part, { id: nativeId });
+          if (isRecord(part.state) && isRecord(nativeState)) part.state = mergeNative({ ...part.state }, nativeState, STATE_NESTED);
+        }
+        const streamed = timing?.streamed ?? at ?? undefined;
+        const relativeChanged = changed
+          .map((file) => relative(cwd, file))
+          .filter((file) => file && !file.startsWith("..") && !isAbsolute(file));
+        const info = mergeNative({
+          id,
+          sessionID: sessionId,
+          role: "assistant",
+          ...(lastUserId ? { parentID: lastUserId } : {}),
+          time: { created: created ?? 0, completed: completed ?? 0 },
+          ...(typeof message.model === "string" ? { modelID: message.model } : {}),
+          ...(typeof message.provider === "string" ? { providerID: message.provider } : {}),
+          ...(typeof message.responseModel === "string" && message.responseModel && message.responseModel !== message.model
+            ? { responseModelID: message.responseModel }
+            : {}),
+          agent,
+          mode: agent,
+          ...(effort ? { variant: effort } : {}),
+          path: { cwd, root: cwd },
+          cost,
+          tokens,
+          ...(finish ? { finish } : {}),
+          ...(error ? { error } : {}),
+        }, defined({
+          // The 2.x assistant record's own fields (engine2/native.ts enrichAssistant).
+          type: "assistant",
+          agent,
+          model: typeof message.model === "string" && typeof message.provider === "string"
+            ? { id: message.model, providerID: message.provider, ...(effort ? { variant: effort } : {}) }
+            : undefined,
+          snapshot: relativeChanged.length > 0 ? { files: relativeChanged } : undefined,
+          finish,
+          rawFinish: nativeRawFinish(message),
+          providerState: responsesApi(api) && typeof message.responseId === "string" && message.responseId ? { responseId: message.responseId } : undefined,
+          cost,
+          tokens: nativeTokens(message.usage),
+          error: nativeError(message),
+          time: defined({ created: created ?? 0, streamed, completed: pending ? undefined : settled ?? undefined }),
+        }), INFO_NESTED);
+        emit({ info, parts: stepParts });
         return;
       }
       default:
@@ -590,6 +941,9 @@ export function engineMessagesFromEntries(entries: readonly PiEntry[], options:
         return;
     }
   });
+  // Changes after the last message: at the end of the last message that takes them.
+  settleSwitch();
+  for (const record of pendingRecords.splice(0)) appendRecord(record);
   return out;
 }
 

package/assets/extensions/omnirush/secret-env.ts

@@ -0,0 +1,15 @@
+/** Environment passed to user-controlled tools and MCP servers. */
+const SECRET_ENV_NAMES = new Set([
+  "OMNIRUSH_TOKEN",
+  "OMNIRUSH_ACCESS_TOKEN",
+  "OMNIRUSH_APPROVAL_TOKEN",
+]);
+
+/** Keep provider credentials in the agent process, never in shell/MCP children. */
+export function sanitizeToolEnvironment(input: NodeJS.ProcessEnv = process.env): NodeJS.ProcessEnv {
+  const output: NodeJS.ProcessEnv = {};
+  for (const [key, value] of Object.entries(input)) {
+    if (!SECRET_ENV_NAMES.has(key) && value !== undefined) output[key] = value;
+  }
+  return output;
+}