pi-t3-pstack 0.0.0-stage → 0.1.0

package/src/accounts.ts

@@ -0,0 +1,527 @@
+import { Type } from "@earendil-works/pi-ai";
+import type { Static } from "@earendil-works/pi-ai";
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+  ExtensionToolContext,
+} from "@earendil-works/pi-coding-agent";
+import { Value } from "typebox/value";
+
+import { family } from "../pstack/skills/setup-pstack/scripts/t3-sheet.mjs";
+import { sessionBrief } from "./brief.ts";
+import { classifyFailure } from "./failures.ts";
+import type { Classified } from "./failures.ts";
+import { scheduleResume } from "./resume.ts";
+import { readFamilies, updateFamily } from "./state.ts";
+import type { FamilyRecord } from "./state.ts";
+import { callT3Tool } from "./t3-http.ts";
+import { t3Json } from "./t3-json.ts";
+
+// Which model families are out of usage. The state lives in one file for the
+// whole machine (state.ts), so every chat routes around a family as soon as
+// any chat or child sees it fail. This session only remembers which of its own
+// children it still has to check.
+
+const ACCOUNT_ENTRY = "t3-pstack-account";
+
+const THREAD_READ_TOOLS = [
+  "mcp__t3-code__t3_thread_read",
+  "mcp__t3__t3_thread_read",
+] as const;
+
+const DELEGATE_TOOLS: ReadonlySet<string> = new Set([
+  "mcp__t3-code__delegate_task",
+  "mcp__t3__delegate_task",
+]);
+
+export const TERMINAL_STATUSES: ReadonlySet<string> = new Set([
+  "cancelled",
+  "completed",
+  "failed",
+  "interrupted",
+  "rolled_back",
+]);
+
+const MAX_READ_PAGES = 20;
+
+// Enough of an error item's text to classify it, short enough that reading a
+// child's activity never pulls its whole result into this session.
+const ERROR_CHARS = 400;
+
+// A probe that has not settled in this long is treated as lost, so a crashed
+// chat cannot hold a family closed forever.
+const PROBE_TIMEOUT_MS = 2 * 60 * 60_000;
+
+const USAGE_LIMIT_TITLE = "Usage limit reached";
+
+const AccountEventSchema = Type.Union([
+  Type.Object({
+    childThreadId: Type.String(),
+    family: Type.String(),
+    kind: Type.Literal("delegated"),
+  }),
+  Type.Object({ childThreadId: Type.String(), kind: Type.Literal("settled") }),
+]);
+
+type AccountEvent = Static<typeof AccountEventSchema>;
+
+export interface Accounts {
+  readonly delegated: readonly string[];
+  readonly families: Readonly<Record<string, FamilyRecord>>;
+  readonly watching: ReadonlyMap<string, string>;
+}
+
+const resetPassed = (record: FamilyRecord | undefined, now: number): boolean =>
+  record?.state === "exhausted" &&
+  record.resetAt !== undefined &&
+  Date.parse(record.resetAt) <= now;
+
+const probeLost = (record: FamilyRecord | undefined, now: number): boolean =>
+  record?.state === "probing" &&
+  Date.parse(record.since) + PROBE_TIMEOUT_MS <= now;
+
+export const readAccounts = (ctx: ExtensionContext): Accounts => {
+  const delegated: string[] = [];
+  const watching = new Map<string, string>();
+
+  for (const entry of ctx.sessionManager.getBranch()) {
+    if (
+      entry.type === "custom" &&
+      entry.customType === ACCOUNT_ENTRY &&
+      Value.Check(AccountEventSchema, entry.data)
+    ) {
+      const event: AccountEvent = entry.data;
+
+      if (event.kind === "delegated") {
+        delegated.push(event.childThreadId);
+        watching.set(event.childThreadId, event.family);
+      } else {
+        watching.delete(event.childThreadId);
+      }
+    }
+  }
+
+  return { delegated, families: readFamilies(), watching };
+};
+
+export const familyAvailable = (
+  accounts: Accounts,
+  modelFamily: string,
+  now = Date.now()
+): boolean => {
+  const record = accounts.families[modelFamily];
+
+  return (
+    record === undefined || resetPassed(record, now) || probeLost(record, now)
+  );
+};
+
+const exhaustedCause = (
+  record: Extract<FamilyRecord, { state: "exhausted" }>,
+  modelFamily: string
+): string => {
+  if (record.source === "user") {
+    return `The user reported the ${modelFamily} account exhausted`;
+  }
+
+  if (record.reason === "unhealthy") {
+    return `${modelFamily} had no working account (overload or a dropped connection)`;
+  }
+
+  return `A ${modelFamily} model hit a usage limit`;
+};
+
+export const unavailableReason = (
+  accounts: Accounts,
+  modelFamily: string,
+  now = Date.now()
+): string | undefined => {
+  const record = accounts.families[modelFamily];
+
+  if (familyAvailable(accounts, modelFamily, now) || record === undefined) {
+    return undefined;
+  }
+
+  if (record.state === "probing") {
+    return `The ${modelFamily} account's reset time has passed and its first attempt is still running. Wait for that task before sending more work to ${modelFamily}.`;
+  }
+
+  const why = exhaustedCause(record, modelFamily);
+
+  const until =
+    record.resetAt === undefined
+      ? "until the user reports recovery with /pstack-recovered"
+      : `until ${record.resetAt}`;
+
+  return `${why}, so ${modelFamily} is unavailable ${until}. Call pstack_model for a target on another family.`;
+};
+
+const shortTime = (iso: string): string =>
+  `${iso.slice(0, 16).replace("T", " ")} UTC`;
+
+// One line per family that is out, for the user rather than the model.
+export const accountLines = (
+  families: Readonly<Record<string, FamilyRecord>>,
+  now = Date.now()
+): string[] =>
+  Object.entries(families)
+    .filter(
+      ([, record]) => !resetPassed(record, now) && !probeLost(record, now)
+    )
+    .map(([name, record]) => {
+      if (record.state === "probing") {
+        return `${name}: its reset time has passed and its first attempt is still running, since ${shortTime(record.since)}.`;
+      }
+
+      const until =
+        record.resetAt === undefined
+          ? `until you run /pstack-recovered ${name}`
+          : `until ${shortTime(record.resetAt)}`;
+
+      return `${name}: ${exhaustedCause(record, name)}, since ${shortTime(record.since)}, ${until}.`;
+    });
+
+export const recordFailure = (
+  modelFamily: string,
+  classified: Classified,
+  source: "delegate" | "user"
+): void => {
+  if (classified.kind === "other") {
+    return;
+  }
+
+  const record: Extract<FamilyRecord, { state: "exhausted" }> = {
+    reason: classified.kind,
+    since: new Date().toISOString(),
+    source,
+    state: "exhausted",
+  };
+
+  if (classified.resetAt !== undefined) {
+    record.resetAt = classified.resetAt;
+  }
+
+  updateFamily(modelFamily, record);
+};
+
+const ThreadReadSchema = Type.Object({
+  hasMore: Type.Boolean(),
+  items: Type.Array(
+    Type.Object({
+      text: Type.Optional(Type.Union([Type.String(), Type.Null()])),
+      title: Type.Union([Type.String(), Type.Null()]),
+      type: Type.String(),
+    })
+  ),
+  nextPosition: Type.Union([Type.Number(), Type.Null()]),
+  recentRuns: Type.Optional(Type.Array(Type.Object({ status: Type.String() }))),
+  thread: Type.Object({ status: Type.String() }),
+});
+
+type ThreadRead = Static<typeof ThreadReadSchema>;
+
+interface PageParams {
+  readonly afterPosition?: number;
+  readonly limit: number;
+  readonly maxCharsPerItem: number;
+  readonly threadId: string;
+  readonly view: string;
+}
+
+type PageReader = (params: PageParams) => Promise<ThreadRead | undefined>;
+
+const bridgeReader =
+  (ctx: ExtensionToolContext, toolName: string): PageReader =>
+  async (params) => {
+    const outcome = await ctx.executeTool(toolName, params);
+
+    return outcome.isError
+      ? undefined
+      : t3Json(ThreadReadSchema, outcome.result);
+  };
+
+// A tool_result handler gets no ctx.executeTool, so it reads through T3 directly.
+export const httpReader: PageReader = async (params) => {
+  const read = await callT3Tool("t3_thread_read", { ...params });
+
+  return Value.Check(ThreadReadSchema, read) ? read : undefined;
+};
+
+type ChildOutcome =
+  | { readonly kind: "running" }
+  | { readonly failure: Classified; readonly kind: "settled" };
+
+const NO_FAILURE: Classified = { kind: "other" };
+
+const worse = (left: Classified, right: Classified): Classified => {
+  const rank = { exhausted: 2, other: 0, unhealthy: 1 };
+
+  return rank[right.kind] > rank[left.kind] ? right : left;
+};
+
+// Pages depend on the previous page's cursor, so they are read one after another.
+const childOutcome = async (
+  read: PageReader,
+  childThreadId: string,
+  cursor: { afterPosition?: number; failure: Classified; page: number }
+): Promise<ChildOutcome> => {
+  const params = {
+    limit: 100,
+    maxCharsPerItem: ERROR_CHARS,
+    threadId: childThreadId,
+    view: "activity",
+  };
+
+  const page = await read(
+    cursor.afterPosition === undefined
+      ? params
+      : { afterPosition: cursor.afterPosition, ...params }
+  );
+
+  if (page === undefined) {
+    return { kind: "running" };
+  }
+
+  const status = page.recentRuns?.[0]?.status ?? page.thread.status;
+
+  if (!TERMINAL_STATUSES.has(status)) {
+    return { kind: "running" };
+  }
+
+  let { failure } = cursor;
+
+  for (const item of page.items) {
+    if (item.type === "error") {
+      failure = worse(
+        failure,
+        classifyFailure({
+          class: item.title === USAGE_LIMIT_TITLE ? "usage_limit" : null,
+          message: item.text ?? "",
+        })
+      );
+    }
+  }
+
+  const lastPage =
+    !page.hasMore ||
+    page.nextPosition === null ||
+    cursor.page + 1 >= MAX_READ_PAGES;
+
+  if (lastPage) {
+    return { failure, kind: "settled" };
+  }
+
+  return childOutcome(read, childThreadId, {
+    afterPosition: page.nextPosition ?? undefined,
+    failure,
+    page: cursor.page + 1,
+  });
+};
+
+export const settleChild = async (
+  pi: ExtensionAPI,
+  read: PageReader,
+  childThreadId: string,
+  childFamily: string
+): Promise<void> => {
+  const outcome = await childOutcome(read, childThreadId, {
+    failure: NO_FAILURE,
+    page: 0,
+  });
+
+  if (outcome.kind !== "settled") {
+    return;
+  }
+
+  pi.appendEntry(ACCOUNT_ENTRY, {
+    childThreadId,
+    kind: "settled",
+  } satisfies AccountEvent);
+
+  const record = readFamilies()[childFamily];
+
+  if (outcome.failure.kind !== "other") {
+    recordFailure(childFamily, outcome.failure, "delegate");
+  } else if (
+    record?.state === "probing" &&
+    record.childThreadId === childThreadId
+  ) {
+    updateFamily(childFamily);
+  }
+};
+
+export const settleChildren = async (
+  pi: ExtensionAPI,
+  ctx: ExtensionToolContext
+): Promise<void> => {
+  const toolName = THREAD_READ_TOOLS.find((name) =>
+    ctx.tools.some((tool) => tool.name === name)
+  );
+
+  if (toolName === undefined) {
+    return;
+  }
+
+  const read = bridgeReader(ctx, toolName);
+
+  await Promise.all(
+    [...readAccounts(ctx).watching].map(([childThreadId, childFamily]) =>
+      settleChild(pi, read, childThreadId, childFamily)
+    )
+  );
+};
+
+const DelegateInputSchema = Type.Object({
+  target: Type.Optional(Type.Object({ model: Type.String() })),
+});
+
+const DelegateResultSchema = Type.Object({ childThreadId: Type.String() });
+
+const sessionFamily = (ctx: ExtensionContext): string | undefined =>
+  ctx.model === undefined ? undefined : family(ctx.model.id);
+
+const AssistantErrorSchema = Type.Object({
+  errorMessage: Type.Optional(Type.String()),
+  role: Type.Literal("assistant"),
+  stopReason: Type.String(),
+});
+
+export const registerAccounts = (
+  pi: ExtensionAPI,
+  knownFamilies: () => readonly string[]
+): void => {
+  pi.on("tool_result", (event, ctx) => {
+    if (!DELEGATE_TOOLS.has(event.toolName) || event.isError) {
+      return;
+    }
+
+    const result = t3Json(DelegateResultSchema, event);
+
+    const input = Value.Check(DelegateInputSchema, event.input)
+      ? event.input
+      : undefined;
+
+    const childFamily =
+      input?.target === undefined
+        ? sessionFamily(ctx)
+        : family(input.target.model);
+
+    if (result === undefined || childFamily === undefined) {
+      return;
+    }
+
+    pi.appendEntry(ACCOUNT_ENTRY, {
+      childThreadId: result.childThreadId,
+      family: childFamily,
+      kind: "delegated",
+    } satisfies AccountEvent);
+
+    const record = readFamilies()[childFamily];
+    const now = Date.now();
+
+    if (resetPassed(record, now) || probeLost(record, now)) {
+      updateFamily(childFamily, {
+        childThreadId: result.childThreadId,
+        since: new Date(now).toISOString(),
+        state: "probing",
+      });
+    }
+  });
+
+  // This session's own model failed: a child or the orchestrator ran out, or
+  // its account went unhealthy. Record it for every chat at once. The
+  // orchestrator also schedules its own resume; a child's parent re-plans.
+  // Pi can retry a failed run after agent_end, so only the error the last run
+  // ended with, read once Pi settles, counts.
+  let lastError: string | undefined;
+
+  pi.on("agent_end", (event) => {
+    // SAFETY: only reads `role`, which every agent message carries.
+    const last = event.messages.findLast(
+      (message) => (message as { role?: string }).role === "assistant"
+    );
+
+    lastError =
+      Value.Check(AssistantErrorSchema, last) && last.stopReason === "error"
+        ? (last.errorMessage ?? "")
+        : undefined;
+  });
+
+  pi.on("agent_settled", async (_event, ctx) => {
+    const message = lastError;
+
+    lastError = undefined;
+
+    const ownFamily = sessionFamily(ctx);
+
+    if (message === undefined || ownFamily === undefined) {
+      return;
+    }
+
+    const classified = classifyFailure({ message });
+
+    recordFailure(ownFamily, classified, "delegate");
+
+    if (classified.kind !== "other" && sessionBrief(ctx) === undefined) {
+      await scheduleResume(pi, ctx, ownFamily, classified);
+    }
+  });
+
+  // Recovery takes any family accounts.json records, such as kimi. Marking one
+  // out takes only a routing family, so a typo cannot add a family.
+  const parseFamily = (
+    word: string | undefined,
+    { recorded }: { readonly recorded: boolean }
+  ): string => {
+    const name = (word ?? "").toLowerCase();
+
+    const names = recorded
+      ? [...new Set([...knownFamilies(), ...Object.keys(readFamilies())])]
+      : knownFamilies();
+
+    if (!names.includes(name)) {
+      throw new Error(`Name a model family: ${names.join(", ")}.`);
+    }
+
+    return name;
+  };
+
+  pi.registerCommand("pstack-exhausted", {
+    description:
+      "Mark a model family's account as out of usage for every chat: /pstack-exhausted <family> [reset time]",
+    handler: async (args, ctx) => {
+      const [word, reset] = args.trim().split(/\s+/u);
+      const exhaustedFamily = parseFamily(word, { recorded: false });
+
+      if (reset !== undefined && Number.isNaN(Date.parse(reset))) {
+        throw new Error(`"${reset}" is not a date and time.`);
+      }
+
+      recordFailure(
+        exhaustedFamily,
+        reset === undefined
+          ? { kind: "exhausted" }
+          : { kind: "exhausted", resetAt: new Date(reset).toISOString() },
+        "user"
+      );
+
+      if (ctx.hasUI) {
+        ctx.ui.notify(`${exhaustedFamily} marked as out of usage.`, "info");
+      }
+    },
+  });
+
+  pi.registerCommand("pstack-recovered", {
+    description:
+      "Mark a model family's account as usable again for every chat: /pstack-recovered <family>",
+    handler: async (args, ctx) => {
+      const recovered = parseFamily(args.trim(), { recorded: true });
+
+      updateFamily(recovered);
+
+      if (ctx.hasUI) {
+        ctx.ui.notify(`${recovered} marked as usable again.`, "info");
+      }
+    },
+  });
+};

