@tmustier/pi-agent-teams 0.2.0 → 0.3.1

package/extensions/teams/leader.ts

@@ -9,36 +9,16 @@ import { TEAM_MAILBOX_NS } from "./protocol.js";
 import { listTasks, unassignTasksForAgent, type TeamTask } from "./task-store.js";
 import { TeammateRpc } from "./teammate-rpc.js";
 import { ensureTeamConfig, loadTeamConfig, setMemberStatus, upsertMember, type TeamConfig } from "./team-config.js";
-import { getTeamDir, getTeamsRootDir } from "./paths.js";
+import { getTeamDir } from "./paths.js";
 import { ensureWorktreeCwd } from "./worktree.js";
 import { ActivityTracker, TranscriptTracker } from "./activity-tracker.js";
 import { openInteractiveWidget } from "./teams-panel.js";
 import { createTeamsWidget } from "./teams-widget.js";
 import { getTeamsStyleFromEnv, type TeamsStyle, formatMemberDisplayName, getTeamsStrings } from "./teams-style.js";
 import { pollLeaderInbox as pollLeaderInboxImpl } from "./leader-inbox.js";
-import { handleTeamTaskCommand } from "./leader-task-commands.js";
-import { handleTeamPlanCommand } from "./leader-plan-commands.js";
-import { handleTeamSpawnCommand } from "./leader-spawn-command.js";
+import { handleTeamCommand } from "./leader-team-command.js";
 import { registerTeamsTool } from "./leader-teams-tool.js";
