@tt-a1i/openpi 0.3.1 → 0.5.0

package/extensions/subagents/src/id-sequence.ts

@@ -0,0 +1,84 @@
+import type { SessionEntry } from "@earendil-works/pi-coding-agent";
+
+export const SUBAGENT_ID_WATERMARK_ENTRY_TYPE = "subagent-id-watermark";
+
+export interface SubagentIdCounters {
+  readonly modelCounter: number;
+  readonly btwCounter: number;
+}
+
+interface IdWatermarkData {
+  readonly version: 1;
+  readonly id: string;
+}
+
+function record(value: unknown): Record<string, unknown> | undefined {
+  return typeof value === "object" && value !== null
+    ? (value as Record<string, unknown>)
+    : undefined;
+}
+
+function idFrom(value: unknown) {
+  const candidate = record(value)?.id;
+  return typeof candidate === "string" ? candidate : undefined;
+}
+
+function idFromEntry(entry: SessionEntry) {
+  if (entry.type === "custom") {
+    const data = record(entry.data);
+    if (
+      entry.customType === SUBAGENT_ID_WATERMARK_ENTRY_TYPE &&
+      data?.version === 1
+    ) {
+      return idFrom(data);
+    }
+    if (
+      entry.customType === "subagent-finished" ||
+      entry.customType === "btw-result"
+    ) {
+      return idFrom(data);
+    }
+    if (entry.customType === "subagent-result") {
+      return idFrom(data?.details) ?? idFrom(data);
+    }
+    return undefined;
+  }
+
+  if (entry.type === "custom_message") {
+    return entry.customType === "subagent-result"
+      ? idFrom(entry.details)
+      : undefined;
+  }
+
+  if (entry.type !== "message") return undefined;
+  const message = record(entry.message);
+  return message?.role === "toolResult" && message.toolName === "subagent_spawn"
+    ? idFrom(message.details)
+    : undefined;
+}
+
+function sequence(id: string | undefined, prefix: "sa" | "btw") {
+  if (!id) return undefined;
+  const match = new RegExp(`^${prefix}-(\\d+)$`, "u").exec(id);
+  if (!match) return undefined;
+  const value = Number(match[1]);
+  return Number.isSafeInteger(value) && value > 0 ? value : undefined;
+}
+
+/** Restore branch-local id high-water marks, including sessions from before watermarks existed. */
+export function restoreSubagentIdCounters(
+  entries: readonly SessionEntry[],
+): SubagentIdCounters {
+  let modelCounter = 0;
+  let btwCounter = 0;
+  for (const entry of entries) {
+    const id = idFromEntry(entry);
+    modelCounter = Math.max(modelCounter, sequence(id, "sa") ?? 0);
+    btwCounter = Math.max(btwCounter, sequence(id, "btw") ?? 0);
+  }
+  return { modelCounter, btwCounter };
+}
+
+export function subagentIdWatermark(id: string): IdWatermarkData {
+  return { version: 1, id };
+}