coder · johnstcn · Apr 11, 2025 · Apr 10, 2025 · Apr 10, 2025 · Apr 11, 2025
diff --git a/cli/exp_mcp.go b/cli/exp_mcp.go
@@ -6,19 +6,19 @@ import (
 	"errors"
 	"os"
 	"path/filepath"
+	"slices"
 	"strings"
 
+	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/mark3labs/mcp-go/server"
 	"github.com/spf13/afero"
 	"golang.org/x/xerrors"
 
-	"cdr.dev/slog"
-	"cdr.dev/slog/sloggers/sloghuman"
 	"github.com/coder/coder/v2/buildinfo"
 	"github.com/coder/coder/v2/cli/cliui"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
-	codermcp "github.com/coder/coder/v2/mcp"
+	"github.com/coder/coder/v2/codersdk/toolsdk"
 	"github.com/coder/serpent"
 )
 
@@ -365,6 +365,8 @@ func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instruct
 	ctx, cancel := context.WithCancel(inv.Context())
 	defer cancel()
 
+	fs := afero.NewOsFs()
+
 	me, err := client.User(ctx, codersdk.Me)
 	if err != nil {
 		cliui.Errorf(inv.Stderr, "Failed to log in to the Coder deployment.")
@@ -397,40 +399,36 @@ func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instruct
 		server.WithInstructions(instructions),
 	)
 
-	// Create a separate logger for the tools.
-	toolLogger := slog.Make(sloghuman.Sink(invStderr))
-
-	toolDeps := codermcp.ToolDeps{
-		Client:        client,
-		Logger:        &toolLogger,
-		AppStatusSlug: appStatusSlug,
-		AgentClient:   agentsdk.New(client.URL),
-	}
-
+	// Create a new context for the tools with all relevant information.
+	clientCtx := toolsdk.WithClient(ctx, client)
 	// Get the workspace agent token from the environment.
