ai_ops/src/agents/provisioning.ts

import { execFile } from "node:child_process";
import { createHash } from "node:crypto";
import { mkdir, open, stat, unlink, writeFile } from "node:fs/promises";
import { dirname, isAbsolute, resolve } from "node:path";
import { promisify } from "node:util";

const execFileAsync = promisify(execFile);

type JsonPrimitive = string | number | boolean | null;
type JsonValue = JsonPrimitive | JsonValue[] | { [key: string]: JsonValue };

export type ResourceRequest = {
  kind: string;
  options?: Record<string, unknown>;
};

export type DiscoverySnapshot = {
  sessionId: string;
  workspaceRoot: string;
  workingDirectory: string;
  hardConstraints: Array<{
    kind: string;
    allocation: Record<string, JsonValue>;
  }>;
  softConstraints: {
    env: Record<string, string>;
    promptSections: string[];
    metadata: Record<string, JsonValue>;
  };
};

export type ChildResourceSuballocationInput = {
  parentSnapshot: DiscoverySnapshot;
  childSessionId: string;
  childIndex: number;
  childCount: number;
};

type ResourceContextPatch = {
  env?: Record<string, string>;
  promptSections?: string[];
  metadata?: Record<string, JsonValue>;
  preferredWorkingDirectory?: string;
};

type ResourceLease = {
  kind: string;
  hard: Record<string, JsonValue>;
  soft?: ResourceContextPatch;
  release: () => Promise<void>;
};

export type ResourceProvider<Options extends Record<string, unknown> = Record<string, unknown>> = {
  kind: string;
  provision: (input: {
    sessionId: string;
    workspaceRoot: string;
    options: Options;
  }) => Promise<ResourceLease>;
};

type RuntimeInjection = {
  workingDirectory: string;
  env: Record<string, string>;
  discoveryFilePath: string;
};

type ProvisionedResourcesState = {
  sessionId: string;
  workspaceRoot: string;
  workingDirectory: string;
  hardConstraints: Array<{
    kind: string;
    allocation: Record<string, JsonValue>;
  }>;
  env: Record<string, string>;
  promptSections: string[];
  metadata: Record<string, JsonValue>;
  releases: Array<{
    kind: string;
    release: () => Promise<void>;
  }>;
};

export class ProvisionedResources {
  private released = false;

  constructor(private readonly state: ProvisionedResourcesState) {}

  getWorkingDirectory(): string {
    return this.state.workingDirectory;
  }

  getInjectedEnv(baseEnv: Record<string, string | undefined> = process.env): Record<string, string> {
    return {
      ...sanitizeEnv(baseEnv),
      ...this.state.env,
    };
  }

  composePrompt(prompt: string, extraPromptSections: string[] = []): string {
    const sections = [...this.state.promptSections, ...extraPromptSections];
    if (sections.length === 0) {
      return prompt;
    }

    return [
      "Runtime resource constraints are already enforced by the orchestrator:",
      ...sections.map((section) => `- ${section}`),
      "",
      prompt,
    ].join("\n");
  }

  toDiscoverySnapshot(): DiscoverySnapshot {
    return {
      sessionId: this.state.sessionId,
      workspaceRoot: this.state.workspaceRoot,
      workingDirectory: this.state.workingDirectory,
      hardConstraints: this.state.hardConstraints.map((entry) => ({
        kind: entry.kind,
        allocation: { ...entry.allocation },
      })),
      softConstraints: {
        env: { ...this.state.env },
        promptSections: [...this.state.promptSections],
        metadata: { ...this.state.metadata },
      },
    };
  }

  async buildRuntimeInjection(input?: {
    discoveryFileRelativePath?: string;
    baseEnv?: Record<string, string | undefined>;
  }): Promise<RuntimeInjection> {
    const relativePath = input?.discoveryFileRelativePath?.trim() || ".agent-context/resources.json";
    const discoveryFilePath = resolve(this.state.workingDirectory, relativePath);
    await mkdir(dirname(discoveryFilePath), { recursive: true });
    await writeFile(
      discoveryFilePath,
      `${JSON.stringify(this.toDiscoverySnapshot(), null, 2)}\n`,
      "utf8",
    );

    const env = this.getInjectedEnv(input?.baseEnv);
    env.AGENT_DISCOVERY_FILE = discoveryFilePath;

    return {
      workingDirectory: this.state.workingDirectory,
      env,
      discoveryFilePath,
    };
  }

