Harden assign_to_agent concurrency: isolate handler state and serialize MCP stdin dispatch (#52034)

Author: Copilot

.changeset/patch-fix-assign-to-agent-concurrency.md

@@ -88,10 +88,10 @@ describe("assign_to_agent", () => {
       const _tempIdMap = loadTemporaryIdMap();
       for (const _item of _items) { await _handler(_item, {}, _tempIdMap); }
     }
-    await writeAssignToAgentSummary();
-    const _errorCount = getAssignToAgentErrorCount();
-    core.setOutput("assigned", getAssignToAgentAssigned());
-    core.setOutput("assignment_errors", getAssignToAgentErrors());
+    await writeAssignToAgentSummary(_handler);
+    const _errorCount = getAssignToAgentErrorCount(_handler);
+    core.setOutput("assigned", getAssignToAgentAssigned(_handler));
+    core.setOutput("assignment_errors", getAssignToAgentErrors(_handler));
     core.setOutput("assignment_error_count", String(_errorCount));
     if (_errorCount > 0) { core.setFailed("Failed to assign " + _errorCount + " agent(s)"); }
   `;
@@ -1359,6 +1359,88 @@ describe("assign_to_agent", () => {
     expect(mockSleep).toHaveBeenCalledWith(10000);
   });
 
+  it("does not consume a max slot for invalid items", async () => {
+    mockGithub.rest.issues.checkUserCanBeAssigned.mockResolvedValue({});
+    mockGithub.rest.users.getByUsername.mockResolvedValue({ data: { id: 99999 } });
+    mockGithub.rest.issues.get.mockImplementation(async ({ issue_number }) => ({
+      data: { id: Number(issue_number) + 1000, number: issue_number, assignees: [], html_url: "", title: "", body: "" },
+    }));
+    mockGithub.request.mockResolvedValue({ data: { id: "task-123" } });
+
+    const result = await eval(`(async () => {
+      ${assignToAgentScript};
+      const _handler = await main({ max: "1", name: "copilot" });
+      const _invalid = await _handler({ type: "assign_to_agent", issue_number: 1, pull_number: 2, agent: "copilot" }, {}, new Map());
+      const _valid = await _handler({ type: "assign_to_agent", issue_number: 3, agent: "copilot" }, {}, new Map());
+      return {
+        invalid: _invalid,
+        valid: _valid,
+        assigned: getAssignToAgentAssigned(_handler),
+      };
+    })()`);
+
+    expect(result.invalid.success).toBe(false);
+    expect(result.valid.success).toBe(true);
+    expect(result.assigned.split("\n").filter(Boolean)).toHaveLength(1);
+  });
+
+  it("atomically reserves the max slot before the inter-assignment delay", async () => {
+    mockGithub.rest.issues.checkUserCanBeAssigned.mockResolvedValue({});
+    mockGithub.rest.users.getByUsername.mockResolvedValue({ data: { id: 99999 } });
+    mockGithub.rest.issues.get.mockImplementation(async ({ issue_number }) => ({
+      data: { id: Number(issue_number) + 1000, number: issue_number, assignees: [], html_url: "", title: "", body: "" },
+    }));
+    mockGithub.request.mockResolvedValue({ data: { id: "task-123" } });
+
+    let releaseSleep;
+    mockSleep.mockImplementationOnce(
+      () =>
+        new Promise(resolve => {
+          releaseSleep = resolve;
+        })
+    );
+
+    const result = await eval(`(async () => {
+      ${assignToAgentScript};
+      const _handler = await main({ max: "2", name: "copilot" });
+      await _handler({ type: "assign_to_agent", issue_number: 1, agent: "copilot" }, {}, new Map());
+      return {
+        second: _handler({ type: "assign_to_agent", issue_number: 2, agent: "copilot" }, {}, new Map()),
+        third: _handler({ type: "assign_to_agent", issue_number: 3, agent: "copilot" }, {}, new Map()),
+      };
+    })()`);
+
+    await vi.waitFor(() => expect(mockSleep).toHaveBeenCalledTimes(1));
+    releaseSleep();
+
+    const [second, third] = await Promise.all([result.second, result.third]);
+    expect(second.success).toBe(true);
+    expect(third.skipped).toBe(true);
+  });
+
+  it("keeps assign_to_agent results isolated per main() invocation", async () => {
+    mockGithub.rest.issues.checkUserCanBeAssigned.mockResolvedValue({});
+    mockGithub.rest.users.getByUsername.mockResolvedValue({ data: { id: 99999 } });
+    mockGithub.rest.issues.get.mockImplementation(async ({ issue_number }) => ({
+      data: { id: Number(issue_number) + 2000, number: issue_number, assignees: [], html_url: "", title: "", body: "" },
+    }));
+    mockGithub.request.mockResolvedValue({ data: { id: "task-123" } });
+
+    const result = await eval(`(async () => {
+      ${assignToAgentScript};
+      const _handlerA = await main({ max: "5", name: "copilot" });
+      const _handlerB = await main({ max: "5", name: "copilot" });
+      await _handlerA({ type: "assign_to_agent", issue_number: 11, agent: "copilot" }, {}, new Map());
+      return {
+        assignedA: getAssignToAgentAssigned(_handlerA),
+        assignedB: getAssignToAgentAssigned(_handlerB),
+      };
+    })()`);
+
+    expect(result.assignedA).toContain("issue:11:copilot");
+    expect(result.assignedB).toBe("");
+  });
+
   describe("Cross-repository allowlist validation", () => {
     it("should reject target repository not in allowlist", async () => {
       process.env.GH_AW_ALLOWED_REPOS = "allowed-owner/allowed-repo";

actions/setup/js/assign_to_agent.cjs

@@ -1043,9 +1043,15 @@ function start(server, options = {}) {
     throw new Error(`${ERR_VALIDATION}: No tools registered`);
   }
 
-  const onData = async chunk => {
+  let processingChain = Promise.resolve();
+  const onData = chunk => {
     server.readBuffer.append(chunk);
-    await processReadBuffer(server, defaultHandler);
+    processingChain = processingChain
+      .then(() => processReadBuffer(server, defaultHandler))
+      .catch(error => {
+        server.debug(`processReadBuffer error: ${getErrorMessage(error)}`);
+      });
+    return processingChain;
   };
 
   process.stdin.on("data", onData);

actions/setup/js/assign_to_agent.test.cjs

@@ -1410,6 +1410,75 @@ echo "result=$INPUT_MY_INPUT" >> "$GITHUB_OUTPUT"
     });
   });
 
+  describe("start", () => {
+    it("serializes pipelined stdin chunks and avoids overlapping tool dispatch", async () => {
+      const { createServer, registerTool, start } = await import("./mcp_server_core.cjs");
+      const server = createServer({ name: "test-server", version: "1.0.0" });
+
+      let activeHandlers = 0;
+      let maxActiveHandlers = 0;
+      let callCount = 0;
+      let releaseFirstCall;
+      const firstCallGate = new Promise(resolve => {
+        releaseFirstCall = resolve;
+      });
+
+      registerTool(server, {
+        name: "slow_tool",
+        description: "A test tool",
+        inputSchema: { type: "object", properties: {}, additionalProperties: false },
+        handler: async () => {
+          callCount++;
+          activeHandlers++;
+          maxActiveHandlers = Math.max(maxActiveHandlers, activeHandlers);
+          if (callCount === 1) {
+            await firstCallGate;
+          }
+          activeHandlers--;
+          return { content: [{ type: "text", text: "ok" }] };
+        },
+      });
+
+      // Avoid writing to real stdout/stderr during this test.
+      server.writeMessage = () => {};
+      server.debug = () => {};
+
+      let dataHandler;
+      const onSpy = vi.spyOn(process.stdin, "on").mockImplementation((event, handler) => {
+        if (event === "data") {
+          dataHandler = handler;
+        }
+        return process.stdin;
+      });
+      const resumeSpy = vi.spyOn(process.stdin, "resume").mockImplementation(() => process.stdin);
+
+      try {
+        start(server);
+        expect(typeof dataHandler).toBe("function");
+
+        const message1 = Buffer.from(`${JSON.stringify({ jsonrpc: "2.0", id: 1, method: "tools/call", params: { name: "slow_tool", arguments: {} } })}\n`);
+        const message2 = Buffer.from(`${JSON.stringify({ jsonrpc: "2.0", id: 2, method: "tools/call", params: { name: "slow_tool", arguments: {} } })}\n`);
+
+        const first = dataHandler(message1);
+        const second = dataHandler(message2);
+
+        // Let queued tasks run before releasing first call.
+        await Promise.resolve();
+        await Promise.resolve();
+        expect(callCount).toBe(1);
+
+        releaseFirstCall();
+        await Promise.all([first, second]);
+
+        expect(callCount).toBe(2);
+        expect(maxActiveHandlers).toBe(1);
+      } finally {
+        onSpy.mockRestore();
+        resumeSpy.mockRestore();
+      }
+    });
+  });
+
   describe("findSimilarTools", () => {
     it("should find tools with typos", async () => {
       const { findSimilarTools } = await import("./mcp_server_core.cjs");

actions/setup/js/mcp_server_core.cjs

@@ -1730,17 +1730,18 @@ async function main() {
 
     // Export assign_to_agent outputs when the handler was loaded
     if (messageHandlers.has("assign_to_agent")) {
-      const assignToAgentAssigned = getAssignToAgentAssigned();
-      const assignToAgentErrors = getAssignToAgentErrors();
-      const assignToAgentErrorCount = getAssignToAgentErrorCount();
+      const assignToAgentHandler = messageHandlers.get("assign_to_agent");
+      const assignToAgentAssigned = getAssignToAgentAssigned(assignToAgentHandler);
+      const assignToAgentErrors = getAssignToAgentErrors(assignToAgentHandler);
+      const assignToAgentErrorCount = getAssignToAgentErrorCount(assignToAgentHandler);
       core.setOutput("assign_to_agent_assigned", assignToAgentAssigned);
       core.setOutput("assign_to_agent_assignment_errors", assignToAgentErrors);
       core.setOutput("assign_to_agent_assignment_error_count", assignToAgentErrorCount.toString());
       if (assignToAgentErrorCount > 0) {
         core.warning(`${assignToAgentErrorCount} agent assignment(s) failed`);
       }
       core.info(`Exported assign_to_agent outputs (${assignToAgentErrorCount} error(s))`);
-      await writeAssignToAgentSummary();
+      await writeAssignToAgentSummary(assignToAgentHandler);
     }
 
     // Export create_agent_session outputs when the handler was loaded