diff --git a/__tests__/security.test.ts b/__tests__/security.test.ts index 782b99da..c57158c2 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 3ceb8551..ce97c365 100644 --- a/src/mcp/tools.ts +++ b/src/mcp/tools.ts @@ -706,11 +706,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)); } /**