Skip to content

Commit 3b2a339

Browse files
zhichlizhichli
committed
demo
1 parent 5b567cd commit 3b2a339

File tree

1 file changed

+22
-114
lines changed

1 file changed

+22
-114
lines changed

‎src/extension/tools/node/searchSubagentTool.ts‎

Lines changed: 22 additions & 114 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,13 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6-
import { JSONTree } from '@vscode/prompt-tsx';
7-
import * as fs from 'fs';
8-
import * as os from 'os';
9-
import * as path from 'path';
106
import type * as vscode from 'vscode';
117
import { ChatFetchResponseType } from '../../../platform/chat/common/commonTypes';
8+
import { CapturingToken } from '../../../platform/requestLogger/common/capturingToken';
9+
import { IRequestLogger } from '../../../platform/requestLogger/node/requestLogger';
1210
import { ChatResponseStreamImpl } from '../../../util/common/chatResponseStreamImpl';
1311
import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
14-
import { ChatPrepareToolInvocationPart, ChatResponseNotebookEditPart, ChatResponseTextEditPart, ExtendedLanguageModelToolResult, LanguageModelDataPart, LanguageModelPromptTsxPart, LanguageModelTextPart } from '../../../vscodeTypes';
12+
import { ChatPrepareToolInvocationPart, ChatResponseNotebookEditPart, ChatResponseTextEditPart, ExtendedLanguageModelToolResult, LanguageModelTextPart } from '../../../vscodeTypes';
1513
import { Conversation, Turn } from '../../prompt/common/conversation';
1614
import { IBuildPromptContext } from '../../prompt/common/intents';
1715
import { SubagentToolCallingLoop } from '../../prompt/node/subagentLoop';
@@ -20,33 +18,6 @@ import { PromptElementCtor } from '../../prompts/node/base/promptElement';
2018
import { ToolName } from '../common/toolNames';
2119
import { CopilotToolMode, ICopilotTool, ToolRegistry } from '../common/toolsRegistry';
2220

