tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

feat: Handle MCP server notification messages (#2613)

Browse Source 
Co-authored-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
Jack Amadeo
2025-05-30 11:50:14 -04:00
committed by GitHub
parent eeb61ace22
commit 03e5549b54
40 changed files with 1186 additions and 443 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/desktop/src/components/ChatView.tsx
+13
View File
@@ -148,6 +148,7 @@ function ChatContent({
handleInputChange: _handleInputChange,
handleSubmit: _submitMessage,
updateMessageStreamBody,
notifications,
} = useMessageStream({
api: getApiUrl('/reply'),
initialMessages: chat.messages,
@@ -492,6 +493,16 @@ function ChatContent({
const handleDragOver = (e: React.DragEvent<HTMLDivElement>) => {
e.preventDefault();
};
const toolCallNotifications = notifications.reduce((map, item) => {
const key = item.request_id;
if (!map.has(key)) {
map.set(key, []);
}
map.get(key).push(item);
return map;
}, new Map());
return (
<div className="flex flex-col w-full h-screen items-center justify-center">
{/* Loader when generating recipe */}
@@ -571,6 +582,7 @@ function ChatContent({
const updatedMessages = [...messages, newMessage];
setMessages(updatedMessages);
}}
toolCallNotifications={toolCallNotifications}
/>
)}
</>
@@ -578,6 +590,7 @@ function ChatContent({
</div>
))}
</SearchView>
{error && (
<div className="flex flex-col items-center justify-center p-4">
<div className="text-red-700 dark:text-red-300 bg-red-400/50 p-3 rounded-lg mb-2">
ui/desktop/src/components/GooseMessage.tsx
+4
View File
@@ -17,6 +17,7 @@ import {
} from '../types/message';
import ToolCallConfirmation from './ToolCallConfirmation';
import MessageCopyLink from './MessageCopyLink';
import { NotificationEvent } from '../hooks/useMessageStream';
interface GooseMessageProps {
// messages up to this index are presumed to be "history" from a resumed session, this is used to track older tool confirmation requests
@@ -25,6 +26,7 @@ interface GooseMessageProps {
message: Message;
messages: Message[];
metadata?: string[];
toolCallNotifications: Map<string, NotificationEvent[]>;
append: (value: string) => void;
appendMessage: (message: Message) => void;
}
@@ -34,6 +36,7 @@ export default function GooseMessage({
message,
metadata,
messages,
toolCallNotifications,
append,
appendMessage,
}: GooseMessageProps) {
@@ -158,6 +161,7 @@ export default function GooseMessage({
}
toolRequest={toolRequest}
toolResponse={toolResponsesMap.get(toolRequest.id)}
notifications={toolCallNotifications.get(toolRequest.id)}
/>
</div>
))}
ui/desktop/src/components/ToolCallWithResponse.tsx
+157 -6
View File
@@ -1,4 +1,4 @@
import React from 'react';
import React, { useEffect, useRef } from 'react';
import { Card } from './ui/card';
import { ToolCallArguments, ToolCallArgumentValue } from './ToolCallArguments';
import MarkdownContent from './MarkdownContent';
@@ -6,17 +6,20 @@ import { Content, ToolRequestMessageContent, ToolResponseMessageContent } from '
import { snakeToTitleCase } from '../utils';
import Dot, { LoadingStatus } from './ui/Dot';
import Expand from './ui/Expand';
import { NotificationEvent } from '../hooks/useMessageStream';
interface ToolCallWithResponseProps {
isCancelledMessage: boolean;
toolRequest: ToolRequestMessageContent;
toolResponse?: ToolResponseMessageContent;
notifications?: NotificationEvent[];
}
export default function ToolCallWithResponse({
isCancelledMessage,
toolRequest,
toolResponse,
notifications,
}: ToolCallWithResponseProps) {
const toolCall = toolRequest.toolCall.status === 'success' ? toolRequest.toolCall.value : null;
if (!toolCall) {
@@ -26,7 +29,7 @@ export default function ToolCallWithResponse({
return (
<div className={'w-full text-textSubtle text-sm'}>
<Card className="">
<ToolCallView {...{ isCancelledMessage, toolCall, toolResponse }} />
<ToolCallView {...{ isCancelledMessage, toolCall, toolResponse, notifications }} />
</Card>
</div>
);
@@ -47,8 +50,9 @@ function ToolCallExpandable({
children,
className = '',
}: ToolCallExpandableProps) {
const [isExpanded, setIsExpanded] = React.useState(isStartExpanded);
const toggleExpand = () => setIsExpanded((prev) => !prev);
const [isExpandedState, setIsExpanded] = React.useState<boolean | null>(null);
const isExpanded = isExpandedState === null ? isStartExpanded : isExpandedState;
const toggleExpand = () => setIsExpanded(!isExpanded);
React.useEffect(() => {
if (isForceExpand) setIsExpanded(true);
}, [isForceExpand]);
@@ -71,9 +75,42 @@ interface ToolCallViewProps {
arguments: Record<string, unknown>;
};
toolResponse?: ToolResponseMessageContent;
notifications?: NotificationEvent[];
}
function ToolCallView({ isCancelledMessage, toolCall, toolResponse }: ToolCallViewProps) {
interface Progress {
progress: number;
progressToken: string;
total?: number;
message?: string;
}
const logToString = (logMessage: NotificationEvent) => {
const params = logMessage.message.params;
// Special case for the developer system shell logs
if (
params &&
params.data &&
typeof params.data === 'object' &&
'output' in params.data &&
'stream' in params.data
) {
return `[${params.data.stream}] ${params.data.output}`;
}
return typeof params.data === 'string' ? params.data : JSON.stringify(params.data);
};
const notificationToProgress = (notification: NotificationEvent): Progress =>
notification.message.params as unknown as Progress;
function ToolCallView({
isCancelledMessage,
toolCall,
toolResponse,
notifications,
}: ToolCallViewProps) {
const responseStyle = localStorage.getItem('response_style');
const isExpandToolDetails = (() => {
switch (responseStyle) {
@@ -103,6 +140,29 @@ function ToolCallView({ isCancelledMessage, toolCall, toolResponse }: ToolCallVi
}))
: [];
const logs = notifications
?.filter((notification) => notification.message.method === 'notifications/message')
.map(logToString);
const progress = notifications
?.filter((notification) => notification.message.method === 'notifications/progress')
.map(notificationToProgress)
.reduce((map, item) => {
const key = item.progressToken;
if (!map.has(key)) {
map.set(key, []);
}
map.get(key)!.push(item);
return map;
}, new Map<string, Progress[]>());
const progressEntries = [...(progress?.values() || [])].map(
(entries) => entries.sort((a, b) => b.progress - a.progress)[0]
);
const isRenderingProgress =
loadingStatus === 'loading' && (progressEntries.length > 0 || (logs || []).length > 0);
const isShouldExpand = isExpandToolDetails || toolResults.some((v) => v.isExpandToolResults);
// Function to create a compact representation of arguments
@@ -136,7 +196,7 @@ function ToolCallView({ isCancelledMessage, toolCall, toolResponse }: ToolCallVi
return (
<ToolCallExpandable
isStartExpanded={isShouldExpand}
isStartExpanded={isShouldExpand || isRenderingProgress}
isForceExpand={isShouldExpand}
label={
<>
@@ -156,6 +216,24 @@ function ToolCallView({ isCancelledMessage, toolCall, toolResponse }: ToolCallVi
</div>
)}
{logs && logs.length > 0 && (
<div className="bg-bgStandard mt-1">
<ToolLogsView
logs={logs}
working={toolResults.length === 0}
isStartExpanded={toolResults.length === 0}
/>
</div>
)}
{toolResults.length === 0 &&
progressEntries.length > 0 &&
progressEntries.map((entry, index) => (
<div className="p-2" key={index}>
<ProgressBar progress={entry.progress} total={entry.total} message={entry.message} />
</div>
))}
{/* Tool Output */}
{!isCancelledMessage && (
<>
@@ -234,3 +312,76 @@ function ToolResultView({ result, isStartExpanded }: ToolResultViewProps) {
</ToolCallExpandable>
);
}
function ToolLogsView({
logs,
working,
isStartExpanded,
}: {
logs: string[];
working: boolean;
isStartExpanded?: boolean;
}) {
const boxRef = useRef(null);
// Whenever logs update, jump to the newest entry
useEffect(() => {
if (boxRef.current) {
boxRef.current.scrollTop = boxRef.current.scrollHeight;
}
}, [logs]);
return (
<ToolCallExpandable
label={
<span className="pl-[19px] py-1">
<span>Logs</span>
{working && (
<div className="mx-2 inline-block">
<span
className="inline-block animate-spin rounded-full border-2 border-t-transparent border-current"
style={{ width: 8, height: 8 }}
role="status"
aria-label="Loading spinner"
/>
</div>
)}
</span>
}
isStartExpanded={isStartExpanded}
>
<div
ref={boxRef}
className={`flex flex-col items-start space-y-2 overflow-y-auto ${working ? 'max-h-[4rem]' : 'max-h-[20rem]'} bg-bgApp`}
>
{logs.map((log, i) => (
<span key={i} className="font-mono text-sm text-textSubtle">
{log}
</span>
))}
</div>
</ToolCallExpandable>
);
}
const ProgressBar = ({ progress, total, message }: Omit<Progress, 'progressToken'>) => {
const isDeterminate = typeof total === 'number';
const percent = isDeterminate ? Math.min((progress / total!) * 100, 100) : 0;
return (
<div className="w-full space-y-2">
{message && <div className="text-sm text-gray-700">{message}</div>}
<div className="w-full bg-gray-200 rounded-full h-4 overflow-hidden relative">
{isDeterminate ? (
<div
className="bg-blue-500 h-full transition-all duration-300"
style={{ width: `${percent}%` }}
/>
) : (
<div className="absolute inset-0 animate-indeterminate bg-blue-500" />
)}
</div>
</div>
);
};
ui/desktop/src/hooks/useMessageStream.ts
+28 -1
View File
@@ -6,11 +6,25 @@ import { Message, createUserMessage, hasCompletedToolCalls } from '../types/mess
// Ensure TextDecoder is available in the global scope
const TextDecoder = globalThis.TextDecoder;
type JsonValue = string | number | boolean | null | JsonValue[] | { [key: string]: JsonValue };
export interface NotificationEvent {
type: 'Notification';
request_id: string;
message: {
method: string;
params: {
[key: string]: JsonValue;
};
};
}
// Event types for SSE stream
type MessageEvent =
| { type: 'Message'; message: Message }
| { type: 'Error'; error: string }
| { type: 'Finish'; reason: string };
| { type: 'Finish'; reason: string }
| NotificationEvent;
export interface UseMessageStreamOptions {
/**
@@ -124,6 +138,8 @@ export interface UseMessageStreamHelpers {
/** Modify body (session id and/or work dir mid-stream) **/
updateMessageStreamBody?: (newBody: object) => void;
notifications: NotificationEvent[];
}
/**
@@ -151,6 +167,8 @@ export function useMessageStream({
fallbackData: initialMessages,
});
const [notifications, setNotifications] = useState<NotificationEvent[]>([]);
// expose a way to update the body so we can update the session id when CLE occurs
const updateMessageStreamBody = useCallback((newBody: object) => {
extraMetadataRef.current.body = {
@@ -247,6 +265,14 @@ export function useMessageStream({
break;
}
case 'Notification': {
const newNotification = {
...parsedEvent,
};
setNotifications((prev) => [...prev, newNotification]);
break;
}
case 'Error':
throw new Error(parsedEvent.error);
@@ -516,5 +542,6 @@ export function useMessageStream({
isLoading: isLoading || false,
addToolResult,
updateMessageStreamBody,
notifications,
};
}