@awak-app/simy-cli 0.2.2 → 0.3.3

package/src/bounded-local-task-subtask-pool.js

@@ -0,0 +1,708 @@
+import { createHash } from "node:crypto";
+import {
+  chmod,
+  lstat,
+  mkdir,
+  readFile,
+  realpath,
+  writeFile,
+} from "node:fs/promises";
+import path from "node:path";
+
+import { prepareLocalTaskWorkspace } from "./local-task.js";
+import {
+  BoundedLocalTaskSubtaskContractError,
+  boundedLocalTaskSubtaskExecution,
+  buildBoundedLocalTaskSubtaskBatchPlan,
+  DEFAULT_LOCAL_TASK_SUBTASK_PARALLELISM,
+  durableLocalTaskSubtaskFence,
+  normalizeBoundedLocalTaskParentAuthority,
+  normalizeDurableLocalTaskSubtaskClaim,
+  normalizeDurableLocalTaskSubtaskResult,
+} from "./bounded-local-task-subtasks-contract.js";
+
+const DEFAULT_LEASE_SECONDS = 120;
+const DEFAULT_HEARTBEAT_INTERVAL_MS = 30_000;
+const TERMINAL_BATCH_STATES = new Set([
+  "waiting_human",
+  "succeeded",
+  "failed",
+  "cancelled",
+]);
+const SUPPORTED_BATCH_STATES = new Set([
+  "active",
+  "running",
+  "cancel_requested",
+  ...TERMINAL_BATCH_STATES,
+]);
+const TERMINAL_POOL_DISPOSITIONS = new Set([
+  "none",
+  "provider_result_unknown",
+  ...TERMINAL_BATCH_STATES,
+]);
+
+/**
+ * Run only Backend-admitted v2 subtasks. The fixed Web route adapter owns
+ * authentication and injects p_device_token_hash; the pool deliberately does
+ * not know the paired device token.
+ */
+export function createBoundedLocalTaskSubtaskPool({
+  api,
+  parentAuthority,
+  batchPlan,
+  storageRoot,
+  workspaceRoot,
+  runSubtask,
+  stopSubtask = async () => {},
+  downloadSource,
+  prepareWorkspace = prepareLocalTaskWorkspace,
+  leaseSeconds = DEFAULT_LEASE_SECONDS,
+  heartbeatIntervalMs = DEFAULT_HEARTBEAT_INTERVAL_MS,
+  setIntervalImpl = setInterval,
+  clearIntervalImpl = clearInterval,
+  onError = () => {},
+} = {}) {
+  const parent = normalizeBoundedLocalTaskParentAuthority(parentAuthority);
+  const plan = buildBoundedLocalTaskSubtaskBatchPlan(batchPlan);
+  requireApi(api);
+  if (typeof runSubtask !== "function") {
+    throw new TypeError("runSubtask is required");
+  }
+  const sourceDownloader =
+    downloadSource ??
+    api.download_local_task_subtask_source_v2?.bind(api);
+  if (typeof sourceDownloader !== "function") {
+    throw new TypeError("downloadSource is required");
+  }
+  if (typeof prepareWorkspace !== "function") {
+    throw new TypeError("prepareWorkspace is required");
+  }
+
+  const controller = new AbortController();
+  const active = new Map();
+  let initialized = false;
+  let closed = false;
+  let terminalDisposition = null;
+  let runPromise = null;
+  let claimCount = 0;
+  let maximumObservedConcurrency = 0;
+  let batchId = null;
+  let batchParallelism = null;
+  const schedulingStopped = () =>
+    TERMINAL_POOL_DISPOSITIONS.has(terminalDisposition);
+
+  const cancel = async (reason = "cancelled") => {
+    if (!controller.signal.aborted) controller.abort(reason);
+    terminalDisposition = "cancelled";
+    await Promise.allSettled(
+      [...active.values()].map(({ normalized }) =>
+        stopSubtask({
+          stepId: normalized.step.step_id,
+          attemptId: normalized.attempt.attempt_id,
+          reason,
+        })),
+    );
+  };
+
+  const runOne = async (claim) => {
+    let normalized;
+    let heartbeatTimer = null;
+    let heartbeatInFlight = null;
+    let completionConfirmed = false;
+    let completionState = null;
+    let providerStarted = false;
+    const startedAt = Date.now();
+    try {
+      normalized = normalizeDurableLocalTaskSubtaskClaim(claim, {
+        parentAuthority: parent,
+        leaseSeconds,
+      });
+      if (normalized.disposition !== "execute") {
+        throw new BoundedLocalTaskSubtaskContractError(
+          "Only an executable server claim can enter the bounded pool.",
+        );
+      }
+      if (batchId !== null && normalized.batch.batch_id !== batchId) {
+        throw new BoundedLocalTaskSubtaskContractError(
+          "The bounded pool received a claim from a different batch.",
+        );
+      }
+      batchId ??= normalized.batch.batch_id;
+
+      const execution = boundedLocalTaskSubtaskExecution(normalized, parent);
+      const runtime = await materializeSubtaskRuntime({
+        normalized,
+        execution,
+        parent,
+        storageRoot,
+        workspaceRoot,
+        prepareWorkspace,
+        downloadSource: sourceDownloader,
+        signal: controller.signal,
+      });
+      if (controller.signal.aborted) return;
+
+      const fence = durableLocalTaskSubtaskFence(normalized, parent);
+      let boundaryConfirmed = false;
+      try {
+        boundaryConfirmed =
+          await api.mark_local_task_subtask_invocation_started_v2(fence);
+      } catch (error) {
+        onError(
+          poolError(
+            "SIMY could not confirm the bounded subtask Provider boundary. The Provider was not started.",
+            "local_task_subtask_boundary_unconfirmed",
+            error,
+          ),
+        );
+        await api.abandon_local_task_subtask_attempt_v2({
+          ...fence,
+          p_reason_code: "invocation_start_unconfirmed",
+        }).catch(onError);
+        return;
+      }
+      if (boundaryConfirmed !== true || controller.signal.aborted) {
+        await api.abandon_local_task_subtask_attempt_v2({
+          ...fence,
+          p_reason_code: controller.signal.aborted
+            ? "cancelled_before_invocation"
+            : "invocation_start_not_confirmed",
+        }).catch(onError);
+        return;
+      }
+      providerStarted = true;
+
+      const heartbeat = async () => {
+        if (controller.signal.aborted || completionConfirmed) return;
+        try {
+          const response =
+            await api.heartbeat_local_task_subtask_attempt_v2({
+              ...fence,
+              p_lease_seconds: leaseSeconds,
+            });
+          if (
+            completionConfirmed &&
+            isTerminalHeartbeatCoveredByCompletion(
+              response,
+              normalized,
+              completionState,
+            )
+          ) {
+            return;
+          }
+          if (
+            response?.ok !== true ||
+            response.cancel_requested === true ||
+            (response.cancel_generation !== undefined &&
+              String(response.cancel_generation) !==
+                normalized.parent.cancel_generation)
+          ) {
+            await cancel("server_cancelled");
+          }
+        } catch (error) {
+          if (completionConfirmed && isKnownSubtaskFenceError(error)) return;
+          onError(
+            poolError(
+              "SIMY could not renew a bounded subtask attempt.",
+              "local_task_subtask_heartbeat_failed",
+              error,
+            ),
+          );
+          if (isHeartbeatLeaseLoss(error)) {
+            await cancel("lease_lost");
+          }
+        }
+      };
+      heartbeatTimer = setIntervalImpl(() => {
+        if (heartbeatInFlight) return;
+        heartbeatInFlight = heartbeat().finally(() => {
+          heartbeatInFlight = null;
+        });
+      }, heartbeatIntervalMs);
+
+      let rawResult;
+      try {
+        rawResult = await runSubtask(runtime, {
+          signal: controller.signal,
+        });
+      } catch (error) {
+        rawResult = {
+          status: controller.signal.aborted ? "cancelled" : "failed",
+          result_summary: controller.signal.aborted
+            ? "The bounded subtask was cancelled."
+            : "The bounded subtask failed without a trusted result.",
+          provider_tokens_used: "0",
+          elapsed_ms: String(Math.max(0, Date.now() - startedAt)),
+          failure_code: controller.signal.aborted
+            ? null
+            : "local_task_subtask_failed",
+        };
+        onError(
+          poolError(
+            "A bounded Local Task subtask failed.",
+            "local_task_subtask_failed",
+            error,
+          ),
+        );
+      }
+      const result = normalizeDurableLocalTaskSubtaskResult(
+        normalizeRunResult(rawResult, startedAt),
+        { normalizedClaim: normalized },
+      );
+      try {
+        // A transport-ambiguous completion is never followed by another
+        // Provider call or an automatic local replay. Backend owns resolution.
+        const completionResponse =
+          await api.complete_local_task_subtask_attempt_v2({
+            ...fence,
+            p_status: result.status,
+            p_result_summary: result.result_summary,
+            p_artifact_refs: [],
+            p_provider_tokens_used: result.provider_tokens_used,
+            p_elapsed_ms: result.elapsed_ms,
+            p_failure_code: result.failure_code,
+          });
+        const completion = normalizeCompletionResponse(
+          completionResponse,
+          fence,
+          result,
+        );
+        completionConfirmed = completion.confirmed;
+        completionState = result.status;
+        if (completion.batchState !== null) {
+          terminalDisposition = completion.batchState;
+        }
+      } catch (error) {
+        terminalDisposition = "provider_result_unknown";
+        onError(
+          poolError(
+            "SIMY could not confirm the bounded subtask's final server state. The Provider will not be replayed.",
+            "local_task_subtask_final_unconfirmed",
+            error,
+          ),
+        );
+      }
+    } catch (error) {
+      onError(
+        poolError(
+          "SIMY rejected an unsafe bounded Local Task subtask.",
+          error?.code || "invalid_local_task_subtask_claim",
+          error,
+        ),
+      );
+      if (providerStarted) {
+        // Once the Provider has started, an invalid or otherwise unreportable
+        // local result is an ambiguous boundary. Stop admitting work and let
+        // Backend/user resolution decide what happens next.
+        terminalDisposition = "provider_result_unknown";
+      }
+      if (normalized?.attempt && normalized?.step && !providerStarted) {
+        const fence = durableLocalTaskSubtaskFence(normalized, parent);
+        await api.abandon_local_task_subtask_attempt_v2({
+          ...fence,
+          p_reason_code: "local_validation_failed",
+        }).catch(onError);
+      }
+    } finally {
+      if (heartbeatTimer) clearIntervalImpl(heartbeatTimer);
+      if (heartbeatInFlight) await heartbeatInFlight;
+    }
+  };
+
+  const scheduler = async () => {
+    if (!initialized) {
+      const initializedBatch =
+        await api.initialize_local_task_subtask_batch_v2({
+          p_turn_id: parent.turnId,
+          p_worker_id: parent.workerId,
+          p_parent_lease_epoch: parent.parentLeaseEpoch,
+          p_plan: plan,
+        });
+      if (
+        initializedBatch?.schema_version !==
+          "local_task_subtask_batch_initialized.v2" ||
+        initializedBatch?.disposition !== "ready"
+      ) {
+        throw poolError(
+          "SIMY did not initialize the bounded Local Task batch.",
+          "local_task_subtask_batch_not_initialized",
+        );
+      }
+      initialized = true;
+    }
+
+    while (
+      !closed &&
+      !controller.signal.aborted &&
+      !schedulingStopped()
+    ) {
+      let admittedInPass = false;
+      while (
+        !closed &&
+        !controller.signal.aborted &&
+        !schedulingStopped() &&
+        active.size <
+          (batchParallelism ?? DEFAULT_LOCAL_TASK_SUBTASK_PARALLELISM)
+      ) {
+        // No prefetch queue: this RPC occurs only when a real execution slot
+        // is available. Therefore the third subtask is not claimed while two
+        // Provider calls are still active.
+        claimCount += 1;
+        const claim =
+          await api.claim_local_task_subtask_attempt_v2({
+            p_turn_id: parent.turnId,
+            p_worker_id: parent.workerId,
+            p_parent_lease_epoch: parent.parentLeaseEpoch,
+            p_cancel_generation: parent.cancelGeneration,
+            p_lease_seconds: leaseSeconds,
+            p_step_id: null,
+          });
+        const normalized = normalizeDurableLocalTaskSubtaskClaim(claim, {
+          parentAuthority: parent,
+          leaseSeconds,
+        });
+        if (normalized.step) assertClaimMatchesPlan(normalized, plan);
+        if (normalized.disposition === "cancelled") {
+          await cancel("server_cancelled");
+          break;
+        }
+        if (normalized.disposition === "provider_result_unknown") {
+          terminalDisposition = "provider_result_unknown";
+          onError(
+            poolError(
+              "A previous bounded subtask Provider result is unknown and cannot be replayed.",
+              "local_task_subtask_provider_result_unknown",
+            ),
+          );
+          break;
+        }
+        if (normalized.disposition === "none") {
+          terminalDisposition = active.size === 0 ? "none" : null;
+          break;
+        }
+        if (
+          batchParallelism !== null &&
+          normalized.batch.parallelism !== batchParallelism
+        ) {
+          throw new BoundedLocalTaskSubtaskContractError(
+            "The bounded pool received inconsistent batch parallelism.",
+          );
+        }
+        batchParallelism ??= normalized.batch.parallelism;
+        admittedInPass = true;
+        const key = normalized.attempt.attempt_id;
+        const operation = runOne(claim).finally(() => {
+          active.delete(key);
+        });
+        active.set(key, { normalized, operation });
+        maximumObservedConcurrency = Math.max(
+          maximumObservedConcurrency,
+          active.size,
+        );
+      }
+
+      if (controller.signal.aborted || closed || schedulingStopped()) break;
+      if (active.size > 0) {
+        await Promise.race([...active.values()].map(({ operation }) => operation));
+        continue;
+      }
+      if (!admittedInPass) break;
+    }
+    await Promise.allSettled(
+      [...active.values()].map(({ operation }) => operation),
+    );
+  };
+
+  return {
+    get signal() {
+      return controller.signal;
+    },
+    snapshot() {
+      return {
+        initialized,
+        closed,
+        cancelled: controller.signal.aborted,
+        terminal_disposition: terminalDisposition,
+        active_count: active.size,
+        claim_count: claimCount,
+        max_observed_concurrency: maximumObservedConcurrency,
+        batch_id: batchId,
+      };
+    },
+    async run() {
+      if (runPromise) return runPromise;
+      runPromise = scheduler().catch((error) => {
+        onError(
+          poolError(
+            "SIMY could not run the bounded Local Task subtask pool.",
+            "bounded_local_task_subtask_pool_failed",
+            error,
+          ),
+        );
+        return false;
+      });
+      await runPromise;
+      return !controller.signal.aborted;
+    },
+    cancel,
+    async waitForIdle() {
+      await runPromise;
+      await Promise.allSettled(
+        [...active.values()].map(({ operation }) => operation),
+      );
+    },
+    async close() {
+      closed = true;
+      await cancel("worker_closed");
+      await runPromise;
+    },
+  };
+}
+
+async function materializeSubtaskRuntime({
+  normalized,
+  execution,
+  parent,
+  storageRoot,
+  workspaceRoot,
+  prepareWorkspace,
+  downloadSource,
+  signal,
+}) {
+  const prepared = await prepareWorkspace({
+    taskId: execution.workspace_id,
+    root: storageRoot,
+    workspaceRoot,
+  });
+  const canonicalInputsRoot = await realpath(prepared.inputs_root);
+  const materializedInputs = [];
+  for (const ref of normalized.step.attachment_refs) {
+    if (signal.aborted) throw abortError();
+    const bytes = await downloadSource({
+      turnId: normalized.turnId,
+      batchId: normalized.batch.batch_id,
+      stepId: normalized.step.step_id,
+      attemptId: normalized.attempt.attempt_id,
+      attemptLeaseEpoch: normalized.attempt.lease_epoch,
+      parentLeaseEpoch: normalized.parent.lease_epoch,
+      cancelGeneration: normalized.parent.cancel_generation,
+      source: { ...ref },
+      signal,
+    });
+    if (!Buffer.isBuffer(bytes)) {
+      throw poolError(
+        "SIMY Web returned an invalid uploaded source body.",
+        "local_task_subtask_source_invalid",
+      );
+    }
+    const expectedSize = BigInt(ref.size_bytes);
+    const digest = createHash("sha256").update(bytes).digest("hex");
+    if (BigInt(bytes.byteLength) !== expectedSize || digest !== ref.sha256) {
+      throw poolError(
+        "A bounded subtask source failed integrity verification.",
+        "local_task_subtask_source_integrity_failed",
+      );
+    }
+    const filename = `${ref.source_id}-${ref.name}`;
+    const target = path.join(canonicalInputsRoot, filename);
+    if (path.dirname(target) !== canonicalInputsRoot) {
+      throw poolError(
+        "A bounded subtask source path escaped its isolated input directory.",
+        "local_task_subtask_source_path_unsafe",
+      );
+    }
+    await mkdir(canonicalInputsRoot, { recursive: true, mode: 0o700 });
+    await writeFile(target, bytes, { flag: "wx", mode: 0o400 });
+    await chmod(target, 0o400);
+    const details = await lstat(target);
+    const canonicalTarget = await realpath(target);
+    if (
+      !details.isFile() ||
+      details.isSymbolicLink() ||
+      path.dirname(canonicalTarget) !== canonicalInputsRoot
+    ) {
+      throw poolError(
+        "A bounded subtask source is not a safe isolated file.",
+        "local_task_subtask_source_path_unsafe",
+      );
+    }
+    const verifiedBytes = await readFile(canonicalTarget);
+    if (
+      createHash("sha256").update(verifiedBytes).digest("hex") !== ref.sha256
+    ) {
+      throw poolError(
+        "A bounded subtask source changed after materialization.",
+        "local_task_subtask_source_integrity_failed",
+      );
+    }
+    materializedInputs.push({
+      source_id: ref.source_id,
+      name: ref.name,
+      local_path: canonicalTarget,
+      mime_type: ref.mime_type,
+      size_bytes: ref.size_bytes,
+      sha256: ref.sha256,
+      verified: true,
+      immutable: true,
+    });
+  }
+  return deepFreeze({
+    ...execution,
+    parent_turn_id: parent.turnId,
+    workspace_path: prepared.workspace_path,
+    task_root: prepared.task_root,
+    inputs_root: canonicalInputsRoot,
+    outputs_root: prepared.outputs_root,
+    outputs_root_expected: prepared.outputs_root_expected,
+    verified_uploaded_inputs: materializedInputs,
+  });
+}
+
+function normalizeRunResult(value, startedAt) {
+  const elapsed = String(Math.max(0, Date.now() - startedAt));
+  if (!value || typeof value !== "object" || Array.isArray(value)) {
+    return {
+      status: "failed",
+      result_summary: "The bounded subtask returned no trusted result.",
+      provider_tokens_used: "0",
+      elapsed_ms: elapsed,
+      failure_code: "local_task_subtask_result_invalid",
+    };
+  }
+  return {
+    status: value.status,
+    result_summary: value.result_summary ?? null,
+    provider_tokens_used: value.provider_tokens_used ?? "0",
+    elapsed_ms: value.elapsed_ms ?? elapsed,
+    failure_code: value.failure_code ?? null,
+  };
+}
+
+function normalizeCompletionResponse(value, fence, result) {
+  // Legacy injected adapters returned `true`; it confirms the transition but
+  // cannot authoritatively stop scheduling because it carries no batch state.
+  if (value === true) {
+    return { confirmed: true, batchState: null };
+  }
+  if (
+    !value ||
+    typeof value !== "object" ||
+    Array.isArray(value) ||
+    value.schema_version !== "local_task_subtask_transition.v2" ||
+    value.batch_id !== fence.p_batch_id ||
+    value.step_id !== fence.p_step_id ||
+    value.attempt_id !== fence.p_attempt_id ||
+    value.state !== result.status ||
+    typeof value.replayed !== "boolean" ||
+    !SUPPORTED_BATCH_STATES.has(value.batch_state)
+  ) {
+    throw poolError(
+      "SIMY Web returned an invalid bounded subtask completion transition.",
+      "local_task_subtask_completion_invalid_response",
+    );
+  }
+  return {
+    confirmed: true,
+    batchState: TERMINAL_BATCH_STATES.has(value.batch_state)
+      ? value.batch_state
+      : null,
+  };
+}
+
+function isKnownSubtaskFenceError(error) {
+  return (
+    error?.status === 409 &&
+    (
+      error?.code === "local_task_lease_fenced" ||
+      error?.code === "local_task_subtask_attempt_fenced"
+    )
+  );
+}
+
+function isHeartbeatLeaseLoss(error) {
+  return error?.status === 409;
+}
+
+function isTerminalHeartbeatCoveredByCompletion(
+  value,
+  normalized,
+  completionState,
+) {
+  return (
+    value?.schema_version === "local_task_subtask_heartbeat.v2" &&
+    value?.terminal === true &&
+    value?.ok === true &&
+    value?.attempt_state === completionState &&
+    SUPPORTED_BATCH_STATES.has(value?.batch_state) &&
+    value?.cancel_requested === false &&
+    String(value?.cancel_generation) ===
+      normalized.parent.cancel_generation &&
+    String(value?.attempt_lease_epoch) ===
+      normalized.attempt.lease_epoch
+  );
+}
+
+function requireApi(api) {
+  const names = [
+    "initialize_local_task_subtask_batch_v2",
+    "claim_local_task_subtask_attempt_v2",
+    "heartbeat_local_task_subtask_attempt_v2",
+    "mark_local_task_subtask_invocation_started_v2",
+    "complete_local_task_subtask_attempt_v2",
+    "abandon_local_task_subtask_attempt_v2",
+  ];
+  if (!api || names.some((name) => typeof api[name] !== "function")) {
+    throw new TypeError(
+      `api must implement ${names.join(", ")}`,
+    );
+  }
+}
+
+function assertClaimMatchesPlan(normalized, plan) {
+  const candidate = plan.candidates.find(
+    (item) => item.candidate_key === normalized.step.stable_key,
+  );
+  const expectedSerial = plan.mode === "serial_uploaded";
+  if (
+    !candidate ||
+    normalized.batch.mode !== plan.mode ||
+    normalized.batch.decision_hash !== plan.decision_hash ||
+    normalized.batch.reason_codes.length !== plan.reason_codes.length ||
+    normalized.batch.reason_codes.some(
+      (code, index) => code !== plan.reason_codes[index],
+    ) ||
+    normalized.batch.fanout !== plan.candidates.length ||
+    normalized.batch.parallelism > (expectedSerial ? 1 : 2) ||
+    candidate.title !== normalized.step.title ||
+    candidate.instruction_fragment !== normalized.step.instruction ||
+    candidate.independence_reason !== normalized.step.independence_reason ||
+    candidate.source_ids.length !== normalized.step.source_ids.length ||
+    candidate.source_ids.some(
+      (sourceId, index) => sourceId !== normalized.step.source_ids[index],
+    )
+  ) {
+    throw new BoundedLocalTaskSubtaskContractError(
+      "The bounded subtask claim does not match its stored batch plan.",
+    );
+  }
+}
+
+function deepFreeze(value) {
+  if (!value || typeof value !== "object" || Object.isFrozen(value)) return value;
+  Object.freeze(value);
+  for (const child of Object.values(value)) deepFreeze(child);
+  return value;
+}
+
+function poolError(message, code, cause = null) {
+  const error = new Error(message, cause ? { cause } : undefined);
+  error.code = code;
+  return error;
+}
+
+function abortError() {
+  return poolError(
+    "The bounded Local Task subtask pool was cancelled.",
+    "local_task_subtask_cancelled",
+  );
+}