clanker-agent/packages/clanker-teams/extensions/index.ts

import type { ExtensionAPI } from "@mariozechner/clanker-coding-agent";
import { Type } from "@sinclair/typebox";
import { StringEnum } from "@mariozechner/clanker-ai";
import * as paths from "../src/utils/paths";
import * as teams from "../src/utils/teams";
import * as tasks from "../src/utils/tasks";
import * as messaging from "../src/utils/messaging";
import { Member } from "../src/utils/models";
import { getTerminalAdapter } from "../src/adapters/terminal-registry";
import { Iterm2Adapter } from "../src/adapters/iterm2-adapter";
import * as path from "node:path";
import * as fs from "node:fs";
import { spawnSync } from "node:child_process";

// Cache for available models
let availableModelsCache: Array<{ provider: string; model: string }> | null =
  null;
let modelsCacheTime = 0;
const MODELS_CACHE_TTL = 60000; // 1 minute

/**
 * Query available models from clanker --list-models
 */
function getAvailableModels(): Array<{ provider: string; model: string }> {
  const now = Date.now();
  if (availableModelsCache && now - modelsCacheTime < MODELS_CACHE_TTL) {
    return availableModelsCache;
  }

  try {
    const result = spawnSync("clanker", ["--list-models"], {
      encoding: "utf-8",
      timeout: 10000,
    });

    if (result.status !== 0 || !result.stdout) {
      return [];
    }

    const models: Array<{ provider: string; model: string }> = [];
    const lines = result.stdout.split("\n");

    for (const line of lines) {
      // Skip header line and empty lines
      if (!line.trim() || line.startsWith("provider")) continue;

      // Parse: provider model context max-out thinking images
      const parts = line.trim().split(/\s+/);
      if (parts.length >= 2) {
        const provider = parts[0];
        const model = parts[1];
        if (provider && model) {
          models.push({ provider, model });
        }
      }
    }

    availableModelsCache = models;
    modelsCacheTime = now;
    return models;
  } catch (e) {
    return [];
  }
}

/**
 * Provider priority list - OAuth/subscription providers first (cheaper), then API-key providers
 */
const PROVIDER_PRIORITY = [
  // OAuth / Subscription providers (typically free/cheaper)
  "google-gemini-cli", // Google Gemini CLI - OAuth, free tier
  "github-copilot", // GitHub Copilot - subscription
  "kimi-sub", // Kimi subscription
  // API key providers
  "anthropic",
  "openai",
  "google",
  "zai",
  "openrouter",
  "azure-openai",
  "amazon-bedrock",
  "mistral",
  "groq",
  "cerebras",
  "xai",
  "vercel-ai-gateway",
];

/**
 * Find the best matching provider for a given model name.
 * Returns the full provider/model string or null if not found.
 */
function resolveModelWithProvider(modelName: string): string | null {
  // If already has provider prefix, return as-is
  if (modelName.includes("/")) {
    return modelName;
  }

  const availableModels = getAvailableModels();
  if (availableModels.length === 0) {
    return null;
  }

  const lowerModelName = modelName.toLowerCase();

  // Find all exact matches (case-insensitive) and sort by provider priority
  const exactMatches = availableModels.filter(
    (m) => m.model.toLowerCase() === lowerModelName,
  );

  if (exactMatches.length > 0) {
    // Sort by provider priority (lower index = higher priority)
    exactMatches.sort((a, b) => {
      const aIndex = PROVIDER_PRIORITY.indexOf(a.provider);
      const bIndex = PROVIDER_PRIORITY.indexOf(b.provider);
      // If provider not in priority list, put it at the end
      const aPriority = aIndex === -1 ? 999 : aIndex;
      const bPriority = bIndex === -1 ? 999 : bIndex;
      return aPriority - bPriority;
    });
    return `${exactMatches[0].provider}/${exactMatches[0].model}`;
  }

  // Try partial match (model name contains the search term)
  const partialMatches = availableModels.filter((m) =>
    m.model.toLowerCase().includes(lowerModelName),
  );

  if (partialMatches.length > 0) {
    for (const preferredProvider of PROVIDER_PRIORITY) {
      const match = partialMatches.find(
        (m) => m.provider === preferredProvider,
      );
      if (match) {
        return `${match.provider}/${match.model}`;
      }
    }
    // Return first match if no preferred provider found
    return `${partialMatches[0].provider}/${partialMatches[0].model}`;
  }

  return null;
}

