@saadnvd1/agent-os 0.3.738 → 0.3.743

package/lib/chat/drivers/pi-mapper.ts

@@ -3,6 +3,12 @@ import type { DriverEvent, UsageTotals } from "../events";
 import { toolDiff, toolTitle } from "../tools";
 import { TurnItems, type ToolStart } from "./turn-items";
 import { logTurnEnd, needsSummary, turnOutcome } from "./turn-outcome";
+import {
+  COMPACT_AT_PERCENT,
+  retryNotice,
+  TurnRetry,
+  type RetryDecision,
+} from "./transient-retry";
 
 type P = Record<string, unknown>;
 const str = (v: unknown) => (typeof v === "string" ? v : "");
@@ -63,10 +69,31 @@ export class PiMapper {
     cacheWriteTokens: 0,
   };
 
-  constructor(private emit: (e: DriverEvent) => void) {
+  // Retrying a turn the model's server dropped, when Pi didn't.
+  readonly retry = new TurnRetry();
+  private agentRetried = false;
+  // A failed turn is waiting for its retry to go out.
+  retrying = false;
+  // The context's fill, as last reported.
+  contextPercent = 0;
+
+  constructor(
+    private emit: (e: DriverEvent) => void,
+    // Called when a failed turn is to be tried again.
+    private onRetry: (
+      d: Extract<RetryDecision, { retry: true }>,
+      error: string
+    ) => void = () => {}
+  ) {
     this.items = new TurnItems(emit);
   }
 
+  // The reader sent a message: a new turn's retries start over.
+  userPrompt(): void {
+    this.retry.reset();
+    this.agentRetried = false;
+  }
+
   // The last answer in the run stopped because the reader stopped it.
   get aborted(): boolean {
     return this.stopReason === "aborted";
@@ -93,12 +120,31 @@ export class PiMapper {
           str(e.toolCallId),
           piTool(str(e.toolName), e.args)
         );
+      case "tool_execution_update":
+        return this.items.progress(
+          str(e.toolCallId),
+          resultText(e.partialResult)
+        );
       case "tool_execution_end":
         return this.items.toolDone(
           str(e.toolCallId),
           e.isError ? "error" : "done",
           resultText(e.result)
         );
+      case "auto_retry_start":
+        this.emit({
+          type: "model_retry",
+          attempt: Number(e.attempt) || 1,
+          message: str(e.errorMessage) || undefined,
+        });
+        this.agentRetried = true;
+        return this.items.note(
+          `Pi is trying again (${str(e.attempt)} of ${str(e.maxAttempts)}): ${str(e.errorMessage).slice(0, 160)}`
+        );
+      case "auto_retry_end":
+        if (e.success === false && e.finalError && !this.error)
+          this.error = str(e.finalError);
+        return;
       case "compaction_start":
         return this.emit({ type: "compacting", active: true });
       case "compaction_end":
@@ -163,6 +209,8 @@ export class PiMapper {
       this.items.retext(`${this.message}:${fixed.index}`, str(fixed.text));
     const u = (m.usage ?? {}) as P;
     const n = (k: string) => Number(u[k]) || 0;
+    const prompt = n("input") + n("cacheRead") + n("cacheWrite");
+    if (prompt > 0) this.emit({ type: "prompt_tokens", tokens: prompt });
     this.totals = {
       costUsd: this.totals.costUsd + (Number((u.cost as P)?.total) || 0),
       inputTokens: this.totals.inputTokens + n("input"),
@@ -184,9 +232,31 @@ export class PiMapper {
     });
   }
 
+  private retryTurn(): boolean {
+    const d = this.retry.next({
+      error: this.error,
+      harnessRetried: this.agentRetried,
+      contextHeavy: this.contextPercent >= COMPACT_AT_PERCENT,
+    });
+    if (!d.retry) return false;
+    this.items.note(retryNotice(d, this.error));
+    this.retrying = true;
+    this.onRetry(d, this.error);
+    return true;
+  }
+
+  // The reader stopped the turn while it waited to retry.
+  stopRetrying(): void {
+    this.retrying = false;
+    this.emit({ type: "usage", totals: { ...this.totals } });
+    this.items.endTurn({ interrupted: true });
+  }
+
   // The run is over: Pi has nothing left to do for it.
   endTurn(): void {
     if (!this.items.inTurn) return;
+    if (this.error && !this.aborted && this.retryTurn()) return;
+    this.retrying = false;
     this.emit({ type: "usage", totals: { ...this.totals } });
     const ended = this.aborted
       ? undefined
@@ -195,6 +265,9 @@ export class PiMapper {
           text: this.last.text,
           toolCalls: this.last.calls,
           error: this.error || undefined,
+          retries: this.retry.tried,
+          agentRetried: this.agentRetried,
+          summary: this.items.summary(),
         });
     if (ended)
       logTurnEnd({

package/lib/chat/drivers/pi.ts

@@ -6,12 +6,17 @@ import { AGENT_DEFAULT_MODEL } from "../../providers/registry";
 import { PiAccess, PiDialogs, writeExtension } from "./pi-approvals";
 import { TEXT_TOOLS_EXTENSION } from "./pi-text-tools";
 import { PiMapper } from "./pi-mapper";
+import { writePiEditGuard } from "./pi-edit-guard";
+import { typescriptPath } from "./edit-guard";
+import { continuePrompt, type RetryDecision } from "./transient-retry";
 import { PiRpc } from "./pi-rpc";
 import { RESULTS_BRIEF, SUMMARY_NUDGE } from "./turn-outcome";
 
 type P = Record<string, unknown>;
 
 const SETTLE_TRIES = 5;
+// Compacting can take a while; the retry waits for it.
+const COMPACT_MS = 300_000;
 
 export function piArgs(o: ChatStartOptions, ...extensions: string[]): string[] {
   const args = extensions.flatMap((e) => ["-e", e]);
@@ -42,6 +47,15 @@ const splitModel = (model: string) => {
     : null;
 };
 
+// The edit guard extension, or none where it can't be written.
+function guardExtension(): string[] {
+  try {
+    return [writePiEditGuard()];
+  } catch {
+    return [];
+  }
+}
+
 export const piDriver: ChatDriver = {
   id: "pi",
   plan: false,
@@ -78,7 +92,11 @@ export const piDriver: ChatDriver = {
   start(options): ChatConversation {
     const out = new InputQueue<DriverEvent>();
     const emit = (e: DriverEvent) => out.push(e);
-    const mapper = new PiMapper(emit);
+    let retryTimer: ReturnType<typeof setTimeout> | null = null;
+    // Bumped when a retry is cancelled, so one already under way stops.
+    let retryGen = 0;
+    let closed = false;
+    const mapper = new PiMapper(emit, (d, error) => scheduleRetry(d, error));
     const dialogs = new PiDialogs(emit);
     const access = new PiAccess(options.access);
     // Sent now: goes once the stopped run has ended.
@@ -87,6 +105,7 @@ export const piDriver: ChatDriver = {
     const rpc: PiRpc = new PiRpc(
       piArgs(
         options,
+        ...guardExtension(),
         writeExtension(),
         writeExtension("agentos-text-tools.mjs", TEXT_TOOLS_EXTENSION)
       ),
@@ -96,6 +115,7 @@ export const piDriver: ChatDriver = {
           ...sessionBaseEnv(),
           ...options.env,
           AGENTOS_PI_ACCESS_FILE: access.file,
+          ...(typescriptPath() && { AGENTOS_TYPESCRIPT: typescriptPath() }),
           PI_OFFLINE: "1",
         },
       },
@@ -110,6 +130,7 @@ export const piDriver: ChatDriver = {
           });
         },
         exit(error) {
+          cancelRetry();
           dialogs.pending.expireAll();
           mapper.items.endTurn({ error });
           access.remove();
@@ -150,6 +171,7 @@ export const piDriver: ChatDriver = {
         .then((s) => {
           const c = s.contextUsage;
           if (c?.tokens == null || !c.contextWindow) return;
+          mapper.contextPercent = Number(c.percent) || 0;
           emit({
             type: "context",
             context: {
@@ -164,8 +186,50 @@ export const piDriver: ChatDriver = {
         .catch(() => {});
     }
 
-    const prompt = (text: string, images?: ChatImage[]) =>
-      rpc
+    // The failed request goes again after its wait, on a compacted context
+    // for the last try. Stopping or closing cancels it.
+    function scheduleRetry(
+      d: Extract<RetryDecision, { retry: true }>,
+      error: string
+    ) {
+      cancelRetry();
+      const gen = retryGen;
+      retryTimer = setTimeout(() => {
+        retryTimer = null;
+        void resume(d, error, gen);
+      }, d.delayMs);
+    }
+
+    // Whether a retry was waiting or under way.
+    function cancelRetry(): boolean {
+      const pending = retryTimer !== null || mapper.retrying;
+      if (retryTimer) clearTimeout(retryTimer);
+      retryTimer = null;
+      retryGen++;
+      return pending;
+    }
+
+    async function resume(
+      d: Extract<RetryDecision, { retry: true }>,
+      error: string,
+      gen: number
+    ) {
+      const live = () => !closed && mapper.retrying && gen === retryGen;
+      if (!live()) return;
+      if (d.compactFirst) {
+        await rpc.request("compact", {}, COMPACT_MS).catch(() => {});
+        if (!live()) return;
+      }
+      mapper.retrying = false;
+      await prompt(continuePrompt(error), undefined, true);
+    }
+
+    const prompt = (text: string, images?: ChatImage[], retried = false) => {
+      if (!retried) {
+        cancelRetry();
+        mapper.userPrompt();
+      }
+      return rpc
         .request("prompt", {
           message: text,
           images: piImages(images),
@@ -175,10 +239,22 @@ export const piDriver: ChatDriver = {
         .catch((error: Error) => {
           mapper.items.error(error.message);
         });
+    };
+
+    // A turn waiting to retry has nothing running to abort: it ends here.
+    function stopWaiting(): boolean {
+      if (!cancelRetry()) return false;
+      mapper.stopRetrying();
+      return true;
+    }
 
     return {
       send(text, images, sendOptions) {
         if (sendOptions?.now && mapper.items.inTurn) {
+          if (stopWaiting()) {
+            void prompt(text, images);
+            return undefined;
+          }
           held.push({ text, images });
           // Abort answers once the run has stopped; what was held goes then.
           void rpc
@@ -193,6 +269,7 @@ export const piDriver: ChatDriver = {
       },
       // Stops the run only: a message already sent still goes after it.
       async interrupt() {
+        if (stopWaiting()) return;
         await rpc.request("abort", {}, 0);
       },
       async setModel(model) {
@@ -211,6 +288,8 @@ export const piDriver: ChatDriver = {
         return { canUndo: false, error: "Pi can't undo from chat", files: [] };
       },
       close() {
+        closed = true;
+        cancelRetry();
         dialogs.pending.expireAll();
         rpc.close();
         out.end();

package/lib/chat/drivers/transient-retry.ts

@@ -0,0 +1,87 @@
+// Retrying a turn the model's server dropped. OpenCode and Pi retry some
+// failures themselves (and say so); this covers the rest, with a hard limit
+// so nothing loops: a few tries with growing waits, the last one on a
+// smaller context, then the turn ends with a report.
+
+export type ErrorKind = "timeout" | "transient";
+
+export const MAX_RETRIES = 3;
+const BASE_MS = 2_000;
+// A context this full (percent of the window) is compacted before the last
+// try: a request that took too long to start is the one that shrinks.
+export const COMPACT_AT_PERCENT = 50;
+
+const NEVER =
+  /context (length|window)|maximum context|too many tokens|prompt is too long|\b40[0-4]\b|unauthori[sz]ed|invalid api key|authentication|forbidden|abort/i;
+const TIMEOUT =
+  /headers timed out|timed out|timeout|UND_ERR_(HEADERS|BODY|CONNECT)_TIMEOUT|ETIMEDOUT|ResponseStreamError/i;
+const TRANSIENT =
+  /fetch failed|socket hang up|ECONNRESET|ECONNREFUSED|EPIPE|terminated|connection (lost|closed|reset)|ended without|overloaded|\b(429|500|502|503|504|524)\b|unavailable|rate limit/i;
+
+// Whether a provider error is worth another try, and which kind it is.
+export function classifyProviderError(message: string): ErrorKind | undefined {
+  if (!message || NEVER.test(message)) return undefined;
+  if (TIMEOUT.test(message)) return "timeout";
+  if (TRANSIENT.test(message)) return "transient";
+  return undefined;
+}
+
+export function backoffMs(attempt: number, jitter = true): number {
+  const base = BASE_MS * 2 ** (attempt - 1);
+  return jitter ? Math.round(base * (0.8 + Math.random() * 0.4)) : base;
+}
+
+export type RetryDecision =
+  | { retry: true; attempt: number; delayMs: number; compactFirst: boolean }
+  | { retry: false; attempts: number; kind?: ErrorKind };
+
+// One per conversation. A turn that sends the user's message starts it over;
+// a retry never does.
+export class TurnRetry {
+  private attempts = 0;
+
+  constructor(private jitter = true) {}
+
+  reset(): void {
+    this.attempts = 0;
+  }
+
+  // Retries made in this turn.
+  get tried(): number {
+    return this.attempts;
+  }
+
+  next(o: {
+    error: string;
+    // The agent already retried this failure itself.
+    harnessRetried?: boolean;
+    // The context is big enough that compacting it is worth a try.
+    contextHeavy?: boolean;
+  }): RetryDecision {
+    const kind = classifyProviderError(o.error);
+    if (!kind || o.harnessRetried || this.attempts >= MAX_RETRIES)
+      return { retry: false, attempts: this.attempts, kind };
+    this.attempts++;
+    return {
+      retry: true,
+      attempt: this.attempts,
+      delayMs: backoffMs(this.attempts, this.jitter),
+      compactFirst: this.attempts === MAX_RETRIES && !!o.contextHeavy,
+    };
+  }
+}
+
+// The message that resumes a turn after a failure: the steps that ran are
+// already in the conversation.
+export function continuePrompt(error: string): string {
+  const why = error.replace(/\s+/g, " ").trim().slice(0, 200);
+  return `The previous request to the model failed (${why}). Continue the task from where it stopped. The tool steps above already ran: don't redo them.`;
+}
+
+export function retryNotice(
+  d: Extract<RetryDecision, { retry: true }>,
+  error: string
+): string {
+  const why = error.replace(/\s+/g, " ").trim().slice(0, 160);
+  return `The model's server didn't answer (${why}). Trying again in ${Math.round(d.delayMs / 1000)}s (${d.attempt} of ${MAX_RETRIES})${d.compactFirst ? ", after compacting the conversation" : ""}.`;
+}

package/lib/chat/drivers/turn-items.ts

@@ -6,12 +6,14 @@ import type {
   ToolStatus,
   TurnEndReason,
 } from "../events";
-import { clipOutput } from "../tools";
+import { clipOutput, tailOutput } from "../tools";
 
 type Tool = Extract<ChatItem, { kind: "tool" }>;
 type Text = Extract<ChatItem, { kind: "assistant" | "reasoning" }>;
 type Todo = Extract<ChatItem, { kind: "todos" }>["todos"][number];
 
+const PROGRESS_MS = 250;
+
 export interface ToolStart {
   name: string;
   title: string;
@@ -19,6 +21,18 @@ export interface ToolStart {
   diff?: FileDiff;
 }
 
+export interface TurnSummary {
+  steps: number;
+  failed: number;
+  files: { path: string; ok: number; failed: number }[];
+}
+
+const newStats = () => ({
+  steps: 0,
+  failed: 0,
+  files: new Map<string, { ok: number; failed: number }>(),
+});
+
 // An API's error as its message: providers wrap JSON in JSON.
 export function readableError(text: string): string {
   let message = text.trim();
@@ -50,6 +64,9 @@ export class TurnItems {
   // Texts that finished in this turn, so one can be corrected afterwards.
   private done = new Map<string, Text>();
   private startedAt = 0;
+  // What the turn has done so far, for the report when it stops short.
+  private stats = newStats();
+  private progressAt = new Map<string, number>();
   inTurn = false;
 
   constructor(private emit: (e: DriverEvent) => void) {}
@@ -58,6 +75,7 @@ export class TurnItems {
     if (this.inTurn) return;
     this.inTurn = true;
     this.startedAt = Date.now();
+    this.stats = newStats();
     this.emit({ type: "turn_start" });
   }
 
@@ -121,6 +139,20 @@ export class TurnItems {
     this.emit({ type: "item", item });
   }
 
+  // A running step's output so far (the whole of it each time, not a
+  // delta). Sent at most every PROGRESS_MS; the end replaces it.
+  progress(id: string, output: string): void {
+    const tool = this.tools.get(id);
+    if (!tool || !output) return;
+    const tail = tailOutput(output);
+    if (tail === tool.output) return;
+    const now = Date.now();
+    if (now - (this.progressAt.get(id) ?? 0) < PROGRESS_MS) return;
+    this.progressAt.set(id, now);
+    tool.output = tail;
+    this.emit({ type: "item", item: { ...tool } });
+  }
+
   hasText(key: string): boolean {
     return this.texts.has(key);
   }
@@ -138,6 +170,8 @@ export class TurnItems {
     const tool = this.tools.get(id);
     if (!tool) return;
     this.tools.delete(id);
+    this.count(tool, update, status);
+    this.progressAt.delete(id);
     this.emit({
       type: "item",
       item: {
@@ -153,6 +187,31 @@ export class TurnItems {
     });
   }
 
+  private count(
+    tool: Tool,
+    update: Partial<ToolStart> | undefined,
+    status: ToolStatus
+  ): void {
+    if (status === "stopped") return;
+    this.stats.steps++;
+    if (status === "error") this.stats.failed++;
+    const input = (update?.input ?? tool.input) as { file_path?: unknown };
+    if (!/^(Edit|Write|MultiEdit)$/.test(tool.name)) return;
+    if (typeof input?.file_path !== "string") return;
+    const file = this.stats.files.get(input.file_path) ?? { ok: 0, failed: 0 };
+    file[status === "error" ? "failed" : "ok"]++;
+    this.stats.files.set(input.file_path, file);
+  }
+
+  // The turn so far: tool steps run, how many failed, files edited.
+  summary(): TurnSummary {
+    return {
+      steps: this.stats.steps,
+      failed: this.stats.failed,
+      files: [...this.stats.files].map(([path, n]) => ({ path, ...n })),
+    };
+  }
+
   todos(todos: Todo[]): void {
     this.emit({
       type: "item",
@@ -160,6 +219,20 @@ export class TurnItems {
     });
   }
 
+  // A line in the conversation that isn't the agent's: a retry under way.
+  note(text: string): void {
+    this.emit({
+      type: "item",
+      item: {
+        id: randomUUID(),
+        kind: "note",
+        text,
+        tone: "note",
+        createdAt: Date.now(),
+      },
+    });
+  }
+
   error(message: string): void {
     this.emit({
       type: "item",

package/lib/chat/drivers/turn-outcome.ts

@@ -1,5 +1,7 @@
 import { redact } from "../../orchestrator/untrusted";
 import type { TurnEndReason } from "../events";
+import type { TurnSummary } from "./turn-items";
+import { classifyProviderError } from "./transient-retry";
 
 export interface TurnEnded {
   reason: TurnEndReason;
@@ -43,11 +45,49 @@ export function hasToolMarkup(text: string): boolean {
   );
 }
 
+const plural = (n: number, one: string, many = `${one}s`) =>
+  `${n} ${n === 1 ? one : many}`;
+
+// Where a turn stopped when the model's server kept failing: what the model
+// got done, what it didn't, and what to do next. The error text is scrubbed.
+export function failureReport(o: {
+  error: string;
+  retries: number;
+  // The agent retried it before AgentOS saw the error.
+  agentRetried?: boolean;
+  summary?: TurnSummary;
+}): string {
+  const kind = classifyProviderError(o.error);
+  const why = redact(o.error.replace(/\s+/g, " ").trim()).slice(0, 240);
+  const tries = o.retries
+    ? ` and ${plural(o.retries, "retry", "retries")} didn't get through`
+    : o.agentRetried
+      ? " even after the agent retried it"
+      : "";
+  const s = o.summary;
+  const done = s?.steps
+    ? `Finished before it stopped: ${plural(s.steps, "tool step")}${s.failed ? ` (${s.failed} failed)` : ""}${
+        s.files.length
+          ? `; edits to ${s.files
+              .map(
+                (f) =>
+                  `${f.path}${f.failed ? ` (${f.failed} failed${f.ok ? `, ${f.ok} applied` : ""})` : ""}`
+              )
+              .join(", ")}`
+          : ""
+      }.`
+    : "Nothing had finished yet.";
+  return `The model's server ${kind === "timeout" ? "timed out" : "kept failing"} (${why})${tries}. ${done} Not finished: the model's next reply, so the task is incomplete. Send "continue" to resume; the finished steps are kept.`;
+}
+
 export function turnOutcome(r: {
   finishReason?: string;
   text: string;
   toolCalls: number;
   error?: string;
+  retries?: number;
+  agentRetried?: boolean;
+  summary?: TurnSummary;
 }): TurnEnded | undefined {
   if (r.finishReason === "length")
     return {
@@ -58,7 +98,14 @@ export function turnOutcome(r: {
   if (r.error)
     return {
       reason: "error",
-      notice: "The model's server returned an error, shown above.",
+      notice: classifyProviderError(r.error)
+        ? failureReport({
+            error: r.error,
+            retries: r.retries ?? 0,
+            agentRetried: r.agentRetried,
+            summary: r.summary,
+          })
+        : "The model's server returned an error, shown above.",
     };
   if (r.toolCalls > 0) return undefined;
   if (hasToolMarkup(r.text))

package/lib/chat/events.ts

@@ -322,6 +322,14 @@ export interface Compacting {
   before?: number;
 }
 
+// A model call the turn is waiting on: since when, the last measured prompt
+// size (a floor: the next prompt is at least that big), and a retry under way.
+export interface ModelWait {
+  since: number;
+  promptTokens?: number;
+  retry?: { attempt: number; message?: string };
+}
+
 export interface UndoPreview {
   canUndo: boolean;
   error?: string;
@@ -347,7 +355,11 @@ export type DriverEvent =
   | { type: "at_rest" } // it has run every message it was given
   | { type: "state"; state: ChatState }
   // The agent started (true) or finished (false) compacting its context.
-  | { type: "compacting"; active: boolean };
+  | { type: "compacting"; active: boolean }
+  // One model response's measured prompt (input plus cache), in tokens.
+  | { type: "prompt_tokens"; tokens: number }
+  // The agent is retrying a failed model request.
+  | { type: "model_retry"; attempt: number; message?: string };
 
 // A tool call's output and diff, loaded when it's opened.
 export interface ToolBody {
@@ -368,6 +380,7 @@ export type ChatServerMessage =
       tasks: ChatItem[];
       state: ChatState;
       compacting?: Compacting | null;
+      modelWait?: ModelWait | null;
       queue: QueuedMessage[];
       suggestion: string | null;
     }
@@ -393,6 +406,7 @@ export type ChatServerMessage =
   | { type: "delta"; id: string; text: string }
   | { type: "state"; state: ChatState }
   | { type: "compacting"; compacting: Compacting | null }
+  | { type: "model_wait"; wait: ModelWait | null }
   | {
       type: "capabilities";
       commands: ChatCommand[];

package/lib/chat/model-wait-live.ts

@@ -0,0 +1,19 @@
+import { emit, type Live } from "./registry";
+import { ModelWaitTracker } from "./model-wait";
+import type { WorkerEvent } from "./worker/protocol";
+
+// Keeps the model call the turn waits on current; the events that only feed
+// it (prompt size, a retry) go no further.
+export function trackModelWait(
+  sessionId: string,
+  live: Live,
+  e: WorkerEvent
+): boolean {
+  const tracker = (live.waitTracker ??= new ModelWaitTracker());
+  const changed = tracker.on(e);
+  if (changed) {
+    live.modelWait = tracker.wait;
+    emit(sessionId, { type: "model_wait", wait: tracker.wait });
+  }
+  return e.type === "prompt_tokens" || e.type === "model_retry";
+}