  async release(): Promise<void> {
    if (this.released) {
      return;
    }
    this.released = true;

    const errors: string[] = [];
    for (let index = this.state.releases.length - 1; index >= 0; index -= 1) {
      const releaser = this.state.releases[index];
      if (!releaser) {
        continue;
      }

      try {
        await releaser.release();
      } catch (error) {
        errors.push(`${releaser.kind}: ${toErrorMessage(error)}`);
      }
    }

    if (errors.length > 0) {
      throw new Error(`Failed to release provisioned resources: ${errors.join(" | ")}`);
    }
  }
}

export class ResourceProvisioningOrchestrator {
  private readonly providers = new Map<string, ResourceProvider>();

  constructor(providers: ResourceProvider[] = []) {
    for (const provider of providers) {
      this.registerProvider(provider);
    }
  }

  registerProvider(provider: ResourceProvider): void {
    if (this.providers.has(provider.kind)) {
      throw new Error(`Resource provider "${provider.kind}" is already registered.`);
    }
    this.providers.set(provider.kind, provider);
  }

  async provisionSession(input: {
    sessionId: string;
    resources: ResourceRequest[];
    workspaceRoot: string;
  }): Promise<ProvisionedResources> {
    const workspaceRoot = resolve(input.workspaceRoot);
    const hardConstraints: ProvisionedResourcesState["hardConstraints"] = [];
    const releases: ProvisionedResourcesState["releases"] = [];
    const env: Record<string, string> = {};
    const promptSections: string[] = [];
    const metadata: Record<string, JsonValue> = {};
    let workingDirectory = workspaceRoot;

    try {
      for (const resource of input.resources) {
        const provider = this.providers.get(resource.kind);
        if (!provider) {
          throw new Error(`No provider registered for resource kind "${resource.kind}".`);
        }

        const lease = await provider.provision({
          sessionId: input.sessionId,
          workspaceRoot,
          options: (resource.options ?? {}) as Record<string, unknown>,
        });

        hardConstraints.push({
          kind: lease.kind,
          allocation: lease.hard,
        });
        releases.push({
          kind: lease.kind,
          release: lease.release,
        });

        if (lease.soft?.env) {
          Object.assign(env, lease.soft.env);
        }
        if (lease.soft?.promptSections) {
          promptSections.push(...lease.soft.promptSections);
        }
        if (lease.soft?.metadata) {
          Object.assign(metadata, lease.soft.metadata);
        }
        if (lease.soft?.preferredWorkingDirectory) {
          workingDirectory = resolve(lease.soft.preferredWorkingDirectory);
        }
      }
    } catch (error) {
      await releaseInReverse(releases);
      throw new Error(`Resource provisioning failed: ${toErrorMessage(error)}`);
    }

    return new ProvisionedResources({
      sessionId: input.sessionId,
      workspaceRoot,
      workingDirectory,
      hardConstraints,
      env,
      promptSections,
      metadata,
      releases,
    });
  }

  async provisionChildSession(input: ChildResourceSuballocationInput): Promise<ProvisionedResources> {
    const childResources = buildChildResourceRequests(input);
    return this.provisionSession({
      sessionId: input.childSessionId,
      resources: childResources,
      workspaceRoot: input.parentSnapshot.workspaceRoot,
    });
  }
}

export type GitWorktreeProviderConfig = {
  rootDirectory: string;
  baseRef: string;
  targetPath?: string;
};

export type PortRangeProviderConfig = {
  basePort: number;
  blockSize: number;
  blockCount: number;
  primaryPortOffset: number;
  lockDirectory: string;
};

