macro-agent 0.2.3 → 0.2.5

package/src/agent/agent-manager-v2.ts

@@ -48,7 +48,7 @@ import type {
   AgentLifecycleEvent,
   AgentConfig,
   ContinueAgentOptions,
-  MCPServerConfig,
+  MCPServerStdioConfig,
 } from "./types.js";
 import { AgentManagerError } from "./types.js";
 import type { RoleRegistry, Capability } from "../roles/types.js";
@@ -220,6 +220,12 @@ export function createAgentManagerV2(
   const activeSessions = new Map<AgentId, ActiveSession>();
   const agentWorkspaces = new Map<AgentId, Workspace>();
   const lifecycleListeners = new Set<AgentLifecycleCallback>();
+  // In-flight fire-and-forget terminate() promises (e.g. the promptUntilDone
+  // completion poller tears the agent down without awaiting it, to keep
+  // completion latency low). close() awaits these so shutdown cannot return
+  // while a group-kill is still racing — which would let process.exit orphan
+  // the subtree.
+  const pendingTerminations = new Set<Promise<unknown>>();
   let spawnInterceptor: SpawnInterceptor | null = null;
   let isShuttingDown = false;
   let mapServerUrl: string | undefined;
@@ -285,7 +291,7 @@ export function createAgentManagerV2(
     lineage?: string[];
     sessionId?: string;
     streamId?: string;
-  }): MCPServerConfig {
+  }): MCPServerStdioConfig {
     const env: Record<string, string> = {
       MACRO_AGENT_ID: opts.agentId,
       MACRO_PARENT_ID: opts.parentId,
@@ -759,15 +765,28 @@ export function createAgentManagerV2(
             value: v,
           })),
         },
-        ...(agentConfig?.mcpServers?.map((s) => ({
-          name: s.name,
-          command: s.command,
-          args: s.args ?? [],
-          env: Object.entries(s.env ?? {}).map(([k, v]) => ({
-            name: k,
-            value: v,
-          })),
-        })) ?? []),
+        ...(agentConfig?.mcpServers?.map((s) => {
+          if ("command" in s) {
+            return {
+              name: s.name,
+              command: s.command,
+              args: s.args ?? [],
+              env: Object.entries(s.env ?? {}).map(([k, v]) => ({
+                name: k,
+                value: v,
+              })),
+            };
+          }
+          return {
+            name: s.name,
+            type: s.type,
+            url: s.url,
+            headers: Object.entries(s.headers ?? {}).map(([k, v]) => ({
+              name: k,
+              value: v,
+            })),
+          };
+        }) ?? []),
       ];
 
       // Always-on subsystem MCP servers (the "trinity"). The macro-agent
