diff --git a/gitnexus/src/mcp/output-budget.ts b/gitnexus/src/mcp/output-budget.ts new file mode 100644 index 0000000000..032d3cd413 --- /dev/null +++ b/gitnexus/src/mcp/output-budget.ts @@ -0,0 +1,58 @@ +const BUDGETED_TOOLS = new Set(['query', 'context', 'impact']); + +export const MCP_TOKEN_ESTIMATE_BYTES = 4; +export const MCP_TRUNCATION_MARKER = '\n…'; + +function parsePositiveInteger(value: unknown, source: string): number { + if (typeof value === 'number' && Number.isSafeInteger(value) && value > 0) return value; + if (typeof value === 'string' && /^[1-9]\d*$/.test(value.trim())) { + const parsed = Number(value.trim()); + if (Number.isSafeInteger(parsed)) return parsed; + } + throw new Error(`${source} must be a positive integer.`); +} + +export function resolveMcpMaxTokens( + toolName: string, + args: Record | undefined, + env: NodeJS.ProcessEnv = process.env, +): number | undefined { + if (!BUDGETED_TOOLS.has(toolName)) return undefined; + if (args?.maxTokens !== undefined) return parsePositiveInteger(args.maxTokens, 'maxTokens'); + + const configured = env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS; + if (configured === undefined || configured.trim() === '') return undefined; + return parsePositiveInteger(configured, 'GITNEXUS_MCP_DEFAULT_MAX_TOKENS'); +} + +function utf8Prefix(text: string, maxBytes: number): string { + let bytes = 0; + const codePoints: string[] = []; + for (const codePoint of text) { + const codePointBytes = Buffer.byteLength(codePoint, 'utf8'); + if (bytes + codePointBytes > maxBytes) break; + codePoints.push(codePoint); + bytes += codePointBytes; + } + return codePoints.join(''); +} + +export function applyMcpMaxTokens(text: string, maxTokens: number | undefined): string { + if (maxTokens === undefined) return text; + + const textBytes = Buffer.byteLength(text, 'utf8'); + if (maxTokens >= Math.ceil(textBytes / MCP_TOKEN_ESTIMATE_BYTES)) return text; + + const maxBytes = maxTokens * MCP_TOKEN_ESTIMATE_BYTES; + const markerBytes = Buffer.byteLength(MCP_TRUNCATION_MARKER, 'utf8'); + return utf8Prefix(text, Math.max(0, maxBytes - markerBytes)) + MCP_TRUNCATION_MARKER; +} + +export function withoutMcpBudgetArg( + args: Record | undefined, +): Record | undefined { + if (!args || !Object.prototype.hasOwnProperty.call(args, 'maxTokens')) return args; + const backendArgs = { ...args }; + delete backendArgs.maxTokens; + return backendArgs; +} diff --git a/gitnexus/src/mcp/server.ts b/gitnexus/src/mcp/server.ts index d4a7c58aa8..bd8b08d03b 100644 --- a/gitnexus/src/mcp/server.ts +++ b/gitnexus/src/mcp/server.ts @@ -27,6 +27,7 @@ import { GITNEXUS_TOOLS } from './tools.js'; import { installGlobalStdoutSentinel } from './stdio-context.js'; import type { LocalBackend } from './local/local-backend.js'; import { getResourceDefinitions, getResourceTemplates, readResource } from './resources.js'; +import { applyMcpMaxTokens, resolveMcpMaxTokens, withoutMcpBudgetArg } from './output-budget.js'; /** * Next-step hints appended to tool responses. @@ -165,9 +166,12 @@ export function createMCPServer(backend: LocalBackend): Server { // Handle tool calls β€” append next-step hints to guide agent workflow server.setRequestHandler(CallToolRequestSchema, async (request) => { const { name, arguments: args } = request.params; + let maxTokens: number | undefined; try { - const result = await backend.callTool(name, args); + const typedArgs = args as Record | undefined; + maxTokens = resolveMcpMaxTokens(name, typedArgs); + const result = await backend.callTool(name, withoutMcpBudgetArg(typedArgs)); const resultText = typeof result === 'string' ? result : JSON.stringify(result, null, 2); const hint = getNextStepHint(name, args as Record | undefined); @@ -175,7 +179,7 @@ export function createMCPServer(backend: LocalBackend): Server { content: [ { type: 'text', - text: resultText + hint, + text: applyMcpMaxTokens(resultText + hint, maxTokens), }, ], }; @@ -185,7 +189,7 @@ export function createMCPServer(backend: LocalBackend): Server { content: [ { type: 'text', - text: `Error: ${message}`, + text: applyMcpMaxTokens(`Error: ${message}`, maxTokens), }, ], isError: true, diff --git a/gitnexus/src/mcp/tools.ts b/gitnexus/src/mcp/tools.ts index 57f71f6752..ec9a96ac57 100644 --- a/gitnexus/src/mcp/tools.ts +++ b/gitnexus/src/mcp/tools.ts @@ -182,6 +182,12 @@ SERVICE: optional monorepo path prefix (POSIX-style, case-sensitive segments). W description: 'Include full symbol source code (default: false)', default: false, }, + maxTokens: { + type: 'integer', + minimum: 1, + description: + 'Maximum estimated tokens in the complete formatted MCP response. Explicit request overrides GITNEXUS_MCP_DEFAULT_MAX_TOKENS.', + }, repo: { type: 'string', description: @@ -311,6 +317,12 @@ SERVICE: optional monorepo path prefix (case-sensitive path segments). When "rep description: 'Include full symbol source code (default: false)', default: false, }, + maxTokens: { + type: 'integer', + minimum: 1, + description: + 'Maximum estimated tokens in the complete formatted MCP response. Explicit request overrides GITNEXUS_MCP_DEFAULT_MAX_TOKENS.', + }, repo: { type: 'string', description: @@ -576,6 +588,12 @@ SERVICE: optional monorepo path prefix (case-sensitive path segments). When "rep 'When true, returns target, summary, risk, byDepthCounts, affected_processes, and affected_modules β€” omits byDepth. Single-repo only; ignored in group mode (@groupName). Use for hub symbols to get actionable signal without output explosion.', default: false, }, + maxTokens: { + type: 'integer', + minimum: 1, + description: + 'Maximum estimated tokens in the complete formatted MCP response. Explicit request overrides GITNEXUS_MCP_DEFAULT_MAX_TOKENS.', + }, timeoutMs: { type: 'number', description: diff --git a/gitnexus/test/unit/mcp-output-budget.test.ts b/gitnexus/test/unit/mcp-output-budget.test.ts new file mode 100644 index 0000000000..8e568f0149 --- /dev/null +++ b/gitnexus/test/unit/mcp-output-budget.test.ts @@ -0,0 +1,61 @@ +import { describe, expect, it } from 'vitest'; +import { + applyMcpMaxTokens, + MCP_TRUNCATION_MARKER, + resolveMcpMaxTokens, + withoutMcpBudgetArg, +} from '../../src/mcp/output-budget.js'; + +describe('MCP output budget helpers', () => { + it('returns the original string byte-for-byte without a configured budget', () => { + const text = 'alphaπŸ˜€omega'; + expect(applyMcpMaxTokens(text, undefined)).toBe(text); + }); + + it('uses the complete marker and stays within a one-token budget', () => { + const text = applyMcpMaxTokens('this response is too long', 1); + expect(text).toBe(MCP_TRUNCATION_MARKER); + expect(Buffer.byteLength(text, 'utf8')).toBe(4); + }); + + it('never splits a multi-byte Unicode code point', () => { + const text = applyMcpMaxTokens('πŸ˜€πŸ˜€πŸ˜€πŸ˜€', 3); + expect(text.endsWith(MCP_TRUNCATION_MARKER)).toBe(true); + expect(text).not.toContain('\uFFFD'); + expect(Buffer.byteLength(text, 'utf8')).toBeLessThanOrEqual(12); + }); + + it('rejects malformed environment defaults for budgeted tools', () => { + expect(() => + resolveMcpMaxTokens('query', undefined, { + GITNEXUS_MCP_DEFAULT_MAX_TOKENS: '1.5', + }), + ).toThrow(/positive integer/i); + }); + + it('lets a valid explicit value override a malformed environment default', () => { + expect( + resolveMcpMaxTokens( + 'impact', + { maxTokens: 17 }, + { + GITNEXUS_MCP_DEFAULT_MAX_TOKENS: 'invalid', + }, + ), + ).toBe(17); + }); + + it('ignores the environment default for tools without output budgets', () => { + expect( + resolveMcpMaxTokens('cypher', undefined, { + GITNEXUS_MCP_DEFAULT_MAX_TOKENS: 'invalid', + }), + ).toBeUndefined(); + }); + + it('removes only the transport-level maxTokens argument', () => { + const args = { search_query: 'auth', maxTokens: 20, repo: 'app' }; + expect(withoutMcpBudgetArg(args)).toEqual({ search_query: 'auth', repo: 'app' }); + expect(args).toEqual({ search_query: 'auth', maxTokens: 20, repo: 'app' }); + }); +}); diff --git a/gitnexus/test/unit/server.test.ts b/gitnexus/test/unit/server.test.ts index f2ac995da6..84850ffa41 100644 --- a/gitnexus/test/unit/server.test.ts +++ b/gitnexus/test/unit/server.test.ts @@ -43,6 +43,27 @@ function createMockBackend(overrides: Record = {}): any { }; } +async function callToolThroughServer( + backend: ReturnType, + name: string, + args: Record, +): Promise<{ text: string; isError: boolean }> { + const server = createMCPServer(backend); + const client = new Client({ name: 'budget-test-client', version: '0.0.0' }); + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + + try { + await Promise.all([server.connect(serverTransport), client.connect(clientTransport)]); + const response = await client.callTool({ name, arguments: args }); + const text = response.content.find((item) => item.type === 'text')?.text; + if (typeof text !== 'string') throw new Error('Expected an MCP text response'); + return { text, isError: response.isError === true }; + } finally { + await client.close(); + await server.close(); + } +} + // ─── createMCPServer ───────────────────────────────────────────────── describe('createMCPServer', () => { @@ -105,6 +126,107 @@ describe('getNextStepHint (via tool call response)', () => { }); }); +describe('MCP output budgets', () => { + it('leaves the complete formatted response unchanged when no budget is configured', async () => { + const previous = process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS; + delete process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS; + try { + const backend = createMockBackend({ + callTool: vi.fn().mockResolvedValue({ payload: 'complete' }), + }); + const { text, isError } = await callToolThroughServer(backend, 'query', { + search_query: 'auth', + }); + expect(isError).toBe(false); + expect(text).toContain('"payload": "complete"'); + expect(text).toContain('**Next:**'); + expect(text.endsWith('\n…')).toBe(false); + } finally { + if (previous === undefined) delete process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS; + else process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS = previous; + } + }); + + it('applies explicit maxTokens to the complete response deterministically and UTF-8 safely', async () => { + const backend = createMockBackend({ + callTool: vi.fn().mockResolvedValue({ payload: 'πŸ˜€'.repeat(100) }), + }); + const args = { search_query: 'auth', maxTokens: 8 }; + + const first = await callToolThroughServer(backend, 'query', args); + const second = await callToolThroughServer(backend, 'query', args); + + expect(first.isError).toBe(false); + expect(first.text).toBe(second.text); + expect(Buffer.byteLength(first.text, 'utf8')).toBeLessThanOrEqual(8 * 4); + expect(first.text.endsWith('\n…')).toBe(true); + expect(first.text).not.toContain('\uFFFD'); + expect(backend.callTool).toHaveBeenCalledWith('query', { search_query: 'auth' }); + }); + + it('uses the environment default when maxTokens is omitted', async () => { + const previous = process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS; + process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS = '8'; + try { + const backend = createMockBackend({ + callTool: vi.fn().mockResolvedValue({ payload: 'x'.repeat(200) }), + }); + const { text } = await callToolThroughServer(backend, 'context', { name: 'auth' }); + expect(Buffer.byteLength(text, 'utf8')).toBeLessThanOrEqual(8 * 4); + expect(text.endsWith('\n…')).toBe(true); + } finally { + if (previous === undefined) delete process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS; + else process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS = previous; + } + }); + + it('lets an explicit request override the environment default', async () => { + const previous = process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS; + process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS = '1'; + try { + const backend = createMockBackend({ + callTool: vi.fn().mockResolvedValue({ payload: 'complete' }), + }); + const { text } = await callToolThroughServer(backend, 'impact', { + target: 'auth', + direction: 'upstream', + maxTokens: 200, + }); + expect(text).toContain('"payload": "complete"'); + expect(text).toContain('**Next:**'); + expect(text.endsWith('\n…')).toBe(false); + } finally { + if (previous === undefined) delete process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS; + else process.env.GITNEXUS_MCP_DEFAULT_MAX_TOKENS = previous; + } + }); + + it('rejects a non-positive explicit maxTokens before backend execution', async () => { + const backend = createMockBackend(); + const { text, isError } = await callToolThroughServer(backend, 'query', { + search_query: 'auth', + maxTokens: 0, + }); + expect(isError).toBe(true); + expect(text).toMatch(/maxTokens.*positive integer/i); + expect(backend.callTool).not.toHaveBeenCalled(); + }); + + it('applies a valid budget to backend error text', async () => { + const backend = createMockBackend({ + callTool: vi.fn().mockRejectedValue(new Error('πŸ˜€'.repeat(100))), + }); + const { text, isError } = await callToolThroughServer(backend, 'context', { + name: 'auth', + maxTokens: 8, + }); + expect(isError).toBe(true); + expect(Buffer.byteLength(text, 'utf8')).toBeLessThanOrEqual(8 * 4); + expect(text.endsWith('\n…')).toBe(true); + expect(text).not.toContain('\uFFFD'); + }); +}); + // ─── Tool handler error handling ────────────────────────────────────── describe('server error handling', () => { diff --git a/gitnexus/test/unit/tools.test.ts b/gitnexus/test/unit/tools.test.ts index 4757fd9f7c..d55a214901 100644 --- a/gitnexus/test/unit/tools.test.ts +++ b/gitnexus/test/unit/tools.test.ts @@ -185,6 +185,16 @@ describe('GITNEXUS_TOOLS', () => { expect(impactTool.description).toContain('truncatedBy'); }); + it.each(['query', 'context', 'impact'])( + '%s advertises an optional positive maxTokens budget', + (name) => { + const tool = GITNEXUS_TOOLS.find((definition) => definition.name === name)!; + const maxTokens = tool.inputSchema.properties.maxTokens; + expect(maxTokens).toMatchObject({ type: 'integer', minimum: 1 }); + expect(tool.inputSchema.required).not.toContain('maxTokens'); + }, + ); + it('rename tool requires new_name', () => { const renameTool = GITNEXUS_TOOLS.find((t) => t.name === 'rename')!; expect(renameTool.inputSchema.required).toContain('new_name');