@tmustier/pi-agent-teams 0.5.3 → 0.5.4

package/CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## 0.5.4
+
+### Fixes
+
+- **Clean branch-mode worker sessions** — teammates spawned with `contextMode=branch` now branch from the last stable pre-turn boundary, replay the active user request, preserve compaction entries, and materialize user-only fallbacks. This prevents workers from inheriting the leader's in-progress tool-use turn or reopening as blank sessions. (#35, #36)
+- **Accurate quality-gate completion notices** — leader completion and batch-complete messages now mention running quality gates only when hooks are actually enabled, instead of inferring that from the presence of an internal enqueue callback. (#37)
+
 ## 0.5.3
 
 ### Fixes

package/extensions/teams/activity-tracker.ts

@@ -15,134 +15,6 @@ export type TranscriptEntry =
 	| { kind: "turn_end"; turnNumber: number; tokens: number; timestamp: number };
 
 const MAX_TRANSCRIPT = 200;
-const MAX_SUMMARY_LENGTH = 120;
-
-// ── Tool content summarization ──
-
-function truncateSummary(text: string): string {
-	if (text.length <= MAX_SUMMARY_LENGTH) return text;
-	return `${text.slice(0, MAX_SUMMARY_LENGTH - 1)}…`;
-}
-
-function summarizeToolArgs(toolName: string, args: unknown): string {
-	if (!isRecord(args)) return "";
-	const key = toolName.toLowerCase();
-
-	if (key === "read" || key === "write") {
-		const path = typeof args.path === "string" ? args.path : null;
-		if (!path) return "";
-		return truncateSummary(path);
-	}
-
-	if (key === "edit") {
-		const path = typeof args.path === "string" ? args.path : null;
-		if (!path) return "";
-		return truncateSummary(path);
-	}
-
-	if (key === "bash") {
-		const cmd = typeof args.command === "string" ? args.command : null;
-		if (!cmd) return "";
-		return truncateSummary(cmd.replace(/\s+/g, " ").trim());
-	}
-
-	if (key === "grep" || key === "glob") {
-		const pattern = typeof args.pattern === "string" ? args.pattern : null;
-		const path = typeof args.path === "string" ? args.path : null;
-		const parts: string[] = [];
-		if (pattern) parts.push(pattern);
-		if (path) parts.push(`in ${path}`);
-		return truncateSummary(parts.join(" "));
-	}
-
-	if (key === "webfetch" || key === "websearch") {
-		const url = typeof args.url === "string" ? args.url : null;
-		const query = typeof args.query === "string" ? args.query : null;
-		return truncateSummary(url ?? query ?? "");
-	}
-
-	if (key === "team_message") {
-		const recipient = typeof args.recipient === "string" ? args.recipient : null;
-		const message = typeof args.message === "string" ? args.message : null;
-		if (recipient && message) return truncateSummary(`→ ${recipient}: ${message.replace(/\s+/g, " ").trim()}`);
-		if (message) return truncateSummary(message.replace(/\s+/g, " ").trim());
-		return recipient ? truncateSummary(`→ ${recipient}`) : "";
-	}
-
-	if (key === "task" || key === "teams") {
-		const action = typeof args.action === "string" ? args.action : null;
-		return action ? truncateSummary(action) : "";
-	}
-
-	// Fallback: try to find a meaningful first-string arg
-	for (const v of Object.values(args)) {
-		if (typeof v === "string" && v.length > 0) return truncateSummary(v);
-	}
-	return "";
-}
-
-/**
- * Extract the first text content from a ToolResultMessage-shaped result.
- * The agent-loop emits `{ content: [{type: "text", text: "..."}], ... }`.
- */
-function extractContentText(result: Record<string, unknown>): string | null {
-	const content = result.content;
-	if (!Array.isArray(content)) return null;
-	for (const item of content) {
-		if (isRecord(item) && item.type === "text" && typeof item.text === "string") {
-			return item.text;
-		}
-	}
-	return null;
-}
-
-function summarizeToolResult(toolName: string, result: unknown, isError: boolean): string {
-	if (isError) {
-		if (typeof result === "string") return truncateSummary(result.replace(/\s+/g, " ").trim());
-		if (isRecord(result)) {
-			// Try content array first (ToolResultMessage shape)
-			const contentText = extractContentText(result);
-			if (contentText !== null) {
-				const trimmed = contentText.replace(/\s+/g, " ").trim();
-				return trimmed.length > 0 ? truncateSummary(trimmed) : "error";
-			}
-			const msg = typeof result.message === "string"
-				? result.message
-				: typeof result.error === "string"
-					? result.error
-					: null;
-			if (msg) return truncateSummary(msg.replace(/\s+/g, " ").trim());
-		}
-		return "error";
-	}
-
-	if (typeof result === "string") {
-		const compact = result.replace(/\s+/g, " ").trim();
-		if (compact.length === 0) return "(empty)";
-		return truncateSummary(compact);
-	}
-
-	if (isRecord(result)) {
-		// Try content array first (ToolResultMessage shape from agent-loop)
-		const contentText = extractContentText(result);
-		if (contentText !== null) {
-			const compact = contentText.replace(/\s+/g, " ").trim();
-			if (compact.length === 0) return "(empty)";
-			return truncateSummary(compact);
-		}
-		// Fallback: check for common structured result shapes
-		const status = typeof result.status === "string" ? result.status : null;
-		if (status) return truncateSummary(status);
-		const output = typeof result.output === "string" ? result.output : null;
-		if (output) return truncateSummary(output.replace(/\s+/g, " ").trim());
-	}
-
-	if (Array.isArray(result)) {
-		return `${String(result.length)} items`;
-	}
-
-	return "";
-}
 
 export class TranscriptLog {
 	private entries: TranscriptEntry[] = [];

package/extensions/teams/leader-inbox.ts

@@ -91,12 +91,14 @@ export async function pollLeaderInbox(opts: {
 	style: TeamsStyle;
 	pendingPlanApprovals: Map<string, { requestId: string; name: string; taskId?: string }>;
 	enqueueHook?: (invocation: TeamsHookInvocation) => void;
+	hooksEnabled?: boolean;
 	sendLeaderLlmMessage?: SendLeaderLlmMessage;
 	/** Batch delegation tracker for all-tasks-complete auto-notify. */
 	delegationTracker?: DelegationTracker;
 }): Promise<void> {
-	const { ctx, teamId, teamDir, taskListId, leadName, style, pendingPlanApprovals, enqueueHook, sendLeaderLlmMessage, delegationTracker } = opts;
+	const { ctx, teamId, teamDir, taskListId, leadName, style, pendingPlanApprovals, enqueueHook, hooksEnabled, sendLeaderLlmMessage, delegationTracker } = opts;
 	const strings = getTeamsStrings(style);
+	const hooksActive = hooksEnabled ?? Boolean(enqueueHook);
 
 	let msgs: Awaited<ReturnType<typeof popUnreadMessages>>;
 	try {
@@ -173,38 +175,42 @@ export async function pollLeaderInbox(opts: {
 			const name = sanitizeName(idle.from);
 
 			// Hook: always emit "idle" (best-effort, non-blocking)
-			try {
-				enqueueHook?.({
-					event: "idle",
-					teamId,
-					teamDir,
-					taskListId,
-					style,
-					memberName: name,
-					timestamp: idle.timestamp,
-					completedTask: null,
-				});
-			} catch {
-				// ignore hook enqueue errors
-			}
-
-			// Hook: task completion / failure
-			if (idle.completedTaskId) {
-				const completedTask = await getTask(teamDir, taskListId, idle.completedTaskId);
+			if (hooksActive) {
 				try {
 					enqueueHook?.({
-						event: idle.completedStatus === "failed" ? "task_failed" : "task_completed",
+						event: "idle",
 						teamId,
 						teamDir,
 						taskListId,
 						style,
 						memberName: name,
 						timestamp: idle.timestamp,
-						completedTask,
+						completedTask: null,
 					});
 				} catch {
 					// ignore hook enqueue errors
 				}
+			}
+
+			// Hook: task completion / failure
+			if (idle.completedTaskId) {
+				const completedTask = await getTask(teamDir, taskListId, idle.completedTaskId);
+				if (hooksActive) {
+					try {
+						enqueueHook?.({
+							event: idle.completedStatus === "failed" ? "task_failed" : "task_completed",
+							teamId,
+							teamDir,
+							taskListId,
+							style,
+							memberName: name,
+							timestamp: idle.timestamp,
+							completedTask,
+						});
+					} catch {
+						// ignore hook enqueue errors
+					}
+				}
 
 				// Event-driven batch tracking: mark this task done and
 				// collect any batches that became fully complete.
@@ -319,7 +325,7 @@ export async function pollLeaderInbox(opts: {
 
 						if (allDone) {
 							lines.push("");
-							if (enqueueHook) {
+							if (hooksActive) {
 								// Hooks run asynchronously and may reopen tasks or create follow-ups.
 								lines.push(`All ${totalTasks} task(s) show completed — quality gates are still running and may change task states.`);
 							} else {
@@ -354,7 +360,7 @@ export async function pollLeaderInbox(opts: {
 	if (sendLeaderLlmMessage) {
 		for (const batch of batchCompletions) {
 			const taskRefs = batch.taskIds.map((id) => `#${id}`).join(", ");
-			const suffix = enqueueHook
+			const suffix = hooksActive
 				? "Quality gates are still running and may change task states."
 				: "Review the results and continue.";
 			const msg = `[Team] All delegated tasks completed (${taskRefs}). ${suffix}`;

package/extensions/teams/leader.ts

@@ -22,6 +22,7 @@ import { getTeamsStyleFromEnv, type TeamsStyle, formatMemberDisplayName, getTeam
 import { DelegationTracker, pollLeaderInbox as pollLeaderInboxImpl } from "./leader-inbox.js";
 import {
 	getHookBaseName,
+	areTeamsHooksEnabled,
 	getTeamsHookFailureAction,
 	getTeamsHookFollowupOwnerPolicy,
 	getTeamsHookMaxReopensPerTask,
@@ -34,6 +35,7 @@ import {
 import { handleTeamCommand } from "./leader-team-command.js";
 import { registerTeamsTool } from "./leader-teams-tool.js";
 import { getParentSessionId, shouldSilenceInheritedParentAttachClaimWarning } from "./session-parent.js";
+import { branchSelectionNote, ensureSessionFileMaterialized, resolveBranchLeafSelection } from "./session-branching.js";
 import type { ContextMode, SpawnTeammateFn, SpawnTeammateResult, WorkspaceMode } from "./spawn-types.js";
 
 function getTeamsExtensionEntryPath(): string | null {
@@ -90,12 +92,17 @@ async function createSessionForTeammate(
 
 	try {
 		const sm = SessionManager.open(parentSessionFile, teamSessionsDir);
-		const branched = sm.createBranchedSession(leafId);
+		const selection = resolveBranchLeafSelection(sm.getBranch(leafId), leafId);
+		const branched = sm.createBranchedSession(selection.leafId);
 		if (!branched) {
 			const fallback = SessionManager.create(ctx.cwd, teamSessionsDir);
 			return { sessionFile: fallback.getSessionFile(), note: "branch(failed->fresh)", warnings };
 		}
-		return { sessionFile: branched, note: "branch", warnings };
+		if (selection.replayUserMessage) {
+			sm.appendMessage(JSON.parse(JSON.stringify(selection.replayUserMessage)) as Parameters<typeof sm.appendMessage>[0]);
+		}
+		await ensureSessionFileMaterialized(sm, branched);
+		return { sessionFile: branched, note: branchSelectionNote(selection), warnings };
 	} catch (err) {
 		const msg = err instanceof Error ? err.message : String(err);
 		if (/Entry .* not found/i.test(msg)) {
@@ -702,6 +709,7 @@ export function runLeader(pi: ExtensionAPI): void {
 			style,
 			pendingPlanApprovals,
 			enqueueHook,
+			hooksEnabled: areTeamsHooksEnabled(process.env),
 			sendLeaderLlmMessage: (content, options) => {
 				pi.sendUserMessage(content, options);
 			},

package/extensions/teams/session-branching.ts

@@ -0,0 +1,133 @@
+import * as fs from "node:fs";
+import * as path from "node:path";
+import type { SessionEntry, SessionManager } from "@mariozechner/pi-coding-agent";
+
+export type BranchLeafSelection = {
+	leafId: string;
+	adjusted: boolean;
+	reason: "requested" | "clean-turn-replay" | "clean-turn-stable" | "clean-turn-user";
+	replayUserMessage?: UserMessageLike;
+};
+
+type MessageLike = Record<string, unknown> & { role: string };
+type UserMessageLike = MessageLike & { role: "user"; content: unknown; timestamp: number };
+
+type MessageEntryLike = SessionEntry & {
+	type: "message";
+	message: MessageLike;
+};
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null;
+}
+
+function isMessageEntry(entry: SessionEntry): entry is MessageEntryLike {
+	if (entry.type !== "message") return false;
+	return isRecord(entry.message) && typeof entry.message.role === "string";
+}
+
+function isUserMessageEntry(entry: SessionEntry): entry is MessageEntryLike & { message: UserMessageLike } {
+	return isMessageEntry(entry) && entry.message.role === "user";
+}
+
+function isAssistantToolUseEntry(entry: SessionEntry): entry is MessageEntryLike & { message: MessageLike & { role: "assistant"; stopReason: "toolUse" } } {
+	return (
+		isMessageEntry(entry) &&
+		entry.message.role === "assistant" &&
+		typeof entry.message.stopReason === "string" &&
+		entry.message.stopReason === "toolUse"
+	);
+}
+
+function isStableAssistantEntry(entry: SessionEntry): entry is MessageEntryLike & { message: MessageLike & { role: "assistant" } } {
+	return (
+		isMessageEntry(entry) &&
+		entry.message.role === "assistant" &&
+		typeof entry.message.stopReason === "string" &&
+		entry.message.stopReason !== "toolUse"
+	);
+}
+
+/**
+ * When the leader is mid-turn, the current leaf may point into an unfinished
+ * assistant/tool-use path. Branching from that leaf causes workers to inherit
+ * the leader's in-progress turn instead of a clean conversation context.
+ *
+ * In that case, branch from the last stable turn boundary instead:
+ * - Prefer the latest completed assistant message (persists as a real branched file)
+ * - Otherwise fall back to the latest user message
+ */
+export function resolveBranchLeafSelection(path: SessionEntry[], requestedLeafId: string): BranchLeafSelection {
+	const lastAssistantIndex = [...path].map((entry, index) => ({ entry, index }))
+		.reverse()
+		.find(({ entry }) => isMessageEntry(entry) && entry.message.role === "assistant")?.index;
+
+	if (lastAssistantIndex === undefined) {
+		return { leafId: requestedLeafId, adjusted: false, reason: "requested" };
+	}
+
+	const lastAssistant = path[lastAssistantIndex];
+	if (!lastAssistant || !isAssistantToolUseEntry(lastAssistant)) {
+		return { leafId: requestedLeafId, adjusted: false, reason: "requested" };
+	}
+
+	let latestUserIndex = -1;
+	let latestUserBeforeActiveTurn: UserMessageLike | undefined;
+	for (let i = lastAssistantIndex - 1; i >= 0; i -= 1) {
+		const candidate = path[i];
+		if (!candidate) continue;
+		if (isUserMessageEntry(candidate)) {
+			latestUserIndex = i;
+			latestUserBeforeActiveTurn = candidate.message;
+			break;
+		}
+	}
+
+	if (latestUserIndex > 0 && latestUserBeforeActiveTurn) {
+		const boundary = path[latestUserIndex - 1];
+		if (boundary) {
+			return {
+				leafId: boundary.id,
+				adjusted: boundary.id !== requestedLeafId,
+				reason: "clean-turn-replay",
+				replayUserMessage: latestUserBeforeActiveTurn,
+			};
+		}
+	}
+
+	for (let i = lastAssistantIndex - 1; i >= 0; i -= 1) {
+		const candidate = path[i];
+		if (!candidate) continue;
+		if (isStableAssistantEntry(candidate)) {
+			return {
+				leafId: candidate.id,
+				adjusted: candidate.id !== requestedLeafId,
+				reason: "clean-turn-stable",
+			};
+		}
+		if (isUserMessageEntry(candidate)) {
+			return { leafId: candidate.id, adjusted: candidate.id !== requestedLeafId, reason: "clean-turn-user" };
+		}
+	}
+
+	return { leafId: requestedLeafId, adjusted: false, reason: "requested" };
+}
+
+export async function ensureSessionFileMaterialized(
+	sm: Pick<SessionManager, "getHeader" | "getEntries">,
+	sessionFile: string,
+): Promise<void> {
+	if (fs.existsSync(sessionFile)) return;
+	const header = sm.getHeader();
+	if (!header) return;
+	await fs.promises.mkdir(path.dirname(sessionFile), { recursive: true });
+	const lines = [JSON.stringify(header), ...sm.getEntries().map((entry) => JSON.stringify(entry))].join("\n") + "\n";
+	await fs.promises.writeFile(sessionFile, lines, "utf8");
+}
+
+export function branchSelectionNote(selection: BranchLeafSelection): string {
+	if (!selection.adjusted) return "branch";
+	if (selection.reason === "clean-turn-replay" || selection.reason === "clean-turn-stable") return "branch(clean-turn)";
+	if (selection.reason === "clean-turn-user") return "branch(clean-turn:user-fallback)";
+	return "branch";
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tmustier/pi-agent-teams",
-  "version": "0.5.3",
+  "version": "0.5.4",
   "description": "Claude Code agent teams style workflow for Pi.",
   "license": "MIT",
   "author": "Thomas Mustier",
@@ -35,7 +35,8 @@
     "integration-spawn-overrides-test": "tsx scripts/integration-spawn-overrides-test.mts",
     "integration-hooks-remediation-test": "tsx scripts/integration-hooks-remediation-test.mts",
     "integration-todo-test": "tsx scripts/integration-todo-test.mts",
-    "integration-cleanup-test": "tsx scripts/integration-cleanup-test.mts"
+    "integration-cleanup-test": "tsx scripts/integration-cleanup-test.mts",
+    "integration-branch-context-test": "tsx scripts/integration-branch-context-test.mts"
   },
   "devDependencies": {
     "@eslint/js": "^9.39.2",

package/scripts/integration-branch-context-test.mts

@@ -0,0 +1,299 @@
+/**
+ * Integration test: branch-mode worker sessions should strip the leader's
+ * in-progress tool-use turn before starting delegated work.
+ *
+ * What this covers:
+ * - Prepare a persisted parent session whose leaf ends inside an unfinished
+ *   assistant/tool-use turn (user -> assistant toolUse -> toolResult)
+ * - Derive a branch session using the same clean-turn selection logic used by
+ *   teammate spawning
+ * - Start a real worker process from that branched session in worktree mode
+ *   context conditions (git repo + persisted session)
+ * - Deliver an assigned task and verify the worker starts and completes it
+ *
+ * Usage:
+ *   npx tsx scripts/integration-branch-context-test.mts
+ *   npx tsx scripts/integration-branch-context-test.mts --timeoutSec 120
+ */
+
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+import { spawn, spawnSync, type ChildProcess } from "node:child_process";
+import { fileURLToPath } from "node:url";
+
+import { SessionManager } from "@mariozechner/pi-coding-agent";
+import type { AssistantMessage } from "@mariozechner/pi-ai";
+import { writeToMailbox } from "../extensions/teams/mailbox.js";
+import { taskAssignmentPayload } from "../extensions/teams/protocol.js";
+import { branchSelectionNote, resolveBranchLeafSelection } from "../extensions/teams/session-branching.js";
+import { createTask, getTask } from "../extensions/teams/task-store.js";
+import { ensureTeamConfig, loadTeamConfig } from "../extensions/teams/team-config.js";
+import { sleep, terminateAll } from "./lib/pi-workers.js";
+
+function parseArgs(argv: readonly string[]): { timeoutSec: number } {
+	let timeoutSec = 120;
+	for (let i = 0; i < argv.length; i += 1) {
+		if (argv[i] === "--timeoutSec") {
+			const v = argv[i + 1];
+			if (v) timeoutSec = Number.parseInt(v, 10);
+			i += 1;
+		}
+	}
+	if (!Number.isFinite(timeoutSec) || timeoutSec < 30) timeoutSec = 120;
+	return { timeoutSec };
+}
+
+function assert(condition: boolean, message: string): void {
+	if (!condition) throw new Error(message);
+}
+
+async function waitFor(
+	fn: () => boolean | Promise<boolean>,
+	opts: { timeoutMs: number; pollMs: number; label: string },
+): Promise<void> {
+	const { timeoutMs, pollMs, label } = opts;
+	const deadline = Date.now() + timeoutMs;
+	while (Date.now() < deadline) {
+		if (await fn()) return;
+		await sleep(pollMs);
+	}
+	throw new Error(`Timeout waiting for ${label}`);
+}
+
+function git(args: string[], cwd: string): void {
+	const res = spawnSync("git", args, {
+		cwd,
+		encoding: "utf8",
+	});
+	if (res.status !== 0) {
+		throw new Error(`git ${args.join(" ")} failed: ${res.stderr || res.stdout}`);
+	}
+}
+
+async function latestMemberStatus(teamDir: string, name: string): Promise<{ status?: string; sessionFile?: string } | null> {
+	const cfg = await loadTeamConfig(teamDir);
+	if (!cfg) return null;
+	const member = cfg.members.find((m) => m.name === name);
+	if (!member) return null;
+	return { status: member.status, sessionFile: member.sessionFile };
+}
+
+const { timeoutSec } = parseArgs(process.argv.slice(2));
+const tmpRoot = fs.mkdtempSync(path.join(os.tmpdir(), "pi-teams-branch-context-"));
+const teamsRootDir = path.join(tmpRoot, "teams-root");
+const repoDir = path.join(tmpRoot, "repo");
+const teamId = "branch-context-team";
+const taskListId = teamId;
+const teamDir = path.join(teamsRootDir, teamId);
+const sessionsDir = path.join(teamDir, "sessions");
+const agentName = "alpha";
+const leadName = "team-lead";
+const procs: ChildProcess[] = [];
+
+try {
+	fs.mkdirSync(repoDir, { recursive: true });
+	fs.mkdirSync(sessionsDir, { recursive: true });
+
+	git(["init", "-b", "main"], repoDir);
+git(["config", "user.name", "Test User"], repoDir);
+git(["config", "user.email", "test@example.com"], repoDir);
+fs.writeFileSync(path.join(repoDir, "README.md"), "branch context integration\n", "utf8");
+git(["add", "README.md"], repoDir);
+git(["commit", "-m", "init"], repoDir);
+
+const parent = SessionManager.create(repoDir, sessionsDir);
+parent.appendModelChange("openai-codex", "gpt-5.4");
+parent.appendThinkingLevelChange("minimal");
+parent.appendMessage({
+	role: "user",
+	content: [{ type: "text", text: "Summarize the repo history." }],
+	timestamp: Date.now(),
+});
+const stableAssistantId = parent.appendMessage({
+	role: "assistant",
+	content: [{ type: "text", text: "The repo history is summarized." }],
+	api: "test",
+	provider: "test",
+	model: "test",
+	usage: {
+		input: 0,
+		output: 0,
+		cacheRead: 0,
+		cacheWrite: 0,
+		totalTokens: 0,
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+	},
+	stopReason: "stop",
+	timestamp: Date.now(),
+});
+const compactionId = parent.appendCompaction("summarized", stableAssistantId, 1234);
+const currentUserId = parent.appendMessage({
+	role: "user",
+	content: [{ type: "text", text: "Investigate the repo, then delegate part of it." }],
+	timestamp: Date.now(),
+});
+const toolUseAssistant: AssistantMessage = {
+	role: "assistant",
+	content: [{ type: "toolCall", id: "call-1", name: "read", arguments: { path: "README.md" } }],
+	api: "test",
+	provider: "test",
+	model: "test",
+	usage: {
+		input: 0,
+		output: 0,
+		cacheRead: 0,
+		cacheWrite: 0,
+		totalTokens: 0,
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+	},
+	stopReason: "toolUse",
+	timestamp: Date.now(),
+};
+parent.appendMessage(toolUseAssistant);
+parent.appendMessage({
+	role: "toolResult",
+	toolCallId: "call-1",
+	toolName: "read",
+	content: [{ type: "text", text: "branch context integration" }],
+	isError: false,
+	timestamp: Date.now(),
+});
+
+const parentLeafId = parent.getLeafId();
+assert(parentLeafId !== null, "expected parent leaf id");
+if (!parentLeafId) {
+	throw new Error("Missing parent leaf id");
+}
+
+const selection = resolveBranchLeafSelection(parent.getBranch(parentLeafId), parentLeafId);
+assert(selection.adjusted, "expected unfinished turn branch selection to adjust away from active leaf");
+assert(selection.leafId === compactionId, `expected branch selection to use the stable pre-user boundary id, got ${selection.leafId}`);
+assert(branchSelectionNote(selection) === "branch(clean-turn)", "expected clean-turn branch note");
+assert(selection.replayUserMessage?.role === "user", "expected the active user request to be replayed into the cleaned child branch");
+
+const branchedSessionFile = parent.createBranchedSession(selection.leafId);
+assert(Boolean(branchedSessionFile), "expected branched session file to be created");
+if (!branchedSessionFile) {
+	throw new Error("Missing branched session file");
+}
+if (selection.replayUserMessage) {
+	parent.appendMessage(JSON.parse(JSON.stringify(selection.replayUserMessage)) as Parameters<typeof parent.appendMessage>[0]);
+}
+
+const childEntries = parent.getEntries();
+assert(childEntries.some((entry) => entry.id === stableAssistantId), "child session should retain the latest completed assistant message");
+assert(childEntries.some((entry) => entry.id === compactionId), "child session should retain the compaction entry before the active user");
+assert(!childEntries.some((entry) => entry.id === currentUserId), "child session should drop the original unfinished-turn user entry");
+assert(
+	childEntries.some(
+		(entry) =>
+			entry.type === "message" &&
+			typeof entry.message === "object" &&
+			entry.message !== null &&
+			(entry.message as { role?: string }).role === "user" &&
+			JSON.stringify((entry.message as { content?: unknown }).content).includes("Investigate the repo, then delegate part of it."),
+	),
+	"child session should replay the active user request onto the cleaned branch",
+);
+assert(
+	!childEntries.some(
+		(entry) =>
+			entry.type === "message" &&
+			typeof entry.message === "object" &&
+			entry.message !== null &&
+			(entry.message as { role?: string }).role === "assistant" &&
+			entry.id !== stableAssistantId,
+	),
+	"child session should exclude the in-progress assistant tool-use message",
+);
+assert(
+	!childEntries.some(
+		(entry) => entry.type === "message" && typeof entry.message === "object" && entry.message !== null && (entry.message as { role?: string }).role === "toolResult",
+	),
+	"child session should exclude trailing tool results from the unfinished turn",
+);
+
+await ensureTeamConfig(teamDir, { teamId, taskListId, leadName, style: "normal" });
+
+const scriptDir = path.dirname(fileURLToPath(import.meta.url));
+const repoRoot = path.resolve(scriptDir, "..");
+const entryPath = path.join(repoRoot, "extensions", "teams", "index.ts");
+assert(fs.existsSync(entryPath), `Missing teams entry path: ${entryPath}`);
+
+const worker = spawn(
+	"pi",
+	[
+		"--mode",
+		"rpc",
+		"--session",
+		branchedSessionFile,
+		"--session-dir",
+		sessionsDir,
+		"--provider",
+		"openai-codex",
+		"--model",
+		"gpt-5.4",
+		"--thinking",
+		"minimal",
+		"--no-extensions",
+		"-e",
+		entryPath,
+		"--append-system-prompt",
+		`You are teammate '${agentName}'. Prefer working from the shared task list.`,
+	],
+	{
+		cwd: repoDir,
+		env: {
+			...process.env,
+			PI_TEAMS_ROOT_DIR: teamsRootDir,
+			PI_TEAMS_WORKER: "1",
+			PI_TEAMS_TEAM_ID: teamId,
+			PI_TEAMS_TASK_LIST_ID: taskListId,
+			PI_TEAMS_AGENT_NAME: agentName,
+			PI_TEAMS_LEAD_NAME: leadName,
+			PI_TEAMS_STYLE: "normal",
+			PI_TEAMS_AUTO_CLAIM: "0",
+		},
+		stdio: ["pipe", "pipe", "pipe"],
+	},
+);
+procs.push(worker);
+
+await waitFor(
+	async () => {
+		const member = await latestMemberStatus(teamDir, agentName);
+		return member?.status === "online";
+	},
+	{ timeoutMs: timeoutSec * 1000, pollMs: 250, label: `${agentName} online` },
+);
+
+const task = await createTask(teamDir, taskListId, {
+	subject: "Branch context integration",
+	description: "Reply with exactly 'branch context integration ok'. Do not edit files.",
+	owner: agentName,
+});
+await writeToMailbox(teamDir, taskListId, agentName, {
+	from: leadName,
+	text: JSON.stringify(taskAssignmentPayload(task, leadName)),
+	timestamp: new Date().toISOString(),
+});
+
+await waitFor(
+	async () => {
+		const current = await getTask(teamDir, taskListId, task.id);
+		const result = current?.metadata?.result;
+		return current?.status === "completed" && typeof result === "string" && result.includes("branch context integration ok");
+	},
+	{ timeoutMs: timeoutSec * 1000, pollMs: 500, label: `task #${task.id} completion` },
+);
+
+	console.log("PASS: branch context integration test passed");
+} finally {
+	await terminateAll(procs);
+	try {
+		fs.rmSync(tmpRoot, { recursive: true, force: true });
+	} catch {
+		// ignore
+	}
+}

package/scripts/smoke-test.mts

@@ -69,8 +69,9 @@ import {
 	isPlanApprovedMessage,
 	isPlanRejectedMessage,
 } from "../extensions/teams/protocol.js";
-import { pollLeaderInbox } from "../extensions/teams/leader-inbox.js";
+import { DelegationTracker, pollLeaderInbox } from "../extensions/teams/leader-inbox.js";
 import { getParentSessionId, shouldSilenceInheritedParentAttachClaimWarning } from "../extensions/teams/session-parent.js";
+import { branchSelectionNote, ensureSessionFileMaterialized, resolveBranchLeafSelection } from "../extensions/teams/session-branching.js";
 import { SessionManager, type ExtensionContext } from "@mariozechner/pi-coding-agent";
 import type { AssistantMessage } from "@mariozechner/pi-ai";
 
@@ -906,6 +907,200 @@ console.log("\n10b. branched sessions + inherited attach claims");
 			);
 		}
 	}
+
+	const branchFromUser = SessionManager.create(tmpRoot, sessionsDir);
+	branchFromUser.appendModelChange("openai-codex", "gpt-5.4");
+	branchFromUser.appendThinkingLevelChange("minimal");
+	branchFromUser.appendMessage({
+		role: "user",
+		content: [{ type: "text", text: "What should we do next?" }],
+		timestamp: Date.now(),
+	});
+	const stableAssistantId = branchFromUser.appendMessage(assistantMessage);
+	const currentUserId = branchFromUser.appendMessage({
+		role: "user",
+		content: [{ type: "text", text: "Investigate something, then delegate it." }],
+		timestamp: Date.now(),
+	});
+	const activeTurnToolUse: AssistantMessage = {
+		role: "assistant",
+		content: [{ type: "toolCall", id: "call-1", name: "read", arguments: { path: "README.md" } }],
+		api: "test",
+		provider: "test",
+		model: "test",
+		usage: {
+			input: 0,
+			output: 0,
+			cacheRead: 0,
+			cacheWrite: 0,
+			totalTokens: 0,
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+		},
+		stopReason: "toolUse",
+		timestamp: Date.now(),
+	};
+	branchFromUser.appendMessage(activeTurnToolUse);
+	branchFromUser.appendMessage({
+		role: "toolResult",
+		toolCallId: "call-1",
+		toolName: "read",
+		content: [{ type: "text", text: "README contents" }],
+		isError: false,
+		timestamp: Date.now(),
+	});
+	const unfinishedLeafId = branchFromUser.getLeafId();
+	assert(unfinishedLeafId !== null, "unfinished branch test leaf exists");
+	if (unfinishedLeafId) {
+		const selection = resolveBranchLeafSelection(branchFromUser.getBranch(unfinishedLeafId), unfinishedLeafId);
+		assert(selection.adjusted, "unfinished turn adjusts branch leaf away from active leaf");
+		assertEq(selection.leafId, stableAssistantId, "unfinished turn branches from latest completed assistant message");
+		assertEq(branchSelectionNote(selection), "branch(clean-turn)", "unfinished turn note marks clean-turn branch");
+		assert(
+			selection.replayUserMessage?.role === "user",
+			"unfinished turn keeps the active user request available for replay into the child branch",
+		);
+
+		const branchedPath = branchFromUser.createBranchedSession(selection.leafId);
+		assert(branchedPath !== null, "clean-turn branch session created");
+		if (selection.replayUserMessage) {
+			branchFromUser.appendMessage(JSON.parse(JSON.stringify(selection.replayUserMessage)) as Parameters<typeof branchFromUser.appendMessage>[0]);
+		}
+		if (branchedPath) await ensureSessionFileMaterialized(branchFromUser, branchedPath);
+		const childEntries = branchFromUser.getEntries();
+		assert(childEntries.some((entry) => entry.id === stableAssistantId), "clean-turn child keeps latest completed assistant");
+		assert(
+			childEntries.some(
+				(entry) =>
+					entry.type === "message" &&
+					isRecord(entry.message) &&
+					entry.message.role === "user" &&
+					JSON.stringify(entry.message.content).includes("Investigate something, then delegate it."),
+			),
+			"clean-turn child replays the active user request onto the cleaned branch",
+		);
+		assert(
+			childEntries.filter((entry) => entry.id === currentUserId).length === 0,
+			"clean-turn child does not keep the original unfinished-turn user entry id",
+		);
+		assert(
+			!childEntries.some((entry) => entry.type === "message" && isRecord(entry.message) && entry.message.role === "toolResult"),
+			"clean-turn child excludes trailing tool results from active turn",
+		);
+		assert(
+			!childEntries.some(
+				(entry) =>
+					entry.type === "message" &&
+					isRecord(entry.message) &&
+					entry.message.role === "assistant" &&
+					entry.id !== stableAssistantId,
+			),
+			"clean-turn child excludes in-progress assistant tool-use turn",
+		);
+	}
+
+	const compactedTurn = SessionManager.create(tmpRoot, sessionsDir);
+	compactedTurn.appendModelChange("openai-codex", "gpt-5.4");
+	compactedTurn.appendThinkingLevelChange("minimal");
+	compactedTurn.appendMessage({
+		role: "user",
+		content: [{ type: "text", text: "Earlier request" }],
+		timestamp: Date.now(),
+	});
+	const compactedAssistantId = compactedTurn.appendMessage(assistantMessage);
+	const compactionId = compactedTurn.appendCompaction("summarized", compactedAssistantId, 1234);
+	compactedTurn.appendMessage({
+		role: "user",
+		content: [{ type: "text", text: "Current request after compaction" }],
+		timestamp: Date.now(),
+	});
+	compactedTurn.appendMessage(activeTurnToolUse);
+	compactedTurn.appendMessage({
+		role: "toolResult",
+		toolCallId: "call-1",
+		toolName: "read",
+		content: [{ type: "text", text: "README contents" }],
+		isError: false,
+		timestamp: Date.now(),
+	});
+	const compactedLeafId = compactedTurn.getLeafId();
+	assert(compactedLeafId !== null, "compacted branch test leaf exists");
+	if (compactedLeafId) {
+		const selection = resolveBranchLeafSelection(compactedTurn.getBranch(compactedLeafId), compactedLeafId);
+		assert(selection.adjusted, "compacted unfinished turn adjusts branch leaf");
+		assertEq(selection.leafId, compactionId, "compacted unfinished turn branches from the entry immediately before the active user");
+		assert(selection.replayUserMessage?.role === "user", "compacted unfinished turn replays the active user request");
+		const branchedPath = compactedTurn.createBranchedSession(selection.leafId);
+		assert(branchedPath !== null, "compacted branch session created");
+		if (selection.replayUserMessage) {
+			compactedTurn.appendMessage(JSON.parse(JSON.stringify(selection.replayUserMessage)) as Parameters<typeof compactedTurn.appendMessage>[0]);
+		}
+		const childEntries = compactedTurn.getEntries();
+		assert(childEntries.some((entry) => entry.id === compactionId), "compacted child keeps the compaction entry before the active user");
+		assert(
+			childEntries.some(
+				(entry) =>
+					entry.type === "message" &&
+					isRecord(entry.message) &&
+					entry.message.role === "user" &&
+					JSON.stringify(entry.message.content).includes("Current request after compaction"),
+			),
+			"compacted child replays the active user after the preserved compaction boundary",
+		);
+		assertEq(branchSelectionNote(selection), "branch(clean-turn)", "compacted unfinished turn keeps clean-turn note");
+	}
+
+	const userOnlyTurn = SessionManager.create(tmpRoot, sessionsDir);
+	userOnlyTurn.appendModelChange("openai-codex", "gpt-5.4");
+	userOnlyTurn.appendThinkingLevelChange("minimal");
+	userOnlyTurn.appendMessage({
+		role: "user",
+		content: [{ type: "text", text: "Only user context so far" }],
+		timestamp: Date.now(),
+	});
+	userOnlyTurn.appendMessage(activeTurnToolUse);
+	userOnlyTurn.appendMessage({
+		role: "toolResult",
+		toolCallId: "call-1",
+		toolName: "read",
+		content: [{ type: "text", text: "README contents" }],
+		isError: false,
+		timestamp: Date.now(),
+	});
+	const userOnlyLeafId = userOnlyTurn.getLeafId();
+	assert(userOnlyLeafId !== null, "user-only fallback test leaf exists");
+	if (userOnlyLeafId) {
+		const selection = resolveBranchLeafSelection(userOnlyTurn.getBranch(userOnlyLeafId), userOnlyLeafId);
+		assert(selection.adjusted, "user-only unfinished turn still adjusts branch leaf");
+		assert(selection.leafId !== userOnlyLeafId, "user-only fallback rewinds away from the active unfinished leaf");
+		assert(selection.replayUserMessage?.role === "user", "user-only fallback keeps the active user message for replay");
+		assertEq(branchSelectionNote(selection), "branch(clean-turn)", "user-only fallback keeps the clean-turn note");
+		const branchedPath = userOnlyTurn.createBranchedSession(selection.leafId);
+		assert(branchedPath !== null, "user-only fallback branch session created");
+		if (selection.replayUserMessage) {
+			userOnlyTurn.appendMessage(JSON.parse(JSON.stringify(selection.replayUserMessage)) as Parameters<typeof userOnlyTurn.appendMessage>[0]);
+		}
+		if (branchedPath) {
+			await ensureSessionFileMaterialized(userOnlyTurn, branchedPath);
+			assert(fs.existsSync(branchedPath), "user-only fallback materializes a real session file");
+		}
+	}
+
+	const completedTurn = SessionManager.create(tmpRoot, sessionsDir);
+	completedTurn.appendMessage({
+		role: "user",
+		content: [{ type: "text", text: "Done already" }],
+		timestamp: Date.now(),
+	});
+	const completedAssistantId = completedTurn.appendMessage(assistantMessage);
+	const completedLeafId = completedTurn.getLeafId();
+	assert(completedLeafId !== null, "completed branch test leaf exists");
+	if (completedLeafId) {
+		const selection = resolveBranchLeafSelection(completedTurn.getBranch(completedLeafId), completedLeafId);
+		assert(!selection.adjusted, "completed turn keeps requested leaf");
+		assertEq(selection.leafId, completedLeafId, "completed turn branches from current leaf");
+		assertEq(completedAssistantId, completedLeafId, "completed leaf stays on assistant reply");
+		assertEq(branchSelectionNote(selection), "branch", "completed turn keeps plain branch note");
+	}
 }
 
 // ── 11. /team done (end-of-run cleanup) ──────────────────────────────
@@ -1114,6 +1309,7 @@ console.log("\n14. leader-inbox LLM message injection");
 		cwd: inboxTeamDir,
 		ui: { notify: () => {} },
 		sessionManager: { getSessionId: () => "inbox-team" },
+		isIdle: () => false,
 	} as unknown as ExtensionContext;
 
 	await pollLeaderInbox({
@@ -1220,7 +1416,8 @@ console.log("\n14. leader-inbox LLM message injection");
 		leadName,
 		style,
 		pendingPlanApprovals: new Map(),
-		enqueueHook: () => {}, // hooks present → should qualify allDone
+		enqueueHook: () => {},
+		hooksEnabled: true,
 		sendLeaderLlmMessage: (content, options) => {
 			llmMessages.push({ content, options });
 		},
@@ -1232,6 +1429,88 @@ console.log("\n14. leader-inbox LLM message injection");
 		assert(hookMsg.content.includes("quality gates are still running"), "allDone qualified when hooks active");
 		assert(!hookMsg.content.includes("Review results and determine next steps"), "no premature wrap-up prompt when hooks active");
 	}
+
+	// Hooks disabled should not qualify all-done messages just because a callback is wired.
+	const t4 = await createTask(inboxTeamDir, inboxTaskListId, { subject: "Post-review cleanup", description: "", owner: "dave" });
+	await completeTask(inboxTeamDir, inboxTaskListId, t4.id, "dave", "Cleanup complete");
+	const ts4 = new Date().toISOString();
+	await writeToMailbox(inboxTeamDir, TEAM_MAILBOX_NS, leadName, {
+		from: "dave",
+		text: JSON.stringify({
+			type: "idle_notification",
+			from: "dave",
+			timestamp: ts4,
+			completedTaskId: t4.id,
+			completedStatus: "completed",
+		}),
+		timestamp: ts4,
+	});
+
+	llmMessages.length = 0;
+	await pollLeaderInbox({
+		ctx: stubCtx,
+		teamId: "inbox-team",
+		teamDir: inboxTeamDir,
+		taskListId: inboxTaskListId,
+		leadName,
+		style,
+		pendingPlanApprovals: new Map(),
+		enqueueHook: () => {},
+		hooksEnabled: false,
+		sendLeaderLlmMessage: (content, options) => {
+			llmMessages.push({ content, options });
+		},
+	});
+
+	assert(llmMessages.length === 1, "one LLM message sent when hooks callback is wired but disabled");
+	const disabledHookMsg = llmMessages[0];
+	if (disabledHookMsg) {
+		assert(!disabledHookMsg.content.includes("quality gates are still running"), "disabled hooks do not qualify the per-task allDone summary");
+		assert(disabledHookMsg.content.includes("Review results and determine next steps"), "disabled hooks keep the normal per-task allDone summary");
+	}
+
+	// Batch-complete auto-wake should use the same hooks-enabled check.
+	const t5 = await createTask(inboxTeamDir, inboxTaskListId, { subject: "Batch wake task", description: "", owner: "erin" });
+	await completeTask(inboxTeamDir, inboxTaskListId, t5.id, "erin", "Batch wake done");
+	const ts5 = new Date().toISOString();
+	await writeToMailbox(inboxTeamDir, TEAM_MAILBOX_NS, leadName, {
+		from: "erin",
+		text: JSON.stringify({
+			type: "idle_notification",
+			from: "erin",
+			timestamp: ts5,
+			completedTaskId: t5.id,
+			completedStatus: "completed",
+		}),
+		timestamp: ts5,
+	});
+
+	const batchTracker = new DelegationTracker();
+	batchTracker.addBatch([t5.id]);
+	llmMessages.length = 0;
+	await pollLeaderInbox({
+		ctx: stubCtx,
+		teamId: "inbox-team",
+		teamDir: inboxTeamDir,
+		taskListId: inboxTaskListId,
+		leadName,
+		style,
+		pendingPlanApprovals: new Map(),
+		enqueueHook: () => {},
+		hooksEnabled: false,
+		delegationTracker: batchTracker,
+		sendLeaderLlmMessage: (content, options) => {
+			llmMessages.push({ content, options });
+		},
+	});
+
+	assert(llmMessages.length === 2, "per-task completion plus batch-complete messages sent when a tracked delegation finishes");
+	const batchMsg = llmMessages.find((entry) => entry.content.includes("All delegated tasks completed"));
+	assert(batchMsg !== undefined, "batch-complete notification sent");
+	if (batchMsg) {
+		assert(!batchMsg.content.includes("Quality gates are still running"), "disabled hooks do not qualify the batch-complete summary");
+		assert(batchMsg.content.includes("Review the results and continue."), "disabled hooks keep the normal batch-complete summary");
+	}
 }
 
 // ── 15. docs/help drift guard ────────────────────────────────────────