fix: branch teammates from the last stable turn

extensions/teams/leader.ts

@@ -34,6 +34,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 +91,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)) {

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.json

@@ -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",