From dda985150d0b0f6ef48be378de0c8cce6349a2a6 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Tue, 24 Mar 2026 16:29:36 -0400 Subject: [PATCH] feat: add MCP server config ID to tool-call message parts (#23522) --- coderd/x/chatd/chatd.go | 28 +++++- coderd/x/chatd/mcpclient/mcpclient.go | 17 +++- coderd/x/chatd/mcpclient/mcpclient_test.go | 85 +++++++++++++++++++ codersdk/chats.go | 39 ++++----- site/src/api/typesGenerated.ts | 2 + site/src/components/ai-elements/tool/Tool.tsx | 31 ++++++- .../components/ai-elements/tool/ToolIcon.tsx | 33 ++++++- .../components/ai-elements/tool/ToolLabel.tsx | 13 ++- .../components/ai-elements/tool/utils.test.ts | 31 +++++++ site/src/components/ai-elements/tool/utils.ts | 20 +++++ .../AgentDetail/ConversationTimeline.tsx | 20 +++++ .../components/AgentDetail/messageParsing.ts | 4 + .../components/AgentDetail/streamState.ts | 6 ++ .../components/AgentDetail/types.ts | 5 ++ .../components/AgentDetailContent.tsx | 6 ++ .../AgentsPage/components/AgentDetailView.tsx | 1 + 16 files changed, 310 insertions(+), 31 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 554bd71256..b036b3a3f1 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -3017,6 +3017,16 @@ func (p *Server) runChat( defer mcpCleanup() } + // Build a lookup from tool name to MCP server config ID + // so we can annotate persisted parts with the originating + // server. + toolNameToConfigID := make(map[string]uuid.UUID) + for _, t := range mcpTools { + if mcp, ok := t.(mcpclient.MCPToolIdentifier); ok { + toolNameToConfigID[t.Info().Name] = mcp.MCPServerConfigID() + } + } + if instruction != "" { prompt = chatprompt.InsertSystem(prompt, instruction) } @@ -3079,7 +3089,13 @@ func (p *Server) runChat( if len(assistantBlocks) > 0 { sdkParts := make([]codersdk.ChatMessagePart, 0, len(assistantBlocks)) for _, block := range assistantBlocks { - sdkParts = append(sdkParts, chatprompt.PartFromContent(block)) + part := chatprompt.PartFromContent(block) + if part.ToolName != "" { + if configID, ok := toolNameToConfigID[part.ToolName]; ok { + part.MCPServerConfigID = uuid.NullUUID{UUID: configID, Valid: true} + } + } + sdkParts = append(sdkParts, part) } finalAssistantText = strings.TrimSpace(contentBlocksToText(sdkParts)) var marshalErr error @@ -3092,6 +3108,11 @@ func (p *Server) runChat( toolResultContents := make([]pqtype.NullRawMessage, len(toolResults)) for i, tr := range toolResults { trPart := chatprompt.PartFromContent(tr) + if trPart.ToolName != "" { + if configID, ok := toolNameToConfigID[trPart.ToolName]; ok { + trPart.MCPServerConfigID = uuid.NullUUID{UUID: configID, Valid: true} + } + } var marshalErr error toolResultContents[i], marshalErr = chatprompt.MarshalParts([]codersdk.ChatMessagePart{trPart}) if marshalErr != nil { @@ -3463,6 +3484,11 @@ func (p *Server) runChat( role codersdk.ChatMessageRole, part codersdk.ChatMessagePart, ) { + if part.ToolName != "" { + if configID, ok := toolNameToConfigID[part.ToolName]; ok { + part.MCPServerConfigID = uuid.NullUUID{UUID: configID, Valid: true} + } + } p.publishMessagePart(chat.ID, role, part) }, Compaction: compactionOptions, diff --git a/coderd/x/chatd/mcpclient/mcpclient.go b/coderd/x/chatd/mcpclient/mcpclient.go index 006f0c749d..481f46c604 100644 --- a/coderd/x/chatd/mcpclient/mcpclient.go +++ b/coderd/x/chatd/mcpclient/mcpclient.go @@ -195,7 +195,7 @@ func connectOne( } tools = append( - tools, newMCPTool(cfg.Slug, mcpTool, mcpClient), + tools, newMCPTool(cfg.ID, cfg.Slug, mcpTool, mcpClient), ) } @@ -383,10 +383,17 @@ func redactErrorURL(err error) string { return err.Error() } +// MCPToolIdentifier is implemented by tools that originate from +// an MCP server config and can report the config's database ID. +type MCPToolIdentifier interface { + MCPServerConfigID() uuid.UUID +} + // mcpToolWrapper adapts a single MCP tool into a // fantasy.AgentTool. It stores the prefixed name for Info() but // strips the prefix when forwarding calls to the remote server. type mcpToolWrapper struct { + configID uuid.UUID prefixedName string originalName string description string @@ -396,14 +403,22 @@ type mcpToolWrapper struct { providerOptions fantasy.ProviderOptions } +// MCPServerConfigID returns the database ID of the MCP server +// config that this tool originates from. +func (t *mcpToolWrapper) MCPServerConfigID() uuid.UUID { + return t.configID +} + // newMCPTool creates an mcpToolWrapper from an mcp.Tool // discovered on a remote server. func newMCPTool( + configID uuid.UUID, serverSlug string, tool mcp.Tool, mcpClient *client.Client, ) *mcpToolWrapper { return &mcpToolWrapper{ + configID: configID, prefixedName: serverSlug + toolNameSep + tool.Name, originalName: tool.Name, description: tool.Description, diff --git a/coderd/x/chatd/mcpclient/mcpclient_test.go b/coderd/x/chatd/mcpclient/mcpclient_test.go index 172f1e4a34..36c3da3cc2 100644 --- a/coderd/x/chatd/mcpclient/mcpclient_test.go +++ b/coderd/x/chatd/mcpclient/mcpclient_test.go @@ -621,6 +621,91 @@ func TestConnectAll_EmptyAccessToken(t *testing.T) { require.NotEmpty(t, tools) } +// TestConnectAll_MCPToolIdentifier verifies that tools returned +// by ConnectAll implement the MCPToolIdentifier interface and +// report the correct server config ID. +func TestConnectAll_MCPToolIdentifier(t *testing.T) { + t.Parallel() + ctx := context.Background() + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + ts := newTestMCPServer(t, echoTool()) + + configID := uuid.New() + cfg := database.MCPServerConfig{ + ID: configID, + Slug: "id-srv", + DisplayName: "ID Server", + Url: ts.URL, + Transport: "streamable_http", + AuthType: "none", + Enabled: true, + } + + tools, cleanup := mcpclient.ConnectAll(ctx, logger, []database.MCPServerConfig{cfg}, nil) + t.Cleanup(cleanup) + + require.Len(t, tools, 1) + + // Assert the tool implements MCPToolIdentifier. + identifier, ok := tools[0].(mcpclient.MCPToolIdentifier) + require.True(t, ok, "tool should implement MCPToolIdentifier") + assert.Equal(t, configID, identifier.MCPServerConfigID()) +} + +// TestConnectAll_MCPToolIdentifier_MultipleServers verifies that +// each tool from a different MCP server carries its own config ID. +func TestConnectAll_MCPToolIdentifier_MultipleServers(t *testing.T) { + t.Parallel() + ctx := context.Background() + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + ts1 := newTestMCPServer(t, echoTool()) + ts2 := newTestMCPServer(t, greetTool()) + + configID1 := uuid.New() + configID2 := uuid.New() + cfg1 := database.MCPServerConfig{ + ID: configID1, + Slug: "srv-a", + DisplayName: "Server A", + Url: ts1.URL, + Transport: "streamable_http", + AuthType: "none", + Enabled: true, + } + cfg2 := database.MCPServerConfig{ + ID: configID2, + Slug: "srv-b", + DisplayName: "Server B", + Url: ts2.URL, + Transport: "streamable_http", + AuthType: "none", + Enabled: true, + } + + tools, cleanup := mcpclient.ConnectAll( + ctx, logger, + []database.MCPServerConfig{cfg1, cfg2}, + nil, + ) + t.Cleanup(cleanup) + + require.Len(t, tools, 2) + + // Map tool name to config ID via the MCPToolIdentifier + // interface. + idByName := make(map[string]uuid.UUID) + for _, tool := range tools { + identifier, ok := tool.(mcpclient.MCPToolIdentifier) + require.True(t, ok, "tool %q should implement MCPToolIdentifier", tool.Info().Name) + idByName[tool.Info().Name] = identifier.MCPServerConfigID() + } + + assert.Equal(t, configID1, idByName["srv-a__echo"]) + assert.Equal(t, configID2, idByName["srv-b__greet"]) +} + func TestConnectAll_CallToolError(t *testing.T) { t.Parallel() ctx := context.Background() diff --git a/codersdk/chats.go b/codersdk/chats.go index fe491a8c35..ecc26d6045 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -146,25 +146,26 @@ func AllChatMessagePartTypes() []ChatMessagePartType { // the frontend does not expect adds noise to the wire format // and wastes space in persisted chat_messages rows. type ChatMessagePart struct { - Type ChatMessagePartType `json:"type"` - Text string `json:"text" variants:"text,reasoning"` - Signature string `json:"signature,omitempty"` - ToolCallID string `json:"tool_call_id,omitempty" variants:"tool-call?,tool-result?"` - ToolName string `json:"tool_name,omitempty" variants:"tool-call?,tool-result?"` - Args json.RawMessage `json:"args,omitempty" variants:"tool-call?"` - ArgsDelta string `json:"args_delta,omitempty" variants:"tool-call?"` - Result json.RawMessage `json:"result,omitempty" variants:"tool-result?"` - ResultDelta string `json:"result_delta,omitempty"` - IsError bool `json:"is_error,omitempty" variants:"tool-result?"` - SourceID string `json:"source_id,omitempty" variants:"source?"` - URL string `json:"url" variants:"source"` - Title string `json:"title,omitempty" variants:"source?"` - MediaType string `json:"media_type" variants:"file"` - Data []byte `json:"data,omitempty" variants:"file?"` - FileID uuid.NullUUID `json:"file_id,omitempty" format:"uuid" variants:"file?"` - FileName string `json:"file_name" variants:"file-reference"` - StartLine int `json:"start_line" variants:"file-reference"` - EndLine int `json:"end_line" variants:"file-reference"` + Type ChatMessagePartType `json:"type"` + Text string `json:"text" variants:"text,reasoning"` + Signature string `json:"signature,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty" variants:"tool-call?,tool-result?"` + ToolName string `json:"tool_name,omitempty" variants:"tool-call?,tool-result?"` + MCPServerConfigID uuid.NullUUID `json:"mcp_server_config_id,omitempty" format:"uuid" variants:"tool-call?,tool-result?"` + Args json.RawMessage `json:"args,omitempty" variants:"tool-call?"` + ArgsDelta string `json:"args_delta,omitempty" variants:"tool-call?"` + Result json.RawMessage `json:"result,omitempty" variants:"tool-result?"` + ResultDelta string `json:"result_delta,omitempty"` + IsError bool `json:"is_error,omitempty" variants:"tool-result?"` + SourceID string `json:"source_id,omitempty" variants:"source?"` + URL string `json:"url" variants:"source"` + Title string `json:"title,omitempty" variants:"source?"` + MediaType string `json:"media_type" variants:"file"` + Data []byte `json:"data,omitempty" variants:"file?"` + FileID uuid.NullUUID `json:"file_id,omitempty" format:"uuid" variants:"file?"` + FileName string `json:"file_name" variants:"file-reference"` + StartLine int `json:"start_line" variants:"file-reference"` + EndLine int `json:"end_line" variants:"file-reference"` // The code content from the diff that was commented on. Content string `json:"content" variants:"file-reference"` // ProviderMetadata holds provider-specific response metadata diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 165f91b146..83b08eb51a 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -1860,6 +1860,7 @@ export interface ChatToolCallPart { readonly type: "tool-call"; readonly tool_call_id?: string; readonly tool_name?: string; + readonly mcp_server_config_id?: string; readonly args?: Record; readonly args_delta?: string; /** @@ -1874,6 +1875,7 @@ export interface ChatToolResultPart { readonly type: "tool-result"; readonly tool_call_id?: string; readonly tool_name?: string; + readonly mcp_server_config_id?: string; readonly result?: Record; readonly is_error?: boolean; /** diff --git a/site/src/components/ai-elements/tool/Tool.tsx b/site/src/components/ai-elements/tool/Tool.tsx index 9ac001842e..a9eac9f18c 100644 --- a/site/src/components/ai-elements/tool/Tool.tsx +++ b/site/src/components/ai-elements/tool/Tool.tsx @@ -1,5 +1,6 @@ import { useTheme } from "@emotion/react"; import { FileDiff, File as FileViewer } from "@pierre/diffs/react"; +import type * as TypesGen from "api/typesGenerated"; import { ScrollArea } from "components/ScrollArea/ScrollArea"; import { type ComponentPropsWithRef, type FC, memo } from "react"; import { cn } from "utils/cn"; @@ -52,6 +53,10 @@ interface ToolProps extends Omit, "children"> { subagentTitles?: Map; /** Maps sub-agent chat IDs to real-time status updates from stream events. */ subagentStatusOverrides?: Map; + /** MCP server config ID associated with this tool call. */ + mcpServerConfigId?: string; + /** Available MCP server configs for icon/name lookup. */ + mcpServers?: readonly TypesGen.MCPServerConfig[]; } // Props passed to each tool-specific renderer function. Each renderer @@ -64,6 +69,8 @@ type ToolRendererProps = { isError: boolean; subagentTitles?: Map; subagentStatusOverrides?: Map; + mcpServerConfigId?: string; + mcpServers?: readonly TypesGen.MCPServerConfig[]; }; // --------------------------------------------------------------------------- @@ -451,6 +458,8 @@ const GenericToolRenderer: FC = ({ args, result, isError, + mcpServerConfigId, + mcpServers, }) => { const theme = useTheme(); const isDark = theme.palette.mode === "dark"; @@ -466,11 +475,25 @@ const GenericToolRenderer: FC = ({ } : fileViewerOpts; + // Look up MCP server config for icon and slug. + const mcpServer = mcpServerConfigId + ? mcpServers?.find((s) => s.id === mcpServerConfigId) + : undefined; + return ( <>
- - + +
{writeFileDiff ? ( { @@ -583,6 +608,8 @@ export const Tool = memo( isError={isError} subagentTitles={subagentTitles} subagentStatusOverrides={subagentStatusOverrides} + mcpServerConfigId={mcpServerConfigId} + mcpServers={mcpServers} /> ); diff --git a/site/src/components/ai-elements/tool/ToolIcon.tsx b/site/src/components/ai-elements/tool/ToolIcon.tsx index 571c5599ed..90005e9a36 100644 --- a/site/src/components/ai-elements/tool/ToolIcon.tsx +++ b/site/src/components/ai-elements/tool/ToolIcon.tsx @@ -1,3 +1,4 @@ +import { ExternalImage } from "components/ExternalImage/ExternalImage"; import { BotIcon, ClipboardListIcon, @@ -9,14 +10,38 @@ import { WrenchIcon, } from "lucide-react"; import type React from "react"; +import { useState } from "react"; import { cn } from "utils/cn"; -export const ToolIcon: React.FC<{ name: string; isError: boolean }> = ({ - name, - isError, -}) => { +export const ToolIcon: React.FC<{ + name: string; + isError: boolean; + iconUrl?: string; +}> = ({ name, isError, iconUrl }) => { + const [imgError, setImgError] = useState(false); const color = isError ? "text-content-destructive" : "text-content-secondary"; const base = cn("h-4 w-4 shrink-0", color); + + // If an MCP icon URL is provided and hasn't failed, render it. + if (iconUrl && !imgError) { + return ( +
+ setImgError(true)} + /> +
+ ); + } + switch (name) { case "execute": case "process_output": diff --git a/site/src/components/ai-elements/tool/ToolLabel.tsx b/site/src/components/ai-elements/tool/ToolLabel.tsx index 0fbe03d7fe..a6ee73908a 100644 --- a/site/src/components/ai-elements/tool/ToolLabel.tsx +++ b/site/src/components/ai-elements/tool/ToolLabel.tsx @@ -1,11 +1,12 @@ import type React from "react"; -import { asRecord, asString, parseArgs } from "./utils"; +import { asRecord, asString, humanizeMCPToolName, parseArgs } from "./utils"; export const ToolLabel: React.FC<{ name: string; args: unknown; result: unknown; -}> = ({ name, args, result }) => { + mcpSlug?: string; +}> = ({ name, args, result, mcpSlug }) => { const parsed = parseArgs(args); const parsedResult = asRecord(result); @@ -169,9 +170,13 @@ export const ToolLabel: React.FC<{ ); } - default: + default: { + const displayName = mcpSlug ? humanizeMCPToolName(mcpSlug, name) : name; return ( - {name} + + {displayName} + ); + } } }; diff --git a/site/src/components/ai-elements/tool/utils.test.ts b/site/src/components/ai-elements/tool/utils.test.ts index a8525f1c49..111479b752 100644 --- a/site/src/components/ai-elements/tool/utils.test.ts +++ b/site/src/components/ai-elements/tool/utils.test.ts @@ -15,6 +15,7 @@ import { getFileViewerOptionsMinimal, getFileViewerOptionsNoHeader, getWriteFileDiff, + humanizeMCPToolName, isSubagentRunningStatus, isSubagentSuccessStatus, mapSubagentStatusToToolStatus, @@ -667,6 +668,36 @@ describe("stripSvnIndexHeaders", () => { }); }); +describe("humanizeMCPToolName", () => { + it("strips slug prefix and humanizes", () => { + expect(humanizeMCPToolName("linear", "linear__list_issues")).toBe( + "List issues", + ); + }); + + it("handles single-word tool name after prefix", () => { + expect(humanizeMCPToolName("github", "github__search")).toBe("Search"); + }); + + it("humanizes entire name when prefix does not match", () => { + expect(humanizeMCPToolName("linear", "github__list_repos")).toBe( + "Github list repos", + ); + }); + + it("falls back to prefixedName when stripping prefix leaves empty string", () => { + expect(humanizeMCPToolName("linear", "linear__")).toBe("linear__"); + }); + + it("collapses consecutive underscores into a single space", () => { + expect(humanizeMCPToolName("srv", "srv__get__data")).toBe("Get data"); + }); + + it("humanizes tool name without slug prefix", () => { + expect(humanizeMCPToolName("linear", "list_issues")).toBe("List issues"); + }); +}); + describe("constants", () => { it("COLLAPSED_OUTPUT_HEIGHT is 54", () => { expect(COLLAPSED_OUTPUT_HEIGHT).toBe(54); diff --git a/site/src/components/ai-elements/tool/utils.ts b/site/src/components/ai-elements/tool/utils.ts index b11901ccac..6bf59c1fb9 100644 --- a/site/src/components/ai-elements/tool/utils.ts +++ b/site/src/components/ai-elements/tool/utils.ts @@ -555,6 +555,26 @@ export const buildEditDiff = ( return parsed[0].files[0]; }; +/** + * Converts an MCP-prefixed tool name into a human-readable label. + * E.g. "linear__list_issues" with slug "linear" → "List issues" + */ +export function humanizeMCPToolName( + slug: string, + prefixedName: string, +): string { + const prefix = `${slug}__`; + const raw = prefixedName.startsWith(prefix) + ? prefixedName.slice(prefix.length) + : prefixedName; + // Replace runs of underscores with a single space, then trim. + const words = raw.replace(/_+/g, " ").trim(); + if (!words) { + return prefixedName; + } + return words.charAt(0).toUpperCase() + words.slice(1); +} + // Re-export runtime type utils used by sub-components so they // can import from a single location. export { asNumber, asRecord, asString } from "../runtimeTypeUtils"; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.tsx b/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.tsx index daefb8a4ea..3f14c5539f 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.tsx @@ -97,6 +97,7 @@ type RenderBlockListParams = { isStreaming?: boolean; subagentTitles?: Map; subagentStatusOverrides?: Map; + mcpServers?: readonly TypesGen.MCPServerConfig[]; onImageClick?: (src: string) => void; onTextFileClick?: (content: string) => void; urlTransform?: UrlTransform; @@ -265,6 +266,7 @@ function renderBlockList({ isStreaming = false, subagentTitles, subagentStatusOverrides, + mcpServers, onImageClick, onTextFileClick, urlTransform, @@ -329,6 +331,7 @@ function renderBlockList({ isError={false} subagentTitles={subagentTitles} subagentStatusOverrides={subagentStatusOverrides} + mcpServers={mcpServers} /> ); } @@ -345,6 +348,8 @@ function renderBlockList({ subagentStatusOverrides={ isStreaming ? subagentStatusOverrides : undefined } + mcpServerConfigId={tool.mcpServerConfigId} + mcpServers={mcpServers} /> ); } @@ -386,6 +391,7 @@ interface ChatMessageItemProps { // overlay to indicate truncated content. fadeFromBottom?: boolean; urlTransform?: UrlTransform; + mcpServers?: readonly TypesGen.MCPServerConfig[]; } const ChatMessageItem: FC = ({ @@ -397,6 +403,7 @@ const ChatMessageItem: FC = ({ isAfterEditingMessage = false, fadeFromBottom = false, urlTransform, + mcpServers, }) => { const isUser = message.role === "user"; const isSavingMessage = savingMessageId === message.id; @@ -462,6 +469,7 @@ const ChatMessageItem: FC = ({ onImageClick: setPreviewImage, onTextFileClick: (content) => setPreviewText(content), urlTransform, + mcpServers, }); const remainingTools = parsed.tools.filter( (tool) => !renderedToolIDs.has(tool.id), @@ -599,6 +607,8 @@ const ChatMessageItem: FC = ({ result={tool.result} status={tool.status} isError={tool.isError} + mcpServerConfigId={tool.mcpServerConfigId} + mcpServers={mcpServers} /> ))} {!hasRenderableContent && ( @@ -635,6 +645,7 @@ export const StreamingOutput: FC<{ showInitialPlaceholder?: boolean; retryState?: { attempt: number; error: string } | null; urlTransform?: UrlTransform; + mcpServers?: readonly TypesGen.MCPServerConfig[]; }> = ({ streamState, streamTools, @@ -643,6 +654,7 @@ export const StreamingOutput: FC<{ showInitialPlaceholder = false, retryState, urlTransform, + mcpServers, }) => { const conversationItemProps = { role: "assistant" as const }; const toolByID = new Map(streamTools.map((tool) => [tool.id, tool])); @@ -655,6 +667,7 @@ export const StreamingOutput: FC<{ subagentTitles, subagentStatusOverrides, urlTransform, + mcpServers, }); const remainingTools = streamTools.filter( (tool) => !renderedToolIDs.has(tool.id), @@ -696,6 +709,8 @@ export const StreamingOutput: FC<{ isError={tool.isError} subagentTitles={subagentTitles} subagentStatusOverrides={subagentStatusOverrides} + mcpServerConfigId={tool.mcpServerConfigId} + mcpServers={mcpServers} /> ))} @@ -989,6 +1004,7 @@ interface ConversationTimelineProps { editingMessageId?: number | null; savingMessageId?: number | null; urlTransform?: UrlTransform; + mcpServers?: readonly TypesGen.MCPServerConfig[]; } export const ConversationTimeline: FC = ({ @@ -1006,6 +1022,7 @@ export const ConversationTimeline: FC = ({ editingMessageId, savingMessageId, urlTransform, + mcpServers, }) => { const shouldRenderStreamAfterMessages = hasStreamOutput && parsedMessages.length > 0; @@ -1054,6 +1071,7 @@ export const ConversationTimeline: FC = ({ savingMessageId={savingMessageId} urlTransform={urlTransform} isAfterEditingMessage={afterEditingMessageIds.has(message.id)} + mcpServers={mcpServers} /> ), )} @@ -1066,6 +1084,7 @@ export const ConversationTimeline: FC = ({ showInitialPlaceholder={isAwaitingFirstStreamChunk} retryState={retryState} urlTransform={urlTransform} + mcpServers={mcpServers} /> )} {hasStreamOutput && parsedMessages.length === 0 && ( @@ -1077,6 +1096,7 @@ export const ConversationTimeline: FC = ({ showInitialPlaceholder={isAwaitingFirstStreamChunk} retryState={retryState} urlTransform={urlTransform} + mcpServers={mcpServers} /> )} diff --git a/site/src/pages/AgentsPage/components/AgentDetail/messageParsing.ts b/site/src/pages/AgentsPage/components/AgentDetail/messageParsing.ts index 15e2dd5e8f..85fde69c34 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/messageParsing.ts +++ b/site/src/pages/AgentsPage/components/AgentDetail/messageParsing.ts @@ -97,6 +97,7 @@ export const mergeTools = ( result: result?.result, isError: result?.isError ?? false, status: result ? (result.isError ? "error" : "completed") : "completed", + mcpServerConfigId: call.mcpServerConfigId || result?.mcpServerConfigId, }); } @@ -108,6 +109,7 @@ export const mergeTools = ( result: result.result, isError: result.isError, status: result.isError ? "error" : "completed", + mcpServerConfigId: result.mcpServerConfigId, }); } } @@ -148,6 +150,7 @@ export const parseMessageContent = ( id, name: part.tool_name || "Tool", args: part.args, + mcpServerConfigId: part.mcp_server_config_id, }); parsed.blocks = ensureToolBlock(parsed.blocks, id); break; @@ -168,6 +171,7 @@ export const parseMessageContent = ( name, result: part.result, isError: parseToolResultIsError(name, part, part.result), + mcpServerConfigId: part.mcp_server_config_id, }); parsed.blocks = ensureToolBlock(parsed.blocks, id); break; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/streamState.ts b/site/src/pages/AgentsPage/components/AgentDetail/streamState.ts index 9272699cf2..5caf2a88ab 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/streamState.ts +++ b/site/src/pages/AgentsPage/components/AgentDetail/streamState.ts @@ -70,6 +70,8 @@ export const applyMessagePartToStreamState = ( name: part.tool_name || existing?.name || "Tool", args: nextArgs.value, argsRaw: nextArgs.rawText, + mcpServerConfigId: + part.mcp_server_config_id || existing?.mcpServerConfigId, }, }, }; @@ -115,6 +117,8 @@ export const applyMessagePartToStreamState = ( result: nextResult.value, resultRaw: nextResult.rawText, isError: nextIsError, + mcpServerConfigId: + part.mcp_server_config_id || existing?.mcpServerConfigId, }, }, }; @@ -192,6 +196,7 @@ export const buildStreamTools = ( result: result?.result, isError: result?.isError ?? false, status: result ? (result.isError ? "error" : "completed") : "running", + mcpServerConfigId: call.mcpServerConfigId || result?.mcpServerConfigId, }); } @@ -203,6 +208,7 @@ export const buildStreamTools = ( result: result.result, isError: result.isError, status: result.isError ? "error" : "completed", + mcpServerConfigId: result.mcpServerConfigId, }); } } diff --git a/site/src/pages/AgentsPage/components/AgentDetail/types.ts b/site/src/pages/AgentsPage/components/AgentDetail/types.ts index aeae99bc7f..7b1b154fdb 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/types.ts +++ b/site/src/pages/AgentsPage/components/AgentDetail/types.ts @@ -4,6 +4,7 @@ export type ParsedToolCall = { id: string; name: string; args?: unknown; + mcpServerConfigId?: string; }; export type ParsedToolResult = { @@ -11,6 +12,7 @@ export type ParsedToolResult = { name: string; result?: unknown; isError: boolean; + mcpServerConfigId?: string; }; export type MergedTool = { @@ -20,6 +22,7 @@ export type MergedTool = { result?: unknown; isError: boolean; status: "completed" | "error" | "running"; + mcpServerConfigId?: string; }; export type RenderBlock = @@ -62,6 +65,7 @@ type StreamToolCall = { name: string; args?: unknown; argsRaw?: string; + mcpServerConfigId?: string; }; type StreamToolResult = { @@ -70,6 +74,7 @@ type StreamToolResult = { result?: unknown; resultRaw?: string; isError: boolean; + mcpServerConfigId?: string; }; export type StreamState = { diff --git a/site/src/pages/AgentsPage/components/AgentDetailContent.tsx b/site/src/pages/AgentsPage/components/AgentDetailContent.tsx index 19e138320b..145e24f5d2 100644 --- a/site/src/pages/AgentsPage/components/AgentDetailContent.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetailContent.tsx @@ -51,6 +51,7 @@ interface AgentDetailTimelineProps { editingMessageId?: number | null; savingMessageId?: number | null; urlTransform?: UrlTransform; + mcpServers?: readonly TypesGen.MCPServerConfig[]; } // Reads only message-related store state (stable during streaming). @@ -63,6 +64,7 @@ const MessageListProvider: FC = ({ editingMessageId, savingMessageId, urlTransform, + mcpServers, }) => { const messagesByID = useChatSelector(store, selectMessagesByID); const orderedMessageIDs = useChatSelector(store, selectOrderedMessageIDs); @@ -105,6 +107,7 @@ const MessageListProvider: FC = ({ editingMessageId={editingMessageId} savingMessageId={savingMessageId} urlTransform={urlTransform} + mcpServers={mcpServers} /> ); }; @@ -129,6 +132,7 @@ const StreamingBridge: FC<{ editingMessageId?: number | null; savingMessageId?: number | null; urlTransform?: UrlTransform; + mcpServers?: readonly TypesGen.MCPServerConfig[]; }> = ({ store, isEmpty, @@ -143,6 +147,7 @@ const StreamingBridge: FC<{ editingMessageId, savingMessageId, urlTransform, + mcpServers, }) => { const streamState = useChatSelector(store, selectStreamState); const streamTools = buildStreamTools(streamState); @@ -170,6 +175,7 @@ const StreamingBridge: FC<{ editingMessageId={editingMessageId} savingMessageId={savingMessageId} urlTransform={urlTransform} + mcpServers={mcpServers} /> ); diff --git a/site/src/pages/AgentsPage/components/AgentDetailView.tsx b/site/src/pages/AgentsPage/components/AgentDetailView.tsx index bafb5fb1b8..b30c6a59a1 100644 --- a/site/src/pages/AgentsPage/components/AgentDetailView.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetailView.tsx @@ -282,6 +282,7 @@ export const AgentDetailView: FC = ({ editingMessageId={editing.editingMessageId} savingMessageId={pendingEditMessageId} urlTransform={urlTransform} + mcpServers={mcpServers} />