@hueyexe/opencode-ensemble 0.3.1 → 0.5.0

package/src/index.ts

@@ -3,10 +3,12 @@ import { tool } from "@opencode-ai/plugin"
 import path from "path"
 import { mkdirSync } from "node:fs"
 import { createDb, getDbPath } from "./db"
-import { recoverStaleMembers, recoverUndeliveredMessages } from "./recovery"
+import { recoverStaleMembers, recoverUndeliveredMessages, recoverOrphanedWorktrees } from "./recovery"
 import { MemberRegistry, DescendantTracker } from "./state"
 import { handleSessionStatusEvent, handleSessionCreatedEvent, checkToolIsolation } from "./hooks"
 import { notifyTeamEvent, notifyWorkingProgress } from "./notify"
+import { buildLeadSystemPrompt, buildTeammateSystemPrompt, buildTeamCompactionContext } from "./system-prompt"
+import { findTeamBySession } from "./types"
 import { executeTeamCreate } from "./tools/team-create"
 import { executeTeamSpawn } from "./tools/team-spawn"
 import { executeTeamMessage } from "./tools/team-message"
@@ -47,7 +49,7 @@ const plugin: Plugin = async (input) => {
 
   // Build shared tool dependencies
   const client = input.client as unknown as PluginClient
-  const deps: ToolDeps = { db, registry, tracker, client }
+  const deps: ToolDeps = { db, registry, tracker, client, directory: input.directory }
 
   // Run recovery on init — mark stale busy members as error + abort orphaned sessions
   const recovery = await recoverStaleMembers(db, client)
@@ -66,6 +68,11 @@ const plugin: Plugin = async (input) => {
     // Best effort — don't block plugin init
   })
 
+  // Clean up orphaned worktrees from crashed teams
+  recoverOrphanedWorktrees(db, client).catch(() => {
+    // Best effort — don't block plugin init
+  })
+
   // Initialize rate limiter
   // OPENCODE_ENSEMBLE_RATE_LIMIT=0 disables it entirely
   const rateLimitEnv = process.env.OPENCODE_ENSEMBLE_RATE_LIMIT
@@ -147,6 +154,42 @@ const plugin: Plugin = async (input) => {
       }
     },
 
