Skip to content
34 changes: 2 additions & 32 deletions ui/src/components/create/SelectToolsDialog.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ import {
isAgentTool,
isAgentResponse,
isMcpTool,
toolResponseToAgentTool,
mergeToolIntoServerEntry,
groupMcpToolsByServer,
serverNamesMatch,
} from "@/lib/toolUtils";
Expand Down Expand Up @@ -326,37 +326,7 @@ export const SelectToolsDialog: React.FC<SelectToolsDialogProps> = ({
setLocalSelectedTools((prev) => [...prev, toolToAdd]);
} else {
const tool = item as ToolsResponse;

const existingServerToolIndex = localSelectedTools.findIndex(
(t) =>
isMcpTool(t) &&
serverNamesMatch(t.mcpServer?.name || "", tool.server_name),
);

if (existingServerToolIndex >= 0) {
const existingTool = localSelectedTools[existingServerToolIndex];

if (existingTool.mcpServer?.toolNames?.includes(tool.id)) {
return;
}

const updatedTool = {
...existingTool,
mcpServer: {
...existingTool.mcpServer!,
toolNames: [...(existingTool.mcpServer!.toolNames || []), tool.id],
},
};

setLocalSelectedTools((prev) =>
prev.map((t, idx) =>
idx === existingServerToolIndex ? updatedTool : t,
),
);
} else {
toolToAdd = toolResponseToAgentTool(tool, tool.server_name);
setLocalSelectedTools((prev) => [...prev, toolToAdd]);
}
setLocalSelectedTools((prev) => mergeToolIntoServerEntry(prev, tool));
}
};

Expand Down
26 changes: 20 additions & 6 deletions ui/src/components/onboarding/steps/ReviewStep.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import { ScrollArea } from "@/components/ui/scroll-area";
import { Separator } from "@/components/ui/separator";
import { Badge } from "@/components/ui/badge";
import { Loader2, FunctionSquare } from 'lucide-react';
import { isMcpTool, isAgentTool } from "@/lib/toolUtils";
import type { Tool } from "@/types";

