@desplega.ai/agent-swarm 1.52.1 → 1.53.0

package/src/oauth/ensure-token.ts

@@ -0,0 +1,50 @@
+import { getOAuthApp, getOAuthTokens, isTokenExpiringSoon } from "../be/db-queries/oauth";
+import { type OAuthProviderConfig, refreshAccessToken } from "./wrapper";
+
+/**
+ * Build an OAuthProviderConfig from the oauth_apps table for any provider.
+ */
+function getOAuthConfig(provider: string): OAuthProviderConfig | null {
+  const app = getOAuthApp(provider);
+  if (!app) return null;
+
+  const metadata = JSON.parse(app.metadata || "{}");
+  return {
+    provider,
+    clientId: app.clientId,
+    clientSecret: app.clientSecret,
+    authorizeUrl: app.authorizeUrl,
+    tokenUrl: app.tokenUrl,
+    redirectUri: app.redirectUri,
+    scopes: app.scopes.split(","),
+    extraParams: metadata.extraParams ?? (metadata.actor ? { actor: metadata.actor } : undefined),
+  };
+}
+
+/**
+ * Ensure a valid OAuth token exists for the given provider.
+ * If the token is expiring soon, attempt to refresh it.
+ * Call this before any API interaction with an OAuth-protected service.
+ */
+export async function ensureToken(provider: string): Promise<void> {
+  if (!isTokenExpiringSoon(provider)) return;
+
+  const config = getOAuthConfig(provider);
+  const tokens = getOAuthTokens(provider);
+  if (!config || !tokens?.refreshToken) {
+    console.warn(
+      `[OAuth] ${provider} token expiring but cannot refresh (missing config or refresh token)`,
+    );
+    return;
+  }
+
+  try {
+    await refreshAccessToken(config, tokens.refreshToken);
+    console.log(`[OAuth] ${provider} token refreshed successfully`);
+  } catch (err) {
+    console.error(
+      `[OAuth] Failed to refresh ${provider} token:`,
+      err instanceof Error ? err.message : err,
+    );
+  }
+}

package/src/prompts/base-prompt.ts

@@ -40,6 +40,8 @@ export type BasePromptArgs = {
   };
   /** Pre-fetched skill summaries for the installed skills section */
   skillsSummary?: { name: string; description: string }[];