export type BuiltInProvisioningConfig = {
  gitWorktree: GitWorktreeProviderConfig;
  portRange: PortRangeProviderConfig;
};

export type BuiltInProvisioningConfigInput = {
  gitWorktree?: Partial<GitWorktreeProviderConfig>;
  portRange?: Partial<PortRangeProviderConfig>;
};

export const DEFAULT_GIT_WORKTREE_CONFIG: GitWorktreeProviderConfig = {
  rootDirectory: ".ai_ops/worktrees",
  baseRef: "HEAD",
};

export const DEFAULT_PORT_RANGE_CONFIG: PortRangeProviderConfig = {
  basePort: 36000,
  blockSize: 32,
  blockCount: 512,
  primaryPortOffset: 0,
  lockDirectory: ".ai_ops/locks/ports",
};

export function createGitWorktreeProvider(
  config: GitWorktreeProviderConfig = DEFAULT_GIT_WORKTREE_CONFIG,
): ResourceProvider {
  return {
    kind: "git-worktree",
    provision: async ({ sessionId, workspaceRoot, options }) => {
      const rootDirectory = readOptionalString(options, "rootDirectory", config.rootDirectory);
      const baseRef = readOptionalString(options, "baseRef", config.baseRef);
      const targetPath = normalizeWorktreeTargetPath(
        readOptionalStringOrUndefined(options, "targetPath") ?? config.targetPath,
        "targetPath",
      );

      const repoRoot = await runGit(["-C", workspaceRoot, "rev-parse", "--show-toplevel"]);
      const worktreeRoot = resolvePath(repoRoot, rootDirectory);
      await mkdir(worktreeRoot, { recursive: true });

      const worktreeName = buildScopedName(sessionId);
      const worktreePath = resolve(worktreeRoot, worktreeName);
      await runGit(["-C", repoRoot, "worktree", "add", "--detach", worktreePath, baseRef]);
      if (targetPath) {
        await runGit(["-C", worktreePath, "sparse-checkout", "init", "--cone"]);
        await runGit(["-C", worktreePath, "sparse-checkout", "set", targetPath]);
      }

      const preferredWorkingDirectory = targetPath ? resolve(worktreePath, targetPath) : worktreePath;
      await assertDirectoryExists(
        preferredWorkingDirectory,
        targetPath
          ? `Configured worktree target path "${targetPath}" is not a directory in ref "${baseRef}".`
          : `Provisioned worktree path "${preferredWorkingDirectory}" does not exist.`,
      );

      return {
        kind: "git-worktree",
        hard: {
          repoRoot,
          worktreeRoot,
          worktreePath,
          baseRef,
          ...(targetPath ? { targetPath } : {}),
        },
        soft: {
          env: {
            AGENT_REPO_ROOT: repoRoot,
            AGENT_WORKTREE_PATH: worktreePath,
            AGENT_WORKTREE_BASE_REF: baseRef,
          },
          promptSections: [
            `Git worktree: ${worktreePath}`,
            `Worktree base ref: ${baseRef}`,
            ...(targetPath ? [`Worktree target path: ${targetPath} (sparse-checkout enabled)`] : []),
          ],
          metadata: {
            git_worktree_path: worktreePath,
            git_worktree_base_ref: baseRef,
            ...(targetPath ? { git_worktree_target_path: targetPath } : {}),
          },
          preferredWorkingDirectory,
        },
        release: async () => {
          await runGit(["-C", repoRoot, "worktree", "remove", "--force", worktreePath]);
          await runGit(["-C", repoRoot, "worktree", "prune"]);
        },
      };
    },
  };
}

