Skip to content

Commit e63c03c

Browse files
debsmita1debsmita1
committed
fix(lightspeed): update how tool call info is fetched after the query response update
1 parent 99aa8ff commit e63c03c

4 files changed

Lines changed: 278 additions & 29 deletions

File tree

‎workspaces/lightspeed/.changeset/clean-tomatoes-clean.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@red-hat-developer-hub/backstage-plugin-lightspeed': patch
3+
---
4+
5+
update how tool call info is retrived

‎workspaces/lightspeed/plugins/lightspeed/src/hooks/__tests__/useConversationMessages.test.tsx‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -559,6 +559,143 @@ data: {"event": "token", "data": {"id": 2, "token": ""}}\n
559559
expect(onComplete).toHaveBeenCalledWith('Hi from conversation 1!');
560560
});
561561

562+
it('should parse MCP-style tool_call and tool_result (name/args and content)', async () => {
563+
const toolCallId = 'mcp_list_67d3a067-e262-4bef-b467-6c82f622bd4a';
564+
const mcpStream = createSSEStream([
565+
{
566+
event: 'start',
567+
data: { conversation_id: 'conv-mcp' },
568+
},
569+
{
570+
event: 'tool_call',
571+
data: {
572+
id: toolCallId,
573+
name: 'mcp_list_tools',
574+
args: { server_label: 'mcp-integration-tools' },
575+
type: 'mcp_list_tools',
576+
},
577+
},
578+
{
579+
event: 'tool_result',
580+
data: {
581+
id: toolCallId,
582+
status: 'success',
583+
content: '{"server_label":"mcp-integration-tools","tools":[]}',
584+
},
585+
},
586+
{
587+
event: 'token',
588+
data: { id: 0, token: 'Done.' },
589+
},
590+
]);
591+
592+
const mockApi = {
593+
createMessage: jest.fn().mockResolvedValue({
594+
read: jest
595+
.fn()
596+
.mockResolvedValueOnce({
597+
done: false,
598+
value: new TextEncoder().encode(mcpStream),
599+
})
600+
.mockResolvedValueOnce({ done: true, value: null }),
601+
}),
602+
};
603+
(useApi as jest.Mock).mockReturnValue(mockApi);
604+
605+
const { result } = renderHook(
606+
() =>
607+
useConversationMessages(
608+
'conv-mcp',
609+
'test-user',
610+
'gpt-4',
611+
'openai',
612+
'user.png',
613+
),
614+
{ wrapper },
615+
);
616+
617+
await act(async () => {
618+
await result.current.handleInputPrompt('List MCP tools');
619+
});
620+
621+
await waitFor(() => {
622+
const msgs = result.current.conversations['conv-mcp'];
623+
const bot = msgs?.[1];
624+
expect(bot?.toolCalls).toHaveLength(1);
625+
expect(bot?.toolCalls?.[0]).toEqual(
626+
expect.objectContaining({
627+
id: toolCallId,
628+
toolName: 'mcp_list_tools',
629+
arguments: { server_label: 'mcp-integration-tools' },
630+
response: '{"server_label":"mcp-integration-tools","tools":[]}',
631+
isLoading: false,
632+
}),
633+
);
634+
expect(bot?.content).toContain('Done.');
635+
});
636+
});
637+
638+
it('should complete legacy tool_result when response field is omitted', async () => {
639+
const legacyStream = createSSEStream([
640+
{
641+
event: 'start',
642+
data: { conversation_id: 'conv-legacy-res' },
643+
},
644+
{
645+
event: 'tool_call',
646+
data: {
647+
id: 1,
648+
token: { tool_name: 'fetch-techdocs', arguments: { owner: 'a' } },
649+
},
650+
},
651+
{
652+
event: 'tool_result',
653+
data: { id: 1, token: { tool_name: 'fetch-techdocs' } },
654+
},
655+
]);
656+
657+
const mockApi = {
658+
createMessage: jest.fn().mockResolvedValue({
659+
read: jest
660+
.fn()
661+
.mockResolvedValueOnce({
662+
done: false,
663+
value: new TextEncoder().encode(legacyStream),
664+
})
665+
.mockResolvedValueOnce({ done: true, value: null }),
666+
}),
667+
};
668+
(useApi as jest.Mock).mockReturnValue(mockApi);
669+
670+
const { result } = renderHook(
671+
() =>
672+
useConversationMessages(
673+
'conv-legacy-res',
674+
'test-user',
675+
'gpt-4',
676+
'openai',
677+
'user.png',
678+
),
679+
{ wrapper },
680+
);
681+
682+
await act(async () => {
683+
await result.current.handleInputPrompt('techdocs');
684+
});
685+
686+
await waitFor(() => {
687+
const msgs = result.current.conversations['conv-legacy-res'];
688+
const bot = msgs?.[1];
689+
expect(bot?.toolCalls?.[0]).toEqual(
690+
expect.objectContaining({
691+
toolName: 'fetch-techdocs',
692+
response: '',
693+
isLoading: false,
694+
}),
695+
);
696+
});
697+
});
698+
562699
it('should resume streaming for the first conversation after switching back and complete', async () => {
563700
const onComplete = jest.fn();
564701

‎workspaces/lightspeed/plugins/lightspeed/src/hooks/useConversationMessages.ts‎

Lines changed: 125 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,44 @@ import {
4242
} from '../utils/lightspeed-chatbox-utils';
4343
import { useCreateConversationMessage } from './useCreateCoversationMessage';
4444

45+
const toolCallIdKey = (id: string | number): string => {
46+
return String(id);
47+
};
48+
49+
const isMcpStyleToolCallPayload = (
50+
data: Record<string, any> | undefined,
51+
): boolean => {
52+
return (
53+
!!data &&
54+
typeof data.name === 'string' &&
55+
data.name.trim().length > 0 &&
56+
data.id !== null
57+
);
58+
};
59+
60+
/** Legacy tool_result uses data.token with at least tool_name and response. */
61+
const isLegacyToolResultToken = (
62+
token: unknown,
63+
): token is { tool_name: string; response?: unknown } => {
64+
return (
65+
!!token &&
66+
typeof token === 'object' &&
67+
!Array.isArray(token) &&
68+
typeof (token as { tool_name?: string }).tool_name === 'string' &&
69+
(token as { tool_name: string }).tool_name.length > 0
70+
);
71+
};
72+
73+
const legacyToolResultToString = (response: unknown): string => {
74+
if (!response) return '';
75+
if (typeof response === 'string') return response;
76+
try {
77+
return JSON.stringify(response);
78+
} catch {
79+
return String(response);
80+
}
81+
};
82+
4583
// Fetch all conversation messages
4684
export const useFetchConversationMessages = (
4785
currentConversation: string,
@@ -117,7 +155,7 @@ export const useConversationMessages = (
117155
});
118156

119157
// Track pending tool calls during streaming
120-
const pendingToolCalls = useRef<{ [id: number]: ToolCall }>({});
158+
const pendingToolCalls = useRef<Record<string, ToolCall>>({});
121159

122160
// Cache tool calls by conversation ID and message index to persist across refetches
123161
// Key format: `${conversationId}-${messageIndex}`
@@ -298,19 +336,49 @@ export const useConversationMessages = (
298336
// Handle tool_call event
299337
if (event === 'tool_call') {
300338
const toolCallData = data?.token;
301-
if (
339+
const legacyObjectCall =
302340
typeof toolCallData === 'object' &&
303-
toolCallData?.tool_name
304-
) {
305-
// Full tool call with arguments - track start time
306-
const toolCall: ToolCall = {
341+
toolCallData !== null &&
342+
!Array.isArray(toolCallData) &&
343+
(toolCallData as { tool_name?: string }).tool_name;
344+
345+
const mcpStyle = isMcpStyleToolCallPayload(data);
346+
const rawArgs = data?.args ?? data?.arguments;
347+
const mcpArgs: Record<string, any> =
348+
rawArgs &&
349+
typeof rawArgs === 'object' &&
350+
!Array.isArray(rawArgs)
351+
? rawArgs
352+
: {};
353+
354+
let toolCall: ToolCall | undefined;
355+
// Prefer legacy token object when present (backward compatible)
356+
if (legacyObjectCall && data.id !== null) {
357+
toolCall = {
307358
id: data.id,
308-
toolName: toolCallData.tool_name,
309-
arguments: toolCallData.arguments || {},
359+
toolName: (toolCallData as { tool_name: string }).tool_name,
360+
arguments:
361+
(toolCallData as { arguments?: Record<string, any> })
362+
.arguments || {},
310363
startTime: Date.now(),
311364
isLoading: true,
312365
};
313-
pendingToolCalls.current[data.id] = toolCall;
366+
} else if (mcpStyle) {
367+
toolCall = {
368+
id: data.id,
369+
toolName: data.name.trim(),
370+
description:
371+
typeof data.type === 'string' && data.type !== data.name
372+
? data.type
373+
: undefined,
374+
arguments: mcpArgs,
375+
startTime: Date.now(),
376+
isLoading: true,
377+
};
378+
}
379+
380+
if (toolCall && data.id !== null) {
381+
pendingToolCalls.current[toolCallIdKey(data.id)] = toolCall;
314382

315383
// Update the bot message with the pending tool call
316384
setConversations(prevConversations => {
@@ -358,10 +426,41 @@ export const useConversationMessages = (
358426

359427
// Handle tool_result event
360428
if (event === 'tool_result') {
361-
const resultData = data?.token;
362-
if (resultData?.tool_name) {
363-
const toolId = data.id;
364-
const pendingCall = pendingToolCalls.current[toolId];
429+
const tokenResult = data?.token;
430+
const legacyResult = isLegacyToolResultToken(tokenResult);
431+
432+
const mcpHasContent =
433+
data?.id !== null &&
434+
data.content !== undefined &&
435+
!legacyResult;
436+
437+
let responsePayload: string | undefined;
438+
let matchToolName: string | undefined;
439+
let toolIdKey: string | undefined;
440+
441+
if (legacyResult) {
442+
responsePayload = legacyToolResultToString(
443+
tokenResult.response,
444+
);
445+
matchToolName = tokenResult.tool_name;
446+
toolIdKey =
447+
data?.id !== null ? toolCallIdKey(data.id) : undefined;
448+
} else if (mcpHasContent) {
449+
toolIdKey = toolCallIdKey(data.id);
450+
responsePayload =
451+
typeof data.content === 'string'
452+
? data.content
453+
: JSON.stringify(data.content);
454+
if (
455+
typeof data.status === 'string' &&
456+
data.status !== 'success'
457+
) {
458+
responsePayload = `[${data.status}] ${responsePayload}`;
459+
}
460+
}
461+
462+
if (responsePayload !== undefined && toolIdKey !== undefined) {
463+
const pendingCall = pendingToolCalls.current[toolIdKey];
365464
const endTime = Date.now();
366465
const executionTime = pendingCall
367466
? (endTime - pendingCall.startTime) / 1000
@@ -380,13 +479,14 @@ export const useConversationMessages = (
380479

381480
// Find and update the matching tool call
382481
const updatedToolCalls = toolCalls.map(tc => {
383-
if (
384-
tc.id === toolId ||
385-
tc.toolName === resultData.tool_name
386-
) {
482+
const idMatches =
483+
toolCallIdKey(tc.id) === toolIdKey ||
484+
(matchToolName !== undefined &&
485+
tc.toolName === matchToolName);
486+
if (idMatches) {
387487
return {
388488
...tc,
389-
response: resultData.response,
489+
response: responsePayload,
390490
endTime,
391491
executionTime,
392492
isLoading: false,
@@ -419,13 +519,14 @@ export const useConversationMessages = (
419519
if (aiMessage) {
420520
const toolCalls = aiMessage.toolCalls || [];
421521
const updatedToolCalls = toolCalls.map(tc => {
422-
if (
423-
tc.id === toolId ||
424-
tc.toolName === resultData.tool_name
425-
) {
522+
const idMatches =
523+
toolCallIdKey(tc.id) === toolIdKey ||
524+
(matchToolName !== undefined &&
525+
tc.toolName === matchToolName);
526+
if (idMatches) {
426527
return {
427528
...tc,
428-
response: resultData.response,
529+
response: responsePayload,
429530
endTime,
430531
executionTime,
431532
isLoading: false,
@@ -440,7 +541,7 @@ export const useConversationMessages = (
440541
}
441542

442543
// Clean up pending tool call
443-
delete pendingToolCalls.current[toolId];
544+
delete pendingToolCalls.current[toolIdKey];
444545
}
445546
}
446547

‎workspaces/lightspeed/plugins/lightspeed/src/types.ts‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,7 @@ export type CaptureFeedback = {
168168

169169
// Tool Calling Types
170170
export interface ToolCall {
171-
id: number;
171+
id: number | string;
172172
toolName: string;
173173
description?: string;
174174
arguments: Record<string, any>;
@@ -180,11 +180,17 @@ export interface ToolCall {
180180
}
181181

182182
export interface ToolCallEvent {
183-
id: number;
184-
token: string | { tool_name: string; arguments: Record<string, any> };
183+
id: number | string;
184+
token?: string | { tool_name: string; arguments: Record<string, any> };
185+
name?: string;
186+
args?: Record<string, any>;
187+
arguments?: Record<string, any>;
188+
type?: string;
185189
}
186190

187191
export interface ToolResultEvent {
188-
id: number;
189-
token: { tool_name: string; response: string };
192+
id: number | string;
193+
token?: { tool_name: string; response: string };
194+
status?: string;
195+
content?: string;
190196
}

0 commit comments

Comments
 (0)