+  /** Pre-fetched MCP server summaries for the installed MCP servers section */
+  mcpServersSummary?: string;
 };
 
 export const getBasePrompt = async (args: BasePromptArgs): Promise<string> => {
@@ -76,6 +78,11 @@ export const getBasePrompt = async (args: BasePromptArgs): Promise<string> => {
     prompt += `\n\n## Installed Skills\n\nThe following skills are available. Use the Skill tool to invoke them by name.\n\n${summaries}\n`;
   }
 
+  // Installed MCP servers section
+  if (args.mcpServersSummary) {
+    prompt += `\n\n## Installed MCP Servers\n\nThe following MCP servers are configured for your use:\n${args.mcpServersSummary}\n`;
+  }
+
   // Repo context (protected, never truncated)
   if (args.repoContext) {
     prompt += "\n\n## Repository Context\n\n";

package/src/providers/claude-adapter.ts

@@ -1,5 +1,6 @@
 import { unlink, writeFile } from "node:fs/promises";
 import { dirname, join } from "node:path";
+import { computeContextUsed, getContextWindowSize } from "../utils/context-window";
 import { validateClaudeCredentials } from "../utils/credentials";
 import {
   parseStderrForErrors,
@@ -40,14 +41,83 @@ async function cleanupTaskFile(pid: number): Promise<void> {
   }
 }
 
+/** Fetch installed MCP servers from the API and return them as .mcp.json-compatible entries */
+async function fetchInstalledMcpServers(
+  apiUrl: string,
+  apiKey: string,
+  agentId: string,
+): Promise<Record<string, Record<string, unknown>> | null> {
+  try {
+    const res = await fetch(`${apiUrl}/api/agents/${agentId}/mcp-servers?resolveSecrets=true`, {
+      headers: {
+        Authorization: `Bearer ${apiKey}`,
+        "X-Agent-ID": agentId,
+      },
+    });
+    if (!res.ok) return null;
+
+    const data = (await res.json()) as {
+      servers: Array<{
+        name: string;
+        transport: string;
+        isActive: boolean;
+        isEnabled: boolean;
+        command?: string;
+        args?: string;
+        url?: string;
+        headers?: string;
+        resolvedEnv?: Record<string, string>;
+        resolvedHeaders?: Record<string, string>;
+      }>;
+    };
+
+    const entries: Record<string, Record<string, unknown>> = {};
+    for (const srv of data.servers.filter((s) => s.isActive && s.isEnabled)) {
+      if (srv.transport === "stdio" && srv.command) {
+        let args: string[] = [];
+        try {
+          args = srv.args ? JSON.parse(srv.args) : [];
+        } catch {
+          // invalid JSON — use empty args
+        }
+        entries[srv.name] = {
+          command: srv.command,
+          args,
+          env: srv.resolvedEnv || {},
+        };
+      } else if ((srv.transport === "http" || srv.transport === "sse") && srv.url) {
+        let parsedHeaders: Record<string, string> = {};
+        try {
+          parsedHeaders = srv.headers ? JSON.parse(srv.headers) : {};
+        } catch {
+          // invalid JSON — use empty headers
+        }
+        entries[srv.name] = {
+          type: srv.transport,
+          url: srv.url,
+          headers: { ...parsedHeaders, ...(srv.resolvedHeaders || {}) },
+        };
+      }
+    }
+    return Object.keys(entries).length > 0 ? entries : null;
+  } catch {
+    return null;
+  }
+}
+
 /**
- * Create a per-session MCP config file with X-Source-Task-Id header injected.
+ * Create a per-session MCP config file with X-Source-Task-Id header injected
+ * and installed MCP servers merged in.
  * Each session gets its own copy to avoid race conditions when multiple concurrent
  * Claude sessions share the same cwd. The session-specific config is passed to
  * Claude CLI via --mcp-config, so the shared .mcp.json is never modified.
  * Returns the path to the per-session config, or null if no config exists.
  */
-async function createSessionMcpConfig(cwd: string, taskId: string): Promise<string | null> {
+async function createSessionMcpConfig(
+  cwd: string,
+  taskId: string,
+  installedServers?: Record<string, Record<string, unknown>> | null,
+): Promise<string | null> {
   // Walk up from cwd to find .mcp.json (mirrors Claude CLI's project-level config discovery).
   // In Docker, .mcp.json lives at /workspace/.mcp.json but tasks often run with cwd set to
   // a subdirectory like /workspace/repos/<repo>, so a single-directory check misses it.
@@ -64,24 +134,42 @@ async function createSessionMcpConfig(cwd: string, taskId: string): Promise<stri
     searchDir = parent;
   }
 
-  if (!mcpJsonPath) return null;
+  if (!mcpJsonPath && !installedServers) return null;
 
   try {
-    const file = Bun.file(mcpJsonPath);
-
-    const config = await file.json();
+    let config: { mcpServers?: Record<string, unknown> } = { mcpServers: {} };
+    if (mcpJsonPath) {
+      const file = Bun.file(mcpJsonPath);
+      config = await file.json();
+    }
     const servers = config?.mcpServers;
-    if (!servers) return null;
+    if (!servers && !installedServers) return null;
+
+    if (!config.mcpServers) config.mcpServers = {};
 
     // Find the agent-swarm server entry (could be named "agent-swarm" or similar)
-    const serverKey = Object.keys(servers).find(
-      (k) => k === "agent-swarm" || servers[k]?.headers?.["X-Agent-ID"],
+    const serverKey = Object.keys(config.mcpServers).find(
+      (k) =>
+        k === "agent-swarm" ||
+        ((config.mcpServers![k] as Record<string, unknown>)?.headers &&
+          ((config.mcpServers![k] as Record<string, Record<string, unknown>>).headers?.[
+            "X-Agent-ID"
+          ] as unknown)),
     );
-    if (!serverKey) return null;
+    if (serverKey) {
+      const server = config.mcpServers[serverKey] as Record<string, unknown>;
+      if (!server.headers) server.headers = {};
+      (server.headers as Record<string, string>)["X-Source-Task-Id"] = taskId;
+    }
 
-    const server = servers[serverKey];
-    if (!server.headers) server.headers = {};
-    server.headers["X-Source-Task-Id"] = taskId;
+    // Merge installed MCP servers (don't overwrite existing entries)
+    if (installedServers) {
+      for (const [name, serverConfig] of Object.entries(installedServers)) {
+        if (!config.mcpServers[name]) {
+          config.mcpServers[name] = serverConfig;
+        }
+      }
+    }
 
     // Write per-session config to /tmp — no race, each session has its own file
     const sessionConfigPath = `/tmp/mcp-${taskId}.json`;
@@ -102,6 +190,7 @@ class ClaudeSession implements ProviderSession {
   private completionPromise: Promise<ProviderResult>;
   private errorTracker = new SessionErrorTracker();
   private taskFilePid: number;
+  private contextWindowSize: number;
 
   constructor(
     private config: ProviderSessionConfig,
@@ -112,6 +201,7 @@ class ClaudeSession implements ProviderSession {
     private claudeBinary: string = "claude",
   ) {
     this.taskFilePid = taskFilePid;
+    this.contextWindowSize = getContextWindowSize(model);
     const cmd = this.buildCommand();
 
     console.log(
@@ -279,6 +369,19 @@ class ClaudeSession implements ProviderSession {
       if (json.type === "system" && json.subtype === "init" && json.session_id) {
         this._sessionId = json.session_id;
         this.emit({ type: "session_init", sessionId: json.session_id });
+        if (json.model) {
+          this.contextWindowSize = getContextWindowSize(json.model);
+        }
+      }
+
+      // Compaction detection
+      if (json.type === "system" && json.subtype === "compact_boundary" && json.compact_metadata) {
+        this.emit({
+          type: "compaction",
+          preCompactTokens: json.compact_metadata.pre_tokens ?? 0,
+          compactTrigger: json.compact_metadata.trigger ?? "auto",
+          contextTotalTokens: this.contextWindowSize,
+        });
       }
 
       // Cost data from result
@@ -312,6 +415,14 @@ class ClaudeSession implements ProviderSession {
           cost,
           isError: json.is_error || false,
         });
+
+        // Update context window size from modelUsage if available
+        if (json.modelUsage) {
+          const modelKey = Object.keys(json.modelUsage)[0];
+          if (modelKey && json.modelUsage[modelKey]?.contextWindow) {
+            this.contextWindowSize = json.modelUsage[modelKey].contextWindow;
+          }
+        }
       }
 
       // Tool use from assistant messages — emit tool_start for auto-progress
@@ -331,6 +442,21 @@ class ClaudeSession implements ProviderSession {
             }
           }
         }
+
+        // Context usage extraction from assistant message usage
+        if (json.message.usage) {
+          const usage = json.message.usage;
+          const contextUsed = computeContextUsed(usage);
+          const contextTotal = this.contextWindowSize;
+
+          this.emit({
+            type: "context_usage",
+            contextUsedTokens: contextUsed,
+            contextTotalTokens: contextTotal,
+            contextPercent: contextTotal > 0 ? (contextUsed / contextTotal) * 100 : 0,
+            outputTokens: usage.output_tokens ?? 0,
+          });
+        }
       }
 
       trackErrorFromJson(json, this.errorTracker);
@@ -435,8 +561,23 @@ export class ClaudeAdapter implements ProviderAdapter {
 
     console.log(`\x1b[2m[${config.role}]\x1b[0m Task file written: ${taskFilePath}`);
 
-    // Create per-session MCP config with X-Source-Task-Id header (no shared-file race condition)
-    const sessionMcpConfig = await createSessionMcpConfig(config.cwd, config.taskId);
+    // Fetch installed MCP servers from API for this agent
+    const installedServers =
+      config.apiUrl && config.apiKey && config.agentId
+        ? await fetchInstalledMcpServers(config.apiUrl, config.apiKey, config.agentId)
+        : null;
+    if (installedServers) {
+      console.log(
+        `\x1b[2m[${config.role}]\x1b[0m Merging ${Object.keys(installedServers).length} installed MCP server(s) into session config`,
+      );
+    }
+
+    // Create per-session MCP config with X-Source-Task-Id header + installed servers (no shared-file race condition)
+    const sessionMcpConfig = await createSessionMcpConfig(
+      config.cwd,
+      config.taskId,
+      installedServers,
+    );
 
     return new ClaudeSession(
       config,

package/src/providers/pi-mono-adapter.ts

@@ -389,6 +389,74 @@ export class PiMonoAdapter implements ProviderAdapter {
       } catch (err) {
         console.warn(`\x1b[33m[${config.role}] Failed to discover MCP tools: ${err}\x1b[0m`);
       }
+
+      // 2b. Discover tools from installed MCP servers (HTTP/SSE transport only)
+      try {
+        const mcpServersRes = await fetch(
+          `${config.apiUrl}/api/agents/${config.agentId}/mcp-servers?resolveSecrets=true`,
+          {
+            headers: {
+              Authorization: `Bearer ${config.apiKey}`,
+              "X-Agent-ID": config.agentId,
+            },
+          },
+        );
+        if (mcpServersRes.ok) {
+          const mcpServersData = (await mcpServersRes.json()) as {
+            servers: Array<{
+              name: string;
+              transport: string;
+              url?: string;
+              headers?: string;
+              isActive: boolean;
+              isEnabled: boolean;
+              resolvedHeaders?: Record<string, string>;
+            }>;
+          };
+          const httpServers = mcpServersData.servers.filter(
+            (s) =>
+              s.isActive &&
+              s.isEnabled &&
+              (s.transport === "http" || s.transport === "sse") &&
+              s.url,
+          );
+
+          for (const srv of httpServers) {
+            try {
+              const srvClient = new McpHttpClient(srv.url!, "", "");
+              srvClient.useRawUrl = true;
+              // Build custom headers from static headers + resolved secret headers
+              let parsedHeaders: Record<string, string> = {};
+              try {
+                parsedHeaders = srv.headers ? JSON.parse(srv.headers) : {};
+              } catch {
+                // invalid JSON
+              }
+              srvClient.customHeaders = {
+                ...parsedHeaders,
+                ...(srv.resolvedHeaders || {}),
+              };
+              await srvClient.initialize();
+              const srvTools = await srvClient.listTools();
+              // Prefix tool names with mcp__<server-name>__ to avoid conflicts
+              const prefixed = mcpToolsToDefinitions(srvClient, srvTools).map((t) => ({
+                ...t,
+                name: `mcp__${srv.name}__${t.name}`,
+              }));
+              customTools.push(...prefixed);
+              console.log(
+                `\x1b[2m[${config.role}]\x1b[0m Discovered ${srvTools.length} tools from MCP server "${srv.name}"`,
+              );
+            } catch (srvErr) {
+              console.warn(
+                `\x1b[33m[${config.role}] Failed to discover tools from MCP server "${srv.name}": ${srvErr}\x1b[0m`,
+              );
+            }
+          }
+        }
+      } catch {
+        // Non-fatal — installed MCP server tool discovery is optional
+      }
     }
 
     // 3. Resolve model

package/src/providers/pi-mono-extension.ts

@@ -383,6 +383,8 @@ ${transcript}`;
  */
 export function createSwarmHooksExtension(config: SwarmHooksConfig): ExtensionFactory {
   return (pi) => {
+    let lastContextPostTime = 0;
+
     // === session_start → SessionStart ===
     pi.on("session_start", async (_event, _ctx) => {
       // Ping server
@@ -479,7 +481,7 @@ export function createSwarmHooksExtension(config: SwarmHooksConfig): ExtensionFa
     });
 
     // === tool_result → PostToolUse ===
-    pi.on("tool_result", async (event, _ctx) => {
+    pi.on("tool_result", async (event, ctx) => {
       // Heartbeat (workers only, fire-and-forget)
       if (!config.isLead && config.taskId) {
         fireAndForget(`${config.apiUrl}/api/active-sessions/heartbeat/${config.taskId}`, {
@@ -494,6 +496,26 @@ export function createSwarmHooksExtension(config: SwarmHooksConfig): ExtensionFa
         headers: apiHeaders(config),
       });
 
+      // Throttled context usage reporting (every 30s)
+      const usage = ctx.getContextUsage?.();
+      if (config.taskId && usage?.tokens != null) {
+        const now = Date.now();
+        if (now - lastContextPostTime >= 30_000) {
+          lastContextPostTime = now;
+          fireAndForget(`${config.apiUrl}/api/tasks/${config.taskId}/context`, {
+            method: "POST",
+            headers: apiHeaders(config),
+            body: JSON.stringify({
+              eventType: "progress",
+              sessionId: `pi-${config.taskId}`,
+              contextUsedTokens: usage.tokens,
+              contextTotalTokens: usage.contextWindow,
+              contextPercent: usage.percent,
+            }),
+          });
+        }
+      }
+
       // Shared disk write failure detection (Archil only — safety net)
       if (process.env.ARCHIL_MOUNT_TOKEN) {
         // Pi-mono uses lowercase tool names
@@ -562,7 +584,7 @@ export function createSwarmHooksExtension(config: SwarmHooksConfig): ExtensionFa
     // === context → PreCompact ===
     // The context event allows injecting messages before compaction.
     // We log the goal reminder to console (it gets captured in context).
-    pi.on("context", async (_event, _ctx) => {
+    pi.on("context", async (_event, ctx) => {
       if (!config.taskId) return undefined;
 
       try {
@@ -583,6 +605,22 @@ export function createSwarmHooksExtension(config: SwarmHooksConfig): ExtensionFa
         /* don't block compaction */
       }
 
+      // Report context usage as a compaction event
+      const usage = ctx.getContextUsage?.();
+      if (usage) {
+        fireAndForget(`${config.apiUrl}/api/tasks/${config.taskId}/context`, {
+          method: "POST",
+          headers: apiHeaders(config),
+          body: JSON.stringify({
+            eventType: "compaction",
+            sessionId: `pi-${config.taskId}`,
+            contextUsedTokens: usage.tokens ?? undefined,
+            contextTotalTokens: usage.contextWindow,
+            contextPercent: usage.percent ?? undefined,
+          }),
+        });
+      }
+
       return undefined;
     });
 
@@ -604,6 +642,22 @@ export function createSwarmHooksExtension(config: SwarmHooksConfig): ExtensionFa
 
     // === session_shutdown → Stop ===
     pi.on("session_shutdown", async (_event, ctx) => {
+      // Post final context usage before shutdown
+      const usage = ctx.getContextUsage?.();
+      if (config.taskId && usage) {
+        await fetch(`${config.apiUrl}/api/tasks/${config.taskId}/context`, {
+          method: "POST",
+          headers: apiHeaders(config),
+          body: JSON.stringify({
+            eventType: "completion",
+            sessionId: `pi-${config.taskId}`,
+            contextTotalTokens: usage.contextWindow,
+            contextPercent: usage.percent ?? undefined,
+            contextUsedTokens: usage.tokens ?? undefined,
+          }),
+        }).catch(() => {});
+      }
+
       // Sync identity files and setup script
       await syncIdentityFilesToServer(config);
       await syncSetupScriptToServer(config);

package/src/providers/pi-mono-mcp-client.ts

@@ -20,6 +20,13 @@ interface McpToolCallResult {
 export class McpHttpClient {
   private sessionId: string | null = null;
   private nextId = 1;
+  /** Additional headers merged into every request (e.g. for installed MCP servers) */
+  public customHeaders: Record<string, string> = {};
+  /**
+   * When true, baseUrl is used as-is for requests (external MCP servers).
+   * When false (default), /mcp is appended (swarm convention).
+   */
+  public useRawUrl = false;
 
   constructor(
     private baseUrl: string,
@@ -34,6 +41,7 @@ export class McpHttpClient {
       Accept: "application/json, text/event-stream",
       Authorization: `Bearer ${this.apiKey}`,
       "X-Agent-ID": this.agentId,
+      ...this.customHeaders,
     };
     if (this.taskId) {
       headers["X-Source-Task-Id"] = this.taskId;
@@ -42,7 +50,8 @@ export class McpHttpClient {
       headers["mcp-session-id"] = this.sessionId;
     }
 
-    const res = await fetch(`${this.baseUrl}/mcp`, {
+    const url = this.useRawUrl ? this.baseUrl : `${this.baseUrl}/mcp`;
+    const res = await fetch(url, {
       method: "POST",
       headers,
       body: JSON.stringify(body),

package/src/providers/types.ts

@@ -24,7 +24,20 @@ export type ProviderEvent =
   | { type: "error"; message: string; category?: string }
   | { type: "raw_log"; content: string }
   | { type: "raw_stderr"; content: string }
-  | { type: "custom"; name: string; data: unknown };
+  | { type: "custom"; name: string; data: unknown }
+  | {
+      type: "context_usage";
+      contextUsedTokens: number;
+      contextTotalTokens: number;
+      contextPercent: number;
+      outputTokens: number;
+    }
+  | {
+      type: "compaction";
+      preCompactTokens: number;
+      compactTrigger: "auto" | "manual";
+      contextTotalTokens: number;
+    };
 
 /** Configuration passed to a provider adapter to create a session. */
 export interface ProviderSessionConfig {

package/src/server.ts

@@ -25,6 +25,16 @@ import { registerJoinSwarmTool } from "./tools/join-swarm";
 // Messaging capability
 import { registerListChannelsTool } from "./tools/list-channels";
 import { registerListServicesTool } from "./tools/list-services";
+// MCP Servers capability
+import {
+  registerMcpServerCreateTool,
+  registerMcpServerDeleteTool,
+  registerMcpServerGetTool,
+  registerMcpServerInstallTool,
+  registerMcpServerListTool,
+  registerMcpServerUninstallTool,
+  registerMcpServerUpdateTool,
+} from "./tools/mcp-servers";
 // Memory capability
 import { registerMemoryGetTool } from "./tools/memory-get";
 import { registerMemorySearchTool } from "./tools/memory-search";
@@ -272,5 +282,14 @@ export function createServer() {
   registerSkillSyncRemoteTool(server);
   registerSkillPublishTool(server);
 
+  // MCP Servers - always registered
+  registerMcpServerCreateTool(server);
+  registerMcpServerUpdateTool(server);
+  registerMcpServerDeleteTool(server);
+  registerMcpServerGetTool(server);
+  registerMcpServerListTool(server);
+  registerMcpServerInstallTool(server);
+  registerMcpServerUninstallTool(server);
+
   return server;
 }

package/src/tests/context-window.test.ts

@@ -0,0 +1,66 @@
+import { describe, expect, test } from "bun:test";
+import { computeContextUsed, getContextWindowSize } from "../utils/context-window";
+
+describe("getContextWindowSize", () => {
+  test("returns 1M for opus models", () => {
+    expect(getContextWindowSize("claude-opus-4-6")).toBe(1_000_000);
+    expect(getContextWindowSize("opus")).toBe(1_000_000);
+  });
+
+  test("returns 1M for sonnet models", () => {
+    expect(getContextWindowSize("claude-sonnet-4-6")).toBe(1_000_000);
+    expect(getContextWindowSize("sonnet")).toBe(1_000_000);
+  });
+
+  test("returns 200K for haiku models", () => {
+    expect(getContextWindowSize("claude-haiku-4-5")).toBe(200_000);
+    expect(getContextWindowSize("haiku")).toBe(200_000);
+  });
+
+  test("returns 200K default for unknown models", () => {
+    expect(getContextWindowSize("gpt-5")).toBe(200_000);
+    expect(getContextWindowSize("unknown-model")).toBe(200_000);
+    expect(getContextWindowSize("")).toBe(200_000);
+  });
+
+  test("returns default entry value", () => {
+    expect(getContextWindowSize("default")).toBe(200_000);
+  });
+});
+
+describe("computeContextUsed", () => {
+  test("sums all token fields", () => {
+    expect(
+      computeContextUsed({
+        input_tokens: 1000,
+        cache_creation_input_tokens: 500,
+        cache_read_input_tokens: 200,
+      }),
+    ).toBe(1700);
+  });
+
+  test("handles missing fields as zero", () => {
+    expect(computeContextUsed({})).toBe(0);
+    expect(computeContextUsed({ input_tokens: 100 })).toBe(100);
+    expect(computeContextUsed({ cache_read_input_tokens: 50 })).toBe(50);
+  });
+
+  test("handles null fields as zero", () => {
+    expect(
+      computeContextUsed({
+        input_tokens: null,
+        cache_creation_input_tokens: null,
+        cache_read_input_tokens: null,
+      }),
+    ).toBe(0);
+  });
+
+  test("handles mix of values, nulls, and missing", () => {
+    expect(
+      computeContextUsed({
+        input_tokens: 5000,
+        cache_creation_input_tokens: null,
+      }),
+    ).toBe(5000);
+  });
+});