@tt-a1i/openpi 0.1.0

package/extensions/workflows/artifacts.ts

@@ -0,0 +1,225 @@
+import type { TranscriptEntry, WorkflowDetails } from "./model.ts";
+import {
+  boundedJournal,
+  parseJournal,
+  type JournalEntry,
+  type WorkflowJournal,
+} from "./journal.ts";
+import {
+  safeStringify,
+  truncateUtf8,
+  writeFileAtomic,
+} from "./serialization.ts";
+import * as fs from "node:fs";
+import * as path from "node:path";
+
+export const JOURNAL_FILE = "journal.json";
+
+const ARTIFACT_TRANSCRIPT_MAX_BYTES = 32 * 1024;
+const ARTIFACT_TRANSCRIPT_ENTRY_MAX_BYTES = 8 * 1024;
+export const WORKFLOW_CHECKPOINT_INTERVAL_MS = 500;
+const ENTRY_TRUNCATION_MARKER = "\n[entry truncated]";
+const TRANSCRIPT_TRUNCATION_MARKER =
+  "[artifact transcript truncated: older entries omitted]";
+
+function textBytes(text: string) {
+  return Buffer.byteLength(text, "utf8");
+}
+
+function boundEntry(entry: TranscriptEntry, maxBytes: number) {
+  if (textBytes(entry.text) <= maxBytes) return { ...entry };
+  const markerBytes = textBytes(ENTRY_TRUNCATION_MARKER);
+  const text =
+    maxBytes > markerBytes
+      ? `${truncateUtf8(entry.text, maxBytes - markerBytes)}${ENTRY_TRUNCATION_MARKER}`
+      : truncateUtf8(ENTRY_TRUNCATION_MARKER, maxBytes);
+  return { ...entry, text };
+}
+
+/** Keep the initial prompt plus the newest useful context within the artifact cap. */
+export function boundedArtifactTranscript(
+  transcript: TranscriptEntry[],
+  options: { maxBytes?: number; entryMaxBytes?: number } = {},
+) {
+  if (transcript.length === 0) return [];
+  const maxBytes = Math.max(
+    256,
+    options.maxBytes ?? ARTIFACT_TRANSCRIPT_MAX_BYTES,
+  );
+  const entryMaxBytes = Math.max(
+    64,
+    Math.min(
+      maxBytes,
+      options.entryMaxBytes ?? ARTIFACT_TRANSCRIPT_ENTRY_MAX_BYTES,
+    ),
+  );
+  const bounded = transcript.map((entry) => boundEntry(entry, entryMaxBytes));
+  if (
+    bounded.reduce((total, entry) => total + textBytes(entry.text), 0) <=
+    maxBytes
+  ) {
+    return bounded;
+  }
+
+  const initialIndex = transcript.findIndex((entry) => entry.role === "user");
+  const initial = boundEntry(
+    transcript[initialIndex >= 0 ? initialIndex : 0],
+    Math.min(entryMaxBytes, maxBytes - textBytes(TRANSCRIPT_TRUNCATION_MARKER)),
+  );
+  const marker: TranscriptEntry = {
+    role: "toolResult",
+    name: "transcript",
+    text: TRANSCRIPT_TRUNCATION_MARKER,
+  };
+  let remaining = maxBytes - textBytes(initial.text) - textBytes(marker.text);
+  const tail: TranscriptEntry[] = [];
+
+  for (
+    let index = transcript.length - 1;
+    index >= 0 && remaining > 0;
+    index--
+  ) {
+    if (index === initialIndex || (initialIndex < 0 && index === 0)) continue;
+    const entry = boundEntry(
+      transcript[index],
+      Math.min(entryMaxBytes, remaining),
+    );
+    tail.push(entry);
+    remaining -= textBytes(entry.text);
+  }
+
+  tail.reverse();
+  return [initial, marker, ...tail];
+}
+
+function writeRunFile(runDir: string, name: string, content: string) {
+  writeFileAtomic(path.join(runDir, name), content);
+}
+
+export function persistWorkflowJson(
+  runDir: string,
+  details: WorkflowDetails,
+  journal?: readonly JournalEntry[],
+) {
+  const transcripts = Object.fromEntries(
+    details.agents.map((agent) => [
+      agent.index,
+      boundedArtifactTranscript(agent.transcript),
+    ]),
+  );
+  writeRunFile(
+    runDir,
+    "transcripts.json",
+    safeStringify(transcripts, { maxBytes: 2 * 1024 * 1024 }),
+  );
+  // Written alongside the rest so it inherits atomic write, 500ms coalescing,
+  // and the final flush. Only present once a call has actually succeeded.
+  // boundedJournal has already brought this under the cap, and plain
+  // JSON.stringify is deliberate: safeStringify would swap an over-cap value
+  // for a preview stub, silently turning the journal into something unusable.
+  if (journal && journal.length > 0) {
+    writeRunFile(
+      runDir,
+      JOURNAL_FILE,
+      JSON.stringify(boundedJournal(journal).journal, null, 2),
+    );
+  }
+  if (details.result !== undefined) {
+    writeRunFile(
+      runDir,
+      "result.json",
+      safeStringify(details.result, { maxBytes: 1024 * 1024 }),
+    );
+  }
+  const compact: WorkflowDetails = {
+    ...details,
+    ...(details.result !== undefined
+      ? { result: "[stored in result.json]", resultArtifact: "result.json" }
+      : {}),
+    transcriptArtifact: "transcripts.json",
+    agents: details.agents.map((agent) => ({ ...agent, transcript: [] })),
+  };
+  writeRunFile(
+    runDir,
+    "workflow.json",
+    safeStringify(compact, { maxBytes: 1024 * 1024 }),
+  );
+}
+
+/** Coalesce live checkpoints while keeping final persistence synchronous. */
+export function createWorkflowPersistence(
+  runDir: string,
+  details: WorkflowDetails,
+  options: {
+    intervalMs?: number;
+    persist?: (
+      runDir: string,
+      details: WorkflowDetails,
+      journal?: readonly JournalEntry[],
+    ) => void;
+    /** Read at write time so callers only have to append to their own array. */
+    journal?: () => readonly JournalEntry[];
+  } = {},
+) {
+  const intervalMs = Math.max(
+    0,
+    options.intervalMs ?? WORKFLOW_CHECKPOINT_INTERVAL_MS,
+  );
+  const persist = options.persist ?? persistWorkflowJson;
+  const readJournal = options.journal;
+  let lastPersistedAt = Date.now();
+  let dirty = false;
+  let timer: ReturnType<typeof setTimeout> | undefined;
+
+  const savePending = () => {
+    timer = undefined;
+    if (!dirty) return;
+    try {
+      persist(runDir, details, readJournal?.());
+      dirty = false;
+      lastPersistedAt = Date.now();
+    } catch {
+      // Final flush retries and reports persistence failures synchronously.
+    }
+  };
+
+  return {
+    checkpoint(options: { immediate?: boolean } = {}) {
+      dirty = true;
+      if (options.immediate) {
+        if (timer) clearTimeout(timer);
+        timer = undefined;
+        savePending();
+        return;
+      }
+      if (timer) return;
+      const delay = Math.max(0, intervalMs - (Date.now() - lastPersistedAt));
+      if (delay === 0) {
+        savePending();
+        return;
+      }
+      timer = setTimeout(savePending, delay);
+    },
+    flush() {
+      if (timer) clearTimeout(timer);
+      timer = undefined;
+      persist(runDir, details, readJournal?.());
+      dirty = false;
+      lastPersistedAt = Date.now();
+    },
+  };
+}
+
+/**
+ * Read a prior run's replay journal. Any failure (missing run, unreadable or
+ * malformed file) yields undefined: resume is an optimization and must never
+ * turn into a new way for a run to fail.
+ */
+export function loadJournal(runDir: string): WorkflowJournal | undefined {
+  try {
+    const raw = fs.readFileSync(path.join(runDir, JOURNAL_FILE), "utf8");
+    return parseJournal(JSON.parse(raw));
+  } catch {
+    return undefined;
+  }
+}

