From 23ad4ea923b31ed7f3aafe313f137ec7a822f91c Mon Sep 17 00:00:00 2001
From: Baijack-star <71923891+Baijack-star@users.noreply.github.com>
Date: Sat, 23 May 2026 02:20:08 +0800
Subject: [PATCH] fix(mcp): cap codegraph_context output to prevent context
 bloat (#296)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Route handleContext's output through the shared truncateOutput cap
(MAX_OUTPUT_LENGTH) so codegraph_context can no longer blow past the context
budget — every sibling MCP tool already truncates; this was the one uncapped
output path.

Closes #296

Co-authored-by: Baijack-star <71923891+Baijack-star@users.noreply.github.com>
---
 __tests__/security.test.ts | 14 ++++++++++++++
 src/mcp/tools.ts           |  4 ++--
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/__tests__/security.test.ts b/__tests__/security.test.ts
index 782b99d..c57158c 100644
--- a/__tests__/security.test.ts
+++ b/__tests__/security.test.ts
@@ -239,6 +239,20 @@ describe('MCP Input Validation', () => {
     expect(result.content[0].text).toContain('non-empty string');
   });
 
+  it('should truncate oversized codegraph_context output', async () => {
+    const oversizedContext = Array.from({ length: 400 }, (_, i) => `line-${i} ${'x'.repeat(80)}`).join('\n');
+    const fakeCg = {
+      buildContext: async () => oversizedContext,
+    };
+    const fakeHandler = new ToolHandler(fakeCg as unknown as CodeGraph);
+
+    const result = await fakeHandler.execute('codegraph_context', { task: 'find example' });
+
+    expect(result.isError).toBeFalsy();
+    expect(result.content[0].text.length).toBeLessThan(oversizedContext.length);
+    expect(result.content[0].text).toContain('... (output truncated)');
+  });
+
   it('should reject non-string symbol in codegraph_impact', async () => {
     const result = await handler.execute('codegraph_impact', { symbol: [] });
     expect(result.isError).toBe(true);
diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts
index f15cdc5..dfd4154 100644
--- a/src/mcp/tools.ts
+++ b/src/mcp/tools.ts
@@ -775,11 +775,11 @@ export class ToolHandler {
 
     // buildContext returns string when format is 'markdown'
     if (typeof context === 'string') {
-      return this.textResult(context + reminder);
+      return this.textResult(this.truncateOutput(context + reminder));
     }
 
     // If it returns TaskContext, format it
-    return this.textResult(this.formatTaskContext(context) + reminder);
+    return this.textResult(this.truncateOutput(this.formatTaskContext(context) + reminder));
   }
 
   /**