package/src/bash-timeout.ts

@@ -0,0 +1,63 @@
+import { Type } from "@earendil-works/pi-ai";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Value } from "typebox/value";
+
+// Every shell command in every session, parent or child, gets a timeout of at
+// most this many seconds. A command with no timeout gets exactly this one, so
+// a hung test suite or a long `sleep` poll stalls a chat for minutes, not for
+// a quarter of an hour or forever. T3_PSTACK_BASH_TIMEOUT_S overrides it.
+const DEFAULT_SECONDS = 300;
+
+export const bashTimeoutSeconds = (): number => {
+  const fromEnv = Number(process.env.T3_PSTACK_BASH_TIMEOUT_S);
+
+  return Number.isFinite(fromEnv) && fromEnv > 0 ? fromEnv : DEFAULT_SECONDS;
+};
+
+const BashInputSchema = Type.Object({
+  command: Type.String(),
+  timeout: Type.Optional(Type.Unknown()),
+});
+
+const TIMED_OUT = /Command timed out after \d+ seconds/u;
+
+export const registerBashTimeout = (pi: ExtensionAPI): void => {
+  pi.on("tool_call", (event) => {
+    const { input } = event;
+
+    if (event.toolName !== "bash" || !Value.Check(BashInputSchema, input)) {
+      return;
+    }
+
+    const cap = bashTimeoutSeconds();
+
+    input.timeout =
+      Value.Check(Type.Number(), input.timeout) && input.timeout > 0
+        ? Math.min(input.timeout, cap)
+        : cap;
+  });
+
+  pi.on("tool_result", (event) => {
+    if (event.toolName !== "bash") {
+      return;
+    }
+
+    const timedOut = event.content.some(
+      (part) => part.type === "text" && TIMED_OUT.test(part.text)
+    );
+
+    if (!timedOut) {
+      return;
+    }
+
+    return {
+      content: [
+        ...event.content,
+        {
+          text: `Shell commands stop after at most ${bashTimeoutSeconds()} seconds. Do not retry with a longer timeout. For a job that needs longer, start it in the background with its output in a log file (nohup <cmd> > "$TMPDIR/job.log" 2>&1 &), then check the log with short commands.`,
+          type: "text",
+        },
+      ],
+    };
+  });
+};