package/extensions/workflows/controller.ts

@@ -0,0 +1,210 @@
+import {
+  DEFAULT_WORKFLOW_CONCURRENCY,
+  DEFAULT_WORKFLOW_MAX_AGENT_CALLS,
+  MAX_WORKFLOW_AGENT_CALLS,
+  MAX_WORKFLOW_CONCURRENCY,
+} from "../shared/setup-config.ts";
+export const RUN_SHUTDOWN_TIMEOUT_MS = 8_000;
+
+function abortError(signal: AbortSignal) {
+  return signal.reason instanceof Error
+    ? signal.reason
+    : new Error("Workflow was aborted");
+}
+
+class Semaphore {
+  private active = 0;
+  private readonly limit: number;
+  private queue: Array<{
+    resolve: () => void;
+    reject: (error: Error) => void;
+    signal: AbortSignal;
+    onAbort: () => void;
+  }> = [];
+
+  constructor(limit: number) {
+    this.limit = limit;
+  }
+
+  acquire(signal: AbortSignal) {
+    if (signal.aborted) return Promise.reject(abortError(signal));
+    if (this.active < this.limit) {
+      this.active++;
+      return Promise.resolve();
+    }
+    return new Promise<void>((resolve, reject) => {
+      const waiter = {
+        resolve: () => {
+          signal.removeEventListener("abort", onAbort);
+          this.active++;
+          resolve();
+        },
+        reject,
+        signal,
+        onAbort: () => {},
+      };
+      const onAbort = () => {
+        const index = this.queue.indexOf(waiter);
+        if (index >= 0) this.queue.splice(index, 1);
+        reject(abortError(signal));
+      };
+      waiter.onAbort = onAbort;
+      this.queue.push(waiter);
+      signal.addEventListener("abort", onAbort, { once: true });
+    });
+  }
+
+  release() {
+    this.active = Math.max(0, this.active - 1);
+    while (this.queue.length > 0) {
+      const waiter = this.queue.shift()!;
+      if (waiter.signal.aborted) {
+        waiter.signal.removeEventListener("abort", waiter.onAbort);
+        waiter.reject(abortError(waiter.signal));
+        continue;
+      }
+      waiter.resolve();
+      return;
+    }
+  }
+
+  clear() {
+    const queued = this.queue;
+    this.queue = [];
+    for (const waiter of queued) {
+      waiter.signal.removeEventListener("abort", waiter.onAbort);
+      waiter.reject(abortError(waiter.signal));
+    }
+  }
+}
+
+/** Owns every agent task and the run-wide fanout/abort budget. */
+export class RunController {
+  private readonly abortController = new AbortController();
+  private readonly semaphore: Semaphore;
+  private readonly maxAgentCalls: number;
+  private readonly tasks = new Set<Promise<unknown>>();
+  private callCount = 0;
+  private sealed = false;
+  private settlePromise?: Promise<boolean>;
+  private parentAbort?: () => void;
+  private parentSignal?: AbortSignal;
+
+  constructor(
+    parentSignal?: AbortSignal,
+    concurrency = DEFAULT_WORKFLOW_CONCURRENCY,
+    maxAgentCalls = DEFAULT_WORKFLOW_MAX_AGENT_CALLS,
+  ) {
+    this.semaphore = new Semaphore(
+      Math.max(1, Math.min(MAX_WORKFLOW_CONCURRENCY, Math.floor(concurrency))),
+    );
+    this.maxAgentCalls = Math.max(
+      1,
+      Math.min(MAX_WORKFLOW_AGENT_CALLS, Math.floor(maxAgentCalls)),
+    );
+    if (parentSignal) {
+      this.parentSignal = parentSignal;
+      this.parentAbort = () => this.abort("Parent operation was aborted");
+      if (parentSignal.aborted) this.parentAbort();
+      else
+        parentSignal.addEventListener("abort", this.parentAbort, {
+          once: true,
+        });
+    }
+  }
+
+  get signal() {
+    return this.abortController.signal;
+  }
+
+  get calls() {
+    return this.callCount;
+  }
+
+  schedule<T>(
+    task: (signal: AbortSignal) => Promise<T>,
+    invocationSignal?: AbortSignal,
+  ): Promise<T> {
+    if (this.sealed) return Promise.reject(new Error("Workflow is settling"));
+    if (this.signal.aborted) return Promise.reject(abortError(this.signal));
+    if (this.callCount >= this.maxAgentCalls) {
+      return Promise.reject(
+        new Error(
+          `Workflow exceeded the limit of ${this.maxAgentCalls} agent calls`,
+        ),
+      );
+    }
+    this.callCount++;
+
+    const running = (async () => {
+      const taskAbort = new AbortController();
+      const onRunAbort = () => taskAbort.abort(this.signal.reason);
+      const onInvocationAbort = () => taskAbort.abort(invocationSignal?.reason);
+      this.signal.addEventListener("abort", onRunAbort, { once: true });
+      invocationSignal?.addEventListener("abort", onInvocationAbort, {
+        once: true,
+      });
+      if (this.signal.aborted) onRunAbort();
+      else if (invocationSignal?.aborted) onInvocationAbort();
+
+      let acquired = false;
+      try {
+        await this.semaphore.acquire(taskAbort.signal);
+        acquired = true;
+        if (taskAbort.signal.aborted) throw abortError(taskAbort.signal);
+        const result = await task(taskAbort.signal);
+        if (invocationSignal?.aborted) throw abortError(invocationSignal);
+        return result;
+      } finally {
+        this.signal.removeEventListener("abort", onRunAbort);
+        invocationSignal?.removeEventListener("abort", onInvocationAbort);
+        if (acquired) this.semaphore.release();
+      }
+    })();
+    this.tasks.add(running);
+    void running.finally(() => this.tasks.delete(running)).catch(() => {});
+    return running;
+  }
+
+  abort(reason = "Workflow was aborted") {
+    if (!this.signal.aborted) this.abortController.abort(new Error(reason));
+    this.semaphore.clear();
+  }
+
+  /** Seal once and wait a bounded time for the same task snapshot on every caller. */
+  settle(options: { abort?: boolean; timeoutMs?: number } = {}) {
+    this.sealed = true;
+    if (options.abort) this.abort();
+    if (this.settlePromise) return this.settlePromise;
+
+    const tasks = [...this.tasks];
+    this.settlePromise = (async () => {
+      if (tasks.length === 0) {
+        this.detachParent();
+        return true;
+      }
+
+      let timer: ReturnType<typeof setTimeout> | undefined;
+      const timeout = new Promise<false>((resolve) => {
+        timer = setTimeout(
+          () => resolve(false),
+          options.timeoutMs ?? RUN_SHUTDOWN_TIMEOUT_MS,
+        );
+      });
+      const settled = Promise.allSettled(tasks).then(() => true as const);
+      const completed = await Promise.race([settled, timeout]);
+      if (timer) clearTimeout(timer);
+      this.detachParent();
+      return completed;
+    })();
+    return this.settlePromise;
+  }
+
+  private detachParent() {
+    if (this.parentAbort) {
+      this.parentSignal?.removeEventListener("abort", this.parentAbort);
+    }
+    this.parentAbort = undefined;
+    this.parentSignal = undefined;
+  }
+}