@tt-a1i/openpi 0.1.0

package/extensions/goal/controller.ts

@@ -0,0 +1,794 @@
+import type { ImageContent } from "@earendil-works/pi-ai";
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+} from "@earendil-works/pi-coding-agent";
+import { sanitizeTerminalText } from "../shared/terminal-text.ts";
+import {
+  budgetLimitPrompt,
+  continuationPrompt,
+  objectiveUpdatedPrompt,
+} from "./prompts.ts";
+import {
+  GOAL_ENTRY_TYPE,
+  GoalRestoreError,
+  acknowledgeGoalCompletion,
+  budgetLimitTransition,
+  canResumeGoal,
+  clearBlockedAudit,
+  editGoalObjective,
+  emergencyLimitTransition,
+  isGoalActive,
+  isGoalVisible,
+  markContinuationDispatched,
+  normalizeGoalReason,
+  recordBlockedAudit,
+  recordGoalProgress,
+  restoreGoalState,
+  resumeGoal,
+  setContinuationDeferred,
+  transitionGoal,
+  validateGoalSnapshot,
+  type GoalSnapshot,
+} from "./state.ts";
+
+export const GOAL_CONTINUATION_TYPE = "goal-continuation";
+
+type AssistantStopReason = "stop" | "length" | "toolUse" | "error" | "aborted";
+
+const BLOCKED_GOAL_TURN_MINIMUM = 3;
+
+export interface GoalControllerOptions {
+  now?: () => number;
+  createId?: () => string;
+}
+
+export interface GoalModelUpdate {
+  goal: GoalSnapshot;
+  message: string;
+  blockedAudit?: {
+    blocker: string;
+    consecutiveTurns: number;
+    requiredTurns: number;
+    accepted: boolean;
+  };
+}
+
+export function countAssistantTokens(messages: readonly unknown[]) {
+  let total = 0;
+  for (const item of messages) {
+    const message = unwrapMessage(item);
+    if (!message || message.role !== "assistant") continue;
+    const tokens = assistantGoalTokens(message);
+    total += tokens;
+    if (!Number.isSafeInteger(total)) return 0;
+  }
+  return total;
+}
+
+export function lastAssistantStopReason(messages: readonly unknown[]) {
+  return lastAssistantMessage(messages)?.stopReason as
+    AssistantStopReason | undefined;
+}
+
+export function isUsageLimitError(messages: readonly unknown[]) {
+  const message = lastAssistantMessage(messages);
+  if (!message || message.stopReason !== "error") return false;
+  const diagnosticText = Array.isArray(message.diagnostics)
+    ? message.diagnostics
+        .flatMap((diagnostic) =>
+          isRecord(diagnostic) && isRecord(diagnostic.error)
+            ? [diagnostic.error.message, diagnostic.error.code]
+            : [],
+        )
+        .join(" ")
+    : "";
+  const text = `${String(message.errorMessage ?? "")} ${diagnosticText}`;
+  return /usage(?:_|\s+)limit|insufficient_quota|quota (?:has been )?exceeded|billing limit/iu.test(
+    text,
+  );
+}
+
+export class GoalController {
+  private readonly pi: ExtensionAPI;
+  private goal: GoalSnapshot | undefined;
+  private lockedReason: string | undefined;
+  private automationEnabled = false;
+  private continuationPending = false;
+  private completionAcknowledgementPending = false;
+  private readonly completionInputMarkers = new Set<string>();
+  private trackedGoalId: string | undefined;
+  private trackedStartedAt: number | undefined;
+  private trackedTokens = 0;
+  private currentRunObservedTokens = 0;
+  private currentRunFlushedTokens = 0;
+  private trackedFromRunStart = false;
+  private accountingGoalId: string | undefined;
+  private elapsedRemainderMs = 0;
+  private lastStopReason: AssistantStopReason | undefined;
+  private lastErrorWasUsageLimit = false;
+  private suppressAbortedStop = false;
+  private readonly now: () => number;
+  readonly createId: () => string;
+
+  constructor(pi: ExtensionAPI, options: GoalControllerOptions = {}) {
+    this.pi = pi;
+    this.now = options.now ?? Date.now;
+    this.createId = options.createId ?? (() => crypto.randomUUID());
+  }
+
+  snapshot() {
+    return isGoalVisible(this.goal) ? structuredClone(this.goal) : undefined;
+  }
+
+  footerSnapshot() {
+    const current = this.snapshot();
+    if (
+      current?.status === "complete" &&
+      this.completionAcknowledgementPending
+    ) {
+      return { ...current, completionAcknowledged: true as const };
+    }
+    return current;
+  }
+
+  revision() {
+    return this.goal?.revision ?? 0;
+  }
+
+  problem() {
+    return this.lockedReason;
+  }
+
+  restore(ctx: ExtensionContext, deferActive = false) {
+    this.resetRuntime();
+    this.automationEnabled = ctx.mode !== "print" && ctx.mode !== "json";
+    try {
+      const restored = restoreGoalState(ctx.sessionManager.getBranch());
+      this.goal = restored.snapshot;
+      this.lockedReason = undefined;
+      if (restored.migrated && this.goal) this.persist(this.goal);
+      if (deferActive && this.goal?.status === "active") {
+        this.persist(setContinuationDeferred(this.goal, true, this.now()));
+      }
+    } catch (error) {
+      this.goal = undefined;
+      this.lockedReason = sanitizeTerminalText(
+        error instanceof GoalRestoreError || error instanceof Error
+          ? error.message
+          : String(error),
+      );
+    }
+  }
+
+  replace(candidate: GoalSnapshot) {
+    this.assertUnlocked();
+    this.continuationPending = false;
+    this.completionInputMarkers.clear();
+    this.completionAcknowledgementPending = false;
+    if (this.goal?.id !== candidate.id) this.resetAccountingClock();
+    this.persist(candidate);
+    return this.snapshot();
+  }
+
+  edit(objective: string, ctx: ExtensionContext) {
+    this.assertUnlocked();
+    const current = this.requireVisibleGoal();
+    this.flushTrackedProgress();
+    this.persist(
+      editGoalObjective(this.goal ?? current, objective, this.now()),
+    );
+    if (this.goal?.status === "active") {
+      this.continuationPending = false;
+      if (!ctx.isIdle()) this.beginTracking(this.goal);
+      this.dispatchPrompt(
+        ctx,
+        ctx.isIdle() ? "continuation" : "objective_updated",
+        true,
+      );
+    }
+    return this.snapshot();
+  }
+
+  pause(reason = "Paused by user.") {
+    this.assertUnlocked();
+    const current = this.requireVisibleGoal();
+    if (current.status === "paused") return this.snapshot();
+    if (current.status !== "active") {
+      throw new Error(`Goal cannot pause from ${current.status}.`);
+    }
+    this.flushTrackedProgress();
+    this.continuationPending = false;
+    this.persist(
+      transitionGoal(this.goal ?? current, "paused", this.now(), reason),
+    );
+    this.resetAccountingClock();
+    return this.snapshot();
+  }
+
+  resume() {
+    this.assertUnlocked();
+    const current = this.requireVisibleGoal();
+    if (!canResumeGoal(current)) {
+      throw new Error(`Goal cannot resume from ${current.status}.`);
+    }
+    this.continuationPending = false;
+    this.resetAccountingClock();
+    this.persist(resumeGoal(current, this.now(), "Resumed by user."));
+    return this.snapshot();
+  }
+
+  updateFromModel(status: "complete" | "blocked", blocker?: unknown) {
+    this.assertUnlocked();
+    const current = this.requireVisibleGoal();
+    if (status === "blocked" && blocker === undefined) {
+      throw new Error("Blocked goal updates require a blocker description.");
+    }
+    const normalizedBlocker =
+      status === "blocked" ? normalizeGoalReason(blocker) : undefined;
+    if (current.status === status) {
+      return this.modelUpdate(current, `Goal is already ${status}.`);
+    }
+    if (current.status === "budget_limited" && status === "blocked") {
+      return this.modelUpdate(
+        current,
+        "Goal remains limited by budget; blocked cannot replace that status.",
+      );
+    }
+    if (status === "blocked") {
+      const recorded = recordBlockedAudit(
+        current,
+        normalizedBlocker,
+        this.now(),
+      );
+      if (recorded.recorded) this.persist(recorded.snapshot);
+      if (recorded.audit.consecutiveTurns < BLOCKED_GOAL_TURN_MINIMUM) {
+        return this.modelUpdate(
+          this.goal ?? recorded.snapshot,
+          `Blocker recorded for goal turn ${recorded.audit.consecutiveTurns} of ${BLOCKED_GOAL_TURN_MINIMUM}; the goal remains active. Continue working until the same blocker recurs on ${BLOCKED_GOAL_TURN_MINIMUM} consecutive distinct goal turns.`,
+          {
+            blocker: recorded.audit.blocker,
+            consecutiveTurns: recorded.audit.consecutiveTurns,
+            requiredTurns: BLOCKED_GOAL_TURN_MINIMUM,
+            accepted: false,
+          },
+        );
+      }
+    }
+    this.flushTrackedProgress();
+    this.continuationPending = false;
+    this.persist(
+      transitionGoal(
+        this.goal ?? current,
+        status,
+        this.now(),
+        status === "complete"
+          ? "Marked complete by the goal agent."
+          : normalizedBlocker,
+      ),
+    );
+    this.resetAccountingClock();
+    return this.modelUpdate(
+      this.goal!,
+      status === "complete" ? "Goal marked complete." : "Goal marked blocked.",
+      status === "blocked"
+        ? {
+            blocker: normalizedBlocker!,
+            consecutiveTurns: BLOCKED_GOAL_TURN_MINIMUM,
+            requiredTurns: BLOCKED_GOAL_TURN_MINIMUM,
+            accepted: true,
+          }
+        : undefined,
+    );
+  }
+
+  clear() {
+    this.assertUnlocked();
+    if (!isGoalVisible(this.goal)) return false;
+    this.flushTrackedProgress();
+    this.continuationPending = false;
+    this.persist(
+      transitionGoal(this.goal, "cleared", this.now(), "Cleared by user."),
+    );
+    this.resetAccountingClock();
+    return true;
+  }
+
+  kickoff(ctx: ExtensionContext) {
+    this.assertUnlocked();
+    const current = this.requireVisibleGoal();
+    if (current.status !== "active") {
+      throw new Error(`Goal cannot start from ${current.status}.`);
+    }
+    if (!ctx.isIdle()) this.beginTracking(current);
+    return this.dispatchPrompt(ctx, "continuation", true);
+  }
+
+  continueWhenIdle(ctx: ExtensionContext) {
+    return this.dispatchPrompt(ctx, "continuation", false);
+  }
+
+  sanitizeCompletionMarkerImages(images: readonly ImageContent[] = []) {
+    let changed = false;
+    const sanitized = images.filter((image) => {
+      if (image.mimeType !== GOAL_COMPLETION_MARKER_MIME) return true;
+      changed = true;
+      this.completionInputMarkers.delete(image.data);
+      return false;
+    });
+    return { images: sanitized, changed };
+  }
+
+  prepareExplicitInput() {
+    let marker: ImageContent | undefined;
+    if (this.goal?.status === "complete" && !this.goal.completionAcknowledged) {
+      const data = crypto.randomUUID();
+      if (this.completionInputMarkers.size >= 32) {
+        const oldest = this.completionInputMarkers.values().next().value;
+        if (oldest !== undefined) this.completionInputMarkers.delete(oldest);
+      }
+      this.completionInputMarkers.add(data);
+      marker = {
+        type: "image",
+        data,
+        mimeType: GOAL_COMPLETION_MARKER_MIME,
+      };
+    }
+    if (this.goal?.status === "active" && this.goal.deferContinuation) {
+      this.persist(setContinuationDeferred(this.goal, false, this.now()));
+    }
+    return marker;
+  }
+
+  releaseDeferredContinuation() {
+    if (this.goal?.status !== "active" || !this.goal.deferContinuation) {
+      return false;
+    }
+    this.persist(setContinuationDeferred(this.goal, false, this.now()));
+    return true;
+  }
+
+  agentStarted() {
+    this.continuationPending = false;
+    this.lastStopReason = undefined;
+    this.lastErrorWasUsageLimit = false;
+    this.suppressAbortedStop = false;
+    this.currentRunObservedTokens = 0;
+    this.currentRunFlushedTokens = 0;
+    if (!this.goal || !isGoalActive(this.goal) || this.goal.deferContinuation) {
+      if (!this.trackedGoalId) this.resetTrackedRun();
+      else this.trackedFromRunStart = true;
+      return;
+    }
+    this.beginTracking(this.goal, true);
+  }
+
+  messageEnded(message: unknown) {
+    const unwrapped = unwrapMessage(message);
+    if (!unwrapped) return undefined;
+    if (unwrapped.role === "user") {
+      const sanitized = stripKnownMarkers(
+        unwrapped,
+        this.completionInputMarkers,
+      );
+      if (!sanitized) return undefined;
+      if (sanitized.matchedTrackedMarker) {
+        this.completionAcknowledgementPending = true;
+      }
+      return {
+        message: sanitized.message,
+        footerChanged: sanitized.matchedTrackedMarker,
+      };
+    }
+    if (unwrapped.role !== "assistant") return undefined;
+    this.currentRunObservedTokens = safeTokenAdd(
+      this.currentRunObservedTokens,
+      assistantGoalTokens(unwrapped),
+    );
+    return undefined;
+  }
+
+  turnEnded() {
+    this.persistPendingCompletionAcknowledgement();
+  }
+
+  settledWithoutAcknowledgement() {
+    this.completionInputMarkers.clear();
+  }
+
+  agentEnded(messages: readonly unknown[]) {
+    if (!this.trackedGoalId) return;
+    const observed = this.trackedFromRunStart
+      ? Math.max(this.currentRunObservedTokens, countAssistantTokens(messages))
+      : this.currentRunObservedTokens;
+    this.trackedTokens = safeTokenAdd(
+      this.trackedTokens,
+      Math.max(0, observed - this.currentRunFlushedTokens),
+    );
+    this.currentRunObservedTokens = 0;
+    this.currentRunFlushedTokens = 0;
+    this.lastStopReason = lastAssistantStopReason(messages);
+    this.lastErrorWasUsageLimit = isUsageLimitError(messages);
+  }
+
+  toolFinished(ctx: ExtensionContext) {
+    if (
+      !this.trackedGoalId ||
+      !this.goal ||
+      this.goal.id !== this.trackedGoalId ||
+      this.goal.status !== "active"
+    ) {
+      return;
+    }
+    this.flushTrackedProgress();
+    if (!this.goal || this.goal.status !== "active") return;
+    const budgetLimited = budgetLimitTransition(this.goal, this.now());
+    if (!budgetLimited) return;
+    this.persist(budgetLimited);
+    this.resetAccountingClock();
+    this.dispatchPrompt(ctx, "budget_limit", true);
+  }
+
+  compacted(willRetry: boolean) {
+    if (willRetry && this.lastStopReason === "aborted") {
+      this.suppressAbortedStop = true;
+    }
+  }
+
+  settled(ctx: ExtensionContext) {
+    const trackedGoalId = this.trackedGoalId;
+    if (trackedGoalId) this.continuationPending = false;
+    const stopReason = this.lastStopReason;
+    const errorWasUsageLimit = this.lastErrorWasUsageLimit;
+    const suppressAbortedStop = this.suppressAbortedStop;
+
+    if (!trackedGoalId || !this.goal || this.goal.id !== trackedGoalId) {
+      this.resetTrackedRun();
+      return;
+    }
+
+    this.flushTrackedProgress();
+    this.resetTrackedRun();
+    if (!this.goal || this.goal.id !== trackedGoalId) return;
+    if (
+      this.goal.status === "budget_limited" &&
+      stopReason === "error" &&
+      errorWasUsageLimit
+    ) {
+      this.persist(
+        transitionGoal(
+          this.goal,
+          "usage_limited",
+          this.now(),
+          "Stopped after the active turn hit a usage limit.",
+        ),
+      );
+      this.resetAccountingClock();
+      return;
+    }
+    if (this.goal.status !== "active") return;
+
+    this.resetUnreportedBlockedAudit();
+    const budgetLimited = budgetLimitTransition(this.goal, this.now());
+    if (budgetLimited) {
+      this.persist(budgetLimited);
+      this.resetAccountingClock();
+      this.dispatchPrompt(ctx, "budget_limit", true);
+      return;
+    }
+    if (stopReason === "aborted" && !suppressAbortedStop) {
+      this.persist(
+        transitionGoal(
+          this.goal,
+          "paused",
+          this.now(),
+          "Paused after the active turn was interrupted.",
+        ),
+      );
+      this.resetAccountingClock();
+      return;
+    }
+    if (stopReason === "error") {
+      this.persist(
+        transitionGoal(
+          this.goal,
+          errorWasUsageLimit ? "usage_limited" : "blocked",
+          this.now(),
+          errorWasUsageLimit
+            ? "Stopped after the active turn hit a usage limit."
+            : "Blocked after the active turn ended with an error.",
+        ),
+      );
+      this.resetAccountingClock();
+      return;
+    }
+
+    const emergencyLimited = emergencyLimitTransition(this.goal, this.now());
+    if (emergencyLimited) {
+      this.persist(emergencyLimited);
+      this.resetAccountingClock();
+      return;
+    }
+    this.continueWhenIdle(ctx);
+  }
+
+  shutdown() {
+    this.automationEnabled = false;
+    this.resetRuntime();
+  }
+
+  private modelUpdate(
+    goal: GoalSnapshot,
+    message: string,
+    blockedAudit?: GoalModelUpdate["blockedAudit"],
+  ) {
+    return {
+      goal: structuredClone(goal),
+      message,
+      ...(blockedAudit === undefined ? {} : { blockedAudit }),
+    };
+  }
+
+  private resetUnreportedBlockedAudit() {
+    if (!this.goal?.blockedAudit) return;
+    if (
+      this.goal.blockedAudit.lastTurn === (this.goal.continuationCount ?? 0)
+    ) {
+      return;
+    }
+    this.persist(clearBlockedAudit(this.goal, this.now()));
+  }
+
+  private persistPendingCompletionAcknowledgement() {
+    if (
+      !this.completionAcknowledgementPending ||
+      this.goal?.status !== "complete" ||
+      this.goal.completionAcknowledged
+    ) {
+      this.completionAcknowledgementPending = false;
+      return;
+    }
+    this.persist(acknowledgeGoalCompletion(this.goal, this.now()));
+    this.completionAcknowledgementPending = false;
+  }
+
+  private dispatchPrompt(
+    ctx: ExtensionContext,
+    kind: "continuation" | "objective_updated" | "budget_limit",
+    allowBusy: boolean,
+  ) {
+    if (!this.automationEnabled || !this.goal || this.continuationPending) {
+      return false;
+    }
+    if (!allowBusy && !ctx.isIdle()) return false;
+    if (kind !== "budget_limit") {
+      if (this.goal.status !== "active" || this.goal.deferContinuation) {
+        return false;
+      }
+      const emergencyLimited = emergencyLimitTransition(this.goal, this.now());
+      if (emergencyLimited) {
+        this.persist(emergencyLimited);
+        return false;
+      }
+    } else if (this.goal.status !== "budget_limited") {
+      return false;
+    }
+
+    if (kind === "budget_limit") this.beginTracking(this.goal);
+    const goal = this.goal;
+    const content =
+      kind === "continuation"
+        ? continuationPrompt(goal)
+        : kind === "objective_updated"
+          ? objectiveUpdatedPrompt(goal)
+          : budgetLimitPrompt(goal);
+    this.continuationPending = true;
+    try {
+      this.pi.sendMessage(
+        {
+          customType: GOAL_CONTINUATION_TYPE,
+          display: true,
+          content,
+          details: { kind, goalId: goal.id, revision: goal.revision },
+        },
+        {
+          triggerTurn: true,
+          deliverAs: ctx.isIdle() ? "followUp" : "steer",
+        },
+      );
+    } catch (error) {
+      this.continuationPending = false;
+      if (this.goal?.status === "active") {
+        this.persist(
+          transitionGoal(
+            this.goal,
+            "paused",
+            this.now(),
+            "Paused because the next autonomous turn could not be dispatched.",
+          ),
+        );
+      }
+      throw error;
+    }
+    // Count the continuation only after the turn was actually dispatched, so a
+    // failed send never erodes the emergency safety budget with zero productive
+    // turns. The in-memory continuationPending guard (set above) is what blocks
+    // a concurrent dispatch in the window before this persist lands.
+    if (kind !== "budget_limit") {
+      this.persist(markContinuationDispatched(this.goal, this.now()));
+    }
+    return true;
+  }
+
+  private persist(candidate: GoalSnapshot) {
+    const checked = validateGoalSnapshot(candidate);
+    this.pi.appendEntry(GOAL_ENTRY_TYPE, checked);
+    this.goal = checked;
+  }
+
+  private beginTracking(goal: GoalSnapshot, fromRunStart = false) {
+    if (this.accountingGoalId !== goal.id) {
+      this.accountingGoalId = goal.id;
+      this.elapsedRemainderMs = 0;
+    }
+    if (this.trackedGoalId !== goal.id) {
+      this.trackedGoalId = goal.id;
+      this.trackedStartedAt = this.now();
+      this.trackedTokens = 0;
+      this.currentRunFlushedTokens = fromRunStart
+        ? 0
+        : this.currentRunObservedTokens;
+      this.trackedFromRunStart = fromRunStart;
+    } else if (fromRunStart) {
+      this.trackedFromRunStart = true;
+    }
+  }
+
+  private flushTrackedProgress() {
+    if (
+      !this.trackedGoalId ||
+      !this.goal ||
+      this.goal.id !== this.trackedGoalId ||
+      this.goal.status === "cleared"
+    ) {
+      return;
+    }
+    const unflushedCurrentTokens = Math.max(
+      0,
+      this.currentRunObservedTokens - this.currentRunFlushedTokens,
+    );
+    const tokens = safeTokenAdd(this.trackedTokens, unflushedCurrentTokens);
+    const now = this.now();
+    const elapsedMs =
+      this.trackedStartedAt === undefined
+        ? 0
+        : Math.max(0, now - this.trackedStartedAt);
+    const totalElapsedMs = this.elapsedRemainderMs + elapsedMs;
+    const elapsedSeconds = Math.floor(totalElapsedMs / 1_000);
+    this.elapsedRemainderMs = totalElapsedMs % 1_000;
+    if (tokens > 0 || elapsedSeconds > 0) {
+      this.persist(recordGoalProgress(this.goal, tokens, elapsedSeconds, now));
+    }
+    this.trackedTokens = 0;
+    this.currentRunFlushedTokens = this.currentRunObservedTokens;
+    this.trackedStartedAt = now;
+  }
+
+  private requireVisibleGoal() {
+    if (!isGoalVisible(this.goal)) throw new Error("No goal is currently set.");
+    return this.goal;
+  }
+
+  private resetRuntime() {
+    this.continuationPending = false;
+    this.completionInputMarkers.clear();
+    this.completionAcknowledgementPending = false;
+    this.resetTrackedRun();
+    this.resetAccountingClock();
+  }
+
+  private resetTrackedRun() {
+    this.trackedGoalId = undefined;
+    this.trackedStartedAt = undefined;
+    this.trackedTokens = 0;
+    this.currentRunObservedTokens = 0;
+    this.currentRunFlushedTokens = 0;
+    this.trackedFromRunStart = false;
+    this.lastStopReason = undefined;
+    this.lastErrorWasUsageLimit = false;
+    this.suppressAbortedStop = false;
+  }
+
+  private resetAccountingClock() {
+    this.accountingGoalId = undefined;
+    this.elapsedRemainderMs = 0;
+  }
+
+  private assertUnlocked() {
+    if (this.lockedReason) {
+      throw new Error(
+        `Session goal is locked by malformed branch history: ${this.lockedReason}`,
+      );
+    }
+  }
+}
+
+function assistantGoalTokens(message: Record<string, unknown>) {
+  const usage = message.usage;
+  if (!isRecord(usage)) return 0;
+  const input = usage.input;
+  const output = usage.output;
+  if (
+    Number.isSafeInteger(input) &&
+    (input as number) >= 0 &&
+    Number.isSafeInteger(output) &&
+    (output as number) >= 0
+  ) {
+    return safeTokenAdd(input as number, output as number);
+  }
+  const total = usage.totalTokens;
+  return Number.isSafeInteger(total) && (total as number) >= 0
+    ? (total as number)
+    : 0;
+}
+
+function lastAssistantMessage(messages: readonly unknown[]) {
+  for (let index = messages.length - 1; index >= 0; index--) {
+    const message = unwrapMessage(messages[index]);
+    if (!message || message.role !== "assistant") continue;
+    if (
+      message.stopReason === "stop" ||
+      message.stopReason === "length" ||
+      message.stopReason === "toolUse" ||
+      message.stopReason === "error" ||
+      message.stopReason === "aborted"
+    ) {
+      return message;
+    }
+  }
+  return undefined;
+}
+
+function unwrapMessage(value: unknown) {
+  if (!isRecord(value)) return undefined;
+  if (value.role) return value;
+  return isRecord(value.message) ? value.message : undefined;
+}
+
+export const GOAL_COMPLETION_MARKER_MIME =
+  "application/x-pi-goal-completion-ack";
+
+function stripKnownMarkers(
+  message: Record<string, unknown>,
+  markers: Set<string>,
+) {
+  if (!Array.isArray(message.content)) return undefined;
+  let matchedTrackedMarker = false;
+  let foundPrivateMarker = false;
+  const content = message.content.filter((part) => {
+    if (
+      !isRecord(part) ||
+      part.type !== "image" ||
+      part.mimeType !== GOAL_COMPLETION_MARKER_MIME
+    ) {
+      return true;
+    }
+    foundPrivateMarker = true;
+    if (typeof part.data === "string" && markers.delete(part.data)) {
+      matchedTrackedMarker = true;
+    }
+    return false;
+  });
+  return foundPrivateMarker
+    ? { message: { ...message, content }, matchedTrackedMarker }
+    : undefined;
+}
+
+function safeTokenAdd(left: number, right: number) {
+  const result = left + right;
+  return Number.isSafeInteger(result) && result >= 0 ? result : left;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}