mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat(coderd/chatd): connect to external MCP servers for chat tool invocation (#23333)
## Summary
Adds a new `coderd/chatd/mcpclient` package that connects to
admin-configured MCP servers and wraps their tools as
`fantasy.AgentTool` values that the chat loop can invoke.
## What changed
### New: `coderd/chatd/mcpclient/mcpclient.go`
The core package with a single entry point:
```go
func ConnectAll(
ctx context.Context,
logger slog.Logger,
configs []database.MCPServerConfig,
tokens []database.MCPServerUserToken,
) (tools []fantasy.AgentTool, cleanup func(), err error)
```
This:
1. Connects to each enabled MCP server using `mark3labs/mcp-go`
(streamable HTTP or SSE transport)
2. Discovers tools via the MCP `tools/list` method
3. Wraps each tool as a `fantasy.AgentTool` with namespaced name
(`serverslug__toolname`)
4. Applies tool allow/deny list filtering from the server config
5. Handles auth: OAuth2 bearer tokens, API keys, and custom headers
6. Skips broken servers with a warning (10s connect timeout per server)
7. Returns a cleanup function to close all MCP connections
### Modified: `coderd/chatd/chatd.go`
In `runChat()`, after loading the model/messages but before assembling
the tool list:
- Reads `chat.MCPServerIDs` from the chat record
- Loads the MCP server configs from the database
- Resolves the user's auth tokens
- Calls `mcpclient.ConnectAll()` to connect and discover tools
- Appends the MCP tools to the chat's tool set
- Defers cleanup to close connections when the chat turn ends
The chat loop (`chatloop.Run`) already handles tools generically —
MCP-backed tools are invoked identically to built-in workspace tools. No
changes needed in `chatloop/`.
### New: `coderd/chatd/mcpclient/mcpclient_test.go`
10 tests covering:
- Tool discovery and namespacing
- Tool call forwarding and result conversion
- Allow/deny list filtering
- Connection failure handling (graceful skip)
- Multi-server support with correct prefixes
- OAuth2 auth header injection
- Disabled server skipping
- Invalid input handling
- Tool info parameter propagation
## Design decisions
- **Tool namespacing**: `slug__toolname` with double underscore
separator. Avoids collisions with tools containing single underscores.
Stripped when forwarding to `tools/call`.
- **Connection lifecycle**: Fresh connections per chat turn, closed via
`defer`. Matches the `turnWorkspaceContext` pattern.
- **Failure isolation**: Each server connects independently. A broken
server doesn't fail the chat — its tools are simply unavailable.
- **No chatloop changes**: The existing `[]fantasy.AgentTool` interface
is already fully generic.
## What's NOT in this PR (follow-ups)
- Frontend MCP server picker UI (selecting servers for a chat)
- System prompt additions describing available MCP tools
- Token refresh on expiry mid-chat
- The deprecated `aibridged` MCP proxy cleanup
This commit is contained in:
@@ -17,6 +17,8 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
||||
mcpserver "github.com/mark3labs/mcp-go/server"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
@@ -3512,3 +3514,197 @@ func (d *panicOnInTxDB) InTx(f func(database.Store) error, opts *database.TxOpti
|
||||
}
|
||||
return d.Store.InTx(f, opts)
|
||||
}
|
||||
|
||||
// TestMCPServerToolInvocation verifies that when a chat has
|
||||
// mcp_server_ids set, the chat loop connects to those MCP servers,
|
||||
// discovers their tools, and the LLM can invoke them.
|
||||
//
|
||||
// NOTE: This test uses a raw database.Store (no dbauthz wrapper).
|
||||
// The chatd RBAC authorization of GetMCPServerConfigsByIDs (which
|
||||
// requires ActionRead on ResourceDeploymentConfig) is covered by
|
||||
// the chatd role definition tests, not here.
|
||||
func TestMCPServerToolInvocation(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
// Start a real MCP server that exposes an "echo" tool.
|
||||
mcpSrv := mcpserver.NewMCPServer("test-mcp", "1.0.0")
|
||||
mcpSrv.AddTools(mcpserver.ServerTool{
|
||||
Tool: mcpgo.NewTool("echo",
|
||||
mcpgo.WithDescription("Echoes the input"),
|
||||
mcpgo.WithString("input",
|
||||
mcpgo.Description("The input string"),
|
||||
mcpgo.Required(),
|
||||
),
|
||||
),
|
||||
Handler: func(_ context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) {
|
||||
input, _ := req.GetArguments()["input"].(string)
|
||||
return mcpgo.NewToolResultText("echo: " + input), nil
|
||||
},
|
||||
})
|
||||
mcpHTTP := mcpserver.NewStreamableHTTPServer(mcpSrv)
|
||||
mcpTS := httptest.NewServer(mcpHTTP)
|
||||
t.Cleanup(mcpTS.Close)
|
||||
|
||||
// Track which tool names are sent to the LLM and capture
|
||||
// whether the MCP tool result appears in the second call.
|
||||
var (
|
||||
callCount atomic.Int32
|
||||
llmToolNames []string
|
||||
llmToolsMu sync.Mutex
|
||||
foundMCPResult atomic.Bool
|
||||
)
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
|
||||
// Record tool names from the first streamed call.
|
||||
if callCount.Add(1) == 1 {
|
||||
names := make([]string, 0, len(req.Tools))
|
||||
for _, tool := range req.Tools {
|
||||
names = append(names, tool.Function.Name)
|
||||
}
|
||||
llmToolsMu.Lock()
|
||||
llmToolNames = names
|
||||
llmToolsMu.Unlock()
|
||||
|
||||
// Ask the LLM to call the MCP echo tool.
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAIToolCallChunk(
|
||||
"test-mcp__echo",
|
||||
`{"input":"hello from LLM"}`,
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
// Second call: verify the tool result was fed back.
|
||||
for _, msg := range req.Messages {
|
||||
if msg.Role == "tool" && strings.Contains(msg.Content, "echo: hello from LLM") {
|
||||
foundMCPResult.Store(true)
|
||||
}
|
||||
}
|
||||
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("Got it!")...,
|
||||
)
|
||||
})
|
||||
|
||||
user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||
|
||||
// Seed the MCP server config in the database. This must
|
||||
// happen after seedChatDependencies so user.ID exists for
|
||||
// the foreign key.
|
||||
mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{
|
||||
DisplayName: "Test MCP",
|
||||
Slug: "test-mcp",
|
||||
Url: mcpTS.URL,
|
||||
Transport: "streamable_http",
|
||||
AuthType: "none",
|
||||
Availability: "default_off",
|
||||
Enabled: true,
|
||||
ToolAllowList: []string{},
|
||||
ToolDenyList: []string{},
|
||||
CreatedBy: user.ID,
|
||||
UpdatedBy: user.ID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
mockConn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes()
|
||||
mockConn.EXPECT().LS(gomock.Any(), gomock.Any(), gomock.Any()).
|
||||
Return(workspacesdk.LSResponse{}, nil).AnyTimes()
|
||||
mockConn.EXPECT().ReadFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
|
||||
Return(io.NopCloser(strings.NewReader("")), "", nil).AnyTimes()
|
||||
|
||||
server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
|
||||
cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||
require.Equal(t, dbAgent.ID, agentID)
|
||||
return mockConn, func() {}, nil
|
||||
}
|
||||
})
|
||||
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "mcp-tool-test",
|
||||
ModelConfigID: model.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
|
||||
MCPServerIDs: []uuid.UUID{mcpConfig.ID},
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("Echo something via MCP."),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify MCPServerIDs were persisted on the chat record.
|
||||
dbChat, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
require.NoError(t, getErr)
|
||||
require.Equal(t, []uuid.UUID{mcpConfig.ID}, dbChat.MCPServerIDs)
|
||||
|
||||
// Wait for the chat to finish processing.
|
||||
var chatResult database.Chat
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
chatResult = got
|
||||
return got.Status == database.ChatStatusWaiting || got.Status == database.ChatStatusError
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
if chatResult.Status == database.ChatStatusError {
|
||||
require.FailNowf(t, "chat failed", "last_error=%q", chatResult.LastError.String)
|
||||
}
|
||||
|
||||
// The MCP tool (test-mcp__echo) should appear in the tool
|
||||
// list sent to the LLM.
|
||||
llmToolsMu.Lock()
|
||||
recordedNames := append([]string(nil), llmToolNames...)
|
||||
llmToolsMu.Unlock()
|
||||
require.Contains(t, recordedNames, "test-mcp__echo",
|
||||
"MCP tool should be in the tool list sent to the LLM")
|
||||
|
||||
// The tool result from the MCP server ("echo: hello from
|
||||
// LLM") should have been fed back to the LLM as a tool
|
||||
// message in the second call.
|
||||
require.True(t, foundMCPResult.Load(),
|
||||
"MCP tool result should appear in the second LLM call")
|
||||
|
||||
// Verify the tool result was persisted in the database.
|
||||
var foundToolMessage bool
|
||||
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
|
||||
messages, dbErr := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
|
||||
ChatID: chat.ID,
|
||||
AfterID: 0,
|
||||
})
|
||||
if dbErr != nil {
|
||||
return false
|
||||
}
|
||||
for _, msg := range messages {
|
||||
if msg.Role != database.ChatMessageRoleTool {
|
||||
continue
|
||||
}
|
||||
parts, parseErr := chatprompt.ParseContent(msg)
|
||||
if parseErr != nil || len(parts) == 0 {
|
||||
continue
|
||||
}
|
||||
for _, part := range parts {
|
||||
if part.Type == codersdk.ChatMessagePartTypeToolResult &&
|
||||
part.ToolName == "test-mcp__echo" &&
|
||||
strings.Contains(string(part.Result), "echo: hello from LLM") {
|
||||
foundToolMessage = true
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}, testutil.IntervalFast)
|
||||
require.True(t, foundToolMessage,
|
||||
"MCP tool result should be persisted as a tool message in the database")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user