23-
// Local function to render PromptTsx parts to strings (simplified to avoid vscode-node import restrictions)
24-
function renderToolResultToStringNoBudget(part: LanguageModelPromptTsxPart): string {
25-
// Extract text content from the prompt-tsx tree structure
26-
const json = part.value as JSONTree.PromptElementJSON;
27-
return extractTextFromPromptTree(json.node);
28-
}
29-
30-
function extractTextFromPromptTree(node: unknown): string {
31-
if (!node || typeof node !== 'object') {
32-
return '';
33-
}
34-
35-
const nodeObj = node as { type?: number; text?: string; children?: unknown[] };
36-
37-
// If this is a text node, return its text
38-
if (nodeObj.type === 2 && typeof nodeObj.text === 'string') {
39-
return nodeObj.text;
40-
}
41-
42-
// If this is an element node with children, recursively extract text from children
43-
if (Array.isArray(nodeObj.children)) {
44-
return nodeObj.children.map(child => extractTextFromPromptTree(child)).join('');
45-
}
46-
47-
return '';
48-
}
49-
5021
export interface ISearchSubagentParams {
5122

5223
/** Natural language query describing what to search for */
@@ -61,6 +32,7 @@ class SearchSubagentTool implements ICopilotTool<ISearchSubagentParams> {
6132

6233
constructor(
6334
@IInstantiationService private readonly instantiationService: IInstantiationService,
35+
@IRequestLogger private readonly requestLogger: IRequestLogger,
6436
) { }
6537
async invoke(options: vscode.LanguageModelToolInvocationOptions<ISearchSubagentParams>, token: vscode.CancellationToken) {
6638
const searchInstruction = [
@@ -107,90 +79,24 @@ class SearchSubagentTool implements ICopilotTool<ISearchSubagentParams> {
10779
part => part instanceof ChatPrepareToolInvocationPart || part instanceof ChatResponseTextEditPart || part instanceof ChatResponseNotebookEditPart
10880
);
10981

110-
const loopResult = await loop.run(stream, token);
111-
112-
// Write trajectory to file in same format as logToolCall
113-
const timestamp = new Date().toISOString().replace(/[:.]/g, '-');
114-
const trajectoryFile = path.join(os.homedir(), `search_trajectory-${timestamp}.json`);
115-
116-
try {
117-
// Build a structured trajectory similar to chat-export-logs format
118-
const rounds: Array<{
119-
role: 'assistant' | 'tool';
120-
content: Array<{ type: 'text' | 'data'; text?: string; data?: unknown }>;
121-
toolCalls?: Array<{ id: string; name: string; arguments: unknown }>;
122-
toolCallId?: string;
123-
thinking?: unknown;
124-
}> = [];
125-
126-
for (const round of loopResult.toolCallRounds) {
127-
// Assistant message with tool calls
128-
if (round.toolCalls.length > 0) {
129-
rounds.push({
130-
role: 'assistant',
131-
content: round.response ? [{ type: 'text', text: round.response }] : [],
132-
toolCalls: round.toolCalls.map(tc => ({
133-
id: tc.id,
134-
name: tc.name,
135-
arguments: JSON.parse(tc.arguments)
136-
})),
137-
thinking: round.thinking
138-
});
139-
140-
// Tool results - render them properly
141-
for (const toolCall of round.toolCalls) {
142-
const result = loopResult.toolCallResults[toolCall.id];
143-
if (result) {
144-
const content: Array<{ type: 'text' | 'data'; text?: string; data?: unknown }> = [];
145-
for (const part of result.content) {
146-
if (part instanceof LanguageModelTextPart) {
147-
content.push({ type: 'text', text: part.value });
148-
} else if (part instanceof LanguageModelDataPart) {
149-
content.push({ type: 'data', data: part.data });
150-
} else if (part instanceof LanguageModelPromptTsxPart) {
151-
// Render prompt-tsx parts to readable text
152-
const rendered = renderToolResultToStringNoBudget(part);
153-
content.push({ type: 'text', text: rendered });
154-
} else {
155-
content.push({ type: 'text', text: String(part) });
156-
}
157-
}
158-
rounds.push({
159-
role: 'tool',
160-
content: content,
161-
toolCallId: toolCall.id
162-
});
163-
}
164-
}
165-
} else {
166-
// Final assistant message without tool calls
167-
rounds.push({
168-
role: 'assistant',
169-
content: [{ type: 'text', text: round.response }]
170-
});
171-
}
172-
}
82+
// Create a new capturing token to group this search subagent and all its nested tool calls
83+
// Similar to how DefaultIntentRequestHandler does it
84+
const searchSubagentToken = new CapturingToken(
85+
`Search: ${options.input.query.substring(0, 50)}${options.input.query.length > 50 ? '...' : ''}`,
86+
'search',
87+
false
88+
);
17389

174-
const trajectory = {
175-
id: `search-subagent-${timestamp}`,
176-
kind: 'toolCall',
177-
tool: ToolName.SearchSubagent,
178-
metadata: {
179-
query: options.input.query,
180-
description: options.input.description,
181-
time: new Date().toISOString(),
182-
responseType: loopResult.response.type,
183-
success: loopResult.response.type === ChatFetchResponseType.Success
184-
},
185-
conversation: rounds,
186-
finalResponse: loopResult.round.response
187-
};
90+
// Wrap the loop execution in captureInvocation with the new token
91+
// All nested tool calls will now be logged under this same CapturingToken
92+
const loopResult = await this.requestLogger.captureInvocation(searchSubagentToken, () => loop.run(stream, token));
18893

189-
fs.writeFileSync(trajectoryFile, JSON.stringify(trajectory, null, 2), 'utf-8');
190-
console.log('[SearchSubagent] Wrote trajectory to:', trajectoryFile);
191-
} catch (error) {
192-
console.error('[SearchSubagent] FAILED to write trajectory to:', trajectoryFile, error);
193-
}
94+
// Build subagent trajectory metadata that will be logged via toolMetadata
95+
// All nested tool calls are already logged by ToolCallingLoop.logToolResult()
96+
const toolMetadata = {
97+
query: options.input.query,
98+
description: options.input.description
99+
};
194100

195101
let subagentResponse = '';
196102
if (loopResult.response.type === ChatFetchResponseType.Success) {
@@ -199,7 +105,9 @@ class SearchSubagentTool implements ICopilotTool<ISearchSubagentParams> {
199105
subagentResponse = `The search subagent request failed with this message:\n${loopResult.response.type}: ${loopResult.response.reason}`;
200106
}
201107

108+
// toolMetadata will be automatically included in exportAllPromptLogsAsJsonCommand
202109
const result = new ExtendedLanguageModelToolResult([new LanguageModelTextPart(subagentResponse)]);
110+
result.toolMetadata = toolMetadata;
203111
return result;
204112
}
205113

0 commit comments

Comments
 (0)