fix(agents): resolve terminal from admitted gateway (#128348)

Co-authored-by: RoboClaw <309084314+roboclaw-bot@users.noreply.github.com>
Co-authored-by: VACInc <3279061+VACInc@users.noreply.github.com>

src/agents/tools/terminal-tool.test.ts

@@ -13,13 +13,15 @@ import {
 import type { spawnTerminalPty } from "../../process/terminal-pty.js";
 import { GATEWAY_OWNER_ONLY_CORE_TOOLS } from "../../security/dangerous-tools.js";
 import { compactToolOutputHint } from "../tool-schema-hints.js";
+import { withGatewayToolCallerIdentity } from "./gateway-caller-context.js";
 import { createTerminalTool } from "./terminal-tool.js";
 
 const callInProcessGatewayTool = vi.hoisted(() => vi.fn(async () => ({ ok: true })));
+const getInProcessGatewayToolContext = vi.hoisted(() => vi.fn());
 
 vi.mock("./in-process-gateway.js", () => ({
   callInProcessGatewayTool,
-  getInProcessGatewayToolContext: vi.fn(),
+  getInProcessGatewayToolContext,
 }));
 
 type TerminalPtyHandle = Awaited<ReturnType<typeof spawnTerminalPty>>;
@@ -85,6 +87,7 @@ describe("terminal tool", () => {
   beforeEach(() => {
     resetAgentRunRegistryForTest();
     callInProcessGatewayTool.mockClear();
+    getInProcessGatewayToolContext.mockReset();
   });
 
   it("uses a flat action enum and the owner-only core gate", () => {
@@ -105,6 +108,176 @@ describe("terminal tool", () => {
     expect(GATEWAY_OWNER_ONLY_CORE_TOOLS).toContain("terminal");
   });
 
+  it("uses the admitted caller Gateway before ambient context", async () => {
+    const callerManager = new TerminalSessionManager({ emit: vi.fn(), spawn: vi.fn() });
+    const ambientManager = new TerminalSessionManager({ emit: vi.fn(), spawn: vi.fn() });
+    const callerList = vi.spyOn(callerManager, "listAgent");
+    const ambientList = vi.spyOn(ambientManager, "listAgent");
+    const gatewayContextResolver = vi.fn();
+    gatewayContextResolver.mockReturnValue(makeContext(callerManager));
+    getInProcessGatewayToolContext.mockReturnValue(makeContext(ambientManager));
+    const tool = createTerminalTool({
+      agentId: "main",
+      agentSessionKey: "agent:main:main",
+      sessionId: "main-session-id",
+    });
+
+    const result = await withGatewayToolCallerIdentity(
+      {
+        agentId: "main",
+        sessionKey: "agent:main:main",
+        gatewayContextResolver,
+      },
+      async () => await tool.execute("list", { action: "list" }),
+    );
+
+    expect(result.details).toEqual({ sessions: [] });
+    expect(callerList).toHaveBeenCalledOnce();
+    expect(ambientList).not.toHaveBeenCalled();
+  });
+
+  it("fails closed when the admitted caller Gateway has retired", async () => {
+    const ambientManager = new TerminalSessionManager({ emit: vi.fn(), spawn: vi.fn() });
+    const ambientList = vi.spyOn(ambientManager, "listAgent");
+    getInProcessGatewayToolContext.mockReturnValue(makeContext(ambientManager));
+    const tool = createTerminalTool({
+      agentId: "main",
+      agentSessionKey: "agent:main:main",
+      sessionId: "main-session-id",
+    });
+
+    await expect(
+      withGatewayToolCallerIdentity(
+        {
+          agentId: "main",
+          sessionKey: "agent:main:main",
+          gatewayContextResolver: () => undefined,
+        },
+        async () => await tool.execute("list", { action: "list" }),
+      ),
+    ).rejects.toThrow("terminal unavailable");
+    expect(ambientList).not.toHaveBeenCalled();
+  });
+
+  it("revalidates the admitted Gateway after task lookup before opening", async () => {
+    const callerSpawn = vi.fn(async () => makeBackend());
+    const ambientSpawn = vi.fn(async () => makeBackend());
+    const callerManager = new TerminalSessionManager({ emit: vi.fn(), spawn: callerSpawn });
+    const ambientManager = new TerminalSessionManager({ emit: vi.fn(), spawn: ambientSpawn });
+    let callerLive = true;
+    const gatewayContextResolver = vi.fn();
+    gatewayContextResolver.mockImplementation(() =>
+      callerLive ? makeContext(callerManager) : undefined,
+    );
+    getInProcessGatewayToolContext.mockReturnValue(makeContext(ambientManager));
+    const lookupTaskByRunIdForChildSession = vi.fn(async () => {
+      callerLive = false;
+      return undefined;
+    });
+    const tool = createTerminalTool({
+      agentId: "main",
+      agentSessionKey: "agent:main:main",
+      sessionId: "main-session-id",
+      runId: "run-1",
+      lookupTaskByRunIdForChildSession,
+    });
+
+    await expect(
+      withGatewayToolCallerIdentity(
+        {
+          agentId: "main",
+          sessionKey: "agent:main:main",
+          gatewayContextResolver,
+        },
+        async () => await tool.execute("open", { action: "open" }),
+      ),
+    ).rejects.toThrow("terminal unavailable");
+    expect(gatewayContextResolver).toHaveBeenCalledTimes(2);
+    expect(callerSpawn).not.toHaveBeenCalled();
+    expect(ambientSpawn).not.toHaveBeenCalled();
+    expect(getInProcessGatewayToolContext).not.toHaveBeenCalled();
+  });
+
+  it("closes a terminal when the admitted Gateway retires during open", async () => {
+    const spawned = deferred<ReturnType<typeof makeBackend>>();
+    const backend = makeBackend();
+    const callerSpawn = vi.fn(() => spawned.promise);
+    const callerManager = new TerminalSessionManager({ emit: vi.fn(), spawn: callerSpawn });
+    const ambientSpawn = vi.fn(async () => makeBackend());
+    const ambientManager = new TerminalSessionManager({ emit: vi.fn(), spawn: ambientSpawn });
+    let callerLive = true;
+    const gatewayContextResolver = vi.fn();
+    gatewayContextResolver.mockImplementation(() =>
+      callerLive ? makeContext(callerManager) : undefined,
+    );
+    getInProcessGatewayToolContext.mockReturnValue(makeContext(ambientManager));
+    const tool = createTerminalTool({
+      agentId: "main",
+      agentSessionKey: "agent:main:main",
+      sessionId: "main-session-id",
+    });
+
+    const opening = withGatewayToolCallerIdentity(
+      {
+        agentId: "main",
+        sessionKey: "agent:main:main",
+        gatewayContextResolver,
+      },
+      async () => await tool.execute("open", { action: "open", command: "echo unsafe" }),
+    );
+    await vi.waitFor(() => expect(callerSpawn).toHaveBeenCalledOnce());
+    callerLive = false;
+    spawned.resolve(backend);
+
+    await expect(opening).rejects.toThrow("terminal unavailable");
+    expect(gatewayContextResolver).toHaveBeenCalledTimes(2);
+    expect(backend.writes).toEqual([]);
+    expect(backend.killed).toBe(true);
+    expect(callerManager.size).toBe(0);
+    expect(ambientSpawn).not.toHaveBeenCalled();
+    expect(getInProcessGatewayToolContext).not.toHaveBeenCalled();
+  });
+
+  it("keeps ambient Gateway context pinned across task lookup", async () => {
+    const firstSpawn = vi.fn(async () => makeBackend());
+    const secondSpawn = vi.fn(async () => makeBackend());
+    const firstManager = new TerminalSessionManager({ emit: vi.fn(), spawn: firstSpawn });
+    const secondManager = new TerminalSessionManager({ emit: vi.fn(), spawn: secondSpawn });
+    getInProcessGatewayToolContext
+      .mockReturnValueOnce(makeContext(firstManager))
+      .mockReturnValue(makeContext(secondManager));
+    const tool = createTerminalTool({
+      agentId: "main",
+      agentSessionKey: "agent:main:main",
+      sessionId: "main-session-id",
+      runId: "run-1",
+      lookupTaskByRunIdForChildSession: vi.fn(async () => undefined),
+    });
+
+    await expect(tool.execute("open", { action: "open" })).resolves.toMatchObject({
+      details: { ok: true },
+    });
+    expect(getInProcessGatewayToolContext).toHaveBeenCalledOnce();
+    expect(firstSpawn).toHaveBeenCalledOnce();
+    expect(secondSpawn).not.toHaveBeenCalled();
+  });
+
+  it("uses ambient Gateway context without an admitted caller", async () => {
+    const manager = new TerminalSessionManager({ emit: vi.fn(), spawn: vi.fn() });
+    const list = vi.spyOn(manager, "listAgent");
+    getInProcessGatewayToolContext.mockReturnValue(makeContext(manager));
+    const tool = createTerminalTool({
+      agentId: "main",
+      agentSessionKey: "agent:main:main",
+      sessionId: "main-session-id",
+    });
+
+    await expect(tool.execute("list", { action: "list" })).resolves.toMatchObject({
+      details: { sessions: [] },
+    });
+    expect(list).toHaveBeenCalledOnce();
+  });
+
   it("opens in the background, reads, writes, resizes, lists, and closes its terminal", async () => {
     const backend = makeBackend();
     const manager = new TerminalSessionManager({ emit: vi.fn(), spawn: async () => backend });

src/agents/tools/terminal-tool.ts

@@ -19,6 +19,7 @@ import {
   readToolStringParam,
   ToolInputError,
 } from "./common.js";
+import { getGatewayToolCallerIdentity } from "./gateway-caller-context.js";
 import { getInProcessGatewayToolContext } from "./in-process-gateway.js";
 
 const ACTIONS = ["open", "read", "input", "resize", "close", "list"] as const;
@@ -166,8 +167,32 @@ function launchBlockMessage(
   return `terminal unavailable: agent sandboxed (${block.mode})`;
 }
 
+function resolveTerminalOpenTarget(params: {
+  agentId: string;
+  context: TerminalToolGatewayContext | undefined;
+  cwd?: string;
+}) {
+  const manager = params.context?.terminalSessions;
+  if (!params.context || !manager) {
+    throw new ToolInputError("terminal unavailable");
+  }
+  if (!params.context.isTerminalEnabled()) {
+    throw new ToolInputError("terminal disabled");
+  }
+  const launch = params.context.resolveTerminalLaunchPolicy(params.agentId);
+  if (!launch.ok) {
+    throw new ToolInputError(launchBlockMessage(launch.block));
+  }
+  return {
+    manager,
+    spawnPlan: resolveTerminalSpawnPlan({
+      ...launch.plan,
+      ...(params.cwd ? { cwdOverride: params.cwd } : {}),
+    }),
+  };
+}
+
 export function createTerminalTool(opts: TerminalToolOptions = {}): AnyAgentTool {
-  const getContext = opts.getGatewayContext ?? getInProcessGatewayToolContext;
   const findOwnerTask = opts.lookupTaskByRunIdForChildSession ?? lookupTaskByRunIdForChildSession;
   return {
     label: "Terminal",
@@ -189,6 +214,11 @@ export function createTerminalTool(opts: TerminalToolOptions = {}): AnyAgentTool
       }
       const agentId = opts.agentId?.trim() || resolveAgentIdFromSessionKey(agentSessionKey);
       const owner = { kind: "agent", agentSessionKey, agentSessionId, agentId } as const;
+      const admittedResolver = opts.getGatewayContext
+        ? undefined
+        : getGatewayToolCallerIdentity()?.gatewayContextResolver;
+      const getContext =
+        opts.getGatewayContext ?? admittedResolver ?? getInProcessGatewayToolContext;
       const context = getContext();
       const manager = context?.terminalSessions;
       if (!context || !manager) {
@@ -204,23 +234,18 @@ export function createTerminalTool(opts: TerminalToolOptions = {}): AnyAgentTool
         const cwd = readOptionalStringParam(params, "cwd");
         const cols = readDimension(params, "cols", DEFAULT_COLS);
         const rows = readDimension(params, "rows", DEFAULT_ROWS);
-        if (!context.isTerminalEnabled()) {
-          throw new ToolInputError("terminal disabled");
-        }
-        const launch = context.resolveTerminalLaunchPolicy(agentId);
-        if (!launch.ok) {
-          throw new ToolInputError(launchBlockMessage(launch.block));
-        }
-        const spawnPlan = resolveTerminalSpawnPlan({
-          ...launch.plan,
-          ...(cwd ? { cwdOverride: cwd } : {}),
-        });
+        const initialTarget = resolveTerminalOpenTarget({ agentId, context, cwd });
         const runId = opts.runId?.trim();
         const taskLookupId = runId ? (getAgentRunTaskRunId(runId) ?? runId) : undefined;
         const task = taskLookupId ? await findOwnerTask(taskLookupId, agentSessionKey) : undefined;
         if (task && isTerminalTaskStatus(task.status)) {
           throw new ToolInputError("terminal task already ended");
         }
+        // Refresh after task lookup so a retired admitted Gateway cannot allocate a new PTY.
+        const { manager: openManager, spawnPlan } =
+          taskLookupId && admittedResolver
+            ? resolveTerminalOpenTarget({ agentId, context: admittedResolver(), cwd })
+            : initialTarget;
         const taskId = task?.taskId;
         const terminalOwner = { ...owner, ...(taskId ? { taskId } : {}) };
         const deadline = createTerminalOpenDeadline();
@@ -234,11 +259,11 @@ export function createTerminalTool(opts: TerminalToolOptions = {}): AnyAgentTool
         } else {
           signal?.addEventListener("abort", cancelOpen, { once: true });
         }
-        let openingTerminal: ReturnType<typeof manager.open> | undefined;
-        let outcome: Awaited<ReturnType<typeof manager.open>>;
+        let openingTerminal: ReturnType<typeof openManager.open> | undefined;
+        let outcome: Awaited<ReturnType<typeof openManager.open>>;
         try {
           outcome = await waitForTerminalOpenDeadline(() => {
-            openingTerminal = manager.open({
+            openingTerminal = openManager.open({
               owner: terminalOwner,
               agentId: spawnPlan.agentId,
               cwd: spawnPlan.cwd,
@@ -256,7 +281,7 @@ export function createTerminalTool(opts: TerminalToolOptions = {}): AnyAgentTool
             void openingTerminal.then(
               (lateOutcome) => {
                 if (lateOutcome.ok) {
-                  manager.closeAgent(owner, lateOutcome.sessionId);
+                  openManager.closeAgent(owner, lateOutcome.sessionId);
                 }
               },
               () => undefined,
@@ -272,10 +297,25 @@ export function createTerminalTool(opts: TerminalToolOptions = {}): AnyAgentTool
         if (!outcome.ok) {
           throw new ToolInputError(outcome.message);
         }
+        if (admittedResolver) {
+          try {
+            const liveManager = resolveTerminalOpenTarget({
+              agentId,
+              context: admittedResolver(),
+              cwd,
+            }).manager;
+            if (liveManager !== openManager) {
+              throw new ToolInputError("terminal unavailable");
+            }
+          } catch (error) {
+            openManager.closeAgent(owner, outcome.sessionId);
+            throw error;
+          }
+        }
         if (command !== undefined) {
-          const commandOutcome = manager.writeAgent(owner, outcome.sessionId, `${command}\r`);
+          const commandOutcome = openManager.writeAgent(owner, outcome.sessionId, `${command}\r`);
           if (!commandOutcome.ok) {
-            manager.closeAgent(owner, outcome.sessionId);
+            openManager.closeAgent(owner, outcome.sessionId);
             terminalActionResult("initial command", commandOutcome);
           }
         }