@@ -1753,63 +1772,189 @@ export function createAgentManagerV2(
     options?: {
       maxFollowUps?: number;
       onUpdate?: (update: ExtendedSessionUpdate) => void;
+      /**
+       * Optional caller-supplied completion predicate. When it returns true the
+       * loop stops consuming, terminates the agent as "completed", and returns
+       * `completedExternally: true` — WITHOUT requiring the macro `done()` tool.
+       *
+       * Checked three ways (all fire terminate + return completedExternally):
+       *   (a) reactively after each streamed update (fast path),
+       *   (b) reactively after each attempt's stream ends (fast path),
+       *   (c) on a CONCURRENT ~2s timer while consuming the stream. (c) is the
+       *       critical path for the live τ rollout: the agent runs the whole
+       *       episode in ONE prompt() call, calls τ's finish, the env writes its
+       *       reward sink (done:true), and then the prompt() stream goes SILENT
+       *       and does NOT close — the async iterator just blocks. With only the
+       *       reactive checks (a)/(b) the predicate is never re-evaluated and
+       *       solve() hangs until the rollout backstop. The concurrent poller
+       *       observes done:true independent of update arrival, flips
+       *       completedExternally, and terminate()s the agent — which ends the
+       *       in-flight prompt() generator (via return or throw, both handled).
+       *
+       * Generic by design: this module has NO knowledge of what the predicate
+       * reads (e.g. an external reward sink). Back-compat: when absent, NO poller
+       * is started and behavior is unchanged — only `done()` detection drives
+       * completion.
+       *
+       * @param pollIntervalMs Override the concurrent poll cadence (default
+       *   2000ms). Primarily a test seam.
+       */
+      isComplete?: () => boolean | Promise<boolean>;
+      pollIntervalMs?: number;
     }
   ): Promise<{
     doneCalled: boolean;
     doneStatus?: string;
+    completedExternally: boolean;
     updates: ExtendedSessionUpdate[];
   }> {
     const maxFollowUps = options?.maxFollowUps ?? 2;
+    const pollIntervalMs = options?.pollIntervalMs ?? 2000;
     const allUpdates: ExtendedSessionUpdate[] = [];
     let doneCalled = false;
     let doneStatus: string | undefined;
+    let completedExternally = false;
 
     let currentMessage = message;
 
     for (let attempt = 0; attempt <= maxFollowUps; attempt++) {
-      for await (const update of prompt(agentId, currentMessage)) {
-        allUpdates.push(update);
-        options?.onUpdate?.(update);
+      // ── Concurrent completion poller ─────────────────────────────────────
+      // Only armed when the caller supplies a predicate (back-compat: no
+      // predicate → no poller, identical behavior). The live failure mode is a
+      // prompt() stream that has gone SILENT but not CLOSED after the episode
+      // externally completed; the reactive per-update check never fires again
+      // because no updates arrive. This timer re-evaluates isComplete()
+      // independent of update arrival. On true it flips completedExternally and
+      // terminate()s the agent, which ends the in-flight prompt() generator
+      // (return or throw) so the for-await below unblocks. The interval is
+      // always cleared in `finally` so no timer leaks across attempts.
+      let pollTimer: ReturnType<typeof setInterval> | undefined;
+      let polling = false;
+
+      // Resolves the instant completion is detected (external predicate OR the
+      // macro done() tool). The attempt is RACED against this signal so it
+      // returns WITHOUT waiting for the prompt() generator to drain or for
+      // terminate() to finish. The prior design relied on terminate() ending
+      // the silent in-flight stream to unblock the for-await; when that teardown
+      // stalled (e.g. a long, busy real episode whose subprocess is slow to die),
+      // promptUntilDone never returned and the macro session stayed "running"
+      // until the multi-hour hard-timer — the refinement "finalization hang".
+      let signalDone: () => void = () => {};
+      const doneSignal = new Promise<void>((resolve) => {
+        signalDone = resolve;
+      });
 
-        // Detect done() tool call from session updates.
-        // acp-factory uses { sessionUpdate: "tool_call", title: "mcp__macro-agent__done" }
-        const uAny = update as any;
+      const startPoller = () => {
+        if (!options?.isComplete) return;
+        pollTimer = setInterval(() => {
+          if (polling || completedExternally) return;
+          polling = true;
+          void (async () => {
+            try {
+              if (!completedExternally && options.isComplete && (await options.isComplete())) {
+                completedExternally = true;
+                signalDone();
+                // Tear down the agent/session in the BACKGROUND — never await it
+                // here, so a slow teardown cannot gate the return. Tracked in
+                // pendingTerminations so close() awaits it before shutdown
+                // returns (else a still-racing group-kill is abandoned by a
+                // subsequent process.exit, orphaning the subtree). Best effort.
+                {
+                  const t = terminate(agentId, "completed" as any).catch(() => {});
+                  pendingTerminations.add(t);
+                  void t.finally(() => pendingTerminations.delete(t));
+                }
+              }
+            } catch {
+              // Predicate errored — leave completion to the reactive checks.
+            } finally {
+              polling = false;
+            }
+          })();
+        }, pollIntervalMs);
+      };
 
-        // Check title field (primary detection)
-        if (
-          (uAny.sessionUpdate === "tool_call" || uAny.sessionUpdate === "tool_call_update") &&
-          typeof uAny.title === "string" &&
-          uAny.title.endsWith("__done")
-        ) {
-          doneCalled = true;
-          // Extract status from rawInput (may arrive across multiple updates —
-          // first update has rawInput={}, subsequent has full input)
-          try {
-            const raw = uAny.rawInput;
-            const input =
-              typeof raw === "string" ? JSON.parse(raw) :
-              typeof raw === "object" ? raw :
-              uAny.input;
-            if (input?.status) {
-              doneStatus = input.status;
+      // Consume the stream as a background task so a silent-but-open generator
+      // can be raced against the completion signal below.
+      const consume = (async () => {
+        for await (const update of prompt(agentId, currentMessage)) {
+          allUpdates.push(update);
+          options?.onUpdate?.(update);
+
+          // External completion predicate (reactive fast path).
+          if (options?.isComplete && (await options.isComplete())) {
+            completedExternally = true;
+            signalDone();
+            break;
+          }
+          if (completedExternally) break;
+
+          // Detect done() tool call from session updates.
+          // acp-factory uses { sessionUpdate: "tool_call", title: "mcp__macro-agent__done" }
+          const uAny = update as any;
+
+          // Check title field (primary detection)
+          if (
+            (uAny.sessionUpdate === "tool_call" || uAny.sessionUpdate === "tool_call_update") &&
+            typeof uAny.title === "string" &&
+            uAny.title.endsWith("__done")
+          ) {
+            doneCalled = true;
+            // Extract status from rawInput (may arrive across multiple updates —
+            // first update has rawInput={}, subsequent has full input)
+            try {
+              const raw = uAny.rawInput;
+              const input =
+                typeof raw === "string" ? JSON.parse(raw) :
+                typeof raw === "object" ? raw :
+                uAny.input;
+              if (input?.status) {
+                doneStatus = input.status;
+              }
+            } catch {
+              // Best effort — rawInput may not be parseable yet
             }
-          } catch {
-            // Best effort — rawInput may not be parseable yet
+            signalDone();
           }
-        }
 
-        // Fallback: check older format
-        if (
-          uAny.type === "result" &&
-          uAny.subtype === "tool_result" &&
-          uAny.toolName === "done"
-        ) {
-          doneCalled = true;
-          doneStatus = uAny.result?.status;
+          // Fallback: check older format
+          if (
+            uAny.type === "result" &&
+            uAny.subtype === "tool_result" &&
+            uAny.toolName === "done"
+          ) {
+            doneCalled = true;
+            doneStatus = uAny.result?.status;
+            signalDone();
+          }
         }
+      })();
+      // Prevent an unhandled rejection if we return via doneSignal before the
+      // now-orphaned stream throws from a background teardown.
+      consume.catch(() => {});
+
+      try {
+        startPoller();
+        // Return as soon as EITHER the stream drains naturally OR completion is
+        // signalled — whichever comes first. This is the key change: completion
+        // no longer depends on the generator unblocking.
+        await Promise.race([consume, doneSignal]);
+      } catch (err) {
+        // A genuine stream error (not a post-completion teardown) re-throws.
+        if (!completedExternally && !doneCalled) throw err;
+      } finally {
+        // Always clear the concurrent poller so no timer leaks across attempts.
+        if (pollTimer) clearInterval(pollTimer);
+      }
+
+      // Re-check after the stream ends in case completion landed exactly as the
+      // attempt's stream drained (the inner break / poller already set the flag
+      // when it fired mid-stream).
+      if (!completedExternally && options?.isComplete && (await options.isComplete())) {
+        completedExternally = true;
       }
 
-      if (doneCalled) break;
+      if (doneCalled || completedExternally) break;
 
       if (attempt < maxFollowUps) {
         currentMessage =
@@ -1817,9 +1962,11 @@ export function createAgentManagerV2(
       }
     }
 
-    // Auto-terminate when done() was called and the handler signaled shouldTerminate.
-    // This closes the lifecycle gap: without this, agents stay in "running" state
-    // after calling done() because nothing triggers terminate().
+    // Auto-terminate when the agent completed — either via the macro done() tool
+    // or via the caller's external completion predicate. This closes the
+    // lifecycle gap: without this, agents stay in "running" state because
+    // nothing triggers terminate(). The external path mirrors the done() path
+    // (terminate as "completed").
     if (doneCalled) {
       const reason = doneStatus === "completed" ? "completed" : (doneStatus ?? "failed");
       try {
@@ -1827,9 +1974,15 @@ export function createAgentManagerV2(
       } catch {
         // Best effort — agent may already be stopping
       }
+    } else if (completedExternally) {
+      try {
+        await terminate(agentId, "completed" as any);
+      } catch {
+        // Best effort — agent may already be stopping
+      }
     }
 
-    return { doneCalled, doneStatus, updates: allUpdates };
+    return { doneCalled, doneStatus, completedExternally, updates: allUpdates };
   }
 
   function getSession(agentId: AgentId): Session | null {
@@ -1974,6 +2127,19 @@ export function createAgentManagerV2(
       });
     }
     activeSessions.clear();
+
+    // Await any in-flight fire-and-forget terminations (e.g. the completion
+    // poller's background teardown) so shutdown does not return while a
+    // group-kill is still racing — otherwise a follow-on process.exit abandons
+    // it and orphans the acp/claude/mcp_env subtree. Bounded by a backstop so a
+    // genuinely stuck close() cannot wedge shutdown.
+    if (pendingTerminations.size > 0) {
+      await Promise.race([
+        Promise.allSettled([...pendingTerminations]),
+        new Promise((resolve) => setTimeout(resolve, 6000)),
+      ]);
+    }
+
     agentWorkspaces.clear();
     lifecycleListeners.clear();
 

package/src/agent/agent-manager.ts

@@ -164,10 +164,27 @@ export interface AgentManager {
       maxFollowUps?: number;
       /** Callback for each update during prompting */
       onUpdate?: (update: ExtendedSessionUpdate) => void;
+      /**
+       * Optional caller-supplied completion predicate. When it returns true the
+       * loop stops early and terminates the agent as "completed" WITHOUT
+       * requiring the macro done() tool. Back-compat: absent → unchanged.
+       *
+       * Evaluated reactively per-update/per-attempt AND on a concurrent timer
+       * (see implementation) so a silent-but-open prompt() stream still
+       * completes. No predicate → no poller, unchanged behavior.
+       */
+      isComplete?: () => boolean | Promise<boolean>;
+      /**
+       * Override the concurrent isComplete() poll cadence in ms (default 2000).
+       * Only meaningful alongside `isComplete`; primarily a test seam.
+       */
+      pollIntervalMs?: number;
     },
   ): Promise<{
     doneCalled: boolean;
     doneStatus?: string;
+    /** True when the caller's isComplete() predicate ended the loop. */
+    completedExternally: boolean;
     updates: ExtendedSessionUpdate[];
   }>;
 

package/src/agent/types.ts

@@ -202,15 +202,36 @@ export interface AgentConfig {
 }
 
 /**
- * MCP server configuration
+ * MCP server configuration.
+ *
+ * A discriminated union over transport. The legacy stdio shape (no `type`,
+ * or `type: "stdio"`) is preserved byte-for-byte; new `http`/`sse` variants
+ * carry a `url` + optional `headers` instead of a `command`. These map
+ * directly onto the ACP wire's `McpServerStdio` / `McpServerHttp` /
+ * `McpServerSse` shapes (see acp-factory `createSession({ mcpServers })`).
  */
-export interface MCPServerConfig {
+export type MCPServerConfig = MCPServerStdioConfig | MCPServerHttpConfig;
+
+/** Stdio transport (legacy default — preserved for M0 compatibility). */
+export interface MCPServerStdioConfig {
   name: string;
+  /** Optional transport discriminant; absent means stdio. */
+  type?: "stdio";
   command: string;
   args?: string[];
   env?: Record<string, string>;
 }
 
+/** HTTP (streamable-http) or SSE transport for a remote MCP server. */
+export interface MCPServerHttpConfig {
+  name: string;
+  type: "http" | "sse";
+  /** URL to the remote MCP server. */
+  url: string;
+  /** HTTP headers to set on requests to the server. */
+  headers?: Record<string, string>;
+}
+
 // ─────────────────────────────────────────────────────────────────
 // Spawned Agent Result
 // ─────────────────────────────────────────────────────────────────

package/src/cognitive/__tests__/macro-agent-backend.test.ts

@@ -6,16 +6,23 @@ import type { RoleRegistry, RoleDefinition } from "../../roles/types.js";
 
 // ── Mock Helpers ─────────────────────────────────────────────────
 
+// Stand-in for the built-in "generic" role the REAL DefaultRoleRegistry falls
+// back to. Used by resolveRole() below so the mock matches production
+// semantics (resolveRole NEVER throws on a miss).
+const GENERIC_FALLBACK = { name: "generic" } as unknown as RoleDefinition;
+
 function createMockRoleRegistry(): RoleRegistry {
   const roles = new Map<string, RoleDefinition>();
   return {
     registerRole: vi.fn((role: RoleDefinition) => {
       roles.set(role.name, role);
     }),
+    // MUST mirror the real DefaultRoleRegistry.resolveRole(): on an unknown
+    // role it does NOT throw — it logs and returns the GenericRole fallback.
+    // A throwing mock (the old behavior) masked the constructor bug where
+    // analyst registration was gated on resolveRole() throwing.
     resolveRole: vi.fn((name: string) => {
-      const role = roles.get(name);
-      if (!role) throw new Error(`Role not found: ${name}`);
-      return role;
+      return roles.get(name) ?? GENERIC_FALLBACK;
     }),
     getRole: vi.fn((name: string) => roles.get(name)),
     hasCapability: vi.fn(() => true),
@@ -92,10 +99,33 @@ describe("MacroAgentBackend", () => {
       );
     });
 
+    it("registers analyst against a real-semantics registry (resolveRole does NOT throw on miss)", () => {
+      // Regression guard for the live Arm-B bug: the constructor must NOT rely
+      // on resolveRole() throwing to detect a missing analyst role. With a
+      // non-throwing registry (production behavior), analyst must still be
+      // registered, and resolveRole("analyst") must then return analyst —
+      // NOT fall back to "generic" (which would give the spawned agent
+      // all built-in tools + own-workspace + persistent lifecycle, exactly
+      // the conditions that produced 0 env-tool calls and a 600s timeout).
+      const registry = createMockRoleRegistry();
+      const am = createMockAgentManager({
+        getRoleRegistry: vi.fn().mockReturnValue(registry),
+      });
+
+      new MacroAgentBackend(am);
+
+      expect(registry.registerRole).toHaveBeenCalledWith(
+        expect.objectContaining({ name: "analyst" }),
+      );
+      // After construction, the role must resolve to analyst, not generic.
+      expect(registry.resolveRole("analyst").name).toBe("analyst");
+    });
+
     it("skips registration if analyst role already exists", () => {
       const registry = createMockRoleRegistry();
-      // Pre-register analyst
-      (registry.resolveRole as ReturnType<typeof vi.fn>).mockReturnValue({
+      // Pre-register analyst via getRole (the exact-match existence check the
+      // constructor now uses), mirroring the real registry's API.
+      (registry.getRole as ReturnType<typeof vi.fn>).mockReturnValue({
         name: "analyst",
       });
 
@@ -203,6 +233,129 @@ describe("MacroAgentBackend", () => {
       );
     });
 
+    it("forwards mcpServers into the raw macro spawn config (Arm B chain)", async () => {
+      const mcpServers = [
+        {
+          name: "tauenv",
+          command: "/venv/bin/python",
+          args: ["-m", "autonomation_tau_bench.mcp_env", "--split", "airline"],
+          env: { TAU_REWARD_FILE: "/tmp/reward.json" },
+        },
+      ];
+
+      await backend.spawn({
+        agentType: "claude-code",
+        task: { description: "handle the customer" },
+        mcpServers,
+      });
+
+      expect(agentManager.spawn).toHaveBeenCalledWith(
+        expect.objectContaining({
+          config: expect.objectContaining({ mcpServers }),
+        }),
+      );
+    });
+
+    it("forwards both env and mcpServers together into the raw spawn config", async () => {
+      const mcpServers = [
+        { name: "tauenv", command: "/venv/bin/python", args: ["-m", "x"] },
+      ];
+
+      await backend.spawn({
+        agentType: "claude-code",
+        task: { description: "test" },
+        env: { TAU_REWARD_FILE: "/tmp/r.json" },
+        mcpServers,
+      });
+
+      expect(agentManager.spawn).toHaveBeenCalledWith(
+        expect.objectContaining({
+          config: {
+            env: { TAU_REWARD_FILE: "/tmp/r.json" },
+            mcpServers,
+          },
+        }),
+      );
+    });
+
+    it("invokes beforeSpawn BEFORE driving the agent (per-attempt reset hook)", async () => {
+      // beforeSpawn must run before agentManager.spawn so the env subprocess /
+      // reward sink is reset for THIS attempt. We assert ordering by recording
+      // the call sequence.
+      const order: string[] = [];
+      const beforeSpawn = vi.fn(() => {
+        order.push("beforeSpawn");
+      });
+      const am = createMockAgentManager({
+        spawn: vi.fn().mockImplementation(async () => {
+          order.push("spawn");
+          return { id: "agent_test123", session_id: "session_test123" };
+        }),
+      });
+      const b = new MacroAgentBackend(am);
+
+      await b.spawn({
+        agentType: "claude-code",
+        task: { description: "tau episode" },
+        beforeSpawn,
+      });
+
+      expect(beforeSpawn).toHaveBeenCalledTimes(1);
+      expect(order).toEqual(["beforeSpawn", "spawn"]);
+    });
+
+    it("passes completionSignal into promptUntilDone (S1 threading)", async () => {
+      const completionSignal = vi.fn(() => false);
+
+      await backend.spawn({
+        agentType: "claude-code",
+        task: { description: "tau episode" },
+        completionSignal,
+      });
+
+      await vi.waitFor(() => {
+        expect(agentManager.promptUntilDone).toHaveBeenCalledWith(
+          expect.any(String),
+          expect.any(String),
+          expect.objectContaining({ isComplete: completionSignal }),
+        );
+      });
+    });
+
+    it("treats completedExternally as a success (session → completed) even without done()", async () => {
+      const am = createMockAgentManager({
+        promptUntilDone: vi.fn().mockResolvedValue({
+          doneCalled: false,
+          completedExternally: true,
+          updates: [],
+        }),
+      });
+      const b = new MacroAgentBackend(am);
+
+      const session = await b.spawn({
+        agentType: "claude-code",
+        task: { description: "tau episode" },
+        completionSignal: () => true,
+      });
+
+      await vi.waitFor(async () => {
+        const s = await b.getSession(session.id);
+        expect(s!.state).toBe("completed");
+        expect(s!.error).toBeUndefined();
+      });
+    });
+
+    it("passes config: undefined when neither env nor mcpServers set (M0 back-compat)", async () => {
+      await backend.spawn({
+        agentType: "claude-code",
+        task: { description: "test" },
+      });
+
+      expect(agentManager.spawn).toHaveBeenCalledWith(
+        expect.objectContaining({ config: undefined }),
+      );
+    });
+
     it("stores macro agent ID in session metadata", async () => {
       const session = await backend.spawn({
         agentType: "claude-code",