export function createPortRangeProvider(
  config: PortRangeProviderConfig = DEFAULT_PORT_RANGE_CONFIG,
): ResourceProvider {
  return {
    kind: "port-range",
    provision: async ({ sessionId, workspaceRoot, options }) => {
      const basePort = readOptionalInteger(options, "basePort", config.basePort, { min: 1 });
      const blockSize = readOptionalInteger(options, "blockSize", config.blockSize, { min: 1 });
      const blockCount = readOptionalInteger(options, "blockCount", config.blockCount, {
        min: 1,
      });
      const primaryPortOffset = readOptionalInteger(
        options,
        "primaryPortOffset",
        config.primaryPortOffset,
        { min: 0 },
      );
      const lockDirectory = readOptionalString(options, "lockDirectory", config.lockDirectory);

      if (primaryPortOffset >= blockSize) {
        throw new Error("primaryPortOffset must be smaller than blockSize.");
      }

      const maxPort = basePort + blockSize * blockCount - 1;
      if (maxPort > 65535) {
        throw new Error(
          `Port range exceeds 65535 (basePort=${basePort}, blockSize=${blockSize}, blockCount=${blockCount}).`,
        );
      }

      const lockRoot = resolvePath(workspaceRoot, lockDirectory);
      await mkdir(lockRoot, { recursive: true });

      const seed = hashToUnsignedInt(sessionId);
      const preferredBlock = seed % blockCount;

      let startPort = -1;
      let endPort = -1;
      let blockIndex = -1;
      let lockPath = "";

      for (let offset = 0; offset < blockCount; offset += 1) {
        const candidateBlock = (preferredBlock + offset) % blockCount;
        const candidateStart = basePort + candidateBlock * blockSize;
        const candidateEnd = candidateStart + blockSize - 1;
        const candidateLockPath = resolve(
          lockRoot,
          `${String(candidateStart)}-${String(candidateEnd)}.lock`,
        );

        const lockHandle = await tryCreateLockFile(candidateLockPath, {
          sessionId,
          allocatedAt: new Date().toISOString(),
          startPort: candidateStart,
          endPort: candidateEnd,
        });
        if (!lockHandle) {
          continue;
        }

        await lockHandle.close();
        startPort = candidateStart;
        endPort = candidateEnd;
        blockIndex = candidateBlock;
        lockPath = candidateLockPath;
        break;
      }

      if (startPort < 0 || endPort < 0 || blockIndex < 0 || !lockPath) {
        throw new Error("No free deterministic port block is available.");
      }

      const primaryPort = startPort + primaryPortOffset;

      return {
        kind: "port-range",
        hard: {
          basePort,
          blockSize,
          blockCount,
          blockIndex,
          startPort,
          endPort,
          primaryPort,
          lockPath,
        },
        soft: {
          env: {
            AGENT_PORT_RANGE_START: String(startPort),
            AGENT_PORT_RANGE_END: String(endPort),
            AGENT_PORT_PRIMARY: String(primaryPort),
          },
          promptSections: [
            `Assigned deterministic port range: ${String(startPort)}-${String(endPort)}`,
            `Primary port: ${String(primaryPort)}`,
          ],
          metadata: {
            port_range_start: startPort,
            port_range_end: endPort,
            port_primary: primaryPort,
            port_block_index: blockIndex,
          },
        },
        release: async () => {
          await unlink(lockPath);
        },
      };
    },
  };
}

export function createDefaultResourceProvisioningOrchestrator(
  input: BuiltInProvisioningConfigInput = {},
): ResourceProvisioningOrchestrator {
  const orchestrator = new ResourceProvisioningOrchestrator();
  orchestrator.registerProvider(
    createGitWorktreeProvider({
      ...DEFAULT_GIT_WORKTREE_CONFIG,
      ...input.gitWorktree,
    }),
  );
  orchestrator.registerProvider(
    createPortRangeProvider({
      ...DEFAULT_PORT_RANGE_CONFIG,
      ...input.portRange,
    }),
  );
  return orchestrator;
}

async function runGit(args: string[]): Promise<string> {
  try {
    const { stdout } = await execFileAsync("git", args, {
      encoding: "utf8",
    });
    return stdout.trim();
  } catch (error) {
    throw new Error(`git ${args.join(" ")} failed: ${toErrorMessage(error)}`);
  }
}

