From 5b019e1179dce76f432c0f5700161de5727b8a96 Mon Sep 17 00:00:00 2001
From: shreyansh1729 <shreyanshiitian1729@gmail.com>
Date: Sat, 17 Oct 2026 07:24:00 +0530
Subject: [PATCH] fix(mcp): retain same-address procedure aliases in completion

Fixes #4623. Check name uniqueness against every procedure row before deriving completion candidates.
---
 src/server/promptCompletion.ts       | 24 ++++++-------
 tests/boundary/mcp/promptMcp.test.ts | 52 ++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 13 deletions(-)

diff --git a/src/server/promptCompletion.ts b/src/server/promptCompletion.ts
index 97e66546f4..21c595d872 100644
--- a/src/server/promptCompletion.ts
+++ b/src/server/promptCompletion.ts
@@ -99,7 +99,8 @@ const procedureCandidates = async (
   session: BinarySessionPort | undefined,
   document: string | undefined,
 ): Promise<readonly string[]> => {
-  const byAddress = new Map<string, string>();
+  const addresses = new Set<string>();
+  const addressesByName = new Map<string, Set<string>>();
   const input =
     document === undefined || document.length === 0 ? {} : { document };
   const result = await completionInventoryLookup(
@@ -114,20 +115,17 @@ const procedureCandidates = async (
       .array(z.object({ address: z.string(), value: z.string() }))
       .safeParse(result.value.result);
     if (parsed.success)
-      for (const item of parsed.data)
-        if (!byAddress.has(item.address))
-          byAddress.set(item.address, item.value);
+      for (const { address, value: name } of parsed.data) {
+        addresses.add(address);
+        const namedAddresses = addressesByName.get(name) ?? new Set<string>();
+        namedAddresses.add(address);
+        addressesByName.set(name, namedAddresses);
+      }
   }
 
-  const addressesByName = new Map<string, Set<string>>();
-  for (const [address, name] of byAddress) {
-    const addresses = addressesByName.get(name) ?? new Set<string>();
-    addresses.add(address);
-    addressesByName.set(name, addresses);
-  }
-  const candidates = [...byAddress.keys()];
-  for (const [name, addresses] of addressesByName)
-    if (addresses.size === 1 && name.length > 0) candidates.push(name);
+  const candidates = [...addresses];
+  for (const [name, namedAddresses] of addressesByName)
+    if (namedAddresses.size === 1 && name.length > 0) candidates.push(name);
   return candidates;
 };
 
diff --git a/tests/boundary/mcp/promptMcp.test.ts b/tests/boundary/mcp/promptMcp.test.ts
index 9138a102dc..8412c5eae8 100644
--- a/tests/boundary/mcp/promptMcp.test.ts
+++ b/tests/boundary/mcp/promptMcp.test.ts
@@ -54,6 +54,58 @@ afterEach(async () => {
   );
 });
 
+it.each([false, true])(
+  "retains same-address procedure aliases (reversed: %s)",
+  async (reversed) => {
+    const rows = [
+      { address: "0x1000", value: "foo" },
+      { address: "0x1000", value: "bar" },
+      { address: "0x1000", value: "ambiguous" },
+      { address: "0x2000", value: "ambiguous" },
+      { address: "0x1000", value: "bar" },
+    ];
+    if (reversed) rows.reverse();
+    const client = await connect(
+      createServer({
+        kind: "fixed",
+        analysis: {
+          execute: async () =>
+            resultOk(
+              createAnalysisExecution(rows, {
+                id: "fixture",
+                name: "Fixture",
+                version: "1",
+              }),
+            ),
+        },
+      }),
+    );
+    const inventory = await client.callTool({
+      name: "list_procedures",
+      arguments: {},
+    });
+    expect(inventory.isError).not.toBe(true);
+    expect(inventory.structuredContent).toMatchObject({
+      normalized_result: rows,
+    });
+    for (const [value, values] of [
+      ["", ["0x1000", "0x2000", "bar", "foo"]],
+      ["BAR", ["bar"]],
+      ["foo", ["foo"]],
+      ["ambiguous", []],
+    ] as const) {
+      expect(
+        await client.complete({
+          ref: { type: "ref/prompt", name: "investigate_feature" },
+          argument: { name: "procedure", value },
+        }),
+      ).toEqual({
+        completion: { values, total: values.length, hasMore: false },
+      });
+    }
+  },
+);
+
 describe("guided prompts over MCP", () => {
   it("lists and renders all workflows without changing the tool inventory", async () => {
     const session = fixtureSession();