Skip to content

Commit db83b7b

Browse files
author
Zehee
committed
feat: add output-size controls to search_context
- Add detail levels: compact / normal / full - Add max_output_chars budget (default 6000 chars) - Reduce default user/agent truncation to 200/400 chars - Trim matches/clusters/members to stay within budget - Add test for compact mode
1 parent 9ab4f40 commit db83b7b

5 files changed

Lines changed: 221 additions & 24 deletions

File tree

‎src/config.ts‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,3 +74,20 @@ export const ESSENCE_SIZE_LIMIT = 15 * 1024; // 15 KB
7474
* 一段流畅讨论的完整性。每次搜索可以通过 cluster_gap_seconds 参数覆盖。
7575
*/
7676
export const DEFAULT_CLUSTER_GAP_SECONDS = 90;
77+
78+
/**
79+
* search_context 默认输出预算（字符数）。
80+
*
81+
* 默认的 normal 模式会尽量把输出控制在这个范围内，避免一次搜索就占满
82+
* 上下文窗口。可通过 max_output_chars 参数覆盖，detail: 'full' 时禁用。
83+
*/
84+
export const DEFAULT_SEARCH_OUTPUT_BUDGET = 6000;
85+
86+
/** search_context 匹配项摘要最大长度。 */
87+
export const SEARCH_SNIPPET_MAX_LEN = 240;
88+
89+
/** search_context normal 模式匹配项 user 文本最大长度。 */
90+
export const SEARCH_USER_MAX_LEN = 200;
91+
92+
/** search_context normal 模式匹配项 agent 文本最大长度。 */
93+
export const SEARCH_AGENT_MAX_LEN = 400;

‎src/context/wire-context.ts‎

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@ import {
1414
getSessionsRoot,
1515
DEFAULT_CONTEXT_WINDOW,
1616
DEFAULT_RECENT_CHANGE_LIMIT,
17+
SEARCH_SNIPPET_MAX_LEN,
18+
SEARCH_USER_MAX_LEN,
19+
SEARCH_AGENT_MAX_LEN,
1720
} from '../config.js';
1821
import { computeWorkspaceHash } from '../utils/paths.js';
1922
import { scoreText, extractSnippet } from '../utils/search.js';
@@ -67,6 +70,9 @@ export interface SearchOptions {
6770
dateFrom?: string;
6871
dateTo?: string;
6972
refinedManager?: RefinedManager;
73+
userMaxLen?: number;
74+
agentMaxLen?: number;
75+
snippetMaxLen?: number;
7076
}
7177

