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, resolveBranchLeafSelection } from "./session-branching.js";
 import type { ContextMode, SpawnTeammateFn, SpawnTeammateResult, WorkspaceMode } from "./spawn-types.js";
 
 function getTeamsExtensionEntryPath(): string | null {
@@ -90,12 +91,13 @@ 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 };
+		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,94 @@
+import type { SessionEntry } from "@mariozechner/pi-coding-agent";
+
+export type BranchLeafSelection = {
+	leafId: string;
+	adjusted: boolean;
+	reason: "requested" | "clean-turn-assistant" | "clean-turn-user";
+};
+
+type MessageLike = Record<string, unknown> & { role: string };
+
+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: MessageLike & { role: "user" } } {
+	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" };
+	}
+
+	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-assistant" };
+		}
+	}
+
+	for (let i = lastAssistantIndex - 1; i >= 0; i -= 1) {
+		const candidate = path[i];
+		if (!candidate) continue;
+		if (isUserMessageEntry(candidate)) {
+			return { leafId: candidate.id, adjusted: candidate.id !== requestedLeafId, reason: "clean-turn-user" };
+		}
+	}
+
+	return { leafId: requestedLeafId, adjusted: false, reason: "requested" };
+}
+
+export function branchSelectionNote(selection: BranchLeafSelection): string {
+	if (!selection.adjusted) return "branch";
+	if (selection.reason === "clean-turn-assistant") 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",