@giladbarnea/pi-user-agents 0.0.4 → 0.0.6

package/attach.ts

@@ -0,0 +1,339 @@
+import { existsSync, readdirSync } from "node:fs";
+import { join } from "node:path";
+import { createAgentSessionFromServices, SessionManager } from "@earendil-works/pi-coding-agent";
+import { conversationFingerprint } from "./rebase.js";
+import {
+	assistantText,
+	buildChildResourceLoaderOptions,
+	createChildServices,
+	DISPATCH_PREAMBLE,
+	parseForwardedArgs,
+	reportAgentFailure,
+	resolveToolOptions,
+	runChildTurns,
+	subscribeToChildSession,
+	waitForInstruction,
+} from "./runner.js";
+import type {
+	AgentMessage,
+	AgentSession,
+	AttachedEntryData,
+	DispatchRecordData,
+	ExtensionAPI,
+	ExtensionCommandContext,
+	Model,
+	RunningAgent,
+} from "./shared.js";
+import {
+	ATTACHED_ENTRY_TYPE,
+	DISPATCH_ENTRY_TYPE,
+	errorMessage,
+	formatModel,
+	formatModelLabel,
+	logSteering,
+} from "./shared.js";
+import { buildAgentResultMessage, reportCommandError } from "./transcript.js";
+import type { UserAgentWidget } from "./widget.js";
+
+/** A dispatch base that can never fingerprint-match a live context, so rebase stays withheld. */
+export const UNKNOWN_DISPATCH_BASE = "unknown-dispatch-base";
+
+export type DispatchBoundary = {
+	/** The main-session context the child inherited at dispatch — the rebase fast-forward base. */
+	base: AgentMessage[];
+	/** The child's own conversation, starting at the preambled dispatch message. */
+	conversation: AgentMessage[];
+	/** The dispatch task, preamble stripped. */
+	task: string;
+};
+
+/**
+ * Split a child's resolved context at its dispatch boundary — the first user message carrying
+ * the dispatch preamble. Undefined when no message carries it (a /fork child, or a child whose
+ * own compaction summarized the boundary away).
+ */
+export function splitAtDispatchBoundary(
+	messages: readonly AgentMessage[],
+): DispatchBoundary | undefined {
+	for (const [index, message] of messages.entries()) {
+		if (message.role !== "user") continue;
+		const text = userMessageText(message);
+		if (!text.startsWith(DISPATCH_PREAMBLE)) continue;
+		return {
+			base: messages.slice(0, index),
+			conversation: messages.slice(index),
+			task: text.slice(DISPATCH_PREAMBLE.length),
+		};
+	}
+	return undefined;
+}
+
+function userMessageText(message: Extract<AgentMessage, { role: "user" }>): string {
+	if (typeof message.content === "string") return message.content;
+	return message.content
+		.flatMap((part) => (part.type === "text" ? [part.text] : []))
+		.join("\n");
+}
+
+export type ChildSessionFileMatch = { sessionId: string; file: string };
+
+/**
+ * The session files in a directory whose id matches the query exactly or by prefix.
+ * Pi names session files `<timestamp>_<session-id>.jsonl`, so matching reads names only.
+ */
+export function matchChildSessionFiles(
+	query: string,
+	sessionDirectory: string,
+): ChildSessionFileMatch[] {
+	const fileNames = existsSync(sessionDirectory) ? readdirSync(sessionDirectory) : [];
+	return fileNames.flatMap((name) => {
+		const sessionId = sessionIdFromFileName(name);
+		if (sessionId === undefined || !sessionId.startsWith(query)) return [];
+		return [{ sessionId, file: join(sessionDirectory, name) }];
+	});
+}
+
+function sessionIdFromFileName(fileName: string): string | undefined {
+	if (!fileName.endsWith(".jsonl")) return undefined;
+	const separator = fileName.indexOf("_");
+	if (separator === -1) return undefined;
+	return fileName.slice(separator + 1, -".jsonl".length);
+}
+
+/** The dispatch record persisted in the child's file, or undefined for pre-record and /fork children. */
+export function readDispatchRecord(
+	childSessionManager: SessionManager,
+): DispatchRecordData | undefined {
+	for (const entry of childSessionManager.getEntries()) {
+		if (entry.type === "custom" && entry.customType === DISPATCH_ENTRY_TYPE)
+			return entry.data as DispatchRecordData | undefined;
+	}
+	return undefined;
+}
+
+/**
+ * Rebuild an idle, steerable RunningAgent from a reopened child session and its resolved model.
+ * The child's own conversation, task, latest response, and rebase base all come from the file;
+ * the dispatch record refines the task and context flag, and with no recognizable dispatch
+ * boundary the whole context attaches and rebase stays withheld.
+ */
+export function buildAttachedAgent(
+	sequenceNumber: number,
+	childSessionManager: SessionManager,
+	model: Model,
+	invocation: string,
+	record?: DispatchRecordData,
+): RunningAgent {
+	const contextMessages = childSessionManager.buildSessionContext().messages;
+	const boundary = splitAtDispatchBoundary(contextMessages);
+	const conversationMessages = structuredClone(boundary?.conversation ?? contextMessages);
+	const firstUserMessage = contextMessages.find((message) => message.role === "user");
+	const lastAssistantMessage = conversationMessages.findLast(
+		(message): message is Extract<AgentMessage, { role: "assistant" }> =>
+			message.role === "assistant",
+	);
+	const responseText = lastAssistantMessage ? assistantText(lastAssistantMessage).trim() : "";
+	const now = Date.now();
+	const agent: RunningAgent = {
+		id: `user-${sequenceNumber.toString(36)}`,
+		sessionId: childSessionManager.getSessionId(),
+		command: "agent",
+		inheritedContext: record
+			? !record.isolate
+			: boundary
+				? boundary.base.length > 0
+				: true,
+		model: formatModel(model),
+		modelLabel: formatModelLabel(model),
+		task:
+			record?.task ??
+			boundary?.task ??
+			(firstUserMessage === undefined ? "(no messages)" : userMessageText(firstUserMessage)),
+		invocation,
+		notifyMainAgent: false,
+		dispatchBaseFingerprint: boundary
+			? conversationFingerprint(boundary.base)
+			: UNKNOWN_DISPATCH_BASE,
+		mainContextState: "separate",
+		status: "idle",
+		startedAt: now,
+		turnStartedAt: now,
+		completedAt: now,
+		activeTools: new Map(),
+		toolUses: 0,
+		turnCount: 0,
+		responseText,
+		conversationMessages,
+		finished: Promise.resolve(),
+	};
+	if (responseText)
+		agent.pendingSquashMessage = buildAgentResultMessage(
+			agent,
+			{ ok: true, response: responseText },
+			{ display: false },
+		);
+	return agent;
+}
+
+/**
+ * The attached agent's lifecycle: bind the child's extensions, park until the user steers or
+ * retires it, then run turns exactly like a dispatched child. Every failure — the bind's
+ * included — settles through the shared error path, and the end shuts the child's extensions
+ * down.
+ */
+export async function runAttachedTurns(
+	pi: ExtensionAPI,
+	isShuttingDown: () => boolean,
+	session: AgentSession,
+	runningAgent: RunningAgent,
+	widget: UserAgentWidget,
+): Promise<void> {
+	let unsubscribe: (() => void) | undefined;
+	try {
+		await session.bindExtensions({ mode: "print" });
+		unsubscribe = subscribeToChildSession(session, runningAgent, widget);
+		// The user can detach during the extension bind, before this lifecycle parks.
+		if (runningAgent.aborted) return;
+		const instruction = await waitForInstruction(runningAgent, widget);
+		if (instruction !== undefined)
+			await runChildTurns(pi, isShuttingDown, instruction, session, runningAgent, widget);
+	} catch (error) {
+		reportAgentFailure(pi, isShuttingDown, runningAgent, widget, error);
+	} finally {
+		unsubscribe?.();
+		logSteering(runningAgent.id, "session-shutdown", { streaming: session.isStreaming });
+		await session.extensionRunner.emit({ type: "session_shutdown", reason: "quit" });
+	}
+}
+
+export async function handleAttachCommand(
+	pi: ExtensionAPI,
+	runningAgents: Set<RunningAgent>,
+	widget: UserAgentWidget,
+	isShuttingDown: () => boolean,
+	nextAgentNumber: () => number,
+	args: string,
+	ctx: ExtensionCommandContext,
+): Promise<void> {
+	try {
+		await attachUserAgent(pi, runningAgents, widget, isShuttingDown, nextAgentNumber, args, ctx);
+	} catch (error) {
+		reportCommandError(pi, "agent-attach", args, ctx, errorMessage(error));
+	}
+}
+
+/**
+ * Reattach a detached child session as a live, parked agent row. Idempotent for a session
+ * that already has a row; errors for anything that is not a child session of this session.
+ * Both success states end by offering to open the agent's overlay.
+ */
+async function attachUserAgent(
+	pi: ExtensionAPI,
+	runningAgents: Set<RunningAgent>,
+	widget: UserAgentWidget,
+	isShuttingDown: () => boolean,
+	nextAgentNumber: () => number,
+	args: string,
+	ctx: ExtensionCommandContext,
+): Promise<void> {
+	const query = args.trim();
+	if (!query || /\s/.test(query)) throw new Error("Usage: /agent-attach <child-session-id>");
+	if (ctx.hasUI) widget.setUI(ctx.ui);
+
+	// One resolution over the whole universe — attached rows and files on disk together — so a
+	// prefix unique among rows cannot mask an ambiguity with a detached session, and a live row
+	// whose file is not written yet still resolves. Dispatch creates children in pi's default
+	// session directory for the cwd; that is where the files are.
+	const attachedIds = widget.matchAttachedSessionIds(query);
+	const childSessionDirectory = SessionManager.create(ctx.cwd).getSessionDir();
+	const fileMatches = matchChildSessionFiles(query, childSessionDirectory);
+	const matchedSessionIds = [
+		...new Set([...attachedIds, ...fileMatches.map((match) => match.sessionId)]),
+	];
+	if (matchedSessionIds.length === 0) throw new Error(`No agent session matches "${query}"`);
+	if (matchedSessionIds.length > 1)
+		throw new Error(`"${query}" matches ${matchedSessionIds.length} sessions; use more characters`);
+	const sessionId = matchedSessionIds[0]!;
+
+	if (attachedIds.includes(sessionId)) {
+		await confirmViewAgent(ctx, widget, sessionId, "Agent already attached");
+		return;
+	}
+
+	const mainSessionFile = ctx.sessionManager.getSessionFile();
+	if (!mainSessionFile)
+		throw new Error("This session has no file; only a persisted session can attach agents");
+	const childSessionFile = fileMatches.find((match) => match.sessionId === sessionId)!.file;
+	const childSessionManager = SessionManager.open(childSessionFile);
+	if (childSessionManager.getHeader()?.parentSession !== mainSessionFile)
+		throw new Error(
+			`Session ${childSessionManager.getSessionId()} is not a child session of this session`,
+		);
+
+	// The record replays through the same projections dispatch uses; a record-less child (an old
+	// dispatch, or a /fork) restores pi defaults.
+	const record = readDispatchRecord(childSessionManager);
+	const dispatchArgs = parseForwardedArgs(record?.forwardedArgs ?? []);
+	const services = await createChildServices(
+		ctx,
+		buildChildResourceLoaderOptions(dispatchArgs, ctx.cwd),
+	);
+	// A non-empty session manager takes pi's resume path: messages, model, and thinking level
+	// are restored from the child's own file; tools and resources come from the dispatch record.
+	const { session, modelFallbackMessage } = await createAgentSessionFromServices({
+		services,
+		sessionManager: childSessionManager,
+		...resolveToolOptions(dispatchArgs),
+	});
+	const model = session.agent.state.model as Model | undefined;
+	if (!model) {
+		session.dispose();
+		throw new Error(modelFallbackMessage ?? "No model is available for the attached session");
+	}
+	if (modelFallbackMessage && ctx.hasUI) ctx.ui.notify(modelFallbackMessage, "warning");
+
+	const runningAgent = buildAttachedAgent(
+		nextAgentNumber(),
+		childSessionManager,
+		model,
+		`/agent-attach ${query}`,
+		record,
+	);
+	runningAgent.session = session;
+	runningAgents.add(runningAgent);
+	logSteering(runningAgent.id, "agent-attached", { sessionId: runningAgent.sessionId });
+	if (ctx.hasUI) {
+		widget.ensureTimer();
+		widget.update();
+	}
+	// No awaits between the add above and this assignment: every later failure settles inside
+	// the lifecycle, so the row can never outlive its own removal hook.
+	runningAgent.finished = runAttachedTurns(
+		pi,
+		isShuttingDown,
+		session,
+		runningAgent,
+		widget,
+	).finally(() => {
+		logSteering(runningAgent.id, "agent-disposed", { status: runningAgent.status });
+		runningAgent.session?.dispose();
+		runningAgents.delete(runningAgent);
+		widget.update();
+	});
+	// The detach breadcrumb's counterpart: after a reload the transcript still tells the truth.
+	pi.appendEntry<AttachedEntryData>(ATTACHED_ENTRY_TYPE, { sessionId: runningAgent.sessionId });
+	if (runningAgent.aborted) return;
+	await confirmViewAgent(ctx, widget, runningAgent.sessionId, "Agent attached");
+}
+
+/** The follow-up both success states share: offer to open the agent's overlay. */
+async function confirmViewAgent(
+	ctx: ExtensionCommandContext,
+	widget: UserAgentWidget,
+	sessionId: string,
+	title: string,
+): Promise<void> {
+	if (!ctx.hasUI) return;
+	const view = await ctx.ui.confirm(title, "View agent?");
+	if (view) widget.openViewer(sessionId);
+}

