From 5b2dc788975e80b44ffc3b724995ecc1ef8d78ff Mon Sep 17 00:00:00 2001
From: liqoingyu <li@libiao.co>
Date: Sun, 18 Jan 2026 20:38:45 +0800
Subject: [PATCH] fix(cli,core): harden MCP resource references

---
 .../src/ui/hooks/atCommandProcessor.test.ts   | 129 ++++++++++++++++++
 .../cli/src/ui/hooks/atCommandProcessor.ts    |  14 +-
 2 files changed, 140 insertions(+), 3 deletions(-)

diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.test.ts b/packages/cli/src/ui/hooks/atCommandProcessor.test.ts
index f159eb385..3b91dd269 100644
--- a/packages/cli/src/ui/hooks/atCommandProcessor.test.ts
+++ b/packages/cli/src/ui/hooks/atCommandProcessor.test.ts
@@ -257,6 +257,135 @@ describe('handleAtCommand', () => {
     );
   });
 
+  it('should expand an MCP resource reference in @server:resource format', async () => {
+    (mockConfig as unknown as { getMcpServers: () => unknown }).getMcpServers =
+      () =>
+        ({
+          github: {},
+        }) as unknown;
+
+    vi.spyOn(registry, 'readMcpResource').mockResolvedValue({
+      contents: [
+        {
+          uri: 'github://repos/owner/repo/issues',
+          mimeType: 'application/json',
+          text: '{"ok":true}',
+        },
+      ],
+    } as unknown as Awaited<ReturnType<ToolRegistry['readMcpResource']>>);
+
+    const query = 'Show me the data from @github:repos/owner/repo/issues';
+
+    const result = await handleAtCommand({
+      query,
+      config: mockConfig,
+      addItem: mockAddItem,
+      onDebugMessage: mockOnDebugMessage,
+      messageId: 1001,
+      signal: abortController.signal,
+    });
+
+    expect(result).toEqual({
+      processedQuery: [
+        { text: query },
+        { text: '\n--- Content from referenced MCP resources ---' },
+        { text: '\nContent from @github:repos/owner/repo/issues:\n' },
+        { text: '{"ok":true}' },
+        { text: '\n--- End of MCP resource content ---' },
+      ],
+      shouldProceed: true,
+    });
+    expect(registry.readMcpResource).toHaveBeenCalledWith(
+      'github',
+      'github://repos/owner/repo/issues',
+    );
+    expect(mockAddItem).toHaveBeenCalledWith(
+      expect.objectContaining({
+        type: 'tool_group',
+        tools: [expect.objectContaining({ status: ToolCallStatus.Success })],
+      }),
+      1001,
+    );
+  });
+
+  it('should expand an MCP resource reference with a leading slash', async () => {
+    (mockConfig as unknown as { getMcpServers: () => unknown }).getMcpServers =
+      () =>
+        ({
+          github: {},
+        }) as unknown;
+
+    vi.spyOn(registry, 'readMcpResource').mockResolvedValue({
+      contents: [
+        {
+          uri: 'github://repos/owner/repo/issues',
+          mimeType: 'application/json',
+          text: '{"ok":true}',
+        },
+      ],
+    } as unknown as Awaited<ReturnType<ToolRegistry['readMcpResource']>>);
+
+    const query = 'Show me the data from @github:/repos/owner/repo/issues';
+
+    const result = await handleAtCommand({
+      query,
+      config: mockConfig,
+      addItem: mockAddItem,
+      onDebugMessage: mockOnDebugMessage,
+      messageId: 1002,
+      signal: abortController.signal,
+    });
+
+    expect(result).toEqual({
+      processedQuery: [
+        { text: 'Show me the data from @github:repos/owner/repo/issues' },
+        { text: '\n--- Content from referenced MCP resources ---' },
+        { text: '\nContent from @github:repos/owner/repo/issues:\n' },
+        { text: '{"ok":true}' },
+        { text: '\n--- End of MCP resource content ---' },
+      ],
+      shouldProceed: true,
+    });
+    expect(registry.readMcpResource).toHaveBeenCalledWith(
+      'github',
+      'github://repos/owner/repo/issues',
+    );
+    expect(mockAddItem).toHaveBeenCalledWith(
+      expect.objectContaining({
+        type: 'tool_group',
+        tools: [expect.objectContaining({ status: ToolCallStatus.Success })],
+      }),
+      1002,
+    );
+  });
+
+  it('should ignore @server: when no MCP resource is provided', async () => {
+    (mockConfig as unknown as { getMcpServers: () => unknown }).getMcpServers =
+      () =>
+        ({
+          github: {},
+        }) as unknown;
+
+    const readMcpResourceSpy = vi.spyOn(registry, 'readMcpResource');
+    const query = 'Show me the data from @github:';
+
+    const result = await handleAtCommand({
+      query,
+      config: mockConfig,
+      addItem: mockAddItem,
+      onDebugMessage: mockOnDebugMessage,
+      messageId: 1003,
+      signal: abortController.signal,
+    });
+
+    expect(result).toEqual({
+      processedQuery: [{ text: query }],
+      shouldProceed: true,
+    });
+    expect(readMcpResourceSpy).not.toHaveBeenCalled();
+    expect(mockAddItem).not.toHaveBeenCalled();
+  });
+
   it('should handle query with text before and after @command', async () => {
     const fileContent = 'Markdown content.';
     const filePath = await createTestFile(
diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.ts b/packages/cli/src/ui/hooks/atCommandProcessor.ts
index d099958da..a5ea3a46d 100644
--- a/packages/cli/src/ui/hooks/atCommandProcessor.ts
+++ b/packages/cli/src/ui/hooks/atCommandProcessor.ts
@@ -225,15 +225,23 @@ function extractMcpResourceAtReferences(
     }
 
     if (!resource) {
-      merged.push(part);
+      // Treat "@server:" without a resource as plain text, rather than falling
+      // through to file resolution for a path like "server:".
+      merged.push({ type: 'text', content: atText });
       continue;
     }
 
-    const normalizedAtCommand = `@${serverName}:${resource}`;
+    const normalizedResource = resource.includes('://')
+      ? resource
+      : resource.startsWith('/')
+        ? resource.slice(1)
+        : resource;
+
+    const normalizedAtCommand = `@${serverName}:${normalizedResource}`;
     refs.push({
       atCommand: normalizedAtCommand,
       serverName,
-      uri: normalizeMcpResourceUri(serverName, resource),
+      uri: normalizeMcpResourceUri(serverName, normalizedResource),
     });
     merged.push({ type: 'atPath', content: normalizedAtCommand });
   }