function buildScopedName(sessionId: string): string {
  const safeSessionId = sessionId.replace(/[^a-zA-Z0-9_-]/g, "-").replace(/-+/g, "-");
  const base = safeSessionId || "session";
  const hash = createHash("sha256").update(sessionId).digest("hex").slice(0, 8);
  return `${base.slice(0, 32)}-${hash}`;
}

function resolvePath(basePath: string, maybeRelativePath: string): string {
  if (isAbsolute(maybeRelativePath)) {
    return maybeRelativePath;
  }
  return resolve(basePath, maybeRelativePath);
}

function sanitizeEnv(source: Record<string, string | undefined>): Record<string, string> {
  const env: Record<string, string> = {};
  for (const [key, value] of Object.entries(source)) {
    if (value !== undefined) {
      env[key] = value;
    }
  }
  return env;
}

async function releaseInReverse(
  releases: Array<{
    kind: string;
    release: () => Promise<void>;
  }>,
): Promise<void> {
  for (let index = releases.length - 1; index >= 0; index -= 1) {
    const releaser = releases[index];
    if (!releaser) {
      continue;
    }

    try {
      await releaser.release();
    } catch {
      // Ignore rollback release errors to preserve the original provisioning error.
    }
  }
}

function hashToUnsignedInt(value: string): number {
  const digest = createHash("sha256").update(value).digest();
  return digest.readUInt32BE(0);
}

async function tryCreateLockFile(
  lockPath: string,
  payload: Record<string, JsonValue>,
): Promise<Awaited<ReturnType<typeof open>> | undefined> {
  try {
    const handle = await open(lockPath, "wx");
    await handle.writeFile(`${JSON.stringify(payload)}\n`, "utf8");
    return handle;
  } catch (error) {
    if ((error as NodeJS.ErrnoException).code === "EEXIST") {
      return undefined;
    }
    throw new Error(`Failed to create lock file "${lockPath}": ${toErrorMessage(error)}`);
  }
}

function readOptionalString(
  options: Record<string, unknown>,
  key: string,
  fallback: string,
): string {
  const value = options[key];
  if (value === undefined) {
    return fallback;
  }
  if (typeof value !== "string" || value.trim().length === 0) {
    throw new Error(`Option "${key}" must be a non-empty string.`);
  }
  return value.trim();
}

function readOptionalStringOrUndefined(
  options: Record<string, unknown>,
  key: string,
): string | undefined {
  const value = options[key];
  if (value === undefined) {
    return undefined;
  }
  if (typeof value !== "string") {
    throw new Error(`Option "${key}" must be a string when provided.`);
  }
  const trimmed = value.trim();
  return trimmed.length > 0 ? trimmed : undefined;
}

function readOptionalInteger(
  options: Record<string, unknown>,
  key: string,
  fallback: number,
  bounds: {
    min: number;
  },
): number {
  const value = options[key];
  if (value === undefined) {
    return fallback;
  }

  if (typeof value !== "number" || !Number.isInteger(value) || value < bounds.min) {
    throw new Error(`Option "${key}" must be an integer >= ${String(bounds.min)}.`);
  }
  return value;
}

function normalizeWorktreeTargetPath(value: string | undefined, key: string): string | undefined {
  if (!value) {
    return undefined;
  }

  const slashNormalized = value.replace(/\\/g, "/");
  if (isAbsolute(slashNormalized) || /^[a-zA-Z]:\//.test(slashNormalized)) {
    throw new Error(`Option "${key}" must be a relative path within the repository worktree.`);
  }

  const normalizedSegments = slashNormalized
    .split("/")
    .map((segment) => segment.trim())
    .filter((segment) => segment.length > 0 && segment !== ".");

  if (normalizedSegments.some((segment) => segment === "..")) {
    throw new Error(`Option "${key}" must not contain ".." path segments.`);
  }

  if (normalizedSegments.length === 0) {
    return undefined;
  }

  return normalizedSegments.join("/");
}