7278
export interface SearchMatch {
@@ -275,18 +281,21 @@ function buildMatchFromRefined(
275281
score: number;
276282
},
277283
terms: string[],
284+
options: SearchOptions = {},
278285
): SearchMatch {
279286
// Refined summaries already condense user intent, tool usage, and agent lead.
280287
// Use them as the primary content when the original wire is unavailable.
281288
const snippetSource = [refined.summary, ...refined.facts, ...refined.notes].join('\n');
289+
const agentMaxLen = options.agentMaxLen ?? SEARCH_AGENT_MAX_LEN;
290+
const snippetMaxLen = options.snippetMaxLen ?? SEARCH_SNIPPET_MAX_LEN;
282291
return {
283292
sessionId: refined.sessionId,
284293
turnId: refined.turnId,
285294
timestamp: refined.timestamp ?? null,
286295
score: refined.score,
287296
user: '',
288-
agent: truncate(refined.summary, 2000),
289-
snippet: extractSnippet(snippetSource, terms, 240),
297+
agent: truncate(refined.summary, agentMaxLen),
298+
snippet: extractSnippet(snippetSource, terms, snippetMaxLen),
290299
actions: [],
291300
};
292301
}
@@ -529,6 +538,9 @@ export async function searchWireContext(query: string, options: SearchOptions =
529538
const dateFrom = options.dateFrom || null;
530539
const dateTo = options.dateTo || null;
531540
const refinedManager = options.refinedManager;
541+
const userMaxLen = options.userMaxLen ?? SEARCH_USER_MAX_LEN;
542+
const agentMaxLen = options.agentMaxLen ?? SEARCH_AGENT_MAX_LEN;
543+
const snippetMaxLen = options.snippetMaxLen ?? SEARCH_SNIPPET_MAX_LEN;
532544
const terms = query
533545
.toLowerCase()
534546
.split(/\s+/)
@@ -568,13 +580,13 @@ export async function searchWireContext(query: string, options: SearchOptions =
568580
if (!turns) {
569581
// The original wire has been removed, but the refined summary still
570582
// contains the key facts. Return it instead of skipping.
571-
matches.push(buildMatchFromRefined(refined, terms));
583+
matches.push(buildMatchFromRefined(refined, terms, options));
572584
continue;
573585
}
574586
const turn = turns.find((t) => parseInt(t.turnId, 10) === refined.turnId);
575587
if (!turn) {
576588
// The turn is gone from the wire; fall back to the refined record.
577-
matches.push(buildMatchFromRefined(refined, terms));
589+
matches.push(buildMatchFromRefined(refined, terms, options));
578590
continue;
579591
}
580592

@@ -588,9 +600,9 @@ export async function searchWireContext(query: string, options: SearchOptions =
588600
turnId: refined.turnId,
589601
timestamp: turn.timestamp || null,
590602
score: refined.score,
591-
user: truncate(turn.user, 1000),
592-
agent: truncate(turn.agentText, 2000),
593-
snippet: extractSnippet(fullText, terms, 240),
603+
user: truncate(turn.user, userMaxLen),
604+
agent: truncate(turn.agentText, agentMaxLen),
605+
snippet: extractSnippet(fullText, terms, snippetMaxLen),
594606
actions: turn.actions.map((a) => a.name).filter(Boolean),
595607
});
596608
hits.push({ sessionId: refined.sessionId, turn });
@@ -628,9 +640,9 @@ export async function searchWireContext(query: string, options: SearchOptions =
628640
turnId: parseInt(turn.turnId, 10),
629641
timestamp: turn.timestamp || null,
630642
score,
631-
user: truncate(turn.user, 1000),
632-
agent: truncate(turn.agentText, 2000),
633-
snippet: extractSnippet(fullText, terms, 240),
643+
user: truncate(turn.user, userMaxLen),
644+
agent: truncate(turn.agentText, agentMaxLen),
645+
snippet: extractSnippet(fullText, terms, snippetMaxLen),
634646
actions: turn.actions.map((a) => a.name).filter(Boolean),
635647
});
636648
hits.push({ sessionId: session.sessionId, turn });

‎src/tools/context-tools.ts‎