+    // System prompt injection — keeps lead aware of team state, reminds teammates of role
+    "experimental.chat.system.transform": async (input, output) => {
+      if (!input.sessionID) return
+      const teamInfo = findTeamBySession(db, registry, input.sessionID)
+      if (!teamInfo) return
+      const prompt = teamInfo.role === "lead"
+        ? buildLeadSystemPrompt(db, teamInfo.teamId)
+        : buildTeammateSystemPrompt(db, teamInfo.teamId, teamInfo.memberName!)
+      output.system.push(prompt)
+    },
+
+    // Compaction safety — preserves team context when sessions get long
+    "experimental.session.compacting": async (input, output) => {
+      const teamInfo = findTeamBySession(db, registry, input.sessionID)
+      if (!teamInfo) return
+      const context = buildTeamCompactionContext(db, teamInfo.teamId, teamInfo.role, teamInfo.memberName)
+      output.context.push(context)
+    },
+
+    // Team-aware shell environment for scripts and hooks
+    "shell.env": async (input, output) => {
+      if (!input.sessionID) return
+      const teamInfo = findTeamBySession(db, registry, input.sessionID)
+      if (!teamInfo) return
+      output.env.ENSEMBLE_TEAM = teamInfo.teamName
+      output.env.ENSEMBLE_ROLE = teamInfo.role
+      if (teamInfo.memberName) {
+        output.env.ENSEMBLE_MEMBER = teamInfo.memberName
+        const member = db.query("SELECT worktree_branch FROM team_member WHERE team_id = ? AND name = ?")
+          .get(teamInfo.teamId, teamInfo.memberName) as { worktree_branch: string | null } | null
+        if (member?.worktree_branch) {
+          output.env.ENSEMBLE_BRANCH = member.worktree_branch
+        }
+      }
+    },
+
     // Register all team tools
     tool: {
       team_create: tool({
@@ -163,6 +206,7 @@ const plugin: Plugin = async (input) => {
 
       team_spawn: tool({
         description: "Spawn a new teammate that works in parallel. The teammate starts immediately with the given prompt. " +
+          "Each teammate gets their own git worktree for file isolation. " +
           "Teammates work asynchronously and will message you when done. Do not poll for their status.",
         args: {
           name: tool.schema.string().describe("Teammate name (lowercase alphanumeric with hyphens)"),
@@ -170,6 +214,8 @@ const plugin: Plugin = async (input) => {
           prompt: tool.schema.string().describe("Task instructions for the teammate"),
           model: tool.schema.string().optional().describe("Model in provider/model format (optional, uses default)"),
           claim_task: tool.schema.string().optional().describe("Task ID to auto-claim for this teammate (optional)"),
+          worktree: tool.schema.boolean().default(true).describe("Create a git worktree for file isolation (default: true, set false for read-only agents)"),
+          plan_approval: tool.schema.boolean().default(false).describe("Require teammate to send a plan for approval before writing files (default: false)"),
         },
         async execute(args, ctx) {
           const result = await executeTeamSpawn(deps, args, ctx.sessionID)
@@ -183,6 +229,8 @@ const plugin: Plugin = async (input) => {
         args: {
           to: tool.schema.string().describe("Recipient name ('lead' or teammate name)"),
           text: tool.schema.string().describe("Message content (max 10KB)"),
+          approve: tool.schema.boolean().optional().describe("Approve a teammate's plan (only when recipient has plan_approval='pending')"),
+          reject: tool.schema.string().optional().describe("Reject a teammate's plan with reason (only when recipient has plan_approval='pending')"),
         },
         async execute(args, ctx) {
           const result = await executeTeamMessage(deps, args, ctx.sessionID)
@@ -267,9 +315,11 @@ const plugin: Plugin = async (input) => {
       }),
 
       team_shutdown: tool({
-        description: "Request a teammate to shut down. The teammate finishes current work then stops.",
+        description: "Request a teammate to shut down. The teammate finishes current work then stops. " +
+          "Pass force: true to abort immediately without waiting.",
         args: {
           member: tool.schema.string().describe("Teammate name to shut down"),
+          force: tool.schema.boolean().default(false).describe("Force immediate abort without waiting for current work to finish"),
         },
         async execute(args, ctx) {
           const result = await executeTeamShutdown(deps, args, ctx.sessionID)

package/src/recovery.ts

@@ -37,6 +37,43 @@ export async function recoverStaleMembers(db: Database, client?: PluginClient):
   return { interrupted: result.changes }
 }
 
+/**
+ * Clean up orphaned worktrees from archived teams or members that no longer exist.
+ * Compares worktrees on disk (via client.worktree.list) against active team members.
+ */
+export async function recoverOrphanedWorktrees(db: Database, client: PluginClient): Promise<{ removed: number }> {
+  let removed = 0
+
+  try {
+    const worktrees = await client.worktree.list()
+    if (!worktrees.data) return { removed: 0 }
+
+    // Get all active worktree directories from the DB
+    const activeWorktrees = new Set(
+      (db.query(
+        `SELECT tm.worktree_dir FROM team_member tm
+         JOIN team t ON tm.team_id = t.id
+         WHERE tm.worktree_dir IS NOT NULL AND t.status = 'active'`
+      ).all() as Array<{ worktree_dir: string }>).map(r => r.worktree_dir)
+    )
+
+    for (const wt of worktrees.data) {
+      // Only clean up worktrees created by ensemble (name starts with "ensemble-")
+      if (!wt.name.startsWith("ensemble-")) continue
+      if (activeWorktrees.has(wt.directory)) continue
+
+      try {
+        await client.worktree.remove({ worktreeRemoveInput: { directory: wt.directory } })
+        removed++
+      } catch { /* best effort */ }
+    }
+  } catch {
+    // worktree.list may not be available — silently ignore
+  }
+
+  return { removed }
+}
+
 /**
  * Redeliver undelivered messages (delivered=0) via promptAsync.
  * Resolves recipient session IDs from the member registry or team lead.

package/src/schema.ts

@@ -70,6 +70,12 @@ export const MIGRATIONS: string[] = [
   // Migration 2: Add read column to team_message for team_results tracking
   `ALTER TABLE team_message ADD COLUMN read INTEGER NOT NULL DEFAULT 0;
    CREATE INDEX IF NOT EXISTS team_message_unread_idx ON team_message(team_id, read) WHERE read = 0;`,
+  // Migration 3: Add worktree columns to team_member for git worktree isolation
+  `ALTER TABLE team_member ADD COLUMN worktree_dir TEXT;
+   ALTER TABLE team_member ADD COLUMN worktree_branch TEXT;`,
+  // Migration 4: Add plan_approval column to team_member for plan-before-build workflow
+  `ALTER TABLE team_member ADD COLUMN plan_approval TEXT NOT NULL DEFAULT 'none'
+     CHECK(plan_approval IN ('none', 'pending', 'approved', 'rejected'));`,
 ]
 
 /**

package/src/system-prompt.ts

@@ -0,0 +1,111 @@
+import type { Database } from "bun:sqlite"
+
+/** Status display mapping from DB status to human-readable label. */
+const STATUS_DISPLAY: Record<string, string> = {
+  busy: "working",
+  ready: "idle",
+  shutdown_requested: "shutting down",
+  shutdown: "shut down",
+  error: "error",
+}
+
+/**
+ * Build the system prompt injected into the lead's session.
+ * Includes team name, member statuses, task counts, and anti-polling guidance.
+ */
+export function buildLeadSystemPrompt(db: Database, teamId: string): string {
+  const team = db.query("SELECT name FROM team WHERE id = ?").get(teamId) as { name: string } | null
+  if (!team) return ""
+
+  const members = db.query("SELECT name, status FROM team_member WHERE team_id = ?").all(teamId) as Array<{ name: string; status: string }>
+
+  const taskCounts = db.query(
+    "SELECT status, COUNT(*) as count FROM team_task WHERE team_id = ? GROUP BY status",
+  ).all(teamId) as Array<{ status: string; count: number }>
+
+  const countMap: Record<string, number> = {}
+  for (const row of taskCounts) {
+    countMap[row.status] = row.count
+  }
+
+  const completed = countMap["completed"] ?? 0
+  const inProgress = countMap["in_progress"] ?? 0
+  const pending = countMap["pending"] ?? 0
+
+  const memberList = members
+    .map((m) => `${m.name} [${STATUS_DISPLAY[m.status] ?? m.status}]`)
+    .join(", ")
+
+  const teammateLine = members.length > 0
+    ? `Teammates: ${memberList}`
+    : "Teammates: none"
+
+  return [
+    `You are leading team "${team.name}" with ${members.length} active teammates.`,
+    teammateLine,
+    `Tasks: ${completed} completed, ${inProgress} in progress, ${pending} pending`,
+    "",
+    "Teammates work asynchronously and message you when done.",
+    "Do NOT poll team_status or team_tasks_list repeatedly — wait for messages.",
+    "After spawning teammates, tell the user what you've set up and wait.",
+    "When all teammates finish, summarize results and suggest next steps.",
+  ].join("\n")
+}
+
+/**
+ * Build the system prompt injected into a teammate's session.
+ * Short role reminder with name, team name, and communication tool.
+ */
+export function buildTeammateSystemPrompt(db: Database, teamId: string, memberName: string): string {
+  const team = db.query("SELECT name FROM team WHERE id = ?").get(teamId) as { name: string } | null
+  if (!team) return ""
+
+  return `You are "${memberName}", a teammate in team "${team.name}". Use team_message to communicate.`
+}
+
+/**
+ * Build a concise context string for compaction.
+ * Includes team name, member statuses, task progress, and role statement.
+ */
+export function buildTeamCompactionContext(
+  db: Database,
+  teamId: string,
+  role: "lead" | "member",
+  memberName?: string,
+): string {
+  const team = db.query("SELECT name FROM team WHERE id = ?").get(teamId) as { name: string } | null
+  if (!team) return ""
+
+  const members = db.query("SELECT name, status FROM team_member WHERE team_id = ?").all(teamId) as Array<{ name: string; status: string }>
+
+  const taskCounts = db.query(
+    "SELECT status, COUNT(*) as count FROM team_task WHERE team_id = ? GROUP BY status",
+  ).all(teamId) as Array<{ status: string; count: number }>
+
+  const countMap: Record<string, number> = {}
+  for (const row of taskCounts) {
+    countMap[row.status] = row.count
+  }
+
+  const completed = countMap["completed"] ?? 0
+  const inProgress = countMap["in_progress"] ?? 0
+  const pending = countMap["pending"] ?? 0
+
+  const roleLine = role === "lead"
+    ? `[Team Context] You are the lead of team "${team.name}".`
+    : `[Team Context] You are a teammate named "${memberName}" in team "${team.name}".`
+
+  const memberList = members
+    .map((m) => `${m.name} (${STATUS_DISPLAY[m.status] ?? m.status})`)
+    .join(", ")
+
+  const membersLine = members.length > 0
+    ? `Members: ${memberList}`
+    : "Members: none"
+
+  return [
+    roleLine,
+    membersLine,
+    `Tasks: ${completed} completed, ${inProgress} in progress, ${pending} pending`,
+  ].join("\n")
+}

package/src/tools/team-cleanup.ts

@@ -3,6 +3,7 @@ import { findTeamBySession } from "../types"
 
 /**
  * Execute the team_cleanup tool. Archives the team and cleans up resources.
+ * Removes worktrees and lists branches for the lead to merge.
  * If force=true, aborts all active member sessions first.
  */
 export async function executeTeamCleanup(
@@ -14,8 +15,8 @@ export async function executeTeamCleanup(
   if (!teamInfo) throw new Error("This session is not in a team.")
   if (teamInfo.role !== "lead") throw new Error("Only the team lead can clean up the team.")
 
-  const members = deps.db.query("SELECT name, session_id, status FROM team_member WHERE team_id = ?")
-    .all(teamInfo.teamId) as Array<{ name: string; session_id: string; status: string }>
+  const members = deps.db.query("SELECT name, session_id, status, worktree_dir, worktree_branch FROM team_member WHERE team_id = ?")
+    .all(teamInfo.teamId) as Array<{ name: string; session_id: string; status: string; worktree_dir: string | null; worktree_branch: string | null }>
 
   const active = members.filter(m => m.status !== "shutdown" && m.status !== "shutdown_requested" && m.status !== "error")
 
@@ -29,9 +30,20 @@ export async function executeTeamCleanup(
     for (const member of active) {
       try {
         await deps.client.session.abort({ path: { id: member.session_id } })
-      } catch {
-        // Best effort
-      }
+      } catch { /* best effort */ }
+    }
+  }
+
+  // Remove worktrees and collect branches for merging
+  const branches: string[] = []
+  for (const member of members) {
+    if (member.worktree_dir) {
+      try {
+        await deps.client.worktree.remove({ worktreeRemoveInput: { directory: member.worktree_dir } })
+      } catch { /* best effort — worktree may already be gone */ }
+    }
+    if (member.worktree_branch) {
+      branches.push(member.worktree_branch)
     }
   }
 
@@ -41,5 +53,10 @@ export async function executeTeamCleanup(
   // Clean up registry
   deps.registry.unregisterTeam(teamInfo.teamId)
 
+  if (branches.length > 0) {
+    const mergeCommands = branches.map(b => `  git merge ${b}`).join("\n")
+    return `Team "${teamInfo.teamName}" cleaned up. Worktrees removed.\n\nMerge teammate branches:\n${mergeCommands}`
+  }
+
   return `Team "${teamInfo.teamName}" cleaned up.`
 }

package/src/tools/team-message.ts

@@ -4,10 +4,11 @@ import { sendMessage, markDelivered } from "../messaging"
 
 /**
  * Execute the team_message tool. Sends a direct message to a teammate or lead.
+ * Optionally approves or rejects a teammate's plan (lead only).
  */
 export async function executeTeamMessage(
   deps: ToolDeps,
-  args: { to: string; text: string },
+  args: { to: string; text: string; approve?: boolean; reject?: string },
   sessionId: string,
 ): Promise<string> {
   const teamInfo = findTeamBySession(deps.db, deps.registry, sessionId)
@@ -18,22 +19,52 @@ export async function executeTeamMessage(
   const recipientSessionId = resolveRecipientSession(deps.db, deps.registry, teamInfo.teamId, args.to)
   if (!recipientSessionId) throw new Error(`Recipient "${args.to}" not found in team "${teamInfo.teamName}"`)
 
+  // Handle plan approval/rejection
+  let messageText = args.text
+  if (args.approve || args.reject) {
+    if (args.approve && args.reject) {
+      throw new Error("Cannot both approve and reject a plan.")
+    }
+    if (teamInfo.role !== "lead") {
+      throw new Error("Only the lead can approve or reject plans.")
+    }
+    const recipient = deps.db.query(
+      "SELECT plan_approval FROM team_member WHERE team_id = ? AND name = ?"
+    ).get(teamInfo.teamId, args.to) as { plan_approval: string } | null
+    if (!recipient || recipient.plan_approval !== "pending") {
+      throw new Error(`Recipient "${args.to}" is not in plan approval mode (plan_approval is not pending).`)
+    }
+    if (args.approve) {
+      deps.db.run(
+        "UPDATE team_member SET plan_approval = 'approved', time_updated = ? WHERE team_id = ? AND name = ?",
+        [Date.now(), teamInfo.teamId, args.to]
+      )
+      messageText = `[Plan Approved] ${args.text}`
+    } else {
+      deps.db.run(
+        "UPDATE team_member SET plan_approval = 'rejected', time_updated = ? WHERE team_id = ? AND name = ?",
+        [Date.now(), teamInfo.teamId, args.to]
+      )
+      messageText = `[Plan Rejected: ${args.reject}] ${args.text}`
+    }
+  }
+
   const msgId = sendMessage(deps.db, {
     teamId: teamInfo.teamId,
     from: senderName,
     to: args.to,
-    content: args.text,
+    content: messageText,
   })
 
   // Deliver via promptAsync — truncate lead-bound messages over 500 chars
   const isToLead = args.to === "lead"
   const MAX_LEAD_MSG = 500
   let deliveryText: string
-  if (isToLead && args.text.length > MAX_LEAD_MSG) {
-    const truncated = args.text.slice(0, MAX_LEAD_MSG)
+  if (isToLead && messageText.length > MAX_LEAD_MSG) {
+    const truncated = messageText.slice(0, MAX_LEAD_MSG)
     deliveryText = `[Team message from ${senderName}]: ${truncated}... (use team_results to read full message)`
   } else {
-    deliveryText = `[Team message from ${senderName}]: ${args.text}`
+    deliveryText = `[Team message from ${senderName}]: ${messageText}`
   }
 
   await deps.client.session.promptAsync({

package/src/tools/team-shutdown.ts

@@ -3,61 +3,93 @@ import { findTeamBySession } from "../types"
 
 /**
  * Execute the team_shutdown tool. Requests a teammate to shut down.
- * Sets status to shutdown_requested, calls abort, then polls session status.
- * If the session is already idle after abort, transitions directly to shutdown.
- * Otherwise the event hook handles the transition when the session goes idle.
+ *
+ * Graceful negotiation flow:
+ * - If member is already shutdown_requested, treat as force (second call).
+ * - If member is idle or force=true, abort immediately and set status='shutdown'.
+ * - If member is busy and force=false, send a shutdown message via promptAsync
+ *   and set status='shutdown_requested'. The member finishes work and reports back.
  */
 export async function executeTeamShutdown(
   deps: ToolDeps,
-  args: { member: string },
+  args: { member: string; force?: boolean },
   sessionId: string,
 ): Promise<string> {
   const teamInfo = findTeamBySession(deps.db, deps.registry, sessionId)
   if (!teamInfo) throw new Error("This session is not in a team.")
   if (teamInfo.role !== "lead") throw new Error("Only the team lead can shut down teammates.")
 
-  const member = deps.db.query("SELECT session_id, status FROM team_member WHERE team_id = ? AND name = ?")
-    .get(teamInfo.teamId, args.member) as { session_id: string; status: string } | null
+  const member = deps.db.query("SELECT session_id, status, worktree_branch FROM team_member WHERE team_id = ? AND name = ?")
+    .get(teamInfo.teamId, args.member) as { session_id: string; status: string; worktree_branch: string | null } | null
   if (!member) throw new Error(`Teammate "${args.member}" not found in team "${teamInfo.teamName}"`)
   if (member.status === "shutdown") throw new Error(`Teammate "${args.member}" is already shut down`)
 
-  // Set to shutdown_requested
-  deps.db.run(
-    "UPDATE team_member SET status = 'shutdown_requested', time_updated = ? WHERE team_id = ? AND name = ?",
-    [Date.now(), teamInfo.teamId, args.member]
-  )
+  const force = args.force ?? false
 
-  // Call abort on the member's session (best effort)
-  try {
-    await deps.client.session.abort({ path: { id: member.session_id } })
-  } catch {
-    // Abort failed — session may already be gone. Fire a warning toast.
-    try {
-      await deps.client.tui.showToast({
-        title: "Team",
-        message: `Failed to abort ${args.member} session — will rely on event hook`,
-        variant: "warning",
-        duration: 4000,
-      })
-    } catch { /* TUI may not be available */ }
+  // Second call on an already-requested member → force abort
+  if (member.status === "shutdown_requested") {
+    await abortAndShutdown(deps, teamInfo.teamId, args.member, member.session_id)
+    const branchInfo = member.worktree_branch ? ` Changes on branch: ${member.worktree_branch}` : ""
+    return `Force shut down "${args.member}".${branchInfo}`
   }
 
-  // Fallback: poll session status after abort. If already idle, transition
-  // directly to shutdown. This handles the case where abort() on an
-  // already-idle session doesn't fire a session.status event.
+  // Determine if member is idle or busy
+  let isIdle = false
   try {
     const statuses = await deps.client.session.status()
     const sessionStatus = statuses.data?.[member.session_id]
-    if (!sessionStatus || sessionStatus.type === "idle") {
-      deps.db.run(
-        "UPDATE team_member SET status = 'shutdown', execution_status = 'idle', time_updated = ? WHERE team_id = ? AND name = ?",
-        [Date.now(), teamInfo.teamId, args.member]
-      )
-      return `Teammate "${args.member}" has been shut down.`
-    }
+    isIdle = !sessionStatus || sessionStatus.type === "idle"
+  } catch {
+    // Status poll failed — assume busy, fall through to graceful path
+  }
+
+  if (isIdle || force) {
+    await abortAndShutdown(deps, teamInfo.teamId, args.member, member.session_id)
+    const branchInfo = member.worktree_branch ? ` Changes on branch: ${member.worktree_branch}` : ""
+    return `Teammate "${args.member}" has been shut down.${branchInfo}`
+  }
+
+  // Busy + not force → graceful: send shutdown message, set shutdown_requested
+  try {
+    await deps.client.session.promptAsync({
+      path: { id: member.session_id },
+      body: {
+        parts: [{
+          type: "text",
+          text: `[Shutdown requested]: The lead has requested you shut down. Finish your current task, send your final findings to the lead via team_message, then stop.`,
+        }],
+      },
+    })
   } catch {
-    // Status poll failed — fall through to eventual consistency
+    // promptAsync failed — best effort
   }
 
-  return `Shutdown requested for ${args.member}. Will complete when current work finishes.`
+  deps.db.run(
+    "UPDATE team_member SET status = 'shutdown_requested', time_updated = ? WHERE team_id = ? AND name = ?",
+    [Date.now(), teamInfo.teamId, args.member],
+  )
+
+  return `Shutdown requested for ${args.member}. They will finish current work and shut down. Call team_shutdown with force: true to abort immediately.`
+}
+
+/**
+ * Abort a member's session and set their status to shutdown.
+ * Best-effort: if abort fails, we still mark them as shutdown.
+ */
+async function abortAndShutdown(
+  deps: ToolDeps,
+  teamId: string,
+  memberName: string,
+  sessionId: string,
+): Promise<void> {
+  try {
+    await deps.client.session.abort({ path: { id: sessionId } })
+  } catch {
+    // Abort failed — session may already be gone
+  }
+
+  deps.db.run(
+    "UPDATE team_member SET status = 'shutdown', execution_status = 'idle', time_updated = ? WHERE team_id = ? AND name = ?",
+    [Date.now(), teamId, memberName],
+  )
 }