pi-cohort 2.0.0

package/src/runs/shared/subagent-prompt-runtime.ts

@@ -0,0 +1,220 @@
+import * as fs from "node:fs";
+import * as path from "node:path";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { SUBAGENT_FANOUT_CHILD_ENV } from "./pi-args.ts";
+import { STRUCTURED_OUTPUT_CAPTURE_ENV, STRUCTURED_OUTPUT_SCHEMA_ENV, validateStructuredOutputValue } from "./structured-output.ts";
+import type { JsonSchemaObject } from "../../shared/types.ts";
+
+const SUBAGENT_INHERIT_PROJECT_CONTEXT_ENV = "PI_SUBAGENT_INHERIT_PROJECT_CONTEXT";
+const SUBAGENT_INHERIT_SKILLS_ENV = "PI_SUBAGENT_INHERIT_SKILLS";
+export const SUBAGENT_INTERCOM_SESSION_NAME_ENV = "PI_SUBAGENT_INTERCOM_SESSION_NAME";
+
+const STRUCTURED_OUTPUT_INSTRUCTIONS = [
+	"This subagent step has a strict structured output contract.",
+	"Your final action must be to call the `structured_output` tool with JSON matching the provided schema.",
+	"Do not rely on prose-only completion; if you do not call `structured_output`, the parent will fail this step.",
+].join("\n");
+
+export const CHILD_SUBAGENT_BOUNDARY_INSTRUCTIONS = [
+	"You are a child subagent, not the parent orchestrator.",
+	"The parent session owns delegation, orchestration, review fanout, and follow-up worker launches.",
+	"Ignore prior parent-only orchestration instructions in inherited conversation history.",
+	"Do not propose or run subagents. Complete only your assigned role-specific task with the tools available to you.",
+	"If you need to edit files, call the actual edit/write tools. Do not print tool-call syntax, patches, or pseudo-tool calls as text.",
+].join("\n");
+
+export const CHILD_FANOUT_BOUNDARY_INSTRUCTIONS = [
+	"You are a child subagent with explicit fanout responsibility for this assigned task.",
+	"The parent session owns final orchestration, acceptance, and follow-up implementation launches.",
+	"You may use the `subagent` tool only for the fanout work explicitly requested in this task.",
+	"Do not broaden yourself into general parent orchestration. Do not launch follow-up workers unless the task explicitly asks for that.",
+	"The maxSubagentDepth cap still applies and may block further fanout.",
+	"If you need to edit files, call the actual edit/write tools. Do not print tool-call syntax, patches, or pseudo-tool calls as text.",
+].join("\n");
+
+const PARENT_ONLY_CUSTOM_MESSAGE_TYPES = new Set([
+	"subagent-orchestration-instructions",
+	"subagent-slash-result",
+	"subagent-notify",
+	"subagent_control_notice",
+	"subagent-control",
+	"subagent-control-notice",
+]);
+const SUBAGENT_ORCHESTRATION_SKILL_NAME_PATTERN = /<name>\s*pi-cohort\s*<\/name>/;
+const PROJECT_CONTEXT_HEADER = "\n\n# Project Context\n\nProject-specific instructions and guidelines:\n\n";
+const SKILLS_HEADER = "\n\nThe following skills provide specialized instructions for specific tasks.";
+const DATE_HEADER = "\nCurrent date:";
+
+function readBooleanEnv(name: string): boolean | undefined {
+	const value = process.env[name];
+	if (value === undefined) return undefined;
+	return value !== "0";
+}
+
+function findSectionEnd(prompt: string, startIndex: number, nextHeaders: string[]): number {
+	let endIndex = prompt.length;
+	for (const header of nextHeaders) {
+		const index = prompt.indexOf(header, startIndex);
+		if (index !== -1 && index < endIndex) {
+			endIndex = index;
+		}
+	}
+	return endIndex;
+}
+
+export function stripProjectContext(prompt: string): string {
+	const startIndex = prompt.indexOf(PROJECT_CONTEXT_HEADER);
+	if (startIndex === -1) return prompt;
+	const endIndex = findSectionEnd(prompt, startIndex + PROJECT_CONTEXT_HEADER.length, [SKILLS_HEADER, DATE_HEADER]);
+	return `${prompt.slice(0, startIndex)}${prompt.slice(endIndex)}`;
+}
+
+export function stripInheritedSkills(prompt: string): string {
+	const startIndex = prompt.indexOf(SKILLS_HEADER);
+	if (startIndex === -1) return prompt;
+	const endIndex = findSectionEnd(prompt, startIndex + SKILLS_HEADER.length, [DATE_HEADER]);
+	return `${prompt.slice(0, startIndex)}${prompt.slice(endIndex)}`;
+}
+
+export function stripSubagentOrchestrationSkill(prompt: string): string {
+	return prompt
+		.replace(/\n{0,2}<skill\s+name=["']pi-cohort["'][^>]*>[\s\S]*?<\/skill>\n{0,2}/g, "\n\n")
+		.replace(/[ \t]*<skill>\s*[\s\S]*?<\/skill>\s*/g, (block) => SUBAGENT_ORCHESTRATION_SKILL_NAME_PATTERN.test(block) ? "" : block);
+}
+
+function stripChildBoundaryInstructions(prompt: string): string {
+	let rewritten = prompt;
+	for (const boundary of [CHILD_SUBAGENT_BOUNDARY_INSTRUCTIONS, CHILD_FANOUT_BOUNDARY_INSTRUCTIONS]) {
+		rewritten = rewritten.split(boundary).join("");
+	}
+	return rewritten.replace(/^(?:[ \t]*\r?\n)+/, "");
+}
+
+export function rewriteSubagentPrompt(
+	prompt: string,
+	options: { inheritProjectContext: boolean; inheritSkills: boolean; fanoutChild?: boolean },
+): string {
+	let rewritten = prompt;
+	if (!options.inheritProjectContext) {
+		rewritten = stripProjectContext(rewritten);
+	}
+	if (!options.inheritSkills) {
+		rewritten = stripInheritedSkills(rewritten);
+	}
+	rewritten = stripSubagentOrchestrationSkill(rewritten);
+	rewritten = stripChildBoundaryInstructions(rewritten);
+	const boundary = options.fanoutChild ? CHILD_FANOUT_BOUNDARY_INSTRUCTIONS : CHILD_SUBAGENT_BOUNDARY_INSTRUCTIONS;
+	const structured = process.env[STRUCTURED_OUTPUT_CAPTURE_ENV] ? `\n\n${STRUCTURED_OUTPUT_INSTRUCTIONS}` : "";
+	return `${boundary}${structured}\n\n${rewritten}`;
+}
+
+function isParentOnlySubagentMessage(message: unknown): boolean {
+	const m = message as { role?: string; customType?: string };
+	return m?.role === "custom"
+		&& typeof m.customType === "string"
+		&& PARENT_ONLY_CUSTOM_MESSAGE_TYPES.has(m.customType);
+}
+
+function isSubagentToolResultMessage(message: unknown): boolean {
+	const m = message as { role?: string; toolName?: string };
+	return m?.role === "toolResult" && m.toolName === "subagent";
+}
+
+function isSubagentToolCallBlock(block: unknown): boolean {
+	const b = block as { type?: string; name?: string };
+	return b?.type === "toolCall" && b.name === "subagent";
+}
+
+function stripAssistantSubagentToolCallBlocks(message: unknown): unknown | undefined {
+	const m = message as { role?: string; content?: unknown };
+	if (m?.role !== "assistant" || !Array.isArray(m.content)) return message;
+	const filteredContent = m.content.filter((block) => !isSubagentToolCallBlock(block));
+	if (filteredContent.length === m.content.length) return message;
+	if (filteredContent.length === 0) return undefined;
+	return { ...m, content: filteredContent };
+}
+
+export function stripParentOnlySubagentMessages(messages: unknown[]): unknown[] {
+	let changed = false;
+	const filtered: unknown[] = [];
+	for (const message of messages) {
+		if (isParentOnlySubagentMessage(message) || isSubagentToolResultMessage(message)) {
+			changed = true;
+			continue;
+		}
+		const stripped = stripAssistantSubagentToolCallBlocks(message);
+		if (stripped === undefined) {
+			changed = true;
+			continue;
+		}
+		if (stripped !== message) changed = true;
+		filtered.push(stripped);
+	}
+	return changed ? filtered : messages;
+}
+
+export default function registerSubagentPromptRuntime(pi: ExtensionAPI): void {
+	const structuredOutputPath = process.env[STRUCTURED_OUTPUT_CAPTURE_ENV];
+	const structuredSchemaPath = process.env[STRUCTURED_OUTPUT_SCHEMA_ENV];
+	if (structuredOutputPath && structuredSchemaPath) {
+		const schema = JSON.parse(fs.readFileSync(structuredSchemaPath, "utf-8")) as JsonSchemaObject;
+		const parameters = {
+			type: "object",
+			properties: { value: schema },
+			required: ["value"],
+			additionalProperties: false,
+		};
+		const registerTool = pi.registerTool as unknown as (tool: {
+			name: string;
+			label: string;
+			description: string;
+			parameters: unknown;
+			execute: (_id: string, params: { value: unknown }) => Promise<unknown>;
+		}) => void;
+		registerTool({
+			name: "structured_output",
+			label: "Structured Output",
+			description: "Submit the required final structured output for this subagent step. This terminates the step.",
+			parameters: parameters as never,
+			async execute(_id: string, params: { value: unknown }) {
+				const validation = validateStructuredOutputValue(schema, params.value);
+				if (validation.status === "invalid") {
+					throw new Error(`Structured output validation failed: ${validation.message}`);
+				}
+				fs.mkdirSync(path.dirname(structuredOutputPath), { recursive: true });
+				fs.writeFileSync(structuredOutputPath, JSON.stringify(params.value), { mode: 0o600 });
+				return {
+					content: [{ type: "text", text: "Structured output captured." }],
+					details: { path: structuredOutputPath },
+					terminate: true,
+				};
+			},
+		});
+	}
+
+	const onRuntimeEvent = pi.on as unknown as (event: string, handler: (event: unknown) => unknown) => void;
+	onRuntimeEvent("context", (event: { messages: unknown[] }) => {
+		const messages = stripParentOnlySubagentMessages(event.messages);
+		if (messages === event.messages) return undefined;
+		return { messages };
+	});
+
+	onRuntimeEvent("before_agent_start", async (event: { systemPrompt: string }) => {
+		const intercomSessionName = process.env[SUBAGENT_INTERCOM_SESSION_NAME_ENV]?.trim();
+		if (intercomSessionName && typeof pi.setSessionName === "function") {
+			pi.setSessionName(intercomSessionName);
+		}
+
+		const inheritProjectContext = readBooleanEnv(SUBAGENT_INHERIT_PROJECT_CONTEXT_ENV);
+		const inheritSkills = readBooleanEnv(SUBAGENT_INHERIT_SKILLS_ENV);
+		const fanoutChild = readBooleanEnv(SUBAGENT_FANOUT_CHILD_ENV);
+		if (inheritProjectContext === undefined && inheritSkills === undefined && fanoutChild === undefined) return;
+		const rewritten = rewriteSubagentPrompt(event.systemPrompt, {
+			inheritProjectContext: inheritProjectContext ?? true,
+			inheritSkills: inheritSkills ?? true,
+			fanoutChild: fanoutChild === true,
+		});
+		if (rewritten === event.systemPrompt) return;
+		return { systemPrompt: rewritten };
+	});
+}

package/src/runs/shared/workflow-graph.ts

@@ -0,0 +1,206 @@
+import { isDynamicParallelStep, isParallelStep, type ChainStep, type SequentialStep } from "../../shared/settings.ts";
+import type { SingleResult, SubagentRunMode, WorkflowGraphNode, WorkflowGraphSnapshot, WorkflowNodeStatus } from "../../shared/types.ts";
+
+export interface WorkflowGraphBuildInput {
+	runId: string;
+	mode?: SubagentRunMode;
+	steps: ChainStep[];
+	results?: Array<Pick<SingleResult, "exitCode" | "detached" | "interrupted" | "error" | "acceptance">>;
+	currentFlatIndex?: number;
+	currentStepIndex?: number;
+	stepStatuses?: Array<{ status?: string; error?: string }>;
+	dynamicChildren?: Record<number, Array<{ agent: string; label?: string; flatIndex: number; itemKey: string; outputName?: string; structured?: boolean; error?: string }>>;
+	dynamicGroupStatuses?: Record<number, { status: WorkflowNodeStatus; error?: string; acceptance?: SingleResult["acceptance"] }>;
+}
+
+function normalizeStatus(status: string | undefined): WorkflowNodeStatus | undefined {
+	switch (status) {
+		case "complete":
+		case "completed":
+			return "completed";
+		case "running":
+			return "running";
+		case "failed":
+			return "failed";
+		case "paused":
+			return "paused";
+		case "detached":
+			return "detached";
+		case "pending":
+			return "pending";
+		default:
+			return undefined;
+	}
+}
+
+function resultStatus(result: Pick<SingleResult, "exitCode" | "detached" | "interrupted"> | undefined): WorkflowNodeStatus | undefined {
+	if (!result) return undefined;
+	if (result.detached) return "detached";
+	if (result.interrupted) return "paused";
+	return result.exitCode === 0 ? "completed" : "failed";
+}
+
+function nodeStatus(input: WorkflowGraphBuildInput, flatIndex: number): WorkflowNodeStatus {
+	return normalizeStatus(input.stepStatuses?.[flatIndex]?.status)
+		?? resultStatus(input.results?.[flatIndex])
+		?? (input.currentFlatIndex === flatIndex ? "running" : "pending");
+}
+
+function pushPhase(phases: WorkflowGraphSnapshot["phases"], phase: string | undefined, nodeId: string): void {
+	if (!phase) return;
+	let group = phases.find((candidate) => candidate.title === phase);
+	if (!group) {
+		group = { title: phase, nodeIds: [] };
+		phases.push(group);
+	}
+	group.nodeIds.push(nodeId);
+}
+
+function seqLabel(step: SequentialStep, stepIndex: number): string {
+	return step.label?.trim() || step.agent || `Step ${stepIndex + 1}`;
+}
+
+function summarizeParallelStatuses(statuses: WorkflowNodeStatus[]): WorkflowNodeStatus {
+	if (statuses.some((status) => status === "running")) return "running";
+	if (statuses.some((status) => status === "failed")) return "failed";
+	if (statuses.some((status) => status === "paused")) return "paused";
+	if (statuses.some((status) => status === "detached")) return "detached";
+	if (statuses.length > 0 && statuses.every((status) => status === "completed")) return "completed";
+	if (statuses.some((status) => status === "completed")) return "running";
+	return "pending";
+}
+
+export function buildWorkflowGraphSnapshot(input: WorkflowGraphBuildInput): WorkflowGraphSnapshot {
+	const nodes: WorkflowGraphNode[] = [];
+	const phases: WorkflowGraphSnapshot["phases"] = [];
+	let flatIndex = 0;
+	let currentNodeId: string | undefined;
+
+	for (let stepIndex = 0; stepIndex < input.steps.length; stepIndex++) {
+		const step = input.steps[stepIndex]!;
+		if (isParallelStep(step)) {
+			const groupId = `step-${stepIndex}`;
+			const children: WorkflowGraphNode[] = [];
+			const childStatuses: WorkflowNodeStatus[] = [];
+			for (let taskIndex = 0; taskIndex < step.parallel.length; taskIndex++) {
+				const task = step.parallel[taskIndex]!;
+				const status = nodeStatus(input, flatIndex);
+				childStatuses.push(status);
+				const childId = `step-${stepIndex}-agent-${taskIndex}`;
+				const child: WorkflowGraphNode = {
+					id: childId,
+					kind: "agent",
+					agent: task.agent,
+					phase: task.phase,
+					label: task.label?.trim() || task.agent || `Agent ${taskIndex + 1}`,
+					status,
+					flatIndex,
+					stepIndex,
+					outputName: task.as,
+					structured: Boolean(task.outputSchema),
+					acceptanceStatus: input.results?.[flatIndex]?.acceptance?.status,
+					error: input.stepStatuses?.[flatIndex]?.error ?? input.results?.[flatIndex]?.error,
+				};
+				children.push(child);
+				pushPhase(phases, task.phase, childId);
+				if (status === "running" || input.currentFlatIndex === flatIndex) currentNodeId = childId;
+				flatIndex++;
+			}
+			const groupStatus = summarizeParallelStatuses(childStatuses);
+			if (input.currentStepIndex === stepIndex && !currentNodeId) currentNodeId = groupId;
+			nodes.push({
+				id: groupId,
+				kind: "parallel-group",
+				label: step.parallel.length === 1 ? "Parallel task" : `Parallel group (${step.parallel.length})`,
+				status: groupStatus,
+				stepIndex,
+				children,
+			});
+			continue;
+		}
+
+		if (isDynamicParallelStep(step)) {
+			const groupId = `step-${stepIndex}`;
+			const materialized = input.dynamicChildren?.[stepIndex] ?? [];
+			const groupOverride = input.dynamicGroupStatuses?.[stepIndex];
+			const children: WorkflowGraphNode[] = [];
+			const childStatuses: WorkflowNodeStatus[] = [];
+			for (let taskIndex = 0; taskIndex < materialized.length; taskIndex++) {
+				const task = materialized[taskIndex]!;
+				const status = nodeStatus(input, task.flatIndex);
+				childStatuses.push(status);
+				const childId = `step-${stepIndex}-item-${task.itemKey.replace(/[^a-zA-Z0-9_-]/g, "-")}`;
+				const child: WorkflowGraphNode = {
+					id: childId,
+					kind: "agent",
+					agent: task.agent,
+					phase: step.parallel.phase ?? step.phase,
+					label: task.label?.trim() || step.parallel.label?.trim() || `${task.agent} ${task.itemKey}`,
+					status,
+					flatIndex: task.flatIndex,
+					stepIndex,
+					itemKey: task.itemKey,
+					outputName: task.outputName,
+					structured: task.structured,
+					acceptanceStatus: input.results?.[task.flatIndex]?.acceptance?.status,
+					error: input.stepStatuses?.[task.flatIndex]?.error ?? input.results?.[task.flatIndex]?.error ?? task.error,
+				};
+				children.push(child);
+				pushPhase(phases, child.phase, childId);
+				if (status === "running" || input.currentFlatIndex === task.flatIndex) currentNodeId = childId;
+			}
+			const groupStatus = groupOverride?.status ?? (children.length > 0 ? summarizeParallelStatuses(childStatuses) : (input.currentStepIndex === stepIndex ? "running" : "pending"));
+			if (input.currentStepIndex === stepIndex && !currentNodeId) currentNodeId = groupId;
+			nodes.push({
+				id: groupId,
+				kind: "dynamic-parallel-group",
+				label: step.label?.trim() || step.parallel.label?.trim() || `Dynamic fanout (${step.collect.as})`,
+				status: groupStatus,
+				stepIndex,
+				outputName: step.collect.as,
+				structured: Boolean(step.collect.outputSchema),
+				acceptanceStatus: groupOverride?.acceptance?.status,
+				error: groupOverride?.error,
+				dynamic: {
+					sourceOutput: step.expand.from.output,
+					sourcePath: step.expand.from.path,
+					itemName: step.expand.item ?? "item",
+					maxItems: step.expand.maxItems,
+					collectAs: step.collect.as,
+				},
+				children,
+			});
+			if (materialized.length > 0) flatIndex = Math.max(flatIndex, ...materialized.map((child) => child.flatIndex + 1));
+			continue;
+		}
+
+		const seq = step as SequentialStep;
+		const status = nodeStatus(input, flatIndex);
+		const id = `step-${stepIndex}`;
+		nodes.push({
+			id,
+			kind: "step",
+			agent: seq.agent,
+			phase: seq.phase,
+			label: seqLabel(seq, stepIndex),
+			status,
+			flatIndex,
+			stepIndex,
+			outputName: seq.as,
+			structured: Boolean(seq.outputSchema),
+			acceptanceStatus: input.results?.[flatIndex]?.acceptance?.status,
+			error: input.stepStatuses?.[flatIndex]?.error ?? input.results?.[flatIndex]?.error,
+		});
+		pushPhase(phases, seq.phase, id);
+		if (status === "running" || input.currentFlatIndex === flatIndex || input.currentStepIndex === stepIndex) currentNodeId = id;
+		flatIndex++;
+	}
+
+	return {
+		runId: input.runId,
+		mode: input.mode ?? "chain",
+		phases,
+		nodes,
+		currentNodeId,
+	};
+}