Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 58 additions & 0 deletions gitnexus/src/mcp/output-budget.ts
Original file line number Diff line number Diff line change
@@ -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<string, unknown> | 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<string, unknown> | undefined,
): Record<string, unknown> | undefined {
if (!args || !Object.prototype.hasOwnProperty.call(args, 'maxTokens')) return args;
const backendArgs = { ...args };
delete backendArgs.maxTokens;
return backendArgs;
}
10 changes: 7 additions & 3 deletions gitnexus/src/mcp/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -165,17 +166,20 @@ 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<string, unknown> | 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<string, any> | undefined);

return {
content: [
{
type: 'text',
text: resultText + hint,
text: applyMcpMaxTokens(resultText + hint, maxTokens),
},
],
};
Expand All @@ -185,7 +189,7 @@ export function createMCPServer(backend: LocalBackend): Server {
content: [
{
type: 'text',
text: `Error: ${message}`,
text: applyMcpMaxTokens(`Error: ${message}`, maxTokens),
},
],
isError: true,
Expand Down
18 changes: 18 additions & 0 deletions gitnexus/src/mcp/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
61 changes: 61 additions & 0 deletions gitnexus/test/unit/mcp-output-budget.test.ts
Original file line number Diff line number Diff line change
@@ -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' });
});
});
122 changes: 122 additions & 0 deletions gitnexus/test/unit/server.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,27 @@ function createMockBackend(overrides: Record<string, any> = {}): any {
};
}

async function callToolThroughServer(
backend: ReturnType<typeof createMockBackend>,
name: string,
args: Record<string, unknown>,
): 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', () => {
Expand Down Expand Up @@ -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', () => {
Expand Down
10 changes: 10 additions & 0 deletions gitnexus/test/unit/tools.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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');
Expand Down
Loading