@oai404iao/pi-subagent 0.2.0 → 0.4.0-alpha.0

package/src/providers.ts

@@ -1,7 +1,22 @@
 import { rm } from "node:fs/promises";
 import type { SessionEntry } from "@earendil-works/pi-coding-agent";
 import { SessionManager } from "@earendil-works/pi-coding-agent";
-import type { SubagentMode, SubagentProviderName } from "./types.ts";
+import type {
+	ContextInheritance,
+	SubagentMode,
+	SubagentProviderName,
+} from "./types.ts";
+
+const INHERITED_COMPACTION_CUSTOM_TYPE =
+	"pi-subagent/inherited-compaction-summary";
+const INHERITED_BRANCH_CUSTOM_TYPE =
+	"pi-subagent/inherited-branch-summary";
+const COMPACTION_SUMMARY_PREFIX =
+	"The conversation history before this point was compacted into the following summary:\n\n<summary>\n";
+const COMPACTION_SUMMARY_SUFFIX = "\n</summary>";
+const BRANCH_SUMMARY_PREFIX =
+	"The following is a summary of a branch that this conversation came back from:\n\n<summary>\n";
+const BRANCH_SUMMARY_SUFFIX = "</summary>";
 
 export interface SessionView {
 	getBranch(): SessionEntry[];
@@ -9,6 +24,9 @@ export interface SessionView {
 	getSessionDir(): string;
 	getSessionFile(): string | undefined;
 	getSessionId(): string;
+	getEntries(): SessionEntry[];
+	buildContextEntries(): SessionEntry[];
+	appendCustomEntry(customType: string, data?: unknown): string;
 }
 
 export interface ProviderParent {
@@ -25,7 +43,11 @@ export interface ChildProvider {
 	name: SubagentProviderName;
 	inheritsParentContext: boolean;
 	supportsContinuable: boolean;
-	prepare(parent: ProviderParent, mode: SubagentMode): Promise<PreparedChildSession>;
+	prepare(
+		parent: ProviderParent,
+		mode: SubagentMode,
+		context?: ContextInheritance,
+	): Promise<PreparedChildSession>;
 }
 
 async function removeOwnedSession(path: string | undefined): Promise<void> {
@@ -66,33 +88,195 @@ export function completedTurnBoundaryId(entries: readonly SessionEntry[]): strin
 	return undefined;
 }
 
-function forkedSession(parent: ProviderParent): PreparedChildSession {
+interface CompletedTurnSpan {
+	start: number;
+	end: number;
+}
+
+function isSummaryEntry(entry: SessionEntry): boolean {
+	return (
+		entry.type === "compaction"
+		|| entry.type === "branch_summary"
+		|| (
+			entry.type === "message"
+			&& (
+				entry.message.role === "compactionSummary"
+				|| entry.message.role === "branchSummary"
+			)
+		)
+	);
+}
+
+function startsModelTurn(entry: SessionEntry): boolean {
+	if (entry.type === "custom_message") return true;
+	if (entry.type !== "message") return false;
+	switch (entry.message.role) {
+		case "user":
+		case "custom":
+			return true;
+		case "bashExecution":
+			return !entry.message.excludeFromContext;
+		default:
+			return false;
+	}
+}
+
+function completedTurnSpans(
+	entries: readonly SessionEntry[],
+): {
+	turns: CompletedTurnSpan[];
+	latestSummary?: number;
+} {
+	const turns: CompletedTurnSpan[] = [];
+	let latestSummary: number | undefined;
+	let currentStart: number | undefined;
+	let currentEnd: number | undefined;
+	const finishCurrent = () => {
+		if (currentStart !== undefined && currentEnd !== undefined) {
+			turns.push({ start: currentStart, end: currentEnd });
+		}
+		currentStart = undefined;
+		currentEnd = undefined;
+	};
+
+	for (let index = 0; index < entries.length; index++) {
+		const entry = entries[index]!;
+		if (isSummaryEntry(entry)) {
+			finishCurrent();
+			latestSummary = index;
+			continue;
+		}
+		if (startsModelTurn(entry)) {
+			finishCurrent();
+			currentStart = index;
+			continue;
+		}
+		if (entry.type !== "message") continue;
+		if (
+			currentStart === undefined
+			&& (
+				entry.message.role === "assistant"
+				|| entry.message.role === "toolResult"
+			)
+		) {
+			currentStart = latestSummary ?? 0;
+		}
+		if (
+			entry.message.role === "assistant"
+			&& entry.message.stopReason !== "toolUse"
+		) {
+			currentEnd = index;
+		}
+	}
+	finishCurrent();
+	return {
+		turns,
+		...(latestSummary !== undefined ? { latestSummary } : {}),
+	};
+}
+
+export function completedContextEntries(
+	entries: readonly SessionEntry[],
+	context: Exclude<ContextInheritance, { mode: "fresh" }>,
+): SessionEntry[] {
+	const { turns, latestSummary } = completedTurnSpans(entries);
+	const end = Math.max(
+		turns.at(-1)?.end ?? -1,
+		latestSummary ?? -1,
+	);
+	if (end < 0) return [];
+	let start = 0;
+	if (context.mode === "last_n_completed") {
+		if (turns.length >= context.completedTurns) {
+			start = turns[turns.length - context.completedTurns]!.start;
+		}
+	}
+	return entries
+		.slice(start, end + 1)
+		.map((entry) => structuredClone(entry));
+}
+
+function appendInheritedEntry(
+	session: SessionManager,
+	entry: SessionEntry,
+): void {
+	if (entry.type === "message") {
+		if (entry.message.role === "compactionSummary") {
+			session.appendCustomMessageEntry(
+				INHERITED_COMPACTION_CUSTOM_TYPE,
+				`${COMPACTION_SUMMARY_PREFIX}${entry.message.summary}${COMPACTION_SUMMARY_SUFFIX}`,
+				false,
+				{ tokensBefore: entry.message.tokensBefore },
+			);
+			return;
+		}
+		if (entry.message.role === "branchSummary") {
+			session.appendCustomMessageEntry(
+				INHERITED_BRANCH_CUSTOM_TYPE,
+				`${BRANCH_SUMMARY_PREFIX}${entry.message.summary}${BRANCH_SUMMARY_SUFFIX}`,
+				false,
+				{ fromId: entry.message.fromId },
+			);
+			return;
+		}
+		session.appendMessage(
+			structuredClone(
+				entry.message,
+			) as Parameters<SessionManager["appendMessage"]>[0],
+		);
+		return;
+	}
+	if (entry.type === "custom_message") {
+		session.appendCustomMessageEntry(
+			entry.customType,
+			structuredClone(entry.content),
+			entry.display,
+			structuredClone(entry.details),
+		);
+		return;
+	}
+	if (entry.type === "compaction") {
+		session.appendCustomMessageEntry(
+			INHERITED_COMPACTION_CUSTOM_TYPE,
+			`${COMPACTION_SUMMARY_PREFIX}${entry.summary}${COMPACTION_SUMMARY_SUFFIX}`,
+			false,
+			{ tokensBefore: entry.tokensBefore },
+		);
+		return;
+	}
+	if (entry.type === "branch_summary") {
+		session.appendCustomMessageEntry(
+			INHERITED_BRANCH_CUSTOM_TYPE,
+			`${BRANCH_SUMMARY_PREFIX}${entry.summary}${BRANCH_SUMMARY_SUFFIX}`,
+			false,
+			{ fromId: entry.fromId },
+		);
+	}
+}
+
+function forkedSession(
+	parent: ProviderParent,
+	context: Exclude<ContextInheritance, { mode: "fresh" }>,
+): PreparedChildSession {
 	const parentFile = parent.sessionManager.getSessionFile();
-	const boundaryId = completedTurnBoundaryId(parent.sessionManager.getBranch());
-	if (!boundaryId) return freshSession(parent);
+	const inherited = completedContextEntries(
+		parent.sessionManager.buildContextEntries(),
+		context,
+	);
+	if (inherited.length === 0) return freshSession(parent);
 	if (!parentFile) {
 		throw new Error(
 			"fork provider cannot copy completed history from an ephemeral parent session; use spawn instead",
 		);
 	}
-
-	const clone = SessionManager.open(
-		parentFile,
-		parent.sessionManager.getSessionDir(),
-		parent.sessionManager.getCwd(),
-	);
-	if (!clone.getEntry(boundaryId)) return freshSession(parent);
-	const childFile = clone.createBranchedSession(boundaryId);
-	if (!childFile) return freshSession(parent);
-	const sessionManager = SessionManager.open(
-		childFile,
-		parent.sessionManager.getSessionDir(),
-		parent.sessionManager.getCwd(),
-	);
+	const prepared = freshSession(parent);
+	for (const entry of inherited) {
+		appendInheritedEntry(prepared.sessionManager, entry);
+	}
 	return {
-		sessionManager,
-		seedMessageCount: sessionManager.buildSessionContext().messages.length,
-		rollback: () => removeOwnedSession(childFile),
+		...prepared,
+		seedMessageCount:
+			prepared.sessionManager.buildSessionContext().messages.length,
 	};
 }
 
@@ -101,7 +285,14 @@ export class SpawnProvider implements ChildProvider {
 	readonly inheritsParentContext = false;
 	readonly supportsContinuable = true;
 
-	prepare(parent: ProviderParent, _mode: SubagentMode): Promise<PreparedChildSession> {
+	prepare(
+		parent: ProviderParent,
+		_mode: SubagentMode,
+		context: ContextInheritance = { mode: "fresh" },
+	): Promise<PreparedChildSession> {
+		if (context.mode !== "fresh") {
+			throw new Error("spawn provider requires fresh context");
+		}
 		return Promise.resolve(freshSession(parent));
 	}
 }
@@ -109,13 +300,17 @@ export class SpawnProvider implements ChildProvider {
 export class ForkProvider implements ChildProvider {
 	readonly name = "fork";
 	readonly inheritsParentContext = true;
-	readonly supportsContinuable = false;
+	readonly supportsContinuable = true;
 
-	async prepare(parent: ProviderParent, mode: SubagentMode): Promise<PreparedChildSession> {
-		if (mode === "continuable") {
-			throw new Error("fork provider is one-shot only; use spawn for continuable background work");
+	async prepare(
+		parent: ProviderParent,
+		_mode: SubagentMode,
+		context: ContextInheritance = { mode: "all_completed" },
+	): Promise<PreparedChildSession> {
+		if (context.mode === "fresh") {
+			throw new Error("fork provider requires inherited context");
 		}
-		return forkedSession(parent);
+		return forkedSession(parent, context);
 	}
 }
 

package/src/render.ts

@@ -4,7 +4,14 @@ import type { DelegationDetails, ParentMessageDetails } from "./types.ts";
 import { formatUsage } from "./result.ts";
 
 export function renderDelegationCall(
-	args: { agent?: string; description?: string; prompt?: string; run_in_background?: boolean },
+	args: {
+		agent?: string;
+		task_name?: string;
+		description?: string;
+		prompt?: string;
+		run_in_background?: boolean;
+		context?: { mode?: string };
+	},
 	theme: {
 		fg(color: any, text: string): string;
 		bold(text: string): string;
@@ -13,17 +20,20 @@ export function renderDelegationCall(
 ): Text {
 	const mode =
 		provider === "fork"
-			? "fork · foreground"
+			? args.run_in_background === true
+				? "fork · background"
+				: "fork · foreground"
 			: args.run_in_background === false
-				? "spawn · foreground"
+				? `${args.context?.mode ?? "fresh"} · foreground`
 				: args.run_in_background === true
-					? "spawn · background"
-					: "spawn · configured default";
+					? `${args.context?.mode ?? "fresh"} · background`
+					: `${args.context?.mode ?? "fresh"} · configured default`;
 	let text =
 		theme.fg("toolTitle", theme.bold(provider === "fork" ? "subagent_fork " : "subagent ")) +
 		theme.fg("accent", args.agent ?? "…") +
 		theme.fg("muted", ` [${mode}]`);
 	if (args.description) text += `\n  ${theme.fg("dim", args.description)}`;
+	if (args.task_name) text += `\n  ${theme.fg("dim", `path: ${args.task_name}`)}`;
 	return new Text(text, 0, 0);
 }
 
@@ -37,6 +47,10 @@ export function renderDelegationResult(
 	},
 ): Text | Container {
 	if (!details) return new Text(content || "(no output)", 0, 0);
+	const taskPath = details.taskPath || details.agentId;
+	const contextMode =
+		details.context?.mode
+		?? (details.provider === "fork" ? "all_completed" : "fresh");
 	const running = options.isPartial || details.status === "starting" || details.status === "running";
 	const icon = running
 		? theme.fg("warning", "◌")
@@ -45,8 +59,8 @@ export function renderDelegationResult(
 			: theme.fg("success", "✓");
 	const header = `${icon} ${theme.fg("toolTitle", theme.bold(details.agent))} ${theme.fg(
 		"muted",
-		`[${details.provider}/${details.mode}]`,
-	)} ${theme.fg("dim", details.id)}`;
+		`[${details.provider}/${details.mode}/${contextMode}]`,
+	)} ${theme.fg("accent", taskPath)} ${theme.fg("dim", details.agentId)}`;
 
 	if (!options.expanded) {
 		const lines = [header, theme.fg("muted", details.label)];
@@ -102,7 +116,10 @@ export function renderParentMessage(
 		theme.fg("accent", icon) +
 			" " +
 			theme.fg("toolTitle", theme.bold(`subagent ${kind}`)) +
-			theme.fg("muted", ` ${details?.childId ?? "unknown"}${label}`),
+			theme.fg(
+				"muted",
+				` ${details?.taskPath ?? details?.childAgentId ?? "unknown"}${label}`,
+			),
 		outputPad,
 		0,
 	);

package/src/scheduler.ts

@@ -0,0 +1,173 @@
+export interface BackgroundRunPermit {
+	release(): void;
+}
+
+export interface BackgroundRunAcquireOptions {
+	signal?: AbortSignal;
+	waitForCapacity?: boolean;
+}
+
+export class BackgroundConcurrencyLimitError extends Error {
+	constructor(readonly limit: number) {
+		super(
+			`background subagent concurrency limit (${limit}) is already in use; wait for an active background run to finish`,
+		);
+		this.name = "BackgroundConcurrencyLimitError";
+	}
+}
+
+interface BackgroundRunWaiter {
+	resolve: (permit: BackgroundRunPermit) => void;
+	reject: (error: Error) => void;
+	signal?: AbortSignal;
+	onAbort?: () => void;
+}
+
+export class BackgroundRunLimiter {
+	private active = 0;
+	private readonly waiters: BackgroundRunWaiter[] = [];
+	private closedError: Error | undefined;
+
+	constructor(private limit = 4) {
+		validateLimit(limit);
+	}
+
+	get activeCount(): number {
+		return this.active;
+	}
+
+	get pendingCount(): number {
+		return this.waiters.length;
+	}
+
+	configure(limit: number): void {
+		validateLimit(limit);
+		if (limit === this.limit) return;
+		this.limit = limit;
+		this.drain();
+	}
+
+	acquire(
+		options: BackgroundRunAcquireOptions = {},
+	): Promise<BackgroundRunPermit> {
+		if (this.closedError) return Promise.reject(this.closedError);
+		if (options.signal?.aborted) {
+			return Promise.reject(abortReason(options.signal));
+		}
+		if (this.waiters.length === 0 && this.active < this.limit) {
+			this.active += 1;
+			return Promise.resolve(this.createPermit());
+		}
+		if (options.waitForCapacity === false) {
+			return Promise.reject(new BackgroundConcurrencyLimitError(this.limit));
+		}
+
+		return new Promise<BackgroundRunPermit>((resolve, reject) => {
+			const waiter: BackgroundRunWaiter = {
+				resolve,
+				reject,
+				...(options.signal ? { signal: options.signal } : {}),
+			};
+			if (options.signal) {
+				waiter.onAbort = () => {
+					const index = this.waiters.indexOf(waiter);
+					if (index >= 0) this.waiters.splice(index, 1);
+					reject(abortReason(options.signal!));
+					this.drain();
+				};
+				options.signal.addEventListener("abort", waiter.onAbort, { once: true });
+			}
+			this.waiters.push(waiter);
+		});
+	}
+
+	close(error = new Error("background subagent scheduler is shutting down")): void {
+		if (this.closedError) return;
+		this.closedError = error;
+		for (const waiter of this.waiters.splice(0)) {
+			this.removeAbortListener(waiter);
+			waiter.reject(error);
+		}
+	}
+
+	private createPermit(): BackgroundRunPermit {
+		let released = false;
+		return {
+			release: () => {
+				if (released) return;
+				released = true;
+				this.active = Math.max(0, this.active - 1);
+				this.drain();
+			},
+		};
+	}
+
+	private drain(): void {
+		while (this.waiters.length > 0) {
+			const waiter = this.waiters[0]!;
+			if (waiter.signal?.aborted) {
+				this.waiters.shift();
+				this.removeAbortListener(waiter);
+				waiter.reject(abortReason(waiter.signal));
+				continue;
+			}
+			if (this.closedError) {
+				this.waiters.shift();
+				this.removeAbortListener(waiter);
+				waiter.reject(this.closedError);
+				continue;
+			}
+			if (this.active >= this.limit) return;
+			this.waiters.shift();
+			this.removeAbortListener(waiter);
+			this.active += 1;
+			waiter.resolve(this.createPermit());
+		}
+	}
+
+	private removeAbortListener(waiter: BackgroundRunWaiter): void {
+		if (waiter.signal && waiter.onAbort) {
+			waiter.signal.removeEventListener("abort", waiter.onAbort);
+		}
+	}
+}
+
+export class AgentOperationQueue {
+	private readonly tails = new Map<string, Promise<void>>();
+
+	async run<T>(agentId: string, operation: () => Promise<T>): Promise<T> {
+		const predecessor = this.tails.get(agentId) ?? Promise.resolve();
+		let release!: () => void;
+		const gate = new Promise<void>((resolve) => {
+			release = resolve;
+		});
+		const tail = predecessor.catch(() => {}).then(() => gate);
+		this.tails.set(agentId, tail);
+
+		await predecessor.catch(() => {});
+		try {
+			return await operation();
+		} finally {
+			release();
+			if (this.tails.get(agentId) === tail) this.tails.delete(agentId);
+		}
+	}
+
+	async waitForIdle(): Promise<void> {
+		while (this.tails.size > 0) {
+			await Promise.all([...this.tails.values()].map((tail) => tail.catch(() => {})));
+		}
+	}
+}
+
+function validateLimit(limit: number): void {
+	if (!Number.isSafeInteger(limit) || limit < 1) {
+		throw new Error("background run limit must be a positive safe integer");
+	}
+}
+
+function abortReason(signal: AbortSignal): Error {
+	return signal.reason instanceof Error
+		? signal.reason
+		: new Error(signal.reason ? String(signal.reason) : "background run scheduling aborted");
+}