From 05585207ba6217c08cb094aee87331bb3b8d6bfb Mon Sep 17 00:00:00 2001 From: iveteamorim Date: Wed, 23 Sep 2026 23:16:44 +0200 Subject: [PATCH 1/2] fix(graph): preserve matched nodes over output budget --- src/graph/__tests__/cli-agent.test.ts | 83 +++++++++++++++++++++++++++ src/graph/cli-agent.ts | 63 +++++++++++++++++++- 2 files changed, 144 insertions(+), 2 deletions(-) diff --git a/src/graph/__tests__/cli-agent.test.ts b/src/graph/__tests__/cli-agent.test.ts index 34fdf75e..5fb90e82 100644 --- a/src/graph/__tests__/cli-agent.test.ts +++ b/src/graph/__tests__/cli-agent.test.ts @@ -1332,6 +1332,89 @@ describe("runGraphGet", () => { expect(records).toHaveLength(1); expect(records[0]).toMatchObject({ type: "error", code: "INVALID_OUTPUT_BUDGET" }); }); + + function oversizedGetFixture(): { + isolated: string; + primary: GraphNode; + getDeps: AgentCommandDeps; + } { + const isolated = mkdtempSync(join(tmpdir(), "mex-graph-get-budget-")); + const sourceLines = Array.from({ length: 185 }, (_, index) => `// primary body ${index + 1}`); + sourceLines[0] = "export function PrimaryProcedure(): void {"; + sourceLines[184] = "}"; + const source = sourceLines.join("\n"); + writeFileSync(join(isolated, "primary.ts"), source); + const primary: GraphNode = { + id: "function:primary-procedure", kind: "function", name: "PrimaryProcedure", + qualifiedName: "PrimaryProcedure", filePath: "primary.ts", language: "typescript", + startLine: 1, endLine: 185, startColumn: 0, endColumn: 1, updatedAt: 1, + }; + const graph = syntheticScopeGraph({ + nodes: [primary], sources: [{ path: primary.filePath, content: source }], + searchNodes: () => [primary], + }); + return { + isolated, primary, + getDeps: { + open: () => ({ graph, db: deps.open!(root).db, close: () => {} }), + write: (line) => lines.push(line), + }, + }; + } + + it("does not report no-match when a matched node's source exceeds the output budget", () => { + const { isolated, primary, getDeps } = oversizedGetFixture(); + try { + const records = capture(() => runGraphGet([primary.id], isolated, getDeps, { + maxOutputTokens: 1500, maxSourceLines: 200, + })); + const summary = records.at(-1)!; + expect(summary.matchedNodes).toBeGreaterThan(0); + expect(records.some((record) => record.type === "fact" && record.id === primary.id)).toBe(true); + const prefixRanges = records.filter((record) => record.type === "source") + .flatMap((record) => record.ranges as Array<{ + startLine: number; endLine: number; truncated: boolean; nodeIds: string[]; + }>); + expect(prefixRanges).toEqual([expect.objectContaining({ + nodeIds: [primary.id], truncated: true, + })]); + expect(prefixRanges[0]!.endLine - prefixRanges[0]!.startLine + 1).toBeLessThan(185); + expect(summary).toMatchObject({ + status: "partial", truncated: true, evidenceStrength: "strong", returnedNodes: 1, + }); + expect(summary.estimatedOutputTokens as number).toBeLessThanOrEqual(1500); + const retry = (summary.suggestedNextCommands as string[])[0] ?? ""; + const budget = Number((/ --max-output-tokens (\d+)/.exec(retry) ?? [])[1]); + expect(retry).toContain(`mex graph get ${primary.id}`); + expect(budget).toBeGreaterThan(1500); + const retryRecords = capture(() => runGraphGet([primary.id], isolated, getDeps, { + maxOutputTokens: budget, maxSourceLines: 200, + })); + expect(retryRecords.filter((record) => record.type === "source") + .flatMap((record) => record.ranges as Array<{ endLine: number; truncated: boolean }>)) + .toEqual([expect.objectContaining({ endLine: 185, truncated: false })]); + expect(retryRecords.at(-1)).toMatchObject({ status: "ok", returnedNodes: 1 }); + } finally { + rmSync(isolated, { recursive: true, force: true }); + } + }); + + it("keeps a tight valid get budget honest when full source cannot fit", () => { + const { isolated, primary, getDeps } = oversizedGetFixture(); + try { + const records = capture(() => runGraphGet([primary.id], isolated, getDeps, { + maxOutputTokens: 800, maxSourceLines: 200, + })); + const summary = records.at(-1)!; + expect(summary.matchedNodes).toBeGreaterThan(0); + expect(summary.status).not.toBe("no-match"); + expect(summary).toMatchObject({ status: "partial", truncated: true }); + expect(summary.estimatedOutputTokens as number).toBeLessThanOrEqual(800); + expect(records.some((record) => record.type === "fact" && record.id === primary.id)).toBe(true); + } finally { + rmSync(isolated, { recursive: true, force: true }); + } + }); }); describe("runGraphQuery", () => { diff --git a/src/graph/cli-agent.ts b/src/graph/cli-agent.ts index 40f6051b..31890e75 100644 --- a/src/graph/cli-agent.ts +++ b/src/graph/cli-agent.ts @@ -820,15 +820,50 @@ export function runGraphGet( sourceRecords.flatMap((record) => (record.ranges as SourceRange[]).flatMap((range) => range.nodeIds)), ); + // Source-backed gets stay unchanged. A matched node whose full source + // cannot fit still exists: reserve its fact, then spill the largest + // whole-line prefix that fits, instead of falling through to no-match. + const omittedNodes = dedupeById(nodes.filter((node) => !sourcedIds.has(node.id))); + const factRecords: Rec[] = []; + let retry: { id: string; full: Rec } | undefined; + if (omittedNodes.length > 0) truncated = true; + const reserve = estimateTokens(summarySkeleton([])) + RESERVE_PAD; + for (const node of omittedNodes) { + const fact = factFor(session, node.id, opts.detail, opts.fingerprint); + if (fact) { + const record: Rec = { type: "fact", ...agentFactFields(fact, opts) }; + if (ledger.tryAdd(record)) factRecords.push(record); + } + const full = sourceRecordForGetNode(session, node, rootDir, opts); + if (!full) continue; + retry ??= { id: node.id, full }; + const available = ctx.effectiveMax - ledger.estimatedTokens - reserve; + const fitted = fitSourceRange(full, available, ledger); + if (fitted && ledger.tryAdd(fitted)) { + sourceRecords.push(fitted); + for (const range of fitted.ranges as SourceRange[]) { + for (const id of range.nodeIds) sourcedIds.add(id); + } + } + } + const suggestions = retry + ? [`mex graph get ${retry.id} --max-output-tokens ${Math.max( + graphGetFullSourceBudget(opts, retry.full), + opts.maxOutputTokens + 1, + )}`] + : []; + // `get` returns declarations and their proven source bytes only, so the // drift declaration appears without any record being marked stale. - emitAll(write, meta, [...configDriftRecords(session), ...errorRecords, ...sourceRecords]); + emitAll(write, meta, [...configDriftRecords(session), ...errorRecords, ...sourceRecords, ...factRecords]); write(JSON.stringify(summaryRecord(ctx, { matchedNodes: ids.length, returnedNodes: sourcedIds.size, returnedEdges: 0, truncated, - suggestedNextCommands: [], + suggestedNextCommands: suggestions, + ...(omittedNodes.length > 0 ? { status: "partial" as const } : {}), + ...(factRecords.length > 0 ? { evidenceStrength: "strong" as const } : {}), }))); }); } @@ -1052,6 +1087,30 @@ function planSource( return sourceRecords; } +/** Single-node source record for get's omit spill. Does not change `planSource`. */ +function sourceRecordForGetNode( + session: AgentGraphSession, + node: GraphNode, + rootDir: string, + opts: AgentOptions, +): Rec | null { + const source = session.readIndexedSource?.(node.filePath); + const range = source === undefined + ? readNodeSource(node, rootDir, opts.maxSourceLines) + : readNodeSourceBuffer(node, source, opts.maxSourceLines); + if (!range) return null; + return { type: "source", filePath: node.filePath, ranges: [range] }; +} + +/** Minimum honest get ceiling that can admit `fullSource` after protocol framing. */ +function graphGetFullSourceBudget(opts: AgentOptions, fullSource: Rec): number { + const reserve = estimateTokens(summarySkeleton([])) + RESERVE_PAD; + const framingMeta = metaRecord("graph get", { + ...opts, maxNodes: 1, maxFlowSteps: 0, maxOutputTokens: SIZE_PROBE, + }); + return estimateTokens(framingMeta) + estimateTokens(fullSource) + reserve; +} + function readNodeSourceBuffer(node: GraphNode, source: string, maxLines: number): SourceRange { const lines = source.split("\n"); const body = lines.slice(node.startLine - 1, node.endLine); From e31c7b348f39225e21987efc6156e5d4c95f9dd0 Mon Sep 17 00:00:00 2001 From: Yashasvi Date: Thu, 24 Sep 2026 10:23:45 +0530 Subject: [PATCH 2/2] fix(graph): keep every oversized node's fact and carry the line cap into the retry When several ids exceed the budget, admit every omitted node's fact before any source prefix spends the budget, so later nodes are not silently dropped. The suggested retry budget is sized for the current --max-source-lines, so a non-default cap is now included in the suggested command. --- src/graph/__tests__/cli-agent.test.ts | 63 ++++++++++++++++++++++----- src/graph/cli-agent.ts | 20 ++++++--- 2 files changed, 64 insertions(+), 19 deletions(-) diff --git a/src/graph/__tests__/cli-agent.test.ts b/src/graph/__tests__/cli-agent.test.ts index 5fb90e82..2eff0f7a 100644 --- a/src/graph/__tests__/cli-agent.test.ts +++ b/src/graph/__tests__/cli-agent.test.ts @@ -1336,25 +1336,36 @@ describe("runGraphGet", () => { function oversizedGetFixture(): { isolated: string; primary: GraphNode; + secondary: GraphNode; getDeps: AgentCommandDeps; } { const isolated = mkdtempSync(join(tmpdir(), "mex-graph-get-budget-")); - const sourceLines = Array.from({ length: 185 }, (_, index) => `// primary body ${index + 1}`); - sourceLines[0] = "export function PrimaryProcedure(): void {"; - sourceLines[184] = "}"; - const source = sourceLines.join("\n"); - writeFileSync(join(isolated, "primary.ts"), source); - const primary: GraphNode = { - id: "function:primary-procedure", kind: "function", name: "PrimaryProcedure", - qualifiedName: "PrimaryProcedure", filePath: "primary.ts", language: "typescript", - startLine: 1, endLine: 185, startColumn: 0, endColumn: 1, updatedAt: 1, + const oversized = (id: string, name: string, filePath: string): { node: GraphNode; source: string } => { + const sourceLines = Array.from({ length: 185 }, (_, index) => `// ${name} body ${index + 1}`); + sourceLines[0] = `export function ${name}(): void {`; + sourceLines[184] = "}"; + const source = sourceLines.join("\n"); + writeFileSync(join(isolated, filePath), source); + const node: GraphNode = { + id, kind: "function", name, qualifiedName: name, filePath, language: "typescript", + startLine: 1, endLine: 185, startColumn: 0, endColumn: 1, updatedAt: 1, + }; + return { node, source }; }; + const first = oversized("function:primary-procedure", "PrimaryProcedure", "primary.ts"); + const second = oversized("function:secondary-procedure", "SecondaryProcedure", "secondary.ts"); + const primary = first.node; + const secondary = second.node; const graph = syntheticScopeGraph({ - nodes: [primary], sources: [{ path: primary.filePath, content: source }], - searchNodes: () => [primary], + nodes: [primary, secondary], + sources: [ + { path: primary.filePath, content: first.source }, + { path: secondary.filePath, content: second.source }, + ], + searchNodes: () => [primary, secondary], }); return { - isolated, primary, + isolated, primary, secondary, getDeps: { open: () => ({ graph, db: deps.open!(root).db, close: () => {} }), write: (line) => lines.push(line), @@ -1387,6 +1398,7 @@ describe("runGraphGet", () => { const budget = Number((/ --max-output-tokens (\d+)/.exec(retry) ?? [])[1]); expect(retry).toContain(`mex graph get ${primary.id}`); expect(budget).toBeGreaterThan(1500); + expect(retry).toContain(" --max-source-lines 200"); const retryRecords = capture(() => runGraphGet([primary.id], isolated, getDeps, { maxOutputTokens: budget, maxSourceLines: 200, })); @@ -1399,6 +1411,33 @@ describe("runGraphGet", () => { } }); + it("keeps a fact for every oversized node before spending budget on source prefixes", () => { + const { isolated, primary, secondary, getDeps } = oversizedGetFixture(); + try { + const records = capture(() => runGraphGet([primary.id, secondary.id], isolated, getDeps, { + maxOutputTokens: 1500, maxSourceLines: 200, + })); + expect(records.filter((record) => record.type === "fact").map((record) => record.id)) + .toEqual([primary.id, secondary.id]); + const summary = records.at(-1)!; + expect(summary).toMatchObject({ status: "partial", truncated: true, matchedNodes: 2 }); + expect(summary.estimatedOutputTokens as number).toBeLessThanOrEqual(1500); + } finally { + rmSync(isolated, { recursive: true, force: true }); + } + }); + + it("leaves --max-source-lines out of the retry suggestion at the default line cap", () => { + const { isolated, primary, getDeps } = oversizedGetFixture(); + try { + const records = capture(() => runGraphGet([primary.id], isolated, getDeps, { maxOutputTokens: 1000 })); + const retry = (records.at(-1)!.suggestedNextCommands as string[])[0] ?? ""; + expect(retry).toMatch(/^mex graph get function:primary-procedure --max-output-tokens \d+$/); + } finally { + rmSync(isolated, { recursive: true, force: true }); + } + }); + it("keeps a tight valid get budget honest when full source cannot fit", () => { const { isolated, primary, getDeps } = oversizedGetFixture(); try { diff --git a/src/graph/cli-agent.ts b/src/graph/cli-agent.ts index 31890e75..9461f60c 100644 --- a/src/graph/cli-agent.ts +++ b/src/graph/cli-agent.ts @@ -17,7 +17,7 @@ import { import { FingerprintStore } from "./fingerprint-store.js"; import { serializeFingerprint } from "./fingerprint.js"; import { - BudgetLedger, estimateTokens, resolveOptions, resolveScopeOptions, SCHEMA_VERSION, type AgentOptions, + BudgetLedger, DEFAULT_OPTIONS, estimateTokens, resolveOptions, resolveScopeOptions, SCHEMA_VERSION, type AgentOptions, } from "./agent-protocol.js"; import { identifierComponents, isLowValueGraphPath, planGraphQuery } from "./retrieval/query.js"; import { @@ -827,13 +827,15 @@ export function runGraphGet( const factRecords: Rec[] = []; let retry: { id: string; full: Rec } | undefined; if (omittedNodes.length > 0) truncated = true; - const reserve = estimateTokens(summarySkeleton([])) + RESERVE_PAD; + // Every omitted node keeps its fact before any prefix spends the budget. for (const node of omittedNodes) { const fact = factFor(session, node.id, opts.detail, opts.fingerprint); - if (fact) { - const record: Rec = { type: "fact", ...agentFactFields(fact, opts) }; - if (ledger.tryAdd(record)) factRecords.push(record); - } + if (!fact) continue; + const record: Rec = { type: "fact", ...agentFactFields(fact, opts) }; + if (ledger.tryAdd(record)) factRecords.push(record); + } + const reserve = estimateTokens(summarySkeleton([])) + RESERVE_PAD; + for (const node of omittedNodes) { const full = sourceRecordForGetNode(session, node, rootDir, opts); if (!full) continue; retry ??= { id: node.id, full }; @@ -846,11 +848,15 @@ export function runGraphGet( } } } + // The retry budget is sized for the current line cap, so a non-default cap travels with it. + const sourceLines = opts.maxSourceLines === DEFAULT_OPTIONS.maxSourceLines + ? "" + : ` --max-source-lines ${opts.maxSourceLines}`; const suggestions = retry ? [`mex graph get ${retry.id} --max-output-tokens ${Math.max( graphGetFullSourceBudget(opts, retry.full), opts.maxOutputTokens + 1, - )}`] + )}${sourceLines}`] : []; // `get` returns declarations and their proven source bytes only, so the