src/agents/subagents/spawn/subagent-spawn-context.ts

import { finiteSecondsToTimerSafeMilliseconds } from "@openclaw/normalization-core/number-coercion";
import { sanitizeForLog } from "../../../../packages/terminal-core/src/ansi.js";
import { resolveThreadBindingSpawnPolicy } from "../../../channels/thread-bindings-policy.js";
import type { SessionEntry } from "../../../config/sessions/types.js";
import type { OpenClawConfig } from "../../../config/types.openclaw.js";
import type { ContextEngine, SubagentSpawnPreparation } from "../../../context-engine/types.js";
import { summarizeSpawnError } from "../../spawn-pipeline.js";
import { getSubagentSpawnDeps } from "./subagent-spawn-deps.js";
import { resolveGatewaySessionStoreTarget } from "./subagent-spawn.runtime.js";
import type { SpawnSubagentContextMode } from "./subagent-spawn.types.js";

type PreparedSpawnContext =
  | {
      status: "ok";
      mode: "isolated";
      parentEntry?: SessionEntry;
      childEntry?: SessionEntry;
      forkFallbackNote?: string;
    }
  | {
      status: "ok";
      mode: "fork";
      parentEntry: SessionEntry;
      childEntry: SessionEntry;
      forked: { sessionId: string; sessionFile: string };
      forkFallbackNote?: never;
    }
  | { status: "error"; error: string };

export async function prepareSubagentSessionContext(params: {
  assertActive?: () => void;
  cfg: OpenClawConfig;
  contextMode: SpawnSubagentContextMode;
  requesterAgentId: string;
  targetAgentId: string;
  requesterInternalKey: string;
  childSessionKey: string;
}): Promise<PreparedSpawnContext> {
  if (params.contextMode === "isolated") {
    return { status: "ok", mode: "isolated" };
  }
  const childTarget = resolveGatewaySessionStoreTarget({
    cfg: params.cfg,
    key: params.childSessionKey,
    agentId: params.targetAgentId,
  });
  const parentTarget = resolveGatewaySessionStoreTarget({
    cfg: params.cfg,
    key: params.requesterInternalKey,
    agentId: params.requesterAgentId,
  });

  try {
    if (params.targetAgentId !== params.requesterAgentId) {
      throw new Error(
        'context="fork" currently requires the same target agent as the requester; use context="isolated" for cross-agent spawns.',
      );
    }

    const forkedResult = await getSubagentSpawnDeps().forkSessionEntryFromParent({
      commitGuard: params.assertActive,
      storePath: childTarget.storePath,
      parentSessionKey: parentTarget.canonicalKey,
      parentStoreKeys: parentTarget.storeKeys,
      sessionKey: childTarget.canonicalKey,
      sessionStoreKeys: childTarget.storeKeys,
      fallbackEntry: { sessionId: "", updatedAt: Date.now() },
      agentId: params.requesterAgentId,
    });
    if (forkedResult.status === "missing-parent") {
      throw new Error(
        'context="fork" requested but the requester session transcript is not available.',
      );
    }
    if (forkedResult.status === "failed" || forkedResult.status === "missing-entry") {
      throw new Error(
        'context="fork" requested but OpenClaw could not fork the requester transcript.',
      );
    }
    if (forkedResult.status === "skipped") {
      const forkFallbackNote =
        forkedResult.decision?.status === "skip" ? forkedResult.decision.message : undefined;
      if (!forkFallbackNote) {
        throw new Error('context="fork" requested but OpenClaw could not prepare forked context.');
      }
      return {
        status: "ok",
        mode: "isolated",
        parentEntry: forkedResult.parentEntry,
        childEntry: forkedResult.sessionEntry,
        forkFallbackNote,
      };
    }
    return {
      status: "ok",
      mode: "fork",
      parentEntry: forkedResult.parentEntry,
      childEntry: forkedResult.sessionEntry,
      forked: forkedResult.fork,
    };
  } catch (err) {
    return { status: "error", error: summarizeSpawnError(err) };
  }
}

export type PreparedContextEngineSubagentSpawn = SubagentSpawnPreparation & {
  dispose(): Promise<void>;
};

export async function prepareContextEngineSubagentSpawn(params: {
  assertActive?: () => void;
  cfg: OpenClawConfig;
  context: PreparedSpawnContext & { status: "ok" };
  requesterInternalKey: string;
  childSessionKey: string;
  runTimeoutSeconds: number;
}): Promise<
  | { status: "ok"; preparation: PreparedContextEngineSubagentSpawn }
  | { status: "error"; error: string }
> {
  let engine: ContextEngine | undefined;
  let disposal: Promise<void> | undefined;
  const dispose = () =>
    (disposal ??= (async () => {
      try {
        await engine?.dispose?.();
      } catch (error) {
        console.warn(
          `[context-engine] Failed subagent preparation cleanup: ${sanitizeForLog(String(error))}`,
        );
        throw error;
      }
    })());
  try {
    const deps = getSubagentSpawnDeps();
    deps.ensureContextEnginesInitialized();
    engine = await deps.resolveContextEngine(params.cfg);
    // Resolution may outlive the caller. Returned preparation must still reach
    // the pipeline rollback owner before its next authority check.
    params.assertActive?.();
    const preparation = await engine.prepareSubagentSpawn?.({
      parentSessionKey: params.requesterInternalKey,
      childSessionKey: params.childSessionKey,
      contextMode: params.context.mode,
      parentSessionId: params.context.parentEntry?.sessionId,
      parentSessionFile: params.requesterInternalKey,
      childSessionId: params.context.childEntry?.sessionId,
      childSessionFile:
        params.context.mode === "fork" ? params.context.forked.sessionFile : params.childSessionKey,
      ttlMs: finiteSecondsToTimerSafeMilliseconds(params.runTimeoutSeconds, {
        floorSeconds: true,
      }),
    });
    let rollback: Promise<void> | undefined;
    return {
      status: "ok",
      preparation: {
        rollback: () =>
          (rollback ??= (async () => {
            try {
              await preparation?.rollback();
            } finally {
              await dispose();
            }
          })()),
        async dispose() {
          // Cancellation may already be rolling back while its caller unwinds.
          await rollback?.catch(() => {});
          await dispose();
        },
      },
    };
  } catch (err) {
    await dispose().catch(() => {});
    return {
      status: "error",
      error: `Context engine subagent preparation failed: ${summarizeSpawnError(err)}`,
    };
  }
}

export async function rollbackPreparedContextEngine(
  preparation?: SubagentSpawnPreparation,
): Promise<boolean> {
  try {
    await preparation?.rollback();
    return true;
  } catch {
    // Best-effort cleanup only.
    return false;
  }
}

export function resolveSubagentContextMode(params: {
  requestedContext?: SpawnSubagentContextMode;
  threadRequested: boolean;
  cfg: OpenClawConfig;
  requester: {
    channel?: string;
    accountId?: string;
  };
}): SpawnSubagentContextMode {
  if (params.requestedContext === "fork" || params.requestedContext === "isolated") {
    return params.requestedContext;
  }
  if (!params.threadRequested || !params.requester.channel) {
    return "isolated";
  }
  return resolveThreadBindingSpawnPolicy({
    cfg: params.cfg,
    channel: params.requester.channel,
    accountId: params.requester.accountId,
    kind: "subagent",
  }).defaultSpawnContext;
}