@giladbarnea/pi-user-agents 0.0.4 → 0.0.6

package/runner.ts

@@ -5,29 +5,43 @@ import {
 	type AgentSessionServices,
 	type Args,
 	buildSessionContext,
+	estimateTokens,
 	createAgentSessionFromServices,
 	createAgentSessionServices,
 	getAgentDir,
 	type ModelRuntime,
 	parseArgs,
 	resolveCliModel,
+	sessionEntryToContextMessages,
 	SessionManager,
 	SettingsManager,
 } from "@earendil-works/pi-coding-agent";
 import { type AgentValueValidator, parseAgentCommand } from "./command-line.js";
+import { conversationFingerprint, mainContextFingerprint } from "./rebase.js";
 import type {
 	AgentCommandName,
 	AgentMessage,
 	AgentResultMessage,
 	AgentSession,
+	ChildCompactionMessage,
 	ExtensionAPI,
 	ExtensionCommandContext,
 	Model,
 	ParsedAgentCommand,
+	RebaseDelivery,
+	RebasedEntryData,
 	RunningAgent,
 	ThinkingLevel,
 } from "./shared.js";
-import { errorMessage, formatModel, formatModelLabel, logSteering } from "./shared.js";
+import type { DispatchRecordData } from "./shared.js";
+import {
+	DISPATCH_ENTRY_TYPE,
+	errorMessage,
+	formatModel,
+	formatModelLabel,
+	logSteering,
+	REBASED_ENTRY_TYPE,
+} from "./shared.js";
 import {
 	buildAgentResultMessage,
 	formatStartNotification,
@@ -78,17 +92,10 @@ async function startUserAgent(
 ): Promise<void> {
 	const parsed = parseAgentCommand(args, command);
 	const parsedForwardedArgs = parseForwardedArgs(parsed.forwardedArgs);
-	const agentDir = getAgentDir();
-	const settingsManager = SettingsManager.create(ctx.cwd, agentDir, {
-		projectTrusted: ctx.isProjectTrusted(),
-	});
-	const services = await createAgentSessionServices({
-		cwd: ctx.cwd,
-		agentDir,
-		settingsManager,
-		modelRuntime: (ctx.modelRegistry as unknown as { runtime: ModelRuntime }).runtime,
-		resourceLoaderOptions: buildChildResourceLoaderOptions(parsedForwardedArgs, ctx.cwd),
-	});
+	const services = await createChildServices(
+		ctx,
+		buildChildResourceLoaderOptions(parsedForwardedArgs, ctx.cwd),
+	);
 	const forwarded = resolveForwardedOptions(parsedForwardedArgs, services.modelRuntime);
 	const selectedModel = forwarded.model ?? ctx.model;
 	if (!selectedModel) throw new Error("No current model is selected; pass -m MODELNAME");
@@ -100,6 +107,12 @@ async function startUserAgent(
 	const childSessionManager = SessionManager.create(services.cwd, undefined, {
 		parentSession: ctx.sessionManager.getSessionFile(),
 	});
+	// The dispatch record: /agent-attach restores tool and resource options from it later.
+	childSessionManager.appendCustomEntry(DISPATCH_ENTRY_TYPE, {
+		forwardedArgs: parsed.forwardedArgs,
+		task: parsed.task,
+		isolate: parsed.isolate,
+	} satisfies DispatchRecordData);
 	const runningAgent = createRunningAgent(
 		nextAgentNumber(),
 		childSessionManager.getSessionId(),
@@ -107,6 +120,7 @@ async function startUserAgent(
 		model,
 		parsed,
 		invocation,
+		conversationFingerprint(inheritedMessages),
 	);
 
 	runningAgents.add(runningAgent);
@@ -146,6 +160,7 @@ function createRunningAgent(
 	model: Model,
 	parsed: ParsedAgentCommand,
 	invocation: string,
+	dispatchBaseFingerprint: string,
 ): RunningAgent {
 	return {
 		id: `user-${sequenceNumber.toString(36)}`,
@@ -156,8 +171,9 @@ function createRunningAgent(
 		modelLabel: formatModelLabel(model),
 		task: parsed.task,
 		invocation,
-		notifyMainAgent: parsed.context,
-		mainContextState: parsed.context ? "will-join" : "separate",
+		notifyMainAgent: parsed.squash,
+		dispatchBaseFingerprint,
+		mainContextState: parsed.squash ? "will-squash" : "separate",
 		status: "starting",
 		startedAt: Date.now(),
 		turnStartedAt: Date.now(),
@@ -187,8 +203,68 @@ function postUserAgentResult(
 	return true;
 }
 
+/**
+ * The rebase mechanism: fast-forward a child conversation onto the session that dispatched it.
+ *
+ * Delivery appends the messages to the main session file, drops a persisted breadcrumb naming
+ * the child session, and asks the host to switch to that same file — the one path that rebuilds
+ * both the live LLM context and the TUI transcript (see INTERNALS.md).
+ */
+export function createRebaseDelivery(
+	pi: ExtensionAPI,
+	getSessionContext: () => ExtensionCommandContext | undefined,
+): RebaseDelivery {
+	return {
+		canDeliver: (agent) => {
+			const ctx = getSessionContext();
+			return (
+				ctx !== undefined &&
+				mainContextFingerprint(ctx.sessionManager) === agent.dispatchBaseFingerprint
+			);
+		},
+		deliver: (agent, messages) => {
+			const ctx = getSessionContext();
+			if (!ctx) throw new Error("No dispatching session context to rebase onto");
+			// The runtime object behind the extension's readonly facade is the writable SessionManager.
+			const sessionManager = ctx.sessionManager as unknown as SessionManager;
+			// A replayed compaction's kept tail may reach back into main's existing context.
+			const priorContextEntryIds = sessionManager
+				.buildContextEntries()
+				.flatMap((entry) => sessionEntryToContextMessages(entry).map(() => entry.id));
+			persistMessages(sessionManager, messages, priorContextEntryIds);
+			pi.appendEntry<RebasedEntryData>(REBASED_ENTRY_TYPE, {
+				sessionId: agent.sessionId,
+				stats: {
+					messageCount: messages.length,
+					tokenEstimate: messages.reduce((total, message) => total + estimateTokens(message), 0),
+					compactionCount: messages.filter((message) => message.role === "compactionSummary")
+						.length,
+				},
+			});
+			const sessionFile = sessionManager.getSessionFile();
+			if (!sessionFile) throw new Error("The main session has no file to switch to");
+			void ctx.switchSession(sessionFile).then(
+				(result) => {
+					if (result.cancelled && ctx.hasUI)
+						ctx.ui.notify(
+							"The rebase is saved to the session file, but the session switch was cancelled — /resume this session to see it",
+							"warning",
+						);
+				},
+				(error) => {
+					if (ctx.hasUI)
+						ctx.ui.notify(
+							`The rebase is saved to the session file, but the session switch failed (${errorMessage(error)}) — /resume this session to see it`,
+							"error",
+						);
+				},
+			);
+		},
+	};
+}
+
 /** Session inputs derived from forwarded pi CLI options, ready to hand to the child session. */
-type ChildResourceLoaderOptions = {
+export type ChildResourceLoaderOptions = {
 	systemPrompt?: string;
 	appendSystemPrompt?: string[];
 	additionalExtensionPaths?: string[];
@@ -201,14 +277,37 @@ type ChildResourceLoaderOptions = {
 	noContextFiles?: boolean;
 };
 
-type ForwardedOptions = {
-	model?: Model;
-	thinkingLevel?: ThinkingLevel;
+export type ChildToolOptions = {
 	tools?: string[];
 	excludeTools?: string[];
 	noTools?: "all" | "builtin";
 };
 
+type ForwardedOptions = ChildToolOptions & {
+	model?: Model;
+	thinkingLevel?: ThinkingLevel;
+};
+
+/** The cwd-bound services for one child session; the parent's ModelRuntime rides along (see INTERNALS). */
+export function createChildServices(
+	ctx: ExtensionCommandContext,
+	resourceLoaderOptions?: ChildResourceLoaderOptions,
+): Promise<AgentSessionServices> {
+	const agentDir = getAgentDir();
+	const settingsManager = SettingsManager.create(ctx.cwd, agentDir, {
+		projectTrusted: ctx.isProjectTrusted(),
+	});
+	return createAgentSessionServices({
+		cwd: ctx.cwd,
+		agentDir,
+		settingsManager,
+		// The extension facade wraps the parent's runtime as a plain property; sharing it keeps
+		// extension-registered providers (their process-global dedup guard) available to the child.
+		modelRuntime: (ctx.modelRegistry as unknown as { runtime: ModelRuntime }).runtime,
+		resourceLoaderOptions,
+	});
+}
+
 /** Parse forwarded pi CLI tokens once before model and session resolution. */
 export function parseForwardedArgs(forwardedArgs: string[]): Args {
 	const parsed = parseArgs(forwardedArgs);
@@ -307,12 +406,17 @@ export function resolveForwardedOptions(
 		}
 	}
 	if (parsed.thinking) forwarded.thinkingLevel = parsed.thinking;
+	return { ...forwarded, ...resolveToolOptions(parsed) };
+}
 
-	if (parsed.noTools) forwarded.noTools = "all";
-	else if (parsed.noBuiltinTools) forwarded.noTools = "builtin";
-	if (parsed.tools) forwarded.tools = parsed.tools;
-	if (parsed.excludeTools) forwarded.excludeTools = parsed.excludeTools;
-	return forwarded;
+/** Map parsed pi tool options onto child-session inputs; dispatch and attach share this. */
+export function resolveToolOptions(parsed: Args): ChildToolOptions {
+	const options: ChildToolOptions = {};
+	if (parsed.noTools) options.noTools = "all";
+	else if (parsed.noBuiltinTools) options.noTools = "builtin";
+	if (parsed.tools) options.tools = parsed.tools;
+	if (parsed.excludeTools) options.excludeTools = parsed.excludeTools;
+	return options;
 }
 
 export function buildChildResourceLoaderOptions(
@@ -374,29 +478,9 @@ async function runAgentLifecycle(
 			widget,
 		);
 		unsubscribe = subscribeToChildSession(session, runningAgent, widget);
-		await runChildTurns(pi, isShuttingDown, task, session, runningAgent, widget);
+		await runChildTurns(pi, isShuttingDown, DISPATCH_PREAMBLE + task, session, runningAgent, widget);
 	} catch (error) {
-		if (runningAgent.aborted) {
-			logSteering(runningAgent.id, "agent-aborted");
-			return;
-		}
-		const message = errorMessage(error);
-		logSteering(runningAgent.id, "agent-error", { error: message });
-		runningAgent.status = "error";
-		runningAgent.completedAt = Date.now();
-		runningAgent.error = message;
-		if (runningAgent.notifyMainAgent && isShuttingDown()) {
-			runningAgent.mainContextState = "separate";
-			return;
-		}
-		const resultMessage = buildAgentResultMessage(
-			runningAgent,
-			{ ok: false, error: message },
-			{ display: runningAgent.notifyMainAgent },
-		);
-		widget.addCompleted(runningAgent, resultMessage, { joinable: !runningAgent.notifyMainAgent });
-		if (postUserAgentResult(pi, isShuttingDown, runningAgent.notifyMainAgent, resultMessage))
-			runningAgent.status = "posted";
+		reportAgentFailure(pi, isShuttingDown, runningAgent, widget, error);
 	} finally {
 		unsubscribe?.();
 		if (session) {
@@ -406,16 +490,51 @@ async function runAgentLifecycle(
 	}
 }
 
-/** Run turns on an initialized child session until it posts, retires, or shuts down. */
+/** Settle a failed lifecycle: record the error and post the error card, unless the user aborted. */
+export function reportAgentFailure(
+	pi: ExtensionAPI,
+	isShuttingDown: () => boolean,
+	runningAgent: RunningAgent,
+	widget: UserAgentWidget,
+	error: unknown,
+): void {
+	if (runningAgent.aborted) {
+		logSteering(runningAgent.id, "agent-aborted");
+		return;
+	}
+	const message = errorMessage(error);
+	logSteering(runningAgent.id, "agent-error", { error: message });
+	runningAgent.status = "error";
+	runningAgent.completedAt = Date.now();
+	runningAgent.error = message;
+	if (runningAgent.notifyMainAgent && isShuttingDown()) {
+		runningAgent.mainContextState = "separate";
+		return;
+	}
+	const resultMessage = buildAgentResultMessage(
+		runningAgent,
+		{ ok: false, error: message },
+		{ display: runningAgent.notifyMainAgent },
+	);
+	widget.addCompleted(runningAgent, resultMessage, { squashable: !runningAgent.notifyMainAgent });
+	if (postUserAgentResult(pi, isShuttingDown, runningAgent.notifyMainAgent, resultMessage))
+		runningAgent.status = "posted";
+}
+
+/** Every dispatch prefixes its task with this; attach recognizes it as the child's dispatch boundary. */
+export const DISPATCH_PREAMBLE =
+	"You are running in an ephemeral, forked background process now, concurrently with the main session. ";
+
+/** Run turns on an initialized child session until it posts, retires, or shuts down. The initial instruction is sent verbatim. */
 export async function runChildTurns(
 	pi: ExtensionAPI,
 	isShuttingDown: () => boolean,
-	task: string,
+	initialInstruction: string,
 	session: AgentSession,
 	runningAgent: RunningAgent,
 	widget: UserAgentWidget,
 ): Promise<void> {
-	let instruction = `You are running in an ephemeral, forked background process now, concurrently with the main session. ${task}`;
+	let instruction = initialInstruction;
 	while (true) {
 		const turnMessageStart = session.agent.state.messages.length;
 		if (!runningAgent.interruptRequested) {
@@ -438,7 +557,7 @@ export async function runChildTurns(
 				{ display: false },
 			);
 			postUserAgentResult(pi, isShuttingDown, false, resultMessage);
-			runningAgent.pendingJoinMessage = runningAgent.notifyMainAgent
+			runningAgent.pendingSquashMessage = runningAgent.notifyMainAgent
 				? undefined
 				: resultMessage;
 			if (isShuttingDown()) return;
@@ -470,13 +589,13 @@ export async function runChildTurns(
 				resultMessage.details.mainContextState = "separate";
 				return;
 			}
-			widget.addCompleted(runningAgent, resultMessage, { joinable: false });
+			widget.addCompleted(runningAgent, resultMessage, { squashable: false });
 			const posted = postUserAgentResult(pi, isShuttingDown, true, resultMessage);
 			runningAgent.status = posted ? "posted" : runningAgent.status;
 			return;
 		}
 		postUserAgentResult(pi, isShuttingDown, false, resultMessage);
-		runningAgent.pendingJoinMessage = resultMessage;
+		runningAgent.pendingSquashMessage = resultMessage;
 		if (isShuttingDown()) return;
 		runningAgent.status = "idle";
 		logSteering(runningAgent.id, "agent-idle", { turnCount: runningAgent.turnCount });
@@ -528,7 +647,7 @@ async function createChildSession(
 	await session.bindExtensions({ mode: "print" });
 	logSteering(runningAgent.id, "session-bound", { streaming: session.isStreaming });
 	if (inheritedMessages.length > 0) {
-		persistInheritedMessages(childSessionManager, inheritedMessages);
+		persistMessages(childSessionManager, inheritedMessages, []);
 		session.agent.state.messages = inheritedMessages;
 		logSteering(runningAgent.id, "context-inherited", { messageCount: inheritedMessages.length });
 	}
@@ -536,27 +655,42 @@ async function createChildSession(
 }
 
 /**
- * Mirror the inherited snapshot into the child's session file so a resumed child sees the
- * same conversation the live one saw. Each role goes through its own append method: the
- * snapshot is already compaction-resolved, so its summary leads the child's own branch.
+ * Mirror a message stream into a session file, each role through its own append method.
+ * Seeds a child with its inherited snapshot (leading summary), and replays a rebase onto the
+ * main session. A compaction with a keptTailCount maps its kept-tail boundary onto this file's
+ * own entry ids, so the rebuilt context keeps the same messages verbatim the child kept — the
+ * tail may reach back into the pre-replay context, whose per-message ids arrive as
+ * priorContextEntryIds (empty when seeding a child).
  */
-export function persistInheritedMessages(
+export function persistMessages(
 	sessionManager: SessionManager,
 	messages: readonly AgentMessage[],
+	priorContextEntryIds: readonly string[],
 ): void {
+	// One id per context message, prior context first; every appended entry projects to one message.
+	const entryIdsPerMessage: string[] = [...priorContextEntryIds];
 	for (const message of messages) {
-		if (message.role === "compactionSummary")
-			sessionManager.appendCompaction(message.summary, "", message.tokensBefore);
-		else if (message.role === "branchSummary")
-			sessionManager.branchWithSummary(sessionManager.getLeafId(), message.summary);
+		if (message.role === "compactionSummary") {
+			const keptTailCount = (message as ChildCompactionMessage).keptTailCount ?? 0;
+			const firstKeptEntryId =
+				keptTailCount > 0 ? (entryIdsPerMessage[entryIdsPerMessage.length - keptTailCount] ?? "") : "";
+			entryIdsPerMessage.push(
+				sessionManager.appendCompaction(message.summary, firstKeptEntryId, message.tokensBefore),
+			);
+		} else if (message.role === "branchSummary")
+			entryIdsPerMessage.push(
+				sessionManager.branchWithSummary(sessionManager.getLeafId(), message.summary),
+			);
 		else if (message.role === "custom")
-			sessionManager.appendCustomMessageEntry(
-				message.customType,
-				message.content,
-				message.display,
-				message.details,
+			entryIdsPerMessage.push(
+				sessionManager.appendCustomMessageEntry(
+					message.customType,
+					message.content,
+					message.display,
+					message.details,
+				),
 			);
-		else sessionManager.appendMessage(message);
+		else entryIdsPerMessage.push(sessionManager.appendMessage(message));
 	}
 }
 
@@ -573,7 +707,7 @@ export function waitForInstruction(
 		};
 		agent.resume = (instruction) => {
 			logSteering(agent.id, "agent-resumed", { instructionLength: instruction.length });
-			agent.mainContextState = agent.notifyMainAgent ? "will-join" : "separate";
+			agent.mainContextState = agent.notifyMainAgent ? "will-squash" : "separate";
 			agent.status = "running";
 			agent.turnStartedAt = Date.now();
 			agent.completedAt = undefined;
@@ -587,12 +721,26 @@ export function waitForInstruction(
 	});
 }
 
-function subscribeToChildSession(
+export function subscribeToChildSession(
 	session: AgentSession,
 	runningAgent: RunningAgent,
 	widget: UserAgentWidget,
 ): () => void {
 	return session.subscribe((event: AgentSessionEvent) => {
+		if (event.type === "compaction_end" && !event.aborted && event.result) {
+			// Compaction never emits message events; capture it here or the rebase stream misses it.
+			// At this instant the live context is exactly [summary, kept tail], so the tail is countable.
+			const keptTailCount = Math.max(0, session.agent.state.messages.length - 1);
+			const compactionMessage: ChildCompactionMessage = {
+				role: "compactionSummary",
+				summary: event.result.summary,
+				tokensBefore: event.result.tokensBefore,
+				timestamp: Date.now(),
+				keptTailCount,
+			};
+			runningAgent.conversationMessages.push(compactionMessage);
+			logSteering(runningAgent.id, "child-compacted", { keptTailCount });
+		}
 		if (event.type === "agent_start")
 			logSteering(runningAgent.id, "child-agent-started", { streaming: session.isStreaming });
 		if (event.type === "agent_end")
@@ -641,7 +789,7 @@ function subscribeToChildSession(
 	});
 }
 
-function assistantText(message: Extract<AgentMessage, { role: "assistant" }>): string {
+export function assistantText(message: Extract<AgentMessage, { role: "assistant" }>): string {
 	return message.content
 		.filter((part) => part.type === "text")
 		.map((part) => part.text)

package/shared.ts

@@ -15,8 +15,8 @@ export type ThinkingLevel = ReturnType<ExtensionAPI["getThinkingLevel"]>;
 export type Theme = ExtensionCommandContext["ui"]["theme"];
 export type UIContext = ExtensionCommandContext["ui"];
 
-export type AgentCommandName = "agent";
-export type MainContextState = "separate" | "will-join" | "joined";
+export type AgentCommandName = "agent" | "agent-attach";
+export type MainContextState = "separate" | "will-squash" | "squashed" | "rebased";
 export type AgentStatus =
 	| "starting"
 	| "running"
@@ -27,7 +27,7 @@ export type AgentStatus =
 
 export type ParsedAgentCommand = {
 	isolate: boolean;
-	context: boolean;
+	squash: boolean;
 	/** Leading pi CLI tokens (minus the extension's own options) to forward to the child, e.g. ["--thinking", "high"]. */
 	forwardedArgs: string[];
 	task: string;
@@ -44,6 +44,39 @@ export type DetachedEntryData = {
 	sessionId: string;
 };
 
+export type AttachedEntryData = {
+	sessionId: string;
+};
+
+/** Persisted in the child session at dispatch: the facts attach cannot recover from messages alone. */
+export type DispatchRecordData = {
+	/** The pi CLI tokens the dispatch forwarded, e.g. ["--tools", "read,grep"]. */
+	forwardedArgs: string[];
+	/** The plain task, without the dispatch preamble. */
+	task: string;
+	isolate: boolean;
+};
+
+export type RebaseStats = {
+	messageCount: number;
+	tokenEstimate: number;
+	compactionCount: number;
+};
+
+export type RebasedEntryData = {
+	sessionId: string;
+	/** Absent on entries persisted before stats existed. */
+	stats?: RebaseStats;
+};
+
+/**
+ * A child compaction captured from its compaction_end event. keptTailCount records how many
+ * messages the compaction kept verbatim, so a rebase replay can restore the exact boundary.
+ */
+export type ChildCompactionMessage = Extract<AgentMessage, { role: "compactionSummary" }> & {
+	keptTailCount?: number;
+};
+
 export type AgentResultMessage = {
 	customType: string;
 	content: string;
@@ -55,7 +88,8 @@ export type AgentCommandDetails = {
 	agentId?: string;
 	command: AgentCommandName;
 	mainContextState?: MainContextState;
-	inheritedContext: boolean;
+	/** Absent on command-error cards: an error has no inherited-or-isolated notion. */
+	inheritedContext?: boolean;
 	model: string;
 	modelLabel: string;
 	task: string;
@@ -78,10 +112,12 @@ export type RunningAgent = {
 	model: string;
 	modelLabel: string;
 	task: string;
-	/** The slash command line as the user typed it, e.g. `/agent -j fix the bug`. */
+	/** The slash command line as the user typed it, e.g. `/agent -s fix the bug`. */
 	invocation: string;
-	/** -j/--join: post invocation+result to the main agent and trigger its turn, instead of waiting quietly for the next user prompt. */
+	/** -s/--squash: post invocation+result to the main agent and trigger its turn, instead of waiting quietly for the next user prompt. */
 	notifyMainAgent: boolean;
+	/** Fingerprint of the main context the child was dispatched from ("[]" for -i); the rebase fast-forward base. */
+	dispatchBaseFingerprint: string;
 	mainContextState: MainContextState;
 	status: AgentStatus;
 	startedAt: number;
@@ -97,8 +133,8 @@ export type RunningAgent = {
 	latestFinalizedMessage?: AgentMessage;
 	error?: string;
 	session?: AgentSession;
-	/** Latest completed turn's result message, deliverable to the main agent via the overlay's join action. */
-	pendingJoinMessage?: AgentResultMessage;
+	/** Latest completed turn's result message, deliverable to the main agent via the overlay's squash action. */
+	pendingSquashMessage?: AgentResultMessage;
 	/** Starts another turn on an idle agent; assigned while the lifecycle loop awaits the next instruction. */
 	resume?: (instruction: string) => void;
 	/** Ends an idle agent's lifecycle without another turn; assigned alongside resume. */
@@ -116,8 +152,9 @@ export type CompletedAgent = {
 	command: AgentCommandName;
 	modelLabel: string;
 	task: string;
+	dispatchBaseFingerprint: string;
 	mainContextState: MainContextState;
-	pendingJoinMessage?: AgentResultMessage;
+	pendingSquashMessage?: AgentResultMessage;
 	ok: boolean;
 	responseText: string;
 	messages: AgentMessage[];
@@ -129,8 +166,63 @@ export type CompletedAgent = {
 	contextPercent?: number;
 };
 
+/** The mechanism that fast-forwards a child conversation onto the main session, injected into the widget. */
+export type RebaseDelivery = {
+	/** Whether the main session's live context still equals this agent's dispatch base. */
+	canDeliver(agent: RunningAgent | CompletedAgent): boolean;
+	/** Append the processed child conversation onto the main session. */
+	deliver(agent: RunningAgent | CompletedAgent, messages: AgentMessage[]): void;
+};
+
+/** How long a confirmation stays armed, and how long a transient footer notice shows. */
+export const CONFIRMATION_WINDOW_MS = 4000;
+
+/**
+ * A two-press confirmation with a timeout: the first press arms it, a repeat press on the
+ * same target within the window confirms. A press on another target re-arms, cancel or
+ * expiry disarms. One instance serves every confirmable action on a surface.
+ */
+export class TimedConfirmation<Target> {
+	private armedOn: Target | undefined;
+	private expiry: ReturnType<typeof setTimeout> | undefined;
+
+	constructor(
+		private readonly onExpire: () => void,
+		private readonly windowMs: number = CONFIRMATION_WINDOW_MS,
+	) {}
+
+	/** One press of the action key: true means confirmed, false means armed and waiting. */
+	press(target: Target): boolean {
+		if (this.armedOn === target) {
+			this.cancel();
+			return true;
+		}
+		this.armedOn = target;
+		if (this.expiry) clearTimeout(this.expiry);
+		this.expiry = setTimeout(() => {
+			this.armedOn = undefined;
+			this.expiry = undefined;
+			this.onExpire();
+		}, this.windowMs);
+		return false;
+	}
+
+	cancel(): void {
+		this.armedOn = undefined;
+		if (this.expiry) clearTimeout(this.expiry);
+		this.expiry = undefined;
+	}
+
+	isArmedOn(target: Target): boolean {
+		return this.armedOn !== undefined && this.armedOn === target;
+	}
+}
+
 export const MESSAGE_TYPE = "pi-user-agents";
 export const DETACHED_ENTRY_TYPE = "pi-user-agents-detached";
+export const REBASED_ENTRY_TYPE = "pi-user-agents-rebased";
+export const ATTACHED_ENTRY_TYPE = "pi-user-agents-attached";
+export const DISPATCH_ENTRY_TYPE = "pi-user-agents-dispatch";
 export const WIDGET_KEY = "pi-user-agents";
 export const STEERING_LOG_PATH = `${tmpdir()}/pi-user-agents-steer.log`;
 
@@ -170,8 +262,9 @@ export function contextLabel(inheritedContext: boolean): string {
 }
 
 export function mainContextLabel(state: MainContextState | undefined): string | undefined {
-	if (state === "will-join") return "will join context";
-	if (state === "joined") return "joined context";
+	if (state === "will-squash") return "will squash into context";
+	if (state === "squashed") return "squashed messages into context";
+	if (state === "rebased") return "rebased into context";
 	return undefined;
 }