-import {
-	handleTeamBroadcastCommand,
-	handleTeamDmCommand,
-	handleTeamSendCommand,
-	handleTeamSteerCommand,
-} from "./leader-messaging-commands.js";
-import { handleTeamEnvCommand, handleTeamIdCommand, handleTeamListCommand } from "./leader-info-commands.js";
-import {
-	handleTeamCleanupCommand,
-	handleTeamDelegateCommand,
-	handleTeamKillCommand,
-	handleTeamShutdownCommand,
-	handleTeamStopCommand,
-	handleTeamStyleCommand,
-} from "./leader-lifecycle-commands.js";
-
-
-type ContextMode = "fresh" | "branch";
-type WorkspaceMode = "shared" | "worktree";
+import type { ContextMode, SpawnTeammateFn, SpawnTeammateResult, WorkspaceMode } from "./spawn-types.js";
 
 function getTeamsExtensionEntryPath(): string | null {
 	// In dev, teammates won't automatically have this extension unless it is installed or discoverable.
@@ -59,15 +39,6 @@ function shellQuote(v: string): string {
 	return "'" + v.replace(/'/g, `"'"'"'`) + "'";
 }
 
-function parseAssigneePrefix(text: string): { assignee?: string; text: string } {
-	const m = text.match(/^([a-zA-Z0-9_-]+):\s*(.+)$/);
-	if (!m) return { text };
-	const assignee = m[1];
-	const rest = m[2];
-	if (!assignee || !rest) return { text };
-	return { assignee, text: rest };
-}
-
 function getTeamSessionsDir(teamDir: string): string {
 	return path.join(teamDir, "sessions");
 }
@@ -121,17 +92,6 @@ async function createSessionForTeammate(
 	}
 }
 
-function taskAssignmentPayload(task: TeamTask, assignedBy: string) {
-	return {
-		type: "task_assignment",
-		taskId: task.id,
-		subject: task.subject,
-		description: task.description,
-		assignedBy,
-		timestamp: new Date().toISOString(),
-	};
-}
-
 // Message parsers are shared with the worker implementation.
 export function runLeader(pi: ExtensionAPI): void {
 	const teammates = new Map<string, TeammateRpc>();
@@ -226,22 +186,7 @@ export function runLeader(pi: ExtensionAPI): void {
 		currentCtx.ui.setWidget("pi-teams", widgetFactory);
 	};
 
-	type SpawnTeammateResult =
-		| {
-				ok: true;
-				name: string;
-				mode: ContextMode;
-				workspaceMode: WorkspaceMode;
-				childCwd: string;
-				note?: string;
-				warnings: string[];
-		  }
-		| { ok: false; error: string };
-
-	const spawnTeammate = async (
-		ctx: ExtensionContext,
-		opts: { name: string; mode?: ContextMode; workspaceMode?: WorkspaceMode; planRequired?: boolean },
-	): Promise<SpawnTeammateResult> => {
+	const spawnTeammate: SpawnTeammateFn = async (ctx, opts): Promise<SpawnTeammateResult> => {
 		const warnings: string[] = [];
 		const mode: ContextMode = opts.mode ?? "fresh";
 		let workspaceMode: WorkspaceMode = opts.workspaceMode ?? "shared";
@@ -516,7 +461,6 @@ export function runLeader(pi: ExtensionAPI): void {
 		teammates,
 		spawnTeammate,
 		getTaskListId: () => taskListId,
-		taskAssignmentPayload,
 		refreshTasks,
 		renderWidget,
 	});
@@ -550,11 +494,11 @@ export function runLeader(pi: ExtensionAPI): void {
 					timestamp: new Date().toISOString(),
 				});
 			},
-			abortComrade(name: string) {
+			abortMember(name: string) {
 				const rpc = teammates.get(name);
 				if (rpc) void rpc.abort();
 			},
-			killComrade(name: string) {
+			killMember(name: string) {
 				const rpc = teammates.get(name);
 				if (!rpc) return;
 
@@ -613,254 +557,34 @@ export function runLeader(pi: ExtensionAPI): void {
 			currentCtx = ctx;
 			currentTeamId = ctx.sessionManager.getSessionId();
 
-			const [sub, ...rest] = args.trim().split(" ");
-			if (!sub || sub === "help") {
-				ctx.ui.notify(
-					[
-						"Usage:",
-						"  /team id",
-						"  /team env <name>",
-						"  /team spawn <name> [fresh|branch] [shared|worktree] [plan]",
-						"  /team panel",
-						"  /team send <name> <msg...>",
-						"  /team dm <name> <msg...>",
-						"  /team broadcast <msg...>",
-						"  /team steer <name> <msg...>",
-						"  /team stop <name> [reason...]",
-						"  /team kill <name>",
-						"  /team shutdown",
-						"  /team shutdown <name> [reason...]",
-						"  /team delegate [on|off]",
-						"  /team plan approve <name>",
-						"  /team plan reject <name> [feedback...]",
-						"  /team cleanup [--force]",
-						"  /team task add <text...>",
-						"  /team task assign <id> <agent>",
-						"  /team task unassign <id>",
-						"  /team task list",
-						"  /team task clear [completed|all] [--force]",
-						"  /team task show <id>",
-						"  /team task dep add <id> <depId>",
-						"  /team task dep rm <id> <depId>",
-						"  /team task dep ls <id>",
-						"  /team task use <taskListId>",
-					].join("\n"),
-					"info",
-				);
-				return;
-			}
-
-			switch (sub) {
-				case "list": {
-					await handleTeamListCommand({
-						ctx,
-						teammates,
-						getTeamConfig: () => teamConfig,
-						style,
-						refreshTasks,
-						renderWidget,
-					});
-					return;
-				}
-
-				case "id": {
-					await handleTeamIdCommand({
-						ctx,
-						taskListId,
-						leadName: teamConfig?.leadName ?? "team-lead",
-						style,
-					});
-					return;
-				}
-
-				case "env": {
-					await handleTeamEnvCommand({
-						ctx,
-						rest,
-						taskListId,
-						leadName: teamConfig?.leadName ?? "team-lead",
-						style,
-						getTeamsExtensionEntryPath,
-						shellQuote,
-					});
-					return;
-				}
-
-				case "cleanup": {
-					await handleTeamCleanupCommand({
-						ctx,
-						rest,
-						teammates,
-						refreshTasks,
-						getTasks: () => tasks,
-						renderWidget,
-						style,
-					});
-					return;
-				}
-
-				case "delegate": {
-					await handleTeamDelegateCommand({
-						ctx,
-						rest,
-						getDelegateMode: () => delegateMode,
-						setDelegateMode: (next) => {
-							delegateMode = next;
-						},
-						renderWidget,
-					});
-					return;
-				}
-
-				case "shutdown": {
-					await handleTeamShutdownCommand({
-						ctx,
-						rest,
-						teammates,
-						leadName: teamConfig?.leadName ?? "team-lead",
-						style,
-						getCurrentCtx: () => currentCtx,
-					});
-					return;
-				}
-
-				case "spawn": {
-					await handleTeamSpawnCommand({ ctx, rest, teammates, style, spawnTeammate });
-					return;
-				}
-
-				case "style": {
-					const teamId = ctx.sessionManager.getSessionId();
-					const teamDir = getTeamDir(teamId);
-					await handleTeamStyleCommand({
-						ctx,
-						rest,
-						teamDir,
-						getStyle: () => style,
-						setStyle: (next) => {
-							style = next;
-						},
-						refreshTasks,
-						renderWidget,
-					});
-					return;
-				}
-
-				case "panel":
-				case "widget": {
-					await openWidget(ctx);
-					return;
-				}
-
-				case "send": {
-					await handleTeamSendCommand({
-						ctx,
-						rest,
-						teammates,
-						style,
-						renderWidget,
-					});
-					return;
-				}
-
-				case "steer": {
-					await handleTeamSteerCommand({
-						ctx,
-						rest,
-						teammates,
-						style,
-						renderWidget,
-					});
-					return;
-				}
-
-				case "stop": {
-					await handleTeamStopCommand({
-						ctx,
-						rest,
-						teammates,
-						leadName: teamConfig?.leadName ?? "team-lead",
-						style,
-						refreshTasks,
-						getTasks: () => tasks,
-						renderWidget,
-					});
-					return;
-				}
-
-				case "kill": {
-					await handleTeamKillCommand({
-						ctx,
-						rest,
-						teammates,
-						leadName: teamConfig?.leadName ?? "team-lead",
-						style,
-						taskListId,
-						refreshTasks,
-						renderWidget,
-					});
-					return;
-				}
-
-				case "dm": {
-					await handleTeamDmCommand({
-						ctx,
-						rest,
-						leadName: teamConfig?.leadName ?? "team-lead",
-						style,
-					});
-					return;
-				}
-
-				case "broadcast": {
-					await handleTeamBroadcastCommand({
-						ctx,
-						rest,
-						teammates,
-						leadName: teamConfig?.leadName ?? "team-lead",
-						style,
-						refreshTasks,
-						getTasks: () => tasks,
-						getTaskListId: () => taskListId,
-					});
-					return;
-				}
-
-				case "task": {
-					await handleTeamTaskCommand({
-						ctx,
-						rest,
-						leadName: teamConfig?.leadName ?? "team-lead",
-						style,
-						getTaskListId: () => taskListId,
-						setTaskListId: (id) => {
-							taskListId = id;
-						},
-						getTasks: () => tasks,
-						refreshTasks,
-						renderWidget,
-						parseAssigneePrefix,
-						taskAssignmentPayload,
-					});
-					return;
-				}
-
-				case "plan": {
-					await handleTeamPlanCommand({
-						ctx,
-						rest,
-						leadName: teamConfig?.leadName ?? "team-lead",
-						style,
-						pendingPlanApprovals,
-					});
-					return;
-				}
-
-				default: {
-					ctx.ui.notify(`Unknown subcommand: ${sub}`, "error");
-					return;
-				}
-			}
+			await handleTeamCommand({
+				args,
+				ctx,
+				teammates,
+				getTeamConfig: () => teamConfig,
+				getTasks: () => tasks,
+				refreshTasks,
+				renderWidget,
+				getTaskListId: () => taskListId,
+				setTaskListId: (id) => {
+					taskListId = id;
+				},
+				pendingPlanApprovals,
+				getDelegateMode: () => delegateMode,
+				setDelegateMode: (next) => {
+					delegateMode = next;
+				},
+				getStyle: () => style,
+				setStyle: (next) => {
+					style = next;
+				},
+				spawnTeammate,
+				openWidget,
+				getTeamsExtensionEntryPath,
+				shellQuote,
+				getCurrentCtx: () => currentCtx,
+				stopAllTeammates,
+			});
 		},
 	});
 }

package/extensions/teams/protocol.ts

@@ -1,5 +1,25 @@
+import type { TeamTask } from "./task-store.js";
+
 export const TEAM_MAILBOX_NS = "team";
 
+export function taskAssignmentPayload(task: TeamTask, assignedBy: string): {
+	type: "task_assignment";
+	taskId: string;
+	subject: string;
+	description: string;
+	assignedBy: string;
+	timestamp: string;
+} {
+	return {
+		type: "task_assignment",
+		taskId: task.id,
+		subject: task.subject,
+		description: task.description,
+		assignedBy,
+		timestamp: new Date().toISOString(),
+	};
+}
+
 function safeParseJson(text: string): unknown | null {
 	try {
 		const parsed: unknown = JSON.parse(text);

package/extensions/teams/spawn-types.ts

@@ -0,0 +1,21 @@
+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
+
+export type ContextMode = "fresh" | "branch";
+export type WorkspaceMode = "shared" | "worktree";
+
+export type SpawnTeammateResult =
+	| {
+			ok: true;
+			name: string;
+			mode: ContextMode;
+			workspaceMode: WorkspaceMode;
+			childCwd?: string;
+			note?: string;
+			warnings: string[];
+	  }
+	| { ok: false; error: string };
+
+export type SpawnTeammateFn = (
+	ctx: ExtensionContext,
+	opts: { name: string; mode?: ContextMode; workspaceMode?: WorkspaceMode; planRequired?: boolean },
+) => Promise<SpawnTeammateResult>;

package/extensions/teams/task-store.ts

@@ -328,6 +328,10 @@ export async function unassignTasksForAgent(
 		if (t.owner !== agentName) continue;
 		if (t.status === "completed") continue;
 		const updated = await updateTask(teamDir, taskListId, t.id, (cur) => {
+			// Re-check ownership under the per-task lock to avoid races with other claimers.
+			if (cur.owner !== agentName) return cur;
+			if (cur.status === "completed") return cur;
+
 			const metadata = { ...(cur.metadata ?? {}) };
 			if (reason) metadata.unassignedReason = reason;
 			metadata.unassignedAt = new Date().toISOString();

package/extensions/teams/teammate-rpc.ts

@@ -4,7 +4,7 @@ import type { AgentEvent } from "@mariozechner/pi-agent-core";
 export type TeammateStatus = "starting" | "idle" | "streaming" | "stopped" | "error";
 
 type RpcCommand =
-	| { id: string; type: "prompt"; message: string; streamingBehavior?: "steer" | "followUp" }
+	| { id: string; type: "prompt"; message: string }
 	| { id: string; type: "steer"; message: string }
 	| { id: string; type: "follow_up"; message: string }
 	| { id: string; type: "abort" }
@@ -48,7 +48,31 @@ function isRpcResponse(v: unknown): v is RpcResponse {
 
 function isAgentEvent(v: unknown): v is AgentEvent {
 	if (!isRecord(v)) return false;
-	return typeof v.type === "string";
+	if (typeof v.type !== "string") return false;
+
+	// Validate the minimal shapes we actually dereference below.
+	if (v.type === "message_update") {
+		const ame = v.assistantMessageEvent;
+		if (!isRecord(ame)) return false;
+		if (typeof ame.type !== "string") return false;
+		if (ame.type === "text_delta" && typeof ame.delta !== "string") return false;
+		return true;
+	}
+
+	if (v.type === "tool_execution_start" || v.type === "tool_execution_update" || v.type === "tool_execution_end") {
+		if (typeof v.toolCallId !== "string") return false;
+		if (typeof v.toolName !== "string") return false;
+		return true;
+	}
+
+	return (
+		v.type === "agent_start" ||
+		v.type === "agent_end" ||
+		v.type === "turn_start" ||
+		v.type === "turn_end" ||
+		v.type === "message_start" ||
+		v.type === "message_end"
+	);
 }
 
 export class TeammateRpc {