open-multi-agent-kit 0.91.0 → 0.92.0

package/examples/extensions/subagent/adaptive-agent-runtime.test.ts

@@ -0,0 +1,234 @@
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+import type { Message } from "omk-ai";
+import { afterEach, describe, expect, it } from "vitest";
+import { type RunAttemptInput, runAdaptiveAgent } from "./adaptive-agent-runtime.ts";
+import { createExecutionBudget } from "./deadline-budget.ts";
+import { DeadlineProfileStore } from "./deadline-profile-store.ts";
+import { emptyUsage, type SingleResult, type SubagentAttemptResult } from "./subagent-runtime-types.ts";
+
+const cleanupPaths: string[] = [];
+
+function assistant(text: string): Message {
+	return { role: "assistant", content: [{ type: "text", text }] } as Message;
+}
+
+function attempt(input: RunAttemptInput, outcome: "completed" | "cutoff" | "failed", text = ""): SubagentAttemptResult {
+	const exitCode = outcome === "completed" ? 0 : outcome === "failed" ? 1 : 124;
+	const result: SingleResult = {
+		agent: "reviewer",
+		agentSource: "user",
+		task: input.logicalTask,
+		exitCode,
+		messages: text === "" ? [] : [assistant(text)],
+		stderr: "",
+		usage: emptyUsage(),
+		output: text,
+		model: "openai-codex/gpt-5.6",
+		...(outcome === "cutoff"
+			? { stopReason: "deadline", errorMessage: "attempt cutoff" }
+			: { stopReason: outcome === "failed" ? "error" : "end" }),
+	};
+	return {
+		result,
+		process: {
+			reason: outcome === "failed" ? "completed" : outcome,
+			elapsedMs: outcome === "completed" ? 10 : input.cutoffMs,
+			cleanup: { termSent: outcome === "cutoff", killSent: false, processGroup: true },
+		},
+	};
+}
+
+afterEach(async () => {
+	await Promise.all(cleanupPaths.splice(0).map((entry) => fs.promises.rm(entry, { recursive: true, force: true })));
+});
+
+describe("adaptive subagent execution", () => {
+	it("resumes only the unfinished shard from bounded checkpoint evidence", async () => {
+		const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "omk-adaptive-resume-"));
+		cleanupPaths.push(dir);
+		const calls: RunAttemptInput[] = [];
+		const outputs = ["completed boundary inspection", "finished remaining cleanup"];
+		const result = await runAdaptiveAgent({
+			agentName: "reviewer",
+			agentSource: "user",
+			logicalTask: "Review and fix the process boundary.",
+			model: "openai-codex/gpt-5.6",
+			budget: createExecutionBudget({
+				hardDeadlineMs: 5_000,
+				cleanupReserveMs: 100,
+				maxResumeAttempts: 1,
+				maxTaskShards: 3,
+			}),
+			pendingSequentialTasks: 1,
+			checkpointPrefix: path.join(dir, "checkpoint-"),
+			profileStore: new DeadlineProfileStore(path.join(dir, "profiles.json")),
+			policy: { fallbackAttemptMs: 1_000, minimumAttemptMs: 1, maxTokensPerShard: 6_000 },
+			runAttempt: async (input) => {
+				calls.push(input);
+				return attempt(input, calls.length === 1 ? "cutoff" : "completed", outputs[calls.length - 1]);
+			},
+		});
+
+		expect(calls).toHaveLength(2);
+		expect(calls[1]?.task).toContain("completed boundary inspection");
+		expect(calls[1]?.task).toContain("Do not repeat completed side effects");
+		expect(result.deadline).toMatchObject({ outcome: "completed", resumeCount: 1, completedShardIds: ["shard-1"] });
+		expect(result.output).toContain("completed boundary inspection");
+		expect(result.output).toContain("finished remaining cleanup");
+	});
+
+	it("learns the provider/model observed from output when the agent inherits its model", async () => {
+		const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "omk-adaptive-observed-model-"));
+		cleanupPaths.push(dir);
+		const store = new DeadlineProfileStore(path.join(dir, "profiles.json"));
+		const result = await runAdaptiveAgent({
+			agentName: "reviewer",
+			agentSource: "user",
+			logicalTask: "Review the process boundary.",
+			budget: createExecutionBudget({
+				hardDeadlineMs: 5_000,
+				cleanupReserveMs: 100,
+				maxResumeAttempts: 1,
+				maxTaskShards: 3,
+			}),
+			pendingSequentialTasks: 1,
+			checkpointPrefix: path.join(dir, "checkpoint-"),
+			profileStore: store,
+			policy: { fallbackAttemptMs: 1_000, minimumAttemptMs: 1, maxTokensPerShard: 6_000 },
+			runAttempt: async (input) => attempt(input, "completed", "review complete"),
+		});
+
+		expect(result.deadline).toMatchObject({ provider: "openai-codex", model: "gpt-5.6" });
+		expect(await store.profileFor("openai-codex", "gpt-5.6")).toMatchObject({ samples: 1, completions: 1 });
+	});
+
+	it("preserves process failures instead of labeling them as deadline cutoffs", async () => {
+		const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "omk-adaptive-failure-"));
+		cleanupPaths.push(dir);
+		const result = await runAdaptiveAgent({
+			agentName: "reviewer",
+			agentSource: "user",
+			logicalTask: "Run the failing fixture.",
+			budget: createExecutionBudget({
+				hardDeadlineMs: 5_000,
+				cleanupReserveMs: 100,
+				maxResumeAttempts: 1,
+				maxTaskShards: 3,
+			}),
+			pendingSequentialTasks: 1,
+			checkpointPrefix: path.join(dir, "checkpoint-"),
+			profileStore: new DeadlineProfileStore(path.join(dir, "profiles.json")),
+			policy: { fallbackAttemptMs: 1_000, minimumAttemptMs: 1, maxTokensPerShard: 6_000 },
+			runAttempt: async (input) => attempt(input, "failed", ""),
+		});
+
+		expect(result).toMatchObject({ exitCode: 1, stopReason: "error", errorMessage: "Subagent process failed" });
+		expect(result.deadline).toMatchObject({ outcome: "failed" });
+	});
+
+	it("does not retry a cutoff that produced no checkpoint or streamed progress", async () => {
+		const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "omk-adaptive-duplicate-"));
+		cleanupPaths.push(dir);
+		let calls = 0;
+		const result = await runAdaptiveAgent({
+			agentName: "reviewer",
+			agentSource: "user",
+			logicalTask: "Review the process boundary.",
+			model: "openai-codex/gpt-5.6",
+			budget: createExecutionBudget({
+				hardDeadlineMs: 5_000,
+				cleanupReserveMs: 100,
+				maxResumeAttempts: 2,
+				maxTaskShards: 3,
+			}),
+			pendingSequentialTasks: 1,
+			checkpointPrefix: path.join(dir, "checkpoint-"),
+			profileStore: new DeadlineProfileStore(path.join(dir, "profiles.json")),
+			policy: { fallbackAttemptMs: 1_000, minimumAttemptMs: 1, maxTokensPerShard: 6_000 },
+			runAttempt: async (input) => {
+				calls += 1;
+				return attempt(input, "cutoff");
+			},
+		});
+
+		expect(calls).toBe(1);
+		expect(result.deadline).toMatchObject({ outcome: "cutoff", resumeCount: 0, duplicateResumeBlocked: true });
+	});
+
+	it("never re-executes shards that already completed", async () => {
+		const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "omk-adaptive-shards-"));
+		cleanupPaths.push(dir);
+		const shardIds: string[] = [];
+		const task = [
+			"Upgrade runtime:",
+			"1. Inspect the process.",
+			"2. Add cleanup.",
+			"3. Implement metadata.",
+			"4. Run tests.",
+		].join("\n");
+		const result = await runAdaptiveAgent({
+			agentName: "worker",
+			agentSource: "user",
+			logicalTask: task,
+			model: "anthropic/claude-sonnet",
+			budget: createExecutionBudget({
+				hardDeadlineMs: 5_000,
+				cleanupReserveMs: 100,
+				maxResumeAttempts: 1,
+				maxTaskShards: 3,
+			}),
+			pendingSequentialTasks: 1,
+			checkpointPrefix: path.join(dir, "checkpoint-"),
+			profileStore: new DeadlineProfileStore(path.join(dir, "profiles.json")),
+			policy: { fallbackAttemptMs: 200, minimumAttemptMs: 1, maxTokensPerShard: 20 },
+			runAttempt: async (input) => {
+				shardIds.push(input.shardId);
+				return attempt(input, "completed", `done ${input.shardId}`);
+			},
+		});
+
+		expect(shardIds.length).toBeGreaterThan(1);
+		expect(new Set(shardIds).size).toBe(shardIds.length);
+		expect(result.deadline?.remainingShardIds).toEqual([]);
+	});
+
+	it("caps resume cost across the whole logical task rather than once per shard", async () => {
+		const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "omk-adaptive-cost-cap-"));
+		cleanupPaths.push(dir);
+		const shardIds: string[] = [];
+		const task = [
+			"Upgrade runtime:",
+			"1. Inspect the process.",
+			"2. Add cleanup.",
+			"3. Implement metadata.",
+			"4. Run tests.",
+		].join("\n");
+		const result = await runAdaptiveAgent({
+			agentName: "worker",
+			agentSource: "user",
+			logicalTask: task,
+			model: "openai-codex/gpt-5.6",
+			budget: createExecutionBudget({
+				hardDeadlineMs: 5_000,
+				cleanupReserveMs: 100,
+				maxResumeAttempts: 1,
+				maxTaskShards: 3,
+			}),
+			pendingSequentialTasks: 1,
+			checkpointPrefix: path.join(dir, "checkpoint-"),
+			profileStore: new DeadlineProfileStore(path.join(dir, "profiles.json")),
+			policy: { fallbackAttemptMs: 200, minimumAttemptMs: 1, maxTokensPerShard: 20 },
+			runAttempt: async (input) => {
+				shardIds.push(input.shardId);
+				const outcome = shardIds.length === 2 ? "completed" : "cutoff";
+				return attempt(input, outcome, `progress ${input.shardId} ${shardIds.length}`);
+			},
+		});
+
+		expect(shardIds).toEqual(["shard-1", "shard-1", "shard-2"]);
+		expect(result.deadline).toMatchObject({ resumeCount: 1, outcome: "cutoff" });
+		expect(result.deadline?.remainingShardIds).toContain("shard-2");
+	});
+});