Lines changed: 162 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import type {
1212
import type { Ctx } from '../types.js';
1313
import type { ToolDefinition } from './types.js';
1414
import { adaptHandler } from './types.js';
15-
import type { TurnReference, WireTurn } from '../context/wire-context.js';
15+
import type { SearchMatch, TurnReference, WireTurn } from '../context/wire-context.js';
1616
import {
1717
getCurrentSessionWirePath,
1818
findAllWorkspaceSessions,
@@ -25,7 +25,13 @@ import {
2525
import fs from 'fs';
2626
import path from 'path';
2727
import crypto from 'crypto';
28-
import { DEFAULT_CLUSTER_GAP_SECONDS } from '../config.js';
28+
import {
29+
DEFAULT_CLUSTER_GAP_SECONDS,
30+
DEFAULT_SEARCH_OUTPUT_BUDGET,
31+
SEARCH_USER_MAX_LEN,
32+
SEARCH_AGENT_MAX_LEN,
33+
SEARCH_SNIPPET_MAX_LEN,
34+
} from '../config.js';
2935
import { toolResult } from '../utils/tools.js';
3036

3137
export async function buildWorkspaceContext(
@@ -94,6 +100,97 @@ export function createContextTools(ctx: Ctx): ToolDefinition[] {
94100
return DEFAULT_CLUSTER_GAP_SECONDS;
95101
}
96102

103+
function resolveDetail(args: SearchContextArgs): 'compact' | 'normal' | 'full' {
104+
const valid: Array<'compact' | 'normal' | 'full'> = ['compact', 'normal', 'full'];
105+
if (typeof args.detail === 'string' && valid.includes(args.detail as 'compact' | 'normal' | 'full')) {
106+
return args.detail as 'compact' | 'normal' | 'full';
107+
}
108+
return 'normal';
109+
}
110+
111+
function resolveMaxOutputChars(args: SearchContextArgs, detail: 'compact' | 'normal' | 'full'): number {
112+
if (detail === 'full') return Number.MAX_SAFE_INTEGER;
113+
if (typeof args.max_output_chars === 'number' && args.max_output_chars > 0) {
114+
return Math.max(500, Math.floor(args.max_output_chars));
115+
}
116+
return DEFAULT_SEARCH_OUTPUT_BUDGET;
117+
}
118+
119+
function estimateOutputSize(obj: unknown): number {
120+
return JSON.stringify(obj).length;
121+
}
122+
123+
interface CompactMatch {
124+
sessionId: string;
125+
turnId: number;
126+
timestamp: string | null;
127+
score: number;
128+
}
129+
130+
interface CompactCluster {
131+
sessionId: string;
132+
hitTurnId: number;
133+
memberCount: number;
134+
refinedCount: number;
135+
}
136+
137+
interface NormalCluster extends CompactCluster {
138+
members: Array<{ sessionId: string; turnId: number }>;
139+
}
140+
141+
function buildCompactMatch(match: SearchMatch): CompactMatch {
142+
return {
143+
sessionId: match.sessionId,
144+
turnId: match.turnId,
145+
timestamp: match.timestamp,
146+
score: match.score,
147+
};
148+
}
149+
150+
function buildNormalCluster(cluster: Cluster, refinedCount: number): NormalCluster {
151+
return {
152+
sessionId: cluster.sessionId,
153+
hitTurnId: cluster.hitTurnId,
154+
memberCount: cluster.members.length,
155+
refinedCount,
156+
members: cluster.members,
157+
};
158+
}
159+
160+
function buildCompactCluster(cluster: Cluster, refinedCount: number): CompactCluster {
161+
return {
162+
sessionId: cluster.sessionId,
163+
hitTurnId: cluster.hitTurnId,
164+
memberCount: cluster.members.length,
165+
refinedCount,
166+
};
167+
}
168+
169+
function trimOutputToBudget(
170+
result: {
171+
matches: SearchMatch[] | CompactMatch[];
172+
clusters: NormalCluster[] | CompactCluster[];
173+
},
174+
maxChars: number,
175+
): void {
176+
// Matches are already sorted by score descending; drop lowest-scoring last.
177+
while (estimateOutputSize(result) > maxChars && result.matches.length > 0) {
178+
result.matches.pop();
179+
}
180+
while (estimateOutputSize(result) > maxChars && result.clusters.length > 0) {
181+
result.clusters.pop();
182+
}
183+
// If still over budget, strip members arrays from clusters.
184+
if (estimateOutputSize(result) > maxChars && result.clusters.length > 0) {
185+
result.clusters = result.clusters.map((c) => ({
186+
sessionId: c.sessionId,
187+
hitTurnId: c.hitTurnId,
188+
memberCount: c.memberCount,
189+
refinedCount: c.refinedCount,
190+
})) as CompactCluster[];
191+
}
192+
}
193+
97194
function groupHitsIntoBlocks(hits: WireTurn[]): WireTurn[][] {
98195
const sorted = [...hits].sort(
99196
(a, b) => parseInt(a.turnId, 10) - parseInt(b.turnId, 10),
@@ -228,12 +325,26 @@ export function createContextTools(ctx: Ctx): ToolDefinition[] {
228325
typeof args.max_cluster_size === 'number' && args.max_cluster_size > 0
229326
? Math.max(2, Math.floor(args.max_cluster_size))
230327
: 15;
231-
const { matches, hits, skippedSessionIds } = await searchWireContext(query, {
328+
const detail = resolveDetail(args);
329+
const maxOutputChars = resolveMaxOutputChars(args, detail);
330+
331+
const searchOptions: Parameters<typeof searchWireContext>[1] = {
232332
limit: args.limit,
233333
dateFrom: args.date_from,
234334
dateTo: args.date_to,
235335
refinedManager,
236-
});
336+
};
337+
if (detail === 'full') {
338+
searchOptions.userMaxLen = Number.MAX_SAFE_INTEGER;
339+
searchOptions.agentMaxLen = Number.MAX_SAFE_INTEGER;
340+
searchOptions.snippetMaxLen = SEARCH_SNIPPET_MAX_LEN;
341+
} else if (detail === 'normal') {
342+
searchOptions.userMaxLen = SEARCH_USER_MAX_LEN;
343+
searchOptions.agentMaxLen = SEARCH_AGENT_MAX_LEN;
344+
searchOptions.snippetMaxLen = SEARCH_SNIPPET_MAX_LEN;
345+
}
346+
347+
const { matches, hits, skippedSessionIds } = await searchWireContext(query, searchOptions);
237348

238349
// Build clusters around each hit and refine any unrefined turns in them.
239350
const hitsBySession = new Map<string, WireTurn[]>();
@@ -327,21 +438,47 @@ export function createContextTools(ctx: Ctx): ToolDefinition[] {
327438
skippedSessions,
328439
);
329440

441+
const normalClusters = clusters.map((c, i) => buildNormalCluster(c, refinedCountByCluster[i] || 0));
442+
const compactClusters = clusters.map((c, i) => buildCompactCluster(c, refinedCountByCluster[i] || 0));
443+
const compactMatches = matches.map(buildCompactMatch);
444+
445+
if (detail === 'normal') {
446+
const normalResult = {
447+
query,
448+
totalMatches: matches.length,
449+
refinedCount,
450+
clusterGapSeconds,
451+
maxClusterSize,
452+
skippedSessions,
453+
matches,
454+
clusters: normalClusters,
455+
};
456+
trimOutputToBudget(normalResult, maxOutputChars);
457+
return toolResult(normalResult);
458+
}
459+
460+
if (detail === 'compact') {
461+
return toolResult({
462+
query,
463+
totalMatches: matches.length,
464+
refinedCount,
465+
clusterGapSeconds,
466+
maxClusterSize,
467+
skippedSessions,
468+
matches: compactMatches,
469+
clusters: compactClusters,
470+
});
471+
}
472+
330473
return toolResult({
331474
query,
332475
totalMatches: matches.length,
333-
matches,
476+
refinedCount,
334477
clusterGapSeconds,
335478
maxClusterSize,
336479
skippedSessions,
337-
clusters: clusters.map((c, i) => ({
338-
sessionId: c.sessionId,
339-
hitTurnId: c.hitTurnId,
340-
memberCount: c.members.length,
341-
refinedCount: refinedCountByCluster[i] || 0,
342-
members: c.members,
343-
})),
344-
refinedCount,
480+
matches,
481+
clusters: normalClusters,
345482
});
346483
}
347484

@@ -478,7 +615,7 @@ export function createContextTools(ctx: Ctx): ToolDefinition[] {
478615
{
479616
name: 'search_context',
480617
description:
481-
'Search conversation rounds across all workspace session wires by keywords and optional date range.',
618+
"Search conversation rounds across all workspace session wires by keywords and optional date range. Default detail: 'normal' keeps output within ~6000 chars. Use detail: 'compact' for a quick overview (no match text, no cluster members). Use detail: 'full' when you need full match text and all cluster members.",
482619
inputSchema: {
483620
type: 'object',
484621
properties: {
@@ -496,6 +633,17 @@ export function createContextTools(ctx: Ctx): ToolDefinition[] {
496633
description:
497634
'单个 cluster 最多包含的 turn 数，防止连续讨论过长时上下文爆炸。默认 15。',
498635
},
636+
detail: {
637+
type: 'string',
638+
enum: ['compact', 'normal', 'full'],
639+
description:
640+
"Output detail level. 'normal' (default) returns truncated text and cluster members within the output budget. 'compact' returns only references/counts. 'full' disables the budget and returns longer text.",
641+
},
642+
max_output_chars: {
643+
type: 'number',
644+
description:
645+
'Maximum output length in characters for normal mode. Default 6000. Ignored in compact/full.',
646+
},
499647
},
500648
required: ['query'],
501649
},

‎src/types.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@ export interface SearchContextArgs {
7070
limit?: number;
7171
cluster_gap_seconds?: number;
7272
max_cluster_size?: number;
73+
detail?: 'compact' | 'normal' | 'full';
74+
max_output_chars?: number;
7375
}
7476

7577
export interface ListSearchViewsArgs {

‎tests/test.ts‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -570,6 +570,23 @@ async function testClusterMaxSize() {
570570
}
571571
}
572572

573+
async function testSearchContextCompact() {
574+
await withClient(async (client) => {
575+
const result = parseJsonResult(
576+
await client.callTool({
577+
name: 'search_context',
578+
arguments: { query: 'nonexistent-xyz', detail: 'compact' },
579+
}),
580+
);
581+
assert(Array.isArray(result.matches));
582+
assert(Array.isArray(result.clusters));
583+
assert.strictEqual(result.refinedCount, 0);
584+
for (const cluster of result.clusters) {
585+
assert(!('members' in cluster), 'compact clusters should not include members array');
586+
}
587+
});
588+
}
589+
573590
async function testSearchContextReturnsRefinedForMissingWire() {
574591
const tmpSessions = fs.mkdtempSync(path.join(os.tmpdir(), 'kimi-code-missing-sessions-'));
575592
const cwd = process.cwd().replace(/\\/g, '/');
@@ -728,6 +745,7 @@ const tests = [
728745
testSearchContextUsesRefined,
729746
testConfigurableSessionsRoot,
730747
testClusterMaxSize,
748+
testSearchContextCompact,
731749
testSearchContextReturnsRefinedForMissingWire,
732750
testListSearchViews,
733751
testLoadWorkspaceContext,

0 commit comments

Comments
 (0)