export default function (clanker: ExtensionAPI) {
  const isTeammate = !!process.env.CLANKER_AGENT_NAME;
  const agentName = process.env.CLANKER_AGENT_NAME || "team-lead";
  const teamName = process.env.CLANKER_TEAM_NAME;

  const terminal = getTerminalAdapter();

  clanker.on("session_start", async (_event, ctx) => {
    paths.ensureDirs();
    if (isTeammate) {
      if (teamName) {
        const pidFile = path.join(paths.teamDir(teamName), `${agentName}.pid`);
        fs.writeFileSync(pidFile, process.pid.toString());
      }
      ctx.ui.notify(`Teammate: ${agentName} (Team: ${teamName})`, "info");
      ctx.ui.setStatus("00-clanker-teams", `[${agentName.toUpperCase()}]`);

      if (terminal) {
        const fullTitle = teamName ? `${teamName}: ${agentName}` : agentName;
        const setIt = () => {
          if ((ctx.ui as any).setTitle) (ctx.ui as any).setTitle(fullTitle);
          terminal.setTitle(fullTitle);
        };
        setIt();
        setTimeout(setIt, 500);
        setTimeout(setIt, 2000);
        setTimeout(setIt, 5000);
      }

      setTimeout(() => {
        clanker.sendUserMessage(
          `I am starting my work as '${agentName}' on team '${teamName}'. Checking my inbox for instructions...`,
        );
      }, 1000);

      setInterval(async () => {
        if (ctx.isIdle() && teamName) {
          const unread = await messaging.readInbox(
            teamName,
            agentName,
            true,
            false,
          );
          if (unread.length > 0) {
            clanker.sendUserMessage(
              `I have ${unread.length} new message(s) in my inbox. Reading them now...`,
            );
          }
        }
      }, 30000);
    } else if (teamName) {
      ctx.ui.setStatus("clanker-teams", `Lead @ ${teamName}`);
    }
  });

  clanker.on("turn_start", async (_event, ctx) => {
    if (isTeammate) {
      const fullTitle = teamName ? `${teamName}: ${agentName}` : agentName;
      if ((ctx.ui as any).setTitle) (ctx.ui as any).setTitle(fullTitle);
      if (terminal) terminal.setTitle(fullTitle);
    }
  });

  let firstTurn = true;
  clanker.on("before_agent_start", async (event, ctx) => {
    if (isTeammate && firstTurn) {
      firstTurn = false;

      let modelInfo = "";
      if (teamName) {
        try {
          const teamConfig = await teams.readConfig(teamName);
          const member = teamConfig.members.find((m) => m.name === agentName);
          if (member && member.model) {
            modelInfo = `\nYou are currently using model: ${member.model}`;
            if (member.thinking) {
              modelInfo += ` with thinking level: ${member.thinking}`;
            }
            modelInfo += `. When reporting your model or thinking level, use these exact values.`;
          }
        } catch (e) {
          // Ignore
        }
      }

      return {
        systemPrompt:
          event.systemPrompt +
          `\n\nYou are teammate '${agentName}' on team '${teamName}'.\nYour lead is 'team-lead'.${modelInfo}\nStart by calling read_inbox(team_name="${teamName}") to get your initial instructions.`,
      };
    }
  });

  async function killTeammate(teamName: string, member: Member) {
    if (member.name === "team-lead") return;

    const pidFile = path.join(paths.teamDir(teamName), `${member.name}.pid`);
    if (fs.existsSync(pidFile)) {
      try {
        const pid = fs.readFileSync(pidFile, "utf-8").trim();
        process.kill(parseInt(pid), "SIGKILL");
        fs.unlinkSync(pidFile);
      } catch (e) {
        // ignore
      }
    }

    if (member.windowId && terminal) {
      terminal.killWindow(member.windowId);
    }

    if (member.tmuxPaneId && terminal) {
      terminal.kill(member.tmuxPaneId);
    }
  }

  // Tools
  clanker.registerTool({
    name: "team_create",
    label: "Create Team",
    description: "Create a new agent team.",
    parameters: Type.Object({
      team_name: Type.String(),
      description: Type.Optional(Type.String()),
      default_model: Type.Optional(Type.String()),
      separate_windows: Type.Optional(
        Type.Boolean({
          default: false,
          description: "Open teammates in separate OS windows instead of panes",
        }),
      ),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const config = teams.createTeam(
        params.team_name,
        "local-session",
        "lead-agent",
        params.description,
        params.default_model,
        params.separate_windows,
      );
      return {
        content: [{ type: "text", text: `Team ${params.team_name} created.` }],
        details: { config },
      };
    },
  });

  clanker.registerTool({
    name: "spawn_teammate",
    label: "Spawn Teammate",
    description: "Spawn a new teammate in a terminal pane or separate window.",
    parameters: Type.Object({
      team_name: Type.String(),
      name: Type.String(),
      prompt: Type.String(),
      cwd: Type.String(),
      model: Type.Optional(Type.String()),
      thinking: Type.Optional(
        StringEnum(["off", "minimal", "low", "medium", "high"]),
      ),
      plan_mode_required: Type.Optional(Type.Boolean({ default: false })),
      separate_window: Type.Optional(Type.Boolean({ default: false })),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const safeName = paths.sanitizeName(params.name);
      const safeTeamName = paths.sanitizeName(params.team_name);

      if (!teams.teamExists(safeTeamName)) {
        throw new Error(`Team ${params.team_name} does not exist`);
      }

      if (!terminal) {
        throw new Error("No terminal adapter detected.");
      }

      const teamConfig = await teams.readConfig(safeTeamName);
      let chosenModel = params.model || teamConfig.defaultModel;

      // Resolve model to provider/model format
      if (chosenModel) {
        if (!chosenModel.includes("/")) {
          // Try to resolve using available models from clanker --list-models
          const resolved = resolveModelWithProvider(chosenModel);
          if (resolved) {
            chosenModel = resolved;
          } else if (
            teamConfig.defaultModel &&
            teamConfig.defaultModel.includes("/")
          ) {
            // Fall back to team default provider
            const [provider] = teamConfig.defaultModel.split("/");
            chosenModel = `${provider}/${chosenModel}`;
          }
        }
      }

      const useSeparateWindow =
        params.separate_window ?? teamConfig.separateWindows ?? false;
      if (useSeparateWindow && !terminal.supportsWindows()) {
        throw new Error(
          `Separate windows mode is not supported in ${terminal.name}.`,
        );
      }

      const member: Member = {
        agentId: `${safeName}@${safeTeamName}`,
        name: safeName,
        agentType: "teammate",
        model: chosenModel,
        joinedAt: Date.now(),
        tmuxPaneId: "",
        cwd: params.cwd,
        subscriptions: [],
        prompt: params.prompt,
        color: "blue",
        thinking: params.thinking,
        planModeRequired: params.plan_mode_required,
      };

      await teams.addMember(safeTeamName, member);
      await messaging.sendPlainMessage(
        safeTeamName,
        "team-lead",
        safeName,
        params.prompt,
        "Initial prompt",
      );

      const piBinary = "clanker";
      let piCmd = piBinary;

      if (chosenModel) {
        // Use the combined --model provider/model:thinking format
        if (params.thinking) {
          piCmd = `${piBinary} --model ${chosenModel}:${params.thinking}`;
        } else {
          piCmd = `${piBinary} --model ${chosenModel}`;
        }
      } else if (params.thinking) {
        piCmd = `${piBinary} --thinking ${params.thinking}`;
      }

      const env: Record<string, string> = {
        ...process.env,
        CLANKER_TEAM_NAME: safeTeamName,
        CLANKER_AGENT_NAME: safeName,
      };

      let terminalId = "";
      let isWindow = false;

      try {
        if (useSeparateWindow) {
          isWindow = true;
          terminalId = terminal.spawnWindow({
            name: safeName,
            cwd: params.cwd,
            command: piCmd,
            env: env,
            teamName: safeTeamName,
          });
          await teams.updateMember(safeTeamName, safeName, {
            windowId: terminalId,
          });
        } else {
          if (terminal instanceof Iterm2Adapter) {
            const teammates = teamConfig.members.filter(
              (m) =>
                m.agentType === "teammate" && m.tmuxPaneId.startsWith("iterm_"),
            );
            const lastTeammate =
              teammates.length > 0 ? teammates[teammates.length - 1] : null;
            if (lastTeammate?.tmuxPaneId) {
              terminal.setSpawnContext({
                lastSessionId: lastTeammate.tmuxPaneId.replace("iterm_", ""),
              });
            } else {
              terminal.setSpawnContext({});
            }
          }

          terminalId = terminal.spawn({
            name: safeName,
            cwd: params.cwd,
            command: piCmd,
            env: env,
          });
          await teams.updateMember(safeTeamName, safeName, {
            tmuxPaneId: terminalId,
          });
        }
      } catch (e) {
        throw new Error(
          `Failed to spawn ${terminal.name} ${isWindow ? "window" : "pane"}: ${e}`,
        );
      }

      return {
        content: [
          {
            type: "text",
            text: `Teammate ${params.name} spawned in ${isWindow ? "window" : "pane"} ${terminalId}.`,
          },
        ],
        details: { agentId: member.agentId, terminalId, isWindow },
      };
    },
  });

  clanker.registerTool({
    name: "spawn_lead_window",
    label: "Spawn Lead Window",
    description: "Open the team lead in a separate OS window.",
    parameters: Type.Object({
      team_name: Type.String(),
      cwd: Type.Optional(Type.String()),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const safeTeamName = paths.sanitizeName(params.team_name);
      if (!teams.teamExists(safeTeamName))
        throw new Error(`Team ${params.team_name} does not exist`);
      if (!terminal || !terminal.supportsWindows())
        throw new Error("Windows mode not supported.");

      const teamConfig = await teams.readConfig(safeTeamName);
      const cwd = params.cwd || process.cwd();
      const piBinary = "clanker";
      let piCmd = piBinary;
      if (teamConfig.defaultModel) {
        // Use the combined --model provider/model format
        piCmd = `${piBinary} --model ${teamConfig.defaultModel}`;
      }

      const env = {
        ...process.env,
        CLANKER_TEAM_NAME: safeTeamName,
        CLANKER_AGENT_NAME: "team-lead",
      };
      try {
        const windowId = terminal.spawnWindow({
          name: "team-lead",
          cwd,
          command: piCmd,
          env,
          teamName: safeTeamName,
        });
        await teams.updateMember(safeTeamName, "team-lead", { windowId });
        return {
          content: [{ type: "text", text: `Lead window spawned: ${windowId}` }],
          details: { windowId },
        };
      } catch (e) {
        throw new Error(`Failed: ${e}`);
      }
    },
  });

  clanker.registerTool({
    name: "send_message",
    label: "Send Message",
    description: "Send a message to a teammate.",
    parameters: Type.Object({
      team_name: Type.String(),
      recipient: Type.String(),
      content: Type.String(),
      summary: Type.String(),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      await messaging.sendPlainMessage(
        params.team_name,
        agentName,
        params.recipient,
        params.content,
        params.summary,
      );
      return {
        content: [
          { type: "text", text: `Message sent to ${params.recipient}.` },
        ],
        details: {},
      };
    },
  });

  clanker.registerTool({
    name: "broadcast_message",
    label: "Broadcast Message",
    description: "Broadcast a message to all team members except the sender.",
    parameters: Type.Object({
      team_name: Type.String(),
      content: Type.String(),
      summary: Type.String(),
      color: Type.Optional(Type.String()),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      await messaging.broadcastMessage(
        params.team_name,
        agentName,
        params.content,
        params.summary,
        params.color,
      );
      return {
        content: [
          { type: "text", text: `Message broadcasted to all team members.` },
        ],
        details: {},
      };
    },
  });

  clanker.registerTool({
    name: "read_inbox",
    label: "Read Inbox",
    description: "Read messages from an agent's inbox.",
    parameters: Type.Object({
      team_name: Type.String(),
      agent_name: Type.Optional(
        Type.String({
          description: "Whose inbox to read. Defaults to your own.",
        }),
      ),
      unread_only: Type.Optional(Type.Boolean({ default: true })),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const targetAgent = params.agent_name || agentName;
      const msgs = await messaging.readInbox(
        params.team_name,
        targetAgent,
        params.unread_only,
      );
      return {
        content: [{ type: "text", text: JSON.stringify(msgs, null, 2) }],
        details: { messages: msgs },
      };
    },
  });

  clanker.registerTool({
    name: "task_create",
    label: "Create Task",
    description: "Create a new team task.",
    parameters: Type.Object({
      team_name: Type.String(),
      subject: Type.String(),
      description: Type.String(),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const task = await tasks.createTask(
        params.team_name,
        params.subject,
        params.description,
      );
      return {
        content: [{ type: "text", text: `Task ${task.id} created.` }],
        details: { task },
      };
    },
  });

  clanker.registerTool({
    name: "task_submit_plan",
    label: "Submit Plan",
    description: "Submit a plan for a task, updating its status to 'planning'.",
    parameters: Type.Object({
      team_name: Type.String(),
      task_id: Type.String(),
      plan: Type.String(),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const updated = await tasks.submitPlan(
        params.team_name,
        params.task_id,
        params.plan,
      );
      return {
        content: [
          { type: "text", text: `Plan submitted for task ${params.task_id}.` },
        ],
        details: { task: updated },
      };
    },
  });

  clanker.registerTool({
    name: "task_evaluate_plan",
    label: "Evaluate Plan",
    description: "Evaluate a submitted plan for a task.",
    parameters: Type.Object({
      team_name: Type.String(),
      task_id: Type.String(),
      action: StringEnum(["approve", "reject"]),
      feedback: Type.Optional(
        Type.String({ description: "Required for rejection" }),
      ),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const updated = await tasks.evaluatePlan(
        params.team_name,
        params.task_id,
        params.action as any,
        params.feedback,
      );
      return {
        content: [
          {
            type: "text",
            text: `Plan for task ${params.task_id} has been ${params.action}d.`,
          },
        ],
        details: { task: updated },
      };
    },
  });

  clanker.registerTool({
    name: "task_list",
    label: "List Tasks",
    description: "List all tasks for a team.",
    parameters: Type.Object({
      team_name: Type.String(),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const taskList = await tasks.listTasks(params.team_name);
      return {
        content: [{ type: "text", text: JSON.stringify(taskList, null, 2) }],
        details: { tasks: taskList },
      };
    },
  });

  clanker.registerTool({
    name: "task_update",
    label: "Update Task",
    description: "Update a task's status or owner.",
    parameters: Type.Object({
      team_name: Type.String(),
      task_id: Type.String(),
      status: Type.Optional(
        StringEnum([
          "pending",
          "planning",
          "in_progress",
          "completed",
          "deleted",
        ]),
      ),
      owner: Type.Optional(Type.String()),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const updated = await tasks.updateTask(params.team_name, params.task_id, {
        status: params.status as any,
        owner: params.owner,
      });
      return {
        content: [{ type: "text", text: `Task ${params.task_id} updated.` }],
        details: { task: updated },
      };
    },
  });

  clanker.registerTool({
    name: "team_shutdown",
    label: "Shutdown Team",
    description: "Shutdown the entire team and close all panes/windows.",
    parameters: Type.Object({
      team_name: Type.String(),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const teamName = params.team_name;
      try {
        const config = await teams.readConfig(teamName);
        for (const member of config.members) {
          await killTeammate(teamName, member);
        }
        const dir = paths.teamDir(teamName);
        const tasksDir = paths.taskDir(teamName);
        if (fs.existsSync(tasksDir)) fs.rmSync(tasksDir, { recursive: true });
        if (fs.existsSync(dir)) fs.rmSync(dir, { recursive: true });
        return {
          content: [{ type: "text", text: `Team ${teamName} shut down.` }],
          details: {},
        };
      } catch (e) {
        throw new Error(`Failed to shutdown team: ${e}`);
      }
    },
  });

  clanker.registerTool({
    name: "task_read",
    label: "Read Task",
    description: "Read details of a specific task.",
    parameters: Type.Object({
      team_name: Type.String(),
      task_id: Type.String(),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const task = await tasks.readTask(params.team_name, params.task_id);
      return {
        content: [{ type: "text", text: JSON.stringify(task, null, 2) }],
        details: { task },
      };
    },
  });

  clanker.registerTool({
    name: "check_teammate",
    label: "Check Teammate",
    description: "Check a single teammate's status.",
    parameters: Type.Object({
      team_name: Type.String(),
      agent_name: Type.String(),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const config = await teams.readConfig(params.team_name);
      const member = config.members.find((m) => m.name === params.agent_name);
      if (!member) throw new Error(`Teammate ${params.agent_name} not found`);

      let alive = false;
      if (member.windowId && terminal) {
        alive = terminal.isWindowAlive(member.windowId);
      } else if (member.tmuxPaneId && terminal) {
        alive = terminal.isAlive(member.tmuxPaneId);
      }

      const unreadCount = (
        await messaging.readInbox(
          params.team_name,
          params.agent_name,
          true,
          false,
        )
      ).length;
      return {
        content: [
          {
            type: "text",
            text: JSON.stringify({ alive, unreadCount }, null, 2),
          },
        ],
        details: { alive, unreadCount },
      };
    },
  });

  clanker.registerTool({
    name: "process_shutdown_approved",
    label: "Process Shutdown Approved",
    description: "Process a teammate's shutdown.",
    parameters: Type.Object({
      team_name: Type.String(),
      agent_name: Type.String(),
    }),
    async execute(toolCallId, params: any, signal, onUpdate, ctx) {
      const config = await teams.readConfig(params.team_name);
      const member = config.members.find((m) => m.name === params.agent_name);
      if (!member) throw new Error(`Teammate ${params.agent_name} not found`);

      await killTeammate(params.team_name, member);
      await teams.removeMember(params.team_name, params.agent_name);
      return {
        content: [
          {
            type: "text",
            text: `Teammate ${params.agent_name} has been shut down.`,
          },
        ],
        details: {},
      };
    },
  });
}