package/examples/extensions/subagent/adaptive-agent-runtime.ts

@@ -0,0 +1,240 @@
+import * as os from "node:os";
+import * as path from "node:path";
+import {
+	attemptOutcome,
+	createBaseResult,
+	finalizeResult,
+	mergeForUpdate,
+	mergeResult,
+	parseProviderModel,
+	resultOutput,
+} from "./adaptive-result.ts";
+import {
+	buildCheckpointTask,
+	checkpointMadeProgress,
+	createCheckpointWorkspace,
+	createEmptyCheckpoint,
+	readCheckpoint,
+	removeCheckpointWorkspace,
+	writeCheckpoint,
+} from "./checkpoint-runtime.ts";
+import {
+	allocateAttemptBudget,
+	type DeadlineProfile,
+	type ExecutionBudget,
+	estimateTaskDemand,
+	planTaskShards,
+	recommendAttemptCutoff,
+	remainingExecutionMs,
+} from "./deadline-budget.ts";
+import type { DeadlineProfileStore } from "./deadline-profile-store.ts";
+import type {
+	DeadlineAttemptMetadata,
+	DeadlineOutcome,
+	SingleResult,
+	SubagentAttemptResult,
+} from "./subagent-runtime-types.ts";
+
+export interface RunAttemptInput {
+	readonly logicalTask: string;
+	readonly task: string;
+	readonly cutoffMs: number;
+	readonly attempt: number;
+	readonly shardId: string;
+	readonly shardIndex: number;
+	readonly shardCount: number;
+	readonly checkpointFilePath: string;
+	readonly onUpdate?: (partial: SingleResult) => void;
+}
+
+interface AdaptiveRuntimePolicy {
+	readonly fallbackAttemptMs: number;
+	readonly minimumAttemptMs: number;
+	readonly maxTokensPerShard: number;
+}
+
+interface RunAdaptiveAgentOptions {
+	readonly agentName: string;
+	readonly agentSource: "user" | "project";
+	readonly logicalTask: string;
+	readonly model?: string;
+	readonly step?: number;
+	readonly signal?: AbortSignal;
+	readonly budget: ExecutionBudget;
+	readonly pendingSequentialTasks: number;
+	readonly checkpointPrefix?: string;
+	readonly profileStore: DeadlineProfileStore;
+	readonly policy: AdaptiveRuntimePolicy;
+	readonly onUpdate?: (partial: SingleResult) => void;
+	readonly runAttempt: (input: RunAttemptInput) => Promise<SubagentAttemptResult>;
+}
+
+export async function runAdaptiveAgent(options: RunAdaptiveAgentOptions): Promise<SingleResult> {
+	const startedAtMs = Date.now();
+	let identity = parseProviderModel(options.model);
+	let profile = await options.profileStore.profileFor(identity.provider, identity.model);
+	const initialCutoff = cutoffFor(profile, options, remainingExecutionMs(options.budget));
+	const demand = estimateTaskDemand(options.logicalTask);
+	const shards = planTaskShards(options.logicalTask, demand, {
+		attemptCutoffMs: initialCutoff,
+		maxTaskShards: options.budget.maxTaskShards,
+		maxTokensPerShard: options.policy.maxTokensPerShard,
+	});
+	const taskAllocation = allocateAttemptBudget({
+		budget: options.budget,
+		pendingSequentialTasks: options.pendingSequentialTasks,
+		pendingShards: 1,
+		recommendedCutoffMs: remainingExecutionMs(options.budget),
+		minimumAttemptMs: 1,
+	});
+	const taskBudget: ExecutionBudget = {
+		...options.budget,
+		hardDeadlineMs: Math.max(1, taskAllocation.taskDeadlineAtMs - startedAtMs),
+		deadlineAtMs: Math.min(options.budget.deadlineAtMs, taskAllocation.taskDeadlineAtMs),
+	};
+	const aggregate = createBaseResult(options);
+	const attempts: DeadlineAttemptMetadata[] = [];
+	const completedShardIds: string[] = [];
+	let resumeCount = 0;
+	let duplicateResumeBlocked = false;
+	let lastCheckpoint: Awaited<ReturnType<typeof readCheckpoint>> | undefined;
+	let terminalOutcome: DeadlineOutcome = "budget-exhausted";
+	const checkpointPrefix = options.checkpointPrefix ?? path.join(os.tmpdir(), "omk-subagent-checkpoint-");
+	const workspace = await createCheckpointWorkspace(checkpointPrefix);
+
+	try {
+		outer: for (const shard of shards) {
+			let shardAttempt = 0;
+			let shardMessageCount = 0;
+			let checkpoint: typeof lastCheckpoint;
+			while (true) {
+				const pendingShards = shards.length - completedShardIds.length;
+				const recommended = cutoffFor(profile, options, remainingExecutionMs(taskBudget));
+				const allocation = allocateAttemptBudget({
+					budget: taskBudget,
+					pendingSequentialTasks: 1,
+					pendingShards,
+					recommendedCutoffMs: recommended,
+					minimumAttemptMs: options.policy.minimumAttemptMs,
+				});
+				if (allocation.cutoffMs === 0 || options.signal?.aborted) {
+					terminalOutcome = options.signal?.aborted ? "aborted" : "budget-exhausted";
+					break outer;
+				}
+
+				const attemptNumber = shardAttempt + 1;
+				const task = buildCheckpointTask({
+					originalTask: options.logicalTask,
+					shardTask: shard.task,
+					shardId: shard.id,
+					shardIndex: shard.index,
+					shardCount: shards.length,
+					attempt: attemptNumber,
+					cutoffMs: allocation.cutoffMs,
+					checkpointFilePath: workspace.filePath,
+					...(checkpoint === undefined ? {} : { checkpoint }),
+				});
+				const attempt = await options.runAttempt({
+					logicalTask: options.logicalTask,
+					task,
+					cutoffMs: allocation.cutoffMs,
+					attempt: attemptNumber,
+					shardId: shard.id,
+					shardIndex: shard.index,
+					shardCount: shards.length,
+					checkpointFilePath: workspace.filePath,
+					onUpdate: options.onUpdate
+						? (partial) => options.onUpdate?.(mergeForUpdate(aggregate, partial))
+						: undefined,
+				});
+				mergeResult(aggregate, attempt.result);
+				shardMessageCount += attempt.result.messages.length;
+				const outcome = attemptOutcome(attempt);
+				const observedIdentity = parseProviderModel(attempt.result.model ?? options.model);
+				if (observedIdentity.provider !== identity.provider || observedIdentity.model !== identity.model) {
+					identity = observedIdentity;
+					profile = await options.profileStore.profileFor(identity.provider, identity.model);
+				}
+				try {
+					profile = await options.profileStore.record({
+						provider: identity.provider,
+						model: identity.model,
+						outcome,
+						elapsedMs: attempt.process.elapsedMs,
+						estimatedTokens: demand.estimatedTokens,
+						workUnits: demand.workUnits,
+					});
+				} catch {
+					profile = await options.profileStore.profileFor(identity.provider, identity.model);
+				}
+				lastCheckpoint = await readCheckpoint(workspace.filePath, {
+					messageCount: shardMessageCount,
+					outputTail: resultOutput(attempt.result),
+					lastEventAtMs: Date.now(),
+				});
+				attempts.push({
+					attempt: attemptNumber,
+					shardId: shard.id,
+					cutoffMs: allocation.cutoffMs,
+					elapsedMs: attempt.process.elapsedMs,
+					outcome,
+					exitCode: attempt.result.exitCode,
+					messageCount: attempt.result.messages.length,
+					processReason: attempt.process.reason,
+					cleanup: attempt.process.cleanup,
+				});
+
+				if (outcome === "completed") {
+					completedShardIds.push(shard.id);
+					terminalOutcome = "completed";
+					checkpoint = undefined;
+					await writeCheckpoint(workspace.filePath, createEmptyCheckpoint());
+					continue outer;
+				}
+				terminalOutcome = outcome;
+				if (outcome !== "cutoff" || resumeCount >= taskBudget.maxResumeAttempts) break outer;
+				if (!checkpointMadeProgress(checkpoint, lastCheckpoint)) {
+					duplicateResumeBlocked = true;
+					break outer;
+				}
+				checkpoint = lastCheckpoint;
+				shardAttempt += 1;
+				resumeCount += 1;
+			}
+		}
+	} finally {
+		await removeCheckpointWorkspace(workspace);
+	}
+
+	const remainingShardIds = shards.map((shard) => shard.id).filter((id) => !completedShardIds.includes(id));
+	if (remainingShardIds.length === 0) terminalOutcome = "completed";
+	finalizeResult(aggregate, terminalOutcome);
+	aggregate.deadline = {
+		outcome: terminalOutcome,
+		startedAtMs,
+		elapsedMs: Date.now() - startedAtMs,
+		hardDeadlineMs: taskBudget.hardDeadlineMs,
+		estimatedTokens: demand.estimatedTokens,
+		workUnits: demand.workUnits,
+		predictedMs: demand.predictedMs,
+		plannedShardIds: shards.map((shard) => shard.id),
+		completedShardIds,
+		remainingShardIds,
+		resumeCount,
+		duplicateResumeBlocked,
+		provider: identity.provider,
+		model: identity.model,
+		profileSamples: profile?.samples ?? 0,
+		attempts,
+		...(lastCheckpoint === undefined ? {} : { checkpoint: lastCheckpoint }),
+	};
+	return aggregate;
+}
+
+function cutoffFor(profile: DeadlineProfile | undefined, options: RunAdaptiveAgentOptions, maximumMs: number): number {
+	return recommendAttemptCutoff(profile, {
+		fallbackMs: options.policy.fallbackAttemptMs,
+		minimumMs: Math.min(options.policy.minimumAttemptMs, maximumMs),
+		maximumMs: Math.max(1, maximumMs),
+	});
+}