async function assertDirectoryExists(path: string, errorMessage: string): Promise<void> {
  try {
    const stats = await stat(path);
    if (!stats.isDirectory()) {
      throw new Error(errorMessage);
    }
  } catch (error) {
    if ((error as NodeJS.ErrnoException).code === "ENOENT") {
      throw new Error(errorMessage);
    }
    throw error;
  }
}

function readNumberFromAllocation(allocation: Record<string, JsonValue>, key: string): number {
  const value = allocation[key];
  if (typeof value !== "number" || !Number.isInteger(value)) {
    throw new Error(`Allocation field "${key}" must be an integer.`);
  }
  return value;
}

function readStringFromAllocation(allocation: Record<string, JsonValue>, key: string): string {
  const value = allocation[key];
  if (typeof value !== "string" || value.trim().length === 0) {
    throw new Error(`Allocation field "${key}" must be a non-empty string.`);
  }
  return value;
}

function getHardConstraint(
  snapshot: DiscoverySnapshot,
  kind: string,
): Record<string, JsonValue> | undefined {
  for (const constraint of snapshot.hardConstraints) {
    if (constraint.kind === kind) {
      return constraint.allocation;
    }
  }
  return undefined;
}

function clamp(value: number, min: number, max: number): number {
  return Math.min(Math.max(value, min), max);
}

export function buildChildResourceRequests(input: ChildResourceSuballocationInput): ResourceRequest[] {
  if (!Number.isInteger(input.childCount) || input.childCount < 1) {
    throw new Error("childCount must be a positive integer.");
  }
  if (!Number.isInteger(input.childIndex) || input.childIndex < 0 || input.childIndex >= input.childCount) {
    throw new Error("childIndex must be an integer in [0, childCount).");
  }

  const requests: ResourceRequest[] = [];

  const parentGit = getHardConstraint(input.parentSnapshot, "git-worktree");
  if (parentGit) {
    const parentWorktreePath = readStringFromAllocation(parentGit, "worktreePath");
    const baseRefRaw = parentGit.baseRef;
    const baseRef = typeof baseRefRaw === "string" && baseRefRaw.trim().length > 0 ? baseRefRaw : "HEAD";
    const targetPathRaw = parentGit.targetPath;
    const targetPath = typeof targetPathRaw === "string" ? targetPathRaw.trim() : "";

    requests.push({
      kind: "git-worktree",
      options: {
        rootDirectory: resolve(
          parentWorktreePath,
          ".ai_ops/child-worktrees",
          buildScopedName(input.parentSnapshot.sessionId),
        ),
        baseRef,
        ...(targetPath ? { targetPath } : {}),
      },
    });
  }

  const parentPortRange = getHardConstraint(input.parentSnapshot, "port-range");
  if (parentPortRange) {
    const parentStart = readNumberFromAllocation(parentPortRange, "startPort");
    const parentEnd = readNumberFromAllocation(parentPortRange, "endPort");
    const parentPrimary = readNumberFromAllocation(parentPortRange, "primaryPort");
    const lockPath = readStringFromAllocation(parentPortRange, "lockPath");

    const parentSize = parentEnd - parentStart + 1;
    const minChildSize = Math.floor(parentSize / input.childCount);
    if (minChildSize < 1) {
      throw new Error(
        `Cannot suballocate ${String(input.childCount)} child port blocks from parent range ${String(parentStart)}-${String(parentEnd)}.`,
      );
    }

    const childStart = parentStart + minChildSize * input.childIndex;
    const childEnd =
      input.childIndex === input.childCount - 1 ? parentEnd : childStart + minChildSize - 1;
    const childBlockSize = childEnd - childStart + 1;
    const primaryOffset = clamp(parentPrimary - childStart, 0, childBlockSize - 1);

    requests.push({
      kind: "port-range",
      options: {
        basePort: childStart,
        blockSize: childBlockSize,
        blockCount: 1,
        primaryPortOffset: primaryOffset,
        lockDirectory: dirname(lockPath),
      },
    });
  }

  return requests;
}

function toErrorMessage(error: unknown): string {
  if (error instanceof Error) {
    return error.message;
  }
  return String(error);
}