package/command-line.ts

@@ -57,7 +57,7 @@ export type AgentSemanticToken = {
 
 export type AgentArgumentScan = {
 	isolate: boolean;
-	context: boolean;
+	squash: boolean;
 	forwardedArgs: string[];
 	/** Ordered, non-overlapping semantic tokens over the argument text. Prose stays untokenized. */
 	tokens: AgentSemanticToken[];
@@ -115,12 +115,12 @@ export const AGENT_OPTIONS: readonly AgentOptionDefinition[] = [
 		description: "Start without parent-session context",
 	},
 	{
-		semanticId: "context",
-		names: ["-j", "--join"],
+		semanticId: "squash",
+		names: ["-s", "--squash"],
 		role: "extension",
 		arity: "boolean",
 		autocomplete: true,
-		description: "Join the completed result into the main agent context",
+		description: "Squash the completed result into the main agent context",
 	},
 	{
 		semanticId: "provider",
@@ -360,7 +360,7 @@ export function scanAgentArguments(
 	const tokens: AgentSemanticToken[] = [];
 	const forwardedArgs: string[] = [];
 	let isolate = false;
-	let context = false;
+	let squash = false;
 	let blocked: AgentSemanticToken | undefined;
 	let consumedOption = false;
 	let position = 0;
@@ -393,8 +393,8 @@ export function scanAgentArguments(
 			}
 		} else if (option.semanticId === "isolate") {
 			isolate = true;
-		} else if (option.semanticId === "context") {
-			context = true;
+		} else if (option.semanticId === "squash") {
+			squash = true;
 		} else if (option.arity === "value") {
 			forwardedArgs.push(option.forwardName ?? token);
 			const read = readValueSpan(args, position);
@@ -427,7 +427,7 @@ export function scanAgentArguments(
 	}
 	const scan: AgentArgumentScan = {
 		isolate,
-		context,
+		squash,
 		forwardedArgs,
 		tokens,
 		proseStart,
@@ -487,11 +487,11 @@ export function parseAgentCommand(args: string, command: AgentCommandName): Pars
 	}
 	if (!scan.prose.trim())
 		throw new Error(
-			`Usage: /${command} [pi options] [-m MODELNAME] [-i|--isolate] [-j|--join] "<task>"`,
+			`Usage: /${command} [pi options] [-m MODELNAME] [-i|--isolate] [-s|--squash] "<task>"`,
 		);
 	return {
 		isolate: scan.isolate,
-		context: scan.context,
+		squash: scan.squash,
 		forwardedArgs: scan.forwardedArgs,
 		task: scan.task,
 		warnings: scan.warnings,

package/index.ts

@@ -2,12 +2,14 @@ import type {
 	MessageEndEvent,
 	MessageEndEventResult,
 } from "@earendil-works/pi-coding-agent";
+import { handleAttachCommand } from "./attach.js";
 import { registerAgentAutocomplete } from "./autocomplete.js";
-import { handleAgentCommand } from "./runner.js";
+import { createRebaseDelivery, handleAgentCommand } from "./runner.js";
 import type {
 	AgentCommandDetails,
 	DetachedEntryData,
 	ExtensionAPI,
+	ExtensionCommandContext,
 	RunningAgent,
 } from "./shared.js";
 import { DETACHED_ENTRY_TYPE, MESSAGE_TYPE } from "./shared.js";
@@ -16,10 +18,13 @@ import { UserAgentWidget } from "./widget.js";
 
 export default function userAgent(pi: ExtensionAPI): void {
 	const runningAgents = new Set<RunningAgent>();
+	// The session every agent was dispatched from; rebase fast-forwards onto it.
+	let mainSessionContext: ExtensionCommandContext | undefined;
 	const widget = new UserAgentWidget(
 		runningAgents,
 		(message) => pi.sendMessage(message, { triggerTurn: true }),
 		(sessionId) => pi.appendEntry<DetachedEntryData>(DETACHED_ENTRY_TYPE, { sessionId }),
+		createRebaseDelivery(pi, () => mainSessionContext),
 	);
 	let shuttingDown = false;
 	let nextAgentNumber = 0;
@@ -27,12 +32,13 @@ export default function userAgent(pi: ExtensionAPI): void {
 	registerUserAgentRenderer(pi);
 	registerAgentAutocomplete(pi);
 
-	pi.on("message_end", (event) => confirmMainContextJoin(event, widget));
+	pi.on("message_end", (event) => confirmMainContextSquash(event, widget));
 
 	pi.registerCommand("agent", {
 		description:
 			"Run a background SDK agent with current session context; forwards pi CLI options like --thinking (-i/--isolate for none)",
 		handler: async (args, ctx) => {
+			mainSessionContext = ctx;
 			await handleAgentCommand(
 				pi,
 				runningAgents,
@@ -47,6 +53,22 @@ export default function userAgent(pi: ExtensionAPI): void {
 		},
 	});
 
+	pi.registerCommand("agent-attach", {
+		description: "Reattach a detached agent session by id, back into the widget",
+		handler: async (args, ctx) => {
+			mainSessionContext = ctx;
+			await handleAttachCommand(
+				pi,
+				runningAgents,
+				widget,
+				() => shuttingDown,
+				() => ++nextAgentNumber,
+				args,
+				ctx,
+			);
+		},
+	});
+
 	pi.on("session_shutdown", async () => {
 		shuttingDown = true;
 		widget.dispose();
@@ -63,19 +85,19 @@ export default function userAgent(pi: ExtensionAPI): void {
 	});
 }
 
-export function confirmMainContextJoin(
+export function confirmMainContextSquash(
 	event: MessageEndEvent,
 	widget: UserAgentWidget,
 ): MessageEndEventResult | undefined {
 	if (event.message.role !== "custom" || event.message.customType !== MESSAGE_TYPE) return;
 	const details = event.message.details as AgentCommandDetails | undefined;
-	if (!details?.agentId || details.mainContextState !== "will-join") return;
-	widget.confirmMainContextJoin(details.agentId);
-	details.mainContextState = "joined";
+	if (!details?.agentId || details.mainContextState !== "will-squash") return;
+	widget.confirmMainContextSquash(details.agentId);
+	details.mainContextState = "squashed";
 	return {
 		message: {
 			...event.message,
-			details: { ...details, mainContextState: "joined" },
+			details: { ...details, mainContextState: "squashed" },
 		},
 	};
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@giladbarnea/pi-user-agents",
-  "version": "0.0.4",
+  "version": "0.0.6",
   "description": "Run, view, and control background Pi agents.",
   "keywords": [
     "pi-package",

package/rebase.ts

@@ -0,0 +1,67 @@
+import { buildSessionContext } from "@earendil-works/pi-coding-agent";
+import type { AgentMessage, ExtensionCommandContext } from "./shared.js";
+import { MESSAGE_TYPE } from "./shared.js";
+
+type MainSessionManager = ExtensionCommandContext["sessionManager"];
+
+/**
+ * Fingerprint a context snapshot for the rebase fast-forward precondition: the main session's
+ * live context must still fingerprint-equal the base the child was dispatched from.
+ *
+ * Compares content, not file bookkeeping: rebuilding custom, compaction, and branch summaries
+ * from a session file re-stamps their timestamps (and a branch summary's fromId), so a base
+ * round-tripped through a child's file would never equal main's live context otherwise.
+ *
+ * @example conversationFingerprint([]) === conversationFingerprint([]) // true
+ */
+export function conversationFingerprint(messages: readonly AgentMessage[]): string {
+	return JSON.stringify(messages.map(fingerprintedMessage));
+}
+
+function fingerprintedMessage(message: AgentMessage): Omit<AgentMessage, "timestamp"> {
+	const { timestamp: _timestamp, ...content } = message;
+	if (content.role !== "branchSummary") return content;
+	const { fromId: _fromId, ...branchContent } = content;
+	return branchContent;
+}
+
+/** Rebuilding a session-sized context string is only worth doing when the branch tip moved. */
+const contextFingerprintCache = new WeakMap<object, { leafId: string | null; fingerprint: string }>();
+
+/** Fingerprint of the main session's current LLM context, cached by branch leaf. */
+export function mainContextFingerprint(sessionManager: MainSessionManager): string {
+	const leafId = sessionManager.getLeafId();
+	const cached = contextFingerprintCache.get(sessionManager);
+	if (cached && cached.leafId === leafId) return cached.fingerprint;
+	const fingerprint = conversationFingerprint(
+		buildSessionContext(sessionManager.getBranch()).messages,
+	);
+	contextFingerprintCache.set(sessionManager, { leafId, fingerprint });
+	return fingerprint;
+}
+
+/**
+ * Select the child conversation messages to rebase onto the main session, so the record reads
+ * as if the user prompted the main session directly. The child's history passes through
+ * verbatim — compactions included, so main inherits the child's context boundary — except the
+ * first user message, which carries the extension's internal background-process prefix and is
+ * replaced by the plain task, and this extension's own custom messages, which are the one
+ * trace of the user-agents machinery.
+ */
+export function selectRebaseMessages(
+	task: string,
+	messages: readonly AgentMessage[],
+): AgentMessage[] {
+	let firstUserMessage = true;
+	const selected: AgentMessage[] = [];
+	for (const message of messages) {
+		if (message.role === "custom" && message.customType === MESSAGE_TYPE) continue;
+		if (message.role === "user" && firstUserMessage) {
+			firstUserMessage = false;
+			selected.push({ ...message, content: task });
+			continue;
+		}
+		selected.push(message);
+	}
+	return selected;
+}