package/src/bridge-trim.ts

@@ -0,0 +1,109 @@
+import {
+  getCurrentSystemMessage,
+  getCurrentSystemPrompt,
+} from "@earendil-works/pi-ai";
+import type {
+  BeforeAgentStartEvent,
+  ExtensionAPI,
+  ExtensionContext,
+} from "@earendil-works/pi-coding-agent";
+
+// T3 shares its orchestration instructions across providers. Pi has neither
+// Claude's Agent tool nor Codex's spawn_agent, so route its subagents through T3.
+export const BRIDGE_CUTS: readonly string[] = [
+  "- Subagents: use your native tool when the child uses your provider, a model it offers, and no settings beyond yours, and it finishes in this turn (Claude `Agent` always with `run_in_background:false`; Codex `spawn_agent`, then wait). Use `delegate_task` for another provider's model, user-asked settings, a separate checkout (`isolation:\"worktree\"`), background work, or `delegate_task` by name. A native type running another provider's model counts too. Workflow only when asked.",
+];
+
+const T3_MARKER = "## T3 Code orchestration";
+
+const REPLACEMENT =
+  "- Subagents: use `delegate_task` for subagent work. Workflow only when asked.";
+
+const trim = (text: string): string => {
+  if (!text.includes(T3_MARKER)) {
+    return text;
+  }
+
+  let trimmed = text;
+
+  for (const cut of BRIDGE_CUTS) {
+    trimmed = trimmed.replaceAll(cut, REPLACEMENT);
+  }
+
+  return trimmed;
+};
+
+export const registerBridgeTrim = (pi: ExtensionAPI): void => {
+  let warned = false;
+  let runOptions: BeforeAgentStartEvent["systemPromptOptions"] | undefined;
+
+  const warnOnDrift = (text: string, ctx: ExtensionContext): void => {
+    if (!(text.includes(T3_MARKER) && ctx.hasUI) || warned) {
+      return;
+    }
+
+    const missed = BRIDGE_CUTS.filter(
+      (cut) => !(text.includes(cut) || text.includes(REPLACEMENT))
+    ).length;
+
+    if (missed > 0) {
+      warned = true;
+      ctx.ui.notify(
+        `t3-pstack: ${missed} of ${BRIDGE_CUTS.length} cuts to T3's bridge text no longer match. T3 likely reworded it; update src/bridge-trim.ts.`,
+        "warning"
+      );
+    }
+  };
+
+  pi.on("before_agent_start", (event) => {
+    runOptions = event.systemPromptOptions;
+
+    if (runOptions.forceSystemPrompt !== undefined) {
+      runOptions.forceSystemPrompt = trim(runOptions.forceSystemPrompt);
+    }
+  });
+
+  pi.on("context_with_system", (event, ctx) => {
+    // Pi projects a forced prompt after this hook, including one set by a
+    // later before_agent_start handler such as T3's bridge.
+    const forced = runOptions?.forceSystemPrompt;
+
+    const section = getCurrentSystemMessage(event.messages)?.sections
+      ?.t3_code_orchestration;
+
+    warnOnDrift(
+      forced ?? section ?? getCurrentSystemPrompt(event.messages),
+      ctx
+    );
+
+    if (runOptions !== undefined && forced !== undefined) {
+      runOptions.forceSystemPrompt = trim(forced);
+    }
+
+    return {
+      messages: event.messages.map((message) => {
+        if (message.role !== "system") {
+          return message;
+        }
+
+        const content = Array.isArray(message.content)
+          ? message.content.map((part) => ({
+              ...part,
+              text: trim(part.text),
+            }))
+          : trim(message.content);
+
+        const sections =
+          message.sections &&
+          Object.fromEntries(
+            Object.entries(message.sections).map(([key, text]) => [
+              key,
+              text === null ? null : trim(text),
+            ])
+          );
+
+        return { ...message, content, sections };
+      }),
+    };
+  });
+};