-	agentToken, ok := os.LookupEnv("CODER_AGENT_TOKEN")
-	if ok && agentToken != "" {
-		toolDeps.AgentClient.SetSessionToken(agentToken)
+	if agentToken, err := getAgentToken(fs); err == nil && agentToken != "" {
+		agentClient := agentsdk.New(client.URL)
+		agentClient.SetSessionToken(agentToken)
+		clientCtx = toolsdk.WithAgentClient(clientCtx, agentClient)
 	} else {
 		cliui.Warnf(inv.Stderr, "CODER_AGENT_TOKEN is not set, task reporting will not be available")
 	}
-	if appStatusSlug == "" {
+	if appStatusSlug != "" {
 		cliui.Warnf(inv.Stderr, "CODER_MCP_APP_STATUS_SLUG is not set, task reporting will not be available.")
+	} else {
+		clientCtx = toolsdk.WithWorkspaceAppStatusSlug(clientCtx, appStatusSlug)
 	}
 
 	// Register tools based on the allowlist (if specified)
-	reg := codermcp.AllTools()
-	if len(allowedTools) > 0 {
-		reg = reg.WithOnlyAllowed(allowedTools...)
+	for _, tool := range toolsdk.All {
+		if len(allowedTools) == 0 || slices.ContainsFunc(allowedTools, func(t string) bool {
+			return t == tool.Tool.Name
+		}) {
+			mcpSrv.AddTools(mcpFromSDK(tool))
+		}
 	}
 
-	reg.Register(mcpSrv, toolDeps)
-
 	srv := server.NewStdioServer(mcpSrv)
 	done := make(chan error)
 	go func() {
 		defer close(done)
-		srvErr := srv.Listen(ctx, invStdin, invStdout)
+		srvErr := srv.Listen(clientCtx, invStdin, invStdout)
 		done <- srvErr
 	}()
 
@@ -527,8 +525,8 @@ func configureClaude(fs afero.Fs, cfg ClaudeConfig) error {
 	if !ok {
 		mcpServers = make(map[string]any)
 	}
-	for name, mcp := range cfg.MCPServers {
-		mcpServers[name] = mcp
+	for name, cfgmcp := range cfg.MCPServers {
+		mcpServers[name] = cfgmcp
 	}
 	project["mcpServers"] = mcpServers
 	// Prevents Claude from asking the user to complete the project onboarding.
@@ -674,7 +672,7 @@ func indexOf(s, substr string) int {
 
 func getAgentToken(fs afero.Fs) (string, error) {
 	token, ok := os.LookupEnv("CODER_AGENT_TOKEN")
-	if ok {
+	if ok && token != "" {
 		return token, nil
 	}
 	tokenFile, ok := os.LookupEnv("CODER_AGENT_TOKEN_FILE")
@@ -687,3 +685,44 @@ func getAgentToken(fs afero.Fs) (string, error) {
 	}
 	return string(bs), nil
 }
+
+// mcpFromSDK adapts a toolsdk.Tool to go-mcp's server.ServerTool.
+// It assumes that the tool responds with a valid JSON object.
+func mcpFromSDK(sdkTool toolsdk.Tool[any]) server.ServerTool {
+	return server.ServerTool{
+		Tool: mcp.Tool{
+			Name:        sdkTool.Tool.Name,
+			Description: sdkTool.Description,
+			InputSchema: mcp.ToolInputSchema{
+				Type:       "object", // Default of mcp.NewTool()
+				Properties: sdkTool.Schema.Properties,
+				Required:   sdkTool.Schema.Required,
+			},
+		},
+		Handler: func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			result, err := sdkTool.Handler(ctx, request.Params.Arguments)
+			if err != nil {
+				return nil, err
+			}
+			var sb strings.Builder
+			if err := json.NewEncoder(&sb).Encode(result); err == nil {
+				return &mcp.CallToolResult{
+					Content: []mcp.Content{
+						mcp.NewTextContent(sb.String()),
+					},
+				}, nil
+			}
+			// If the result is not JSON, return it as a string.
+			// This is a fallback for tools that return non-JSON data.
+			resultStr, ok := result.(string)
+			if !ok {
+				return nil, xerrors.Errorf("tool call result is neither valid JSON or a string, got: %T", result)
+			}
+			return &mcp.CallToolResult{
+				Content: []mcp.Content{
+					mcp.NewTextContent(resultStr),
+				},
+			}, nil
+		},
+	}
+}
diff --git a/cli/exp_mcp_test.go b/cli/exp_mcp_test.go
@@ -39,12 +39,13 @@ func TestExpMcpServer(t *testing.T) {
 		_ = coderdtest.CreateFirstUser(t, client)
 
 		// Given: we run the exp mcp command with allowed tools set
-		inv, root := clitest.New(t, "exp", "mcp", "server", "--allowed-tools=coder_whoami,coder_list_templates")
+		inv, root := clitest.New(t, "exp", "mcp", "server", "--allowed-tools=coder_get_authenticated_user")
 		inv = inv.WithContext(cancelCtx)
 
 		pty := ptytest.New(t)
 		inv.Stdin = pty.Input()
 		inv.Stdout = pty.Output()
+		// nolint: gocritic // not the focus of this test
 		clitest.SetupConfig(t, client, root)
 
 		cmdDone := make(chan struct{})
@@ -73,13 +74,13 @@ func TestExpMcpServer(t *testing.T) {
 		}
 		err := json.Unmarshal([]byte(output), &toolsResponse)
 		require.NoError(t, err)
-		require.Len(t, toolsResponse.Result.Tools, 2, "should have exactly 2 tools")
+		require.Len(t, toolsResponse.Result.Tools, 1, "should have exactly 1 tool")
 		foundTools := make([]string, 0, 2)
 		for _, tool := range toolsResponse.Result.Tools {
 			foundTools = append(foundTools, tool.Name)
 		}
 		slices.Sort(foundTools)
-		require.Equal(t, []string{"coder_list_templates", "coder_whoami"}, foundTools)
+		require.Equal(t, []string{"coder_get_authenticated_user"}, foundTools)
 	})
 
 	t.Run("OK", func(t *testing.T) {

diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -287,23 +287,25 @@ type TemplateVersionResponse struct {
 }
 
 type TemplateVersionBuilder struct {
-	t         testing.TB
-	db        database.Store
-	seed      database.TemplateVersion
-	fileID    uuid.UUID
-	ps        pubsub.Pubsub
-	resources []*sdkproto.Resource
-	params    []database.TemplateVersionParameter
-	promote   bool
+	t                  testing.TB
+	db                 database.Store
+	seed               database.TemplateVersion
+	fileID             uuid.UUID
+	ps                 pubsub.Pubsub
+	resources          []*sdkproto.Resource
+	params             []database.TemplateVersionParameter
+	promote            bool
+	autoCreateTemplate bool
 }
 
 // TemplateVersion generates a template version and optionally a parent
 // template if no template ID is set on the seed.
 func TemplateVersion(t testing.TB, db database.Store) TemplateVersionBuilder {
 	return TemplateVersionBuilder{
-		t:       t,
-		db:      db,
-		promote: true,
+		t:                  t,
+		db:                 db,
+		promote:            true,
+		autoCreateTemplate: true,
 	}
 }
 
@@ -337,6 +339,13 @@ func (t TemplateVersionBuilder) Params(ps ...database.TemplateVersionParameter)
 	return t
 }
 
+func (t TemplateVersionBuilder) SkipCreateTemplate() TemplateVersionBuilder {
+	// nolint: revive // returns modified struct
+	t.autoCreateTemplate = false
+	t.promote = false
+	return t
+}
+
 func (t TemplateVersionBuilder) Do() TemplateVersionResponse {
 	t.t.Helper()
 
@@ -347,7 +356,7 @@ func (t TemplateVersionBuilder) Do() TemplateVersionResponse {
 	t.fileID = takeFirst(t.fileID, uuid.New())
 
 	var resp TemplateVersionResponse
-	if t.seed.TemplateID.UUID == uuid.Nil {
+	if t.seed.TemplateID.UUID == uuid.Nil && t.autoCreateTemplate {
 		resp.Template = dbgen.Template(t.t, t.db, database.Template{
 			ActiveVersionID: t.seed.ID,
 			OrganizationID:  t.seed.OrganizationID,
@@ -360,16 +369,14 @@ func (t TemplateVersionBuilder) Do() TemplateVersionResponse {
 	}
 
 	version := dbgen.TemplateVersion(t.t, t.db, t.seed)
-
-	// Always make this version the active version. We can easily
-	// add a conditional to the builder to opt out of this when
-	// necessary.
-	err := t.db.UpdateTemplateActiveVersionByID(ownerCtx, database.UpdateTemplateActiveVersionByIDParams{
-		ID:              t.seed.TemplateID.UUID,
-		ActiveVersionID: t.seed.ID,
-		UpdatedAt:       dbtime.Now(),
-	})
-	require.NoError(t.t, err)
+	if t.promote {
+		err := t.db.UpdateTemplateActiveVersionByID(ownerCtx, database.UpdateTemplateActiveVersionByIDParams{
+			ID:              t.seed.TemplateID.UUID,
+			ActiveVersionID: t.seed.ID,
+			UpdatedAt:       dbtime.Now(),
+		})
+		require.NoError(t.t, err)
+	}
 
 	payload, err := json.Marshal(provisionerdserver.TemplateVersionImportJob{
 		TemplateVersionID: t.seed.ID,