package/examples/extensions/subagent/adaptive-result.ts

@@ -0,0 +1,112 @@
+import type { Message } from "omk-ai";
+import {
+	type DeadlineOutcome,
+	emptyUsage,
+	type SingleResult,
+	type SubagentAttemptResult,
+	type UsageStats,
+} from "./subagent-runtime-types.ts";
+
+interface BaseResultOptions {
+	readonly agentName: string;
+	readonly agentSource: "user" | "project";
+	readonly logicalTask: string;
+	readonly model?: string;
+	readonly step?: number;
+}
+
+export function createBaseResult(options: BaseResultOptions): SingleResult {
+	return {
+		agent: options.agentName,
+		agentSource: options.agentSource,
+		task: options.logicalTask,
+		exitCode: 0,
+		messages: [],
+		stderr: "",
+		usage: emptyUsage(),
+		model: options.model,
+		step: options.step,
+	};
+}
+
+export function attemptOutcome(attempt: SubagentAttemptResult): Exclude<DeadlineOutcome, "budget-exhausted"> {
+	if (attempt.process.reason === "cutoff") return "cutoff";
+	if (attempt.process.reason === "aborted") return "aborted";
+	return attempt.result.exitCode === 0 && attempt.result.stopReason !== "error" ? "completed" : "failed";
+}
+
+export function mergeResult(target: SingleResult, source: SingleResult): void {
+	target.messages.push(...source.messages);
+	target.stderr = [target.stderr, source.stderr].filter(Boolean).join("\n");
+	target.usage = mergeUsage(target.usage, source.usage);
+	target.output = mergeOutput(target.output ?? "", resultOutput(source));
+	target.exitCode = source.exitCode;
+	target.stopReason = source.stopReason;
+	target.errorMessage = source.errorMessage;
+	target.model = source.model ?? target.model;
+}
+
+export function mergeForUpdate(base: SingleResult, partial: SingleResult): SingleResult {
+	const merged: SingleResult = { ...base, messages: [...base.messages], usage: { ...base.usage } };
+	mergeResult(merged, partial);
+	return merged;
+}
+
+export function resultOutput(result: SingleResult): string {
+	if (result.output !== undefined) return result.output;
+	for (let index = result.messages.length - 1; index >= 0; index--) {
+		const message: Message = result.messages[index];
+		if (message.role !== "assistant") continue;
+		for (const part of message.content) if (part.type === "text") return part.text;
+	}
+	return "";
+}
+
+export function finalizeResult(result: SingleResult, outcome: DeadlineOutcome): void {
+	if (outcome === "completed") {
+		result.exitCode = 0;
+		result.stopReason = "end";
+		result.errorMessage = undefined;
+		return;
+	}
+	if (outcome === "failed") {
+		result.exitCode = Math.max(1, result.exitCode);
+		result.stopReason = "error";
+		result.errorMessage ??= "Subagent process failed";
+		return;
+	}
+	result.exitCode = outcome === "aborted" ? 130 : 124;
+	result.stopReason = outcome === "aborted" ? "aborted" : "deadline";
+	result.errorMessage ??=
+		outcome === "budget-exhausted" ? "Subagent execution budget exhausted" : "Subagent deadline cutoff";
+}
+
+export function parseProviderModel(value: string | undefined): { provider: string; model: string } {
+	if (!value) return { provider: "unknown", model: "default" };
+	const separator = value.indexOf("/");
+	return separator < 1
+		? { provider: "unknown", model: value }
+		: { provider: value.slice(0, separator), model: value.slice(separator + 1) || "default" };
+}
+
+function mergeUsage(left: UsageStats, right: UsageStats): UsageStats {
+	return {
+		input: left.input + right.input,
+		output: left.output + right.output,
+		cacheRead: left.cacheRead + right.cacheRead,
+		cacheWrite: left.cacheWrite + right.cacheWrite,
+		cost: left.cost + right.cost,
+		contextTokens: Math.max(left.contextTokens, right.contextTokens),
+		turns: left.turns + right.turns,
+	};
+}
+
+function mergeOutput(left: string, right: string): string {
+	if (right === "" || left === right || left.endsWith(right)) return left;
+	if (left === "") return right;
+	const maxOverlap = Math.min(left.length, right.length, 4_096);
+	for (let overlap = maxOverlap; overlap > 0; overlap--) {
+		if (left.endsWith(right.slice(0, overlap))) return `${left}${right.slice(overlap)}`;
+	}
+	return `${left}\n\n${right}`;
+}