interface OnboardingDataForReview {
Expand Down Expand Up @@ -74,12 +75,25 @@ export function ReviewStep({ onboardingData, isLoading, onBack, onSubmit }: Revi
{onboardingData.selectedTools && onboardingData.selectedTools.length > 0 ? (
<ScrollArea className="h-[100px] w-full rounded-md border p-3 bg-muted/50">
<div className="flex flex-wrap gap-2">
{onboardingData.selectedTools.map((tool, index) => (
<Badge variant="secondary" key={`${tool.mcpServer?.name}-${index}`} className="flex items-center gap-1">
<FunctionSquare className="h-3 w-3" />
{tool.mcpServer?.name}
</Badge>
))}
{onboardingData.selectedTools.flatMap((tool) => {
if (isMcpTool(tool)) {
return tool.mcpServer.toolNames.map((toolName) => (
<Badge variant="secondary" key={`${tool.mcpServer.name}-${toolName}`} className="flex items-center gap-1">
<FunctionSquare className="h-3 w-3" />
{toolName}
</Badge>
));
}
if (isAgentTool(tool)) {
return [(
<Badge variant="secondary" key={`${tool.agent.namespace}-${tool.agent.name}`} className="flex items-center gap-1">
<FunctionSquare className="h-3 w-3" />
{tool.agent.name}
</Badge>
)];
}
return [];
})}
</div>
</ScrollArea>
) : (
Expand Down
56 changes: 45 additions & 11 deletions ui/src/components/onboarding/steps/ToolSelectionStep.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert";
import { Info, ChevronDown, ChevronRight, FunctionSquare, Search } from 'lucide-react';
import { LoadingState } from "@/components/LoadingState";
import { ErrorState } from "@/components/ErrorState";
import { getToolResponseDisplayName, getToolResponseDescription, getToolResponseIdentifier, getToolResponseCategory, toolResponseToAgentTool } from "@/lib/toolUtils";
import { getToolResponseDisplayName, getToolResponseDescription, getToolResponseIdentifier, getToolResponseCategory, isMcpTool, serverNamesMatch, mergeToolIntoServerEntry } from "@/lib/toolUtils";
import type { Tool, ToolsResponse } from "@/types";
import { Input } from "@/components/ui/input";

Expand Down Expand Up @@ -39,7 +39,7 @@ export function ToolSelectionStep({
if (tool.type === "Agent" && tool.agent) {
return false; // Agents don't match ToolResponse objects
} else if (tool.type === "McpServer" && tool.mcpServer) {
return tool.mcpServer.name === toolResponse.server_name &&
return serverNamesMatch(tool.mcpServer.name, toolResponse.server_name) &&
tool.mcpServer.toolNames.includes(toolResponse.id);
}
return false;
Expand Down Expand Up @@ -125,12 +125,16 @@ export function ToolSelectionStep({
"k8s_get_resources",
];

const initialSelection: Tool[] = [];
// Merge desired tools that share a server into one entry instead of
// pushing one per tool, so preselection doesn't itself produce
// duplicate-looking rows for the same server.
let initialSelection: Tool[] = [];
availableTools.forEach((tool) => {
const toolId = getToolResponseDisplayName(tool);
if (desiredIds.includes(toolId)) {
initialSelection.push(toolResponseToAgentTool(tool, tool.server_name));
if (!desiredIds.includes(toolId)) {
return;
}
initialSelection = mergeToolIntoServerEntry(initialSelection, tool);
});

if (initialSelection.length > 0) {
Expand All @@ -141,10 +145,30 @@ export function ToolSelectionStep({
}, [availableTools, initialSelectedTools, selectedTools.length]);

const handleToolToggle = (toolResponse: ToolsResponse) => {
const agentTool = toolResponseToAgentTool(toolResponse, toolResponse.server_name);
setSelectedTools(prev => {
const isSelected = prev.some(t => toolResponseMatchesTool(toolResponse, t));
return isSelected ? prev.filter(t => !toolResponseMatchesTool(toolResponse, t)) : [...prev, agentTool];

if (isSelected) {
// Drop just this tool id from whichever entry holds it, and drop
// the entry entirely once it references no tools.
return prev
.map(t => {
if (!isMcpTool(t) || !serverNamesMatch(t.mcpServer.name, toolResponse.server_name)) {
return t;
}
const remainingToolNames = t.mcpServer.toolNames.filter((id) => id !== toolResponse.id);
return remainingToolNames.length > 0
? { ...t, mcpServer: { ...t.mcpServer, toolNames: remainingToolNames } }
: null;
})
.filter((t): t is Tool => t !== null);
}

// Merge into the existing entry for this server instead of adding a
// second entry for the same server, which otherwise produced
// duplicate-looking rows downstream (Review step, agent edit page)
// and duplicate live MCP connections at agent runtime.
return mergeToolIntoServerEntry(prev, toolResponse);
});
};

Expand Down Expand Up @@ -217,11 +241,20 @@ export function ToolSelectionStep({
</div>
{expandedCategories[category] && (
<div className="pl-6 pt-2 space-y-2">
{categoryTools.map((tool: ToolsResponse) => (
<div key={getToolResponseIdentifier(tool)} className="flex items-start space-x-3">
{categoryTools.map((tool: ToolsResponse) => {
const selected = isToolSelected(tool);
return (
<div
key={getToolResponseIdentifier(tool)}
className={`flex items-start space-x-3 rounded-md border p-2 transition-colors ${
selected
? "border-primary/40 bg-primary/5"
: "border-transparent hover:bg-muted/50"
}`}
>
<Checkbox
id={getToolResponseIdentifier(tool)}
checked={isToolSelected(tool)}
checked={selected}
onCheckedChange={() => handleToolToggle(tool)}
className="mt-1"
/>
Expand All @@ -233,7 +266,8 @@ export function ToolSelectionStep({
<p className="text-xs text-muted-foreground">{getToolResponseDescription(tool)}</p>
</div>
</div>
))}
);
})}
</div>
)}
</div>
Expand Down
100 changes: 100 additions & 0 deletions ui/src/components/onboarding/steps/__tests__/ReviewStep.test.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
/**
* @jest-environment jsdom
*
* Bug: the "Selected Tools" badges rendered `tool.mcpServer?.name` (the MCP
* server name, e.g. "kagent-tool-server") - one badge per array entry. Now
* that ToolSelectionStep merges same-server picks into a single Tool entry's
* mcpServer.toolNames array, this collapsed to a single, unhelpful badge no
* matter how many/which tools were actually selected.
*
* Fix: flatMap over each entry's toolNames (for MCP tools) so Review shows
* one badge per selected tool name instead of per array entry.
*/
import React from "react";
import { describe, it, expect } from "@jest/globals";
import { render, screen } from "@testing-library/react";
import { ReviewStep } from "@/components/onboarding/steps/ReviewStep";
import type { Tool } from "@/types";

describe("ReviewStep Selected Tools", () => {
it("shows one badge per tool name, not one per (merged) server entry", () => {
const mergedServerTool: Tool = {
type: "McpServer",
mcpServer: {
kind: "RemoteMCPServer",
apiGroup: "kagent.dev",
name: "kagent-tool-server",
namespace: "kagent",
toolNames: ["k8s_get_pods", "k8s_get_events"],
},
};

render(
<ReviewStep
onboardingData={{ selectedTools: [mergedServerTool] }}
isLoading={false}
onBack={() => {}}
onSubmit={() => {}}
/>,
);

expect(screen.getByText("k8s_get_pods")).toBeInTheDocument();
expect(screen.getByText("k8s_get_events")).toBeInTheDocument();
expect(screen.queryByText("kagent-tool-server")).not.toBeInTheDocument();
});

it("shows tools from two different servers as separate entries, not merged", () => {
const serverATool: Tool = {
type: "McpServer",
mcpServer: {
kind: "RemoteMCPServer",
apiGroup: "kagent.dev",
name: "kagent-tool-server",
namespace: "kagent",
toolNames: ["k8s_get_pods"],
},
};
const serverBTool: Tool = {
type: "McpServer",
mcpServer: {
kind: "RemoteMCPServer",
apiGroup: "kagent.dev",
name: "context-forge",
namespace: "kagent",
toolNames: ["argocd-get-application"],
},
};

render(
<ReviewStep
onboardingData={{ selectedTools: [serverATool, serverBTool] }}
isLoading={false}
onBack={() => {}}
onSubmit={() => {}}
/>,
);

expect(screen.getByText("k8s_get_pods")).toBeInTheDocument();
expect(screen.getByText("argocd-get-application")).toBeInTheDocument();
expect(screen.queryByText("kagent-tool-server")).not.toBeInTheDocument();
expect(screen.queryByText("context-forge")).not.toBeInTheDocument();
});

it("shows the agent name for Agent-type selections", () => {
const agentTool: Tool = {
type: "Agent",
agent: { name: "researcher", namespace: "kagent" },
};

render(
<ReviewStep
onboardingData={{ selectedTools: [agentTool] }}
isLoading={false}
onBack={() => {}}
onSubmit={() => {}}
/>,
);

expect(screen.getByText("researcher")).toBeInTheDocument();
});
});
Loading
Loading