@giladbarnea/pi-user-agents 0.0.4 → 0.0.6

package/transcript.ts

@@ -6,13 +6,16 @@ import type {
 	AgentEntryData,
 	AgentMessage,
 	AgentResultMessage,
+	AttachedEntryData,
 	DetachedEntryData,
 	ExtensionAPI,
 	ExtensionCommandContext,
+	RebasedEntryData,
 	RunningAgent,
 	Theme,
 } from "./shared.js";
 import {
+	ATTACHED_ENTRY_TYPE,
 	contextLabel,
 	DETACHED_ENTRY_TYPE,
 	escapeAttribute,
@@ -22,6 +25,7 @@ import {
 	formatTurns,
 	mainContextLabel,
 	MESSAGE_TYPE,
+	REBASED_ENTRY_TYPE,
 	truncatePlain,
 } from "./shared.js";
 
@@ -42,7 +46,6 @@ export function reportCommandError(
 		display: true,
 		details: {
 			command,
-			inheritedContext: command === "agent",
 			model: "",
 			modelLabel: "",
 			task: args.trim(),
@@ -53,14 +56,14 @@ export function reportCommandError(
 }
 
 /**
- * Select the joined role messages from a child transcript.
+ * Select the squashed role messages from a child transcript.
  *
- * @example selectJoinedMessages([{ role: "user", content: "hello", timestamp: 0 }]).length // 1
+ * @example selectSquashedMessages([{ role: "user", content: "hello", timestamp: 0 }]).length // 1
  */
-type JoinedMessage = Extract<AgentMessage, { role: "user" | "assistant" }>;
+type SquashedMessage = Extract<AgentMessage, { role: "user" | "assistant" }>;
 
-export function selectJoinedMessages(messages: readonly AgentMessage[]): JoinedMessage[] {
-	const selected: JoinedMessage[] = [];
+export function selectSquashedMessages(messages: readonly AgentMessage[]): SquashedMessage[] {
+	const selected: SquashedMessage[] = [];
 	let latestAssistant: Extract<AgentMessage, { role: "assistant" }> | undefined;
 	for (const message of messages) {
 		if (message.role === "assistant" && roleMessageText(message).trim()) {
@@ -110,13 +113,13 @@ export function buildAgentResultMessage(
 	};
 }
 
-const JOINED_CONVERSATION_PREFACE =
+const SQUASHED_CONVERSATION_PREFACE =
 	"The user has dispatched a background sub-agent with a task. The sub-agent is done. The following is the back and forth between them:";
 
 export function formatResultMessage(agent: RunningAgent): string {
-	const messages = selectJoinedMessages(agent.conversationMessages);
+	const messages = selectSquashedMessages(agent.conversationMessages);
 	const lines = [
-		JOINED_CONVERSATION_PREFACE,
+		SQUASHED_CONVERSATION_PREFACE,
 		`<user_agent model="${escapeAttribute(agent.model)}" inherited_context="${agent.inheritedContext}">`,
 	];
 	let firstUserMessage = true;
@@ -135,7 +138,7 @@ export function formatResultMessage(agent: RunningAgent): string {
 	return lines.join("\n");
 }
 
-function roleMessageText(message: JoinedMessage): string {
+function roleMessageText(message: SquashedMessage): string {
 	if (typeof message.content === "string") return message.content;
 	return message.content
 		.flatMap((part) => (part.type === "text" ? [part.text] : []))
@@ -214,11 +217,42 @@ export function registerUserAgentRenderer(pi: ExtensionAPI): void {
 					)
 				: undefined,
 	);
+	// Pi's chat container does not pad renderer components; its own chat lines self-pad with paddingX 1.
 	pi.registerEntryRenderer<DetachedEntryData>(DETACHED_ENTRY_TYPE, (entry, _state, theme) =>
 		entry.data
-			? new Text(theme.fg("dim", `Detached session ${entry.data.sessionId}`), 0, 0)
+			? new Text(theme.fg("dim", `Detached session ${entry.data.sessionId}`), 1, 0)
 			: undefined,
 	);
+	pi.registerEntryRenderer<AttachedEntryData>(ATTACHED_ENTRY_TYPE, (entry, _state, theme) =>
+		entry.data
+			? new Text(theme.fg("dim", `Attached session ${entry.data.sessionId}`), 1, 0)
+			: undefined,
+	);
+	pi.registerEntryRenderer<RebasedEntryData>(REBASED_ENTRY_TYPE, (entry, _state, theme) =>
+		entry.data ? new Text(theme.fg("dim", formatRebasedLine(entry.data)), 1, 0) : undefined,
+	);
+}
+
+function formatRebasedLine(data: RebasedEntryData): string {
+	const base = `Rebased session ${data.sessionId} into this conversation`;
+	if (!data.stats) return base;
+	const parts = [
+		`added ${data.stats.messageCount} message${data.stats.messageCount === 1 ? "" : "s"}`,
+		`~${formatTokenCount(data.stats.tokenEstimate)} tokens`,
+	];
+	if (data.stats.compactionCount > 0)
+		parts.push(
+			`${data.stats.compactionCount} compaction event${data.stats.compactionCount === 1 ? "" : "s"}`,
+		);
+	return `${base} (${parts.join(", ")})`;
+}
+
+/**
+ * @example formatTokenCount(135_264) // "135K"
+ * @example formatTokenCount(12) // "12"
+ */
+function formatTokenCount(tokens: number): string {
+	return tokens < 1000 ? `${tokens}` : `${Math.round(tokens / 1000)}K`;
 }
 
 function liveAgentCard(
@@ -251,7 +285,7 @@ function renderAgentCard(
 	if (parts.length > 0) text += ` ${theme.fg("dim", "·")} ${theme.fg("dim", parts.join(" · "))}`;
 	if (task) text += `\n  ${theme.fg("dim", `task: ${truncatePlain(task, 88)}`)}`;
 	const card = new Container();
-	card.addChild(new Text(text, 0, 0));
+	card.addChild(new Text(text, 1, 0));
 	card.addChild(body);
 	return card;
 }
@@ -270,12 +304,12 @@ function renderAgentCardBody(
 		"No output.";
 	if (expanded)
 		return ok
-			? new Markdown(output, 2, 0, getMarkdownTheme())
-			: new Text(theme.fg("error", output), 2, 0);
+			? new Markdown(output, 3, 0, getMarkdownTheme())
+			: new Text(theme.fg("error", output), 3, 0);
 
 	const previewText = `⎿  ${truncatePlain(details?.responsePreview ?? output, 110)}`;
-	if (!ok) return new Text(`  ${theme.fg("error", previewText)}`, 0, 0);
-	return new Markdown(previewText, 2, 0, getMarkdownTheme(), {
+	if (!ok) return new Text(`  ${theme.fg("error", previewText)}`, 1, 0);
+	return new Markdown(previewText, 3, 0, getMarkdownTheme(), {
 		color: (value) => theme.fg("dim", value),
 	});
 }
@@ -296,7 +330,7 @@ function buildRendererParts(content: string, details: AgentCommandDetails | unde
 	const parts: string[] = [];
 	if (details?.modelLabel) parts.push(details.modelLabel);
 	else if (details?.model) parts.push(details.model);
-	if (details) parts.push(contextLabel(details.inheritedContext));
+	if (details?.inheritedContext !== undefined) parts.push(contextLabel(details.inheritedContext));
 	const mainContext = mainContextLabel(details?.mainContextState);
 	if (mainContext) parts.push(mainContext);
 	if (details?.turnCount) parts.push(formatTurns(details.turnCount));

package/widget.ts

@@ -15,9 +15,11 @@ import {
 	type ViewableAgent,
 } from "./agent-viewer.js";
 import { renderAgentContextMeter } from "./context-meter.js";
+import { selectRebaseMessages } from "./rebase.js";
 import type {
 	AgentResultMessage,
 	CompletedAgent,
+	RebaseDelivery,
 	RunningAgent,
 	Theme,
 	UIContext,
@@ -25,6 +27,7 @@ import type {
 import {
 	contextLabel,
 	formatMs,
+	TimedConfirmation,
 	formatToolUses,
 	formatTurns,
 	logSteering,
@@ -49,13 +52,15 @@ export class UserAgentWidget {
 	private active = false;
 	private selectedIndex = 0;
 	private viewerOpen = false;
-	private detachConfirmationTarget: ViewableAgent | undefined;
+	/** One confirmation slot for detaching the selected agent. */
+	private readonly detachConfirmation = new TimedConfirmation<ViewableAgent>(() => this.update());
 	private readonly completedAgents: CompletedAgent[] = [];
 
 	constructor(
 		private readonly runningAgents: Set<RunningAgent>,
-		private readonly sendJoinedResult: (message: AgentResultMessage) => void,
+		private readonly sendSquashedResult: (message: AgentResultMessage) => void,
 		private readonly announceDetached: (sessionId: string) => void,
+		private readonly rebaseDelivery: RebaseDelivery,
 	) {}
 
 	setUI(ui: UIContext): void {
@@ -64,14 +69,14 @@ export class UserAgentWidget {
 		this.ui = ui;
 		this.widgetRegistered = false;
 		this.tui = undefined;
-		this.detachConfirmationTarget = undefined;
+		this.detachConfirmation.cancel();
 		this.inputUnsub = ui.onTerminalInput((data) => this.handleKey(data));
 	}
 
 	addCompleted(
 		agent: RunningAgent,
 		resultMessage: AgentResultMessage,
-		options: { joinable: boolean },
+		options: { squashable: boolean },
 	): void {
 		this.completedAgents.push({
 			id: agent.id,
@@ -79,8 +84,9 @@ export class UserAgentWidget {
 			command: agent.command,
 			modelLabel: agent.modelLabel,
 			task: agent.task,
+			dispatchBaseFingerprint: agent.dispatchBaseFingerprint,
 			mainContextState: agent.mainContextState,
-			pendingJoinMessage: options.joinable ? resultMessage : undefined,
+			pendingSquashMessage: options.squashable ? resultMessage : undefined,
 			ok: resultMessage.details.ok,
 			responseText: agent.responseText,
 			messages: structuredClone(transcriptMessages(agent)),
@@ -142,7 +148,7 @@ export class UserAgentWidget {
 		this.widgetRegistered = false;
 		this.tui = undefined;
 		this.active = false;
-		this.detachConfirmationTarget = undefined;
+		this.detachConfirmation.cancel();
 	}
 
 	private clear(): void {
@@ -157,7 +163,7 @@ export class UserAgentWidget {
 		}
 		this.active = false;
 		this.selectedIndex = 0;
-		this.detachConfirmationTarget = undefined;
+		this.detachConfirmation.cancel();
 	}
 
 	private runningAgentsForWidget(): RunningAgent[] {
@@ -201,13 +207,13 @@ export class UserAgentWidget {
 		}
 
 		if (matchesKey(data, "down")) {
-			this.detachConfirmationTarget = undefined;
+			this.detachConfirmation.cancel();
 			this.selectedIndex = Math.min(entries.length - 1, this.selectedIndex + 1);
 			this.update();
 			return { consume: true };
 		}
 		if (matchesKey(data, "up")) {
-			this.detachConfirmationTarget = undefined;
+			this.detachConfirmation.cancel();
 			if (this.selectedIndex === 0) {
 				this.deactivate();
 				return { consume: true };
@@ -227,12 +233,12 @@ export class UserAgentWidget {
 			selected?.kind === "running" &&
 			isLiveAgent(selected.agent)
 		) {
-			this.detachConfirmationTarget = undefined;
+			this.detachConfirmation.cancel();
 			this.interruptRunning(selected.agent);
 			return { consume: true };
 		}
 		if (matchesKey(data, "enter")) {
-			this.detachConfirmationTarget = undefined;
+			this.detachConfirmation.cancel();
 			if (selected) this.openSelected(selected.agent);
 			return { consume: true };
 		}
@@ -248,22 +254,33 @@ export class UserAgentWidget {
 	}
 
 	private confirmDetach(target: ViewableAgent): boolean {
-		if (this.detachConfirmationTarget === target) {
-			this.detachConfirmationTarget = undefined;
-			return true;
-		}
-		this.detachConfirmationTarget = target;
-		this.update();
-		return false;
+		const confirmed = this.detachConfirmation.press(target);
+		if (!confirmed) this.update();
+		return confirmed;
 	}
 
 	private deactivate(): void {
 		this.active = false;
 		this.selectedIndex = 0;
-		this.detachConfirmationTarget = undefined;
+		this.detachConfirmation.cancel();
 		this.update();
 	}
 
+	/** Session ids with a live row or completed card, matched by exact id or prefix. */
+	matchAttachedSessionIds(query: string): string[] {
+		return this.entries()
+			.map((entry) => entry.agent.sessionId)
+			.filter((sessionId) => sessionId.startsWith(query));
+	}
+
+	/** Open the overlay of the row holding this session, as if the user selected it. */
+	openViewer(sessionId: string): boolean {
+		const entry = this.entries().find((candidate) => candidate.agent.sessionId === sessionId);
+		if (!entry) return false;
+		this.openSelected(entry.agent);
+		return true;
+	}
+
 	private clampSelection(): void {
 		const total = this.entries().length;
 		if (total === 0) {
@@ -286,8 +303,12 @@ export class UserAgentWidget {
 						agent,
 						theme,
 						done,
-						() => this.canJoinMainContext(agent.id),
-						() => this.joinMainContext(agent.id),
+						() => this.canSquashMainContext(agent.id),
+						() => this.squashMainContext(agent.id),
+						() => this.canRebaseMainContext(agent.id),
+						() => this.rebaseMainContext(agent.id),
+						() => this.rebaseBlockReason(agent.id),
+						() => this.rebaseDetachCount(agent.id),
 						() => {
 							if (isRunningAgent(agent)) this.interruptRunning(agent);
 						},
@@ -319,54 +340,107 @@ export class UserAgentWidget {
 		return [...this.runningAgents].find((agent) => agent.id === agentId && agent.status === "idle");
 	}
 
-	private canJoinMainContext(agentId: string): boolean {
-		if (this.idleAgent(agentId)?.pendingJoinMessage) return true;
-		return this.completedAgents.some(
-			(agent) => agent.id === agentId && agent.pendingJoinMessage !== undefined,
-		);
+	private canSquashMainContext(agentId: string): boolean {
+		return this.deliverableAgent(agentId) !== undefined;
 	}
 
-	confirmMainContextJoin(agentId: string): boolean {
+	confirmMainContextSquash(agentId: string): boolean {
 		const agents = [
 			...[...this.runningAgents].filter((candidate) => candidate.id === agentId),
 			...this.completedAgents.filter((candidate) => candidate.id === agentId),
-		].filter((candidate) => candidate.mainContextState === "will-join");
+		].filter((candidate) => candidate.mainContextState === "will-squash");
 		if (agents.length === 0) return false;
-		for (const agent of agents) agent.mainContextState = "joined";
+		for (const agent of agents) agent.mainContextState = "squashed";
 		this.update();
 		return true;
 	}
 
-	private joinMainContext(agentId: string): void {
+	private squashMainContext(agentId: string): void {
 		const idleAgent = this.idleAgent(agentId);
 		if (idleAgent) {
-			this.joinIdleAgent(idleAgent);
+			this.squashIdleAgent(idleAgent);
 			return;
 		}
 		const completedAgent = this.completedAgents.find((agent) => agent.id === agentId);
-		const message = completedAgent?.pendingJoinMessage;
+		const message = completedAgent?.pendingSquashMessage;
 		if (!completedAgent || !message) return;
 
-		completedAgent.pendingJoinMessage = undefined;
-		completedAgent.mainContextState = "will-join";
-		message.details.mainContextState = "will-join";
-		this.sendJoinedResult(message);
-		logSteering(agentId, "joined-main-context", { display: message.display });
+		completedAgent.pendingSquashMessage = undefined;
+		completedAgent.mainContextState = "will-squash";
+		message.details.mainContextState = "will-squash";
+		this.sendSquashedResult(message);
+		logSteering(agentId, "squashed-main-context", { display: message.display });
+		this.update();
+	}
+
+	/** An agent whose latest turn result is still deliverable: idle live, or completed and squashable. */
+	private deliverableAgent(agentId: string): RunningAgent | CompletedAgent | undefined {
+		const idleAgent = this.idleAgent(agentId);
+		if (idleAgent?.pendingSquashMessage) return idleAgent;
+		return this.completedAgents.find(
+			(agent) => agent.id === agentId && agent.pendingSquashMessage !== undefined,
+		);
+	}
+
+	private canRebaseMainContext(agentId: string): boolean {
+		return this.deliverableAgent(agentId) !== undefined && this.rebaseBlockReason(agentId) === undefined;
+	}
+
+	/** How many other agent sessions the rebase cascade would detach. */
+	private rebaseDetachCount(agentId: string): number {
+		return (
+			[...this.runningAgents].filter((agent) => agent.id !== agentId).length +
+			this.completedAgents.filter((agent) => agent.id !== agentId).length
+		);
+	}
+
+	/** Why a deliverable result cannot rebase right now; undefined when it can, or when nothing is deliverable. */
+	private rebaseBlockReason(agentId: string): string | undefined {
+		const agent = this.deliverableAgent(agentId);
+		if (!agent) return undefined;
+		// The session switch tears the extension runtime down, which would abort a mid-turn sibling.
+		if ([...this.runningAgents].some(isLiveAgent)) return "another agent is mid-turn";
+		if (!this.rebaseDelivery.canDeliver(agent))
+			return "the main session has drifted since dispatch";
+		return undefined;
+	}
+
+	/** Fast-forward the raw child conversation onto the main session, then retire the agent. */
+	private rebaseMainContext(agentId: string): void {
+		if (!this.canRebaseMainContext(agentId)) return;
+		const agent = this.deliverableAgent(agentId);
+		if (!agent) return;
+		// The reload will not preserve live rows; give every sibling its resumable-session breadcrumb.
+		for (const other of [...this.runningAgents]) if (other !== agent) this.closeRunning(other);
+		for (const other of [...this.completedAgents]) if (other !== agent) this.removeCompleted(other);
+		const message = agent.pendingSquashMessage;
+		this.rebaseDelivery.deliver(
+			agent,
+			selectRebaseMessages(agent.task, structuredClone(transcriptMessages(agent))),
+		);
+		agent.pendingSquashMessage = undefined;
+		agent.mainContextState = "rebased";
+		if (isRunningAgent(agent)) {
+			agent.status = "posted";
+			if (message) this.addCompleted(agent, message, { squashable: false });
+			agent.retire?.();
+		}
+		logSteering(agentId, "rebased-main-context");
 		this.update();
 	}
 
-	/** Joining is one of the two terminal actions on a live idle agent: deliver the result, snapshot it, retire the session. */
-	private joinIdleAgent(agent: RunningAgent): void {
-		const message = agent.pendingJoinMessage;
+	/** Squashing is one of the two terminal actions on a live idle agent: deliver the result, snapshot it, retire the session. */
+	private squashIdleAgent(agent: RunningAgent): void {
+		const message = agent.pendingSquashMessage;
 		if (!message) return;
-		agent.pendingJoinMessage = undefined;
-		agent.mainContextState = "will-join";
-		message.details.mainContextState = "will-join";
+		agent.pendingSquashMessage = undefined;
+		agent.mainContextState = "will-squash";
+		message.details.mainContextState = "will-squash";
 		agent.status = "posted";
-		this.addCompleted(agent, message, { joinable: false });
-		this.sendJoinedResult(message);
+		this.addCompleted(agent, message, { squashable: false });
+		this.sendSquashedResult(message);
 		agent.retire?.();
-		logSteering(agent.id, "joined-main-context", { display: message.display, live: true });
+		logSteering(agent.id, "squashed-main-context", { display: message.display, live: true });
 		this.update();
 	}
 
@@ -418,7 +492,8 @@ export class UserAgentWidget {
 			lines.push(truncateToWidth(`  ${theme.fg("dim", "←/↓ select agent")}`, width));
 		} else {
 			const selected = entries[this.selectedIndex];
-			const confirmingDetach = this.detachConfirmationTarget === selected?.agent;
+			const confirmingDetach =
+				selected !== undefined && this.detachConfirmation.isArmedOn(selected.agent);
 			const segments = [theme.fg("accent", "↑↓ select"), theme.fg("accent", "Enter view")];
 			if (
 				!confirmingDetach &&