Check for pending tool calls before appending IDE Context (#6317)
This commit is contained in:
Evan Otero
GitHub
parent
80763f5629
commit
5246aa11f4
|
|
@ -1479,6 +1479,349 @@ ${JSON.stringify(
|
|||
expect(contextJson.activeFile.path).toBe('/path/to/active/file.ts');
|
||||
});
|
||||
});
|
||||
|
||||
describe('IDE context with pending tool calls', () => {
|
||||
let mockChat: Partial<GeminiChat>;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.spyOn(client, 'tryCompressChat').mockResolvedValue(null);
|
||||
|
||||
const mockStream = (async function* () {
|
||||
yield { type: 'content', value: 'response' };
|
||||
})();
|
||||
mockTurnRunFn.mockReturnValue(mockStream);
|
||||
|
||||
mockChat = {
|
||||
addHistory: vi.fn(),
|
||||
getHistory: vi.fn().mockReturnValue([]), // Default empty history
|
||||
setHistory: vi.fn(),
|
||||
sendMessage: vi.fn().mockResolvedValue({ text: 'summary' }),
|
||||
};
|
||||
client['chat'] = mockChat as GeminiChat;
|
||||
|
||||
const mockGenerator: Partial<ContentGenerator> = {
|
||||
countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
|
||||
};
|
||||
client['contentGenerator'] = mockGenerator as ContentGenerator;
|
||||
|
||||
vi.spyOn(client['config'], 'getIdeMode').mockReturnValue(true);
|
||||
vi.mocked(ideContext.getIdeContext).mockReturnValue({
|
||||
workspaceState: {
|
||||
openFiles: [{ path: '/path/to/file.ts', timestamp: Date.now() }],
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it('should NOT add IDE context when a tool call is pending', async () => {
|
||||
// Arrange: History ends with a functionCall from the model
|
||||
const historyWithPendingCall: Content[] = [
|
||||
{ role: 'user', parts: [{ text: 'Please use a tool.' }] },
|
||||
{
|
||||
role: 'model',
|
||||
parts: [{ functionCall: { name: 'some_tool', args: {} } }],
|
||||
},
|
||||
];
|
||||
vi.mocked(mockChat.getHistory!).mockReturnValue(historyWithPendingCall);
|
||||
|
||||
// Act: Simulate sending the tool's response back
|
||||
const stream = client.sendMessageStream(
|
||||
[
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'some_tool',
|
||||
response: { success: true },
|
||||
},
|
||||
},
|
||||
],
|
||||
new AbortController().signal,
|
||||
'prompt-id-tool-response',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
// consume stream to complete the call
|
||||
}
|
||||
|
||||
// Assert: The IDE context message should NOT have been added to the history.
|
||||
expect(mockChat.addHistory).not.toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
parts: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
text: expect.stringContaining("user's editor context"),
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should add IDE context when no tool call is pending', async () => {
|
||||
// Arrange: History is normal, no pending calls
|
||||
const normalHistory: Content[] = [
|
||||
{ role: 'user', parts: [{ text: 'A normal message.' }] },
|
||||
{ role: 'model', parts: [{ text: 'A normal response.' }] },
|
||||
];
|
||||
vi.mocked(mockChat.getHistory!).mockReturnValue(normalHistory);
|
||||
|
||||
// Act
|
||||
const stream = client.sendMessageStream(
|
||||
[{ text: 'Another normal message' }],
|
||||
new AbortController().signal,
|
||||
'prompt-id-normal',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
// consume stream
|
||||
}
|
||||
|
||||
// Assert: The IDE context message SHOULD have been added.
|
||||
expect(mockChat.addHistory).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
role: 'user',
|
||||
parts: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
text: expect.stringContaining("user's editor context"),
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should send the latest IDE context on the next message after a skipped context', async () => {
|
||||
// --- Step 1: A tool call is pending, context should be skipped ---
|
||||
|
||||
// Arrange: History ends with a functionCall
|
||||
const historyWithPendingCall: Content[] = [
|
||||
{ role: 'user', parts: [{ text: 'Please use a tool.' }] },
|
||||
{
|
||||
role: 'model',
|
||||
parts: [{ functionCall: { name: 'some_tool', args: {} } }],
|
||||
},
|
||||
];
|
||||
vi.mocked(mockChat.getHistory!).mockReturnValue(historyWithPendingCall);
|
||||
|
||||
// Arrange: Set the initial IDE context
|
||||
const initialIdeContext = {
|
||||
workspaceState: {
|
||||
openFiles: [{ path: '/path/to/fileA.ts', timestamp: Date.now() }],
|
||||
},
|
||||
};
|
||||
vi.mocked(ideContext.getIdeContext).mockReturnValue(initialIdeContext);
|
||||
|
||||
// Act: Send the tool response
|
||||
let stream = client.sendMessageStream(
|
||||
[
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'some_tool',
|
||||
response: { success: true },
|
||||
},
|
||||
},
|
||||
],
|
||||
new AbortController().signal,
|
||||
'prompt-id-tool-response',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
|
||||
// Assert: The initial context was NOT sent
|
||||
expect(mockChat.addHistory).not.toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
parts: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
text: expect.stringContaining("user's editor context"),
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
);
|
||||
|
||||
// --- Step 2: A new message is sent, latest context should be included ---
|
||||
|
||||
// Arrange: The model has responded to the tool, and the user is sending a new message.
|
||||
const historyAfterToolResponse: Content[] = [
|
||||
...historyWithPendingCall,
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'some_tool',
|
||||
response: { success: true },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{ role: 'model', parts: [{ text: 'The tool ran successfully.' }] },
|
||||
];
|
||||
vi.mocked(mockChat.getHistory!).mockReturnValue(
|
||||
historyAfterToolResponse,
|
||||
);
|
||||
vi.mocked(mockChat.addHistory!).mockClear(); // Clear previous calls for the next assertion
|
||||
|
||||
// Arrange: The IDE context has now changed
|
||||
const newIdeContext = {
|
||||
workspaceState: {
|
||||
openFiles: [{ path: '/path/to/fileB.ts', timestamp: Date.now() }],
|
||||
},
|
||||
};
|
||||
vi.mocked(ideContext.getIdeContext).mockReturnValue(newIdeContext);
|
||||
|
||||
// Act: Send a new, regular user message
|
||||
stream = client.sendMessageStream(
|
||||
[{ text: 'Thanks!' }],
|
||||
new AbortController().signal,
|
||||
'prompt-id-final',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
|
||||
// Assert: The NEW context was sent as a FULL context because there was no previously sent context.
|
||||
const addHistoryCalls = vi.mocked(mockChat.addHistory!).mock.calls;
|
||||
const contextCall = addHistoryCalls.find((call) =>
|
||||
JSON.stringify(call[0]).includes("user's editor context"),
|
||||
);
|
||||
expect(contextCall).toBeDefined();
|
||||
expect(JSON.stringify(contextCall![0])).toContain(
|
||||
"Here is the user's editor context as a JSON object",
|
||||
);
|
||||
// Check that the sent context is the new one (fileB.ts)
|
||||
expect(JSON.stringify(contextCall![0])).toContain('fileB.ts');
|
||||
// Check that the sent context is NOT the old one (fileA.ts)
|
||||
expect(JSON.stringify(contextCall![0])).not.toContain('fileA.ts');
|
||||
});
|
||||
|
||||
it('should send a context DELTA on the next message after a skipped context', async () => {
|
||||
// --- Step 0: Establish an initial context ---
|
||||
vi.mocked(mockChat.getHistory!).mockReturnValue([]); // Start with empty history
|
||||
const contextA = {
|
||||
workspaceState: {
|
||||
openFiles: [
|
||||
{
|
||||
path: '/path/to/fileA.ts',
|
||||
isActive: true,
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
vi.mocked(ideContext.getIdeContext).mockReturnValue(contextA);
|
||||
|
||||
// Act: Send a regular message to establish the initial context
|
||||
let stream = client.sendMessageStream(
|
||||
[{ text: 'Initial message' }],
|
||||
new AbortController().signal,
|
||||
'prompt-id-initial',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
|
||||
// Assert: Full context for fileA.ts was sent and stored.
|
||||
const initialCall = vi.mocked(mockChat.addHistory!).mock.calls[0][0];
|
||||
expect(JSON.stringify(initialCall)).toContain(
|
||||
"user's editor context as a JSON object",
|
||||
);
|
||||
expect(JSON.stringify(initialCall)).toContain('fileA.ts');
|
||||
// This implicitly tests that `lastSentIdeContext` is now set internally by the client.
|
||||
vi.mocked(mockChat.addHistory!).mockClear();
|
||||
|
||||
// --- Step 1: A tool call is pending, context should be skipped ---
|
||||
const historyWithPendingCall: Content[] = [
|
||||
{ role: 'user', parts: [{ text: 'Please use a tool.' }] },
|
||||
{
|
||||
role: 'model',
|
||||
parts: [{ functionCall: { name: 'some_tool', args: {} } }],
|
||||
},
|
||||
];
|
||||
vi.mocked(mockChat.getHistory!).mockReturnValue(historyWithPendingCall);
|
||||
|
||||
// Arrange: IDE context changes, but this should be skipped
|
||||
const contextB = {
|
||||
workspaceState: {
|
||||
openFiles: [
|
||||
{
|
||||
path: '/path/to/fileB.ts',
|
||||
isActive: true,
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
vi.mocked(ideContext.getIdeContext).mockReturnValue(contextB);
|
||||
|
||||
// Act: Send the tool response
|
||||
stream = client.sendMessageStream(
|
||||
[
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'some_tool',
|
||||
response: { success: true },
|
||||
},
|
||||
},
|
||||
],
|
||||
new AbortController().signal,
|
||||
'prompt-id-tool-response',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
|
||||
// Assert: No context was sent
|
||||
expect(mockChat.addHistory).not.toHaveBeenCalled();
|
||||
|
||||
// --- Step 2: A new message is sent, latest context DELTA should be included ---
|
||||
const historyAfterToolResponse: Content[] = [
|
||||
...historyWithPendingCall,
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'some_tool',
|
||||
response: { success: true },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{ role: 'model', parts: [{ text: 'The tool ran successfully.' }] },
|
||||
];
|
||||
vi.mocked(mockChat.getHistory!).mockReturnValue(
|
||||
historyAfterToolResponse,
|
||||
);
|
||||
|
||||
// Arrange: The IDE context has changed again
|
||||
const contextC = {
|
||||
workspaceState: {
|
||||
openFiles: [
|
||||
// fileA is now closed, fileC is open
|
||||
{
|
||||
path: '/path/to/fileC.ts',
|
||||
isActive: true,
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
vi.mocked(ideContext.getIdeContext).mockReturnValue(contextC);
|
||||
|
||||
// Act: Send a new, regular user message
|
||||
stream = client.sendMessageStream(
|
||||
[{ text: 'Thanks!' }],
|
||||
new AbortController().signal,
|
||||
'prompt-id-final',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
/* consume */
|
||||
}
|
||||
|
||||
// Assert: The DELTA context was sent
|
||||
const finalCall = vi.mocked(mockChat.addHistory!).mock.calls[0][0];
|
||||
expect(JSON.stringify(finalCall)).toContain('summary of changes');
|
||||
// The delta should reflect fileA being closed and fileC being opened.
|
||||
expect(JSON.stringify(finalCall)).toContain('filesClosed');
|
||||
expect(JSON.stringify(finalCall)).toContain('fileA.ts');
|
||||
expect(JSON.stringify(finalCall)).toContain('activeFileChanged');
|
||||
expect(JSON.stringify(finalCall)).toContain('fileC.ts');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('generateContent', () => {
|
||||
|
|
|
|||
|
|
@ -468,9 +468,22 @@ export class GeminiClient {
|
|||
yield { type: GeminiEventType.ChatCompressed, value: compressed };
|
||||
}
|
||||
|
||||
if (this.config.getIdeMode()) {
|
||||
// Prevent context updates from being sent while a tool call is
|
||||
// waiting for a response. The Gemini API requires that a functionResponse
|
||||
// part from the user immediately follows a functionCall part from the model
|
||||
// in the conversation history . The IDE context is not discarded; it will
|
||||
// be included in the next regular message sent to the model.
|
||||
const history = this.getHistory();
|
||||
const lastMessage =
|
||||
history.length > 0 ? history[history.length - 1] : undefined;
|
||||
const hasPendingToolCall =
|
||||
!!lastMessage &&
|
||||
lastMessage.role === 'model' &&
|
||||
(lastMessage.parts?.some((p) => 'functionCall' in p) || false);
|
||||
|
||||
if (this.config.getIdeMode() && !hasPendingToolCall) {
|
||||
const { contextParts, newIdeContext } = this.getIdeContextParts(
|
||||
this.forceFullIdeContext || this.getHistory().length === 0,
|
||||
this.forceFullIdeContext || history.length === 0,
|
||||
);
|
||||
if (contextParts.length > 0) {
|
||||
this.getChat().addHistory({
|
||||
|
|
|
|||
Loading…
Reference in New Issue