chore(codersdk/toolsdk): improve static analyzability of toolsdk.Tools (#17562)

* Refactors toolsdk.Tools to remove opaque `map[string]any` argument in
favour of typed args structs.
* Refactors toolsdk.Tools to remove opaque passing of dependencies via
`context.Context` in favour of a tool dependencies struct.
* Adds panic recovery and clean context middleware to all tools.
* Adds `GenericTool` implementation to allow keeping `toolsdk.All` with
uniform type signature while maintaining type information in handlers.
* Adds stricter checks to `patchWorkspaceAgentAppStatus` handler.
This commit is contained in:
Cian Johnston
2025-04-29 16:05:23 +01:00
committed by GitHub
parent 1fc74f629e
commit 2acf0adcf2
6 changed files with 1139 additions and 865 deletions
+21 -25
View File
@@ -1,6 +1,7 @@
package cli
import (
"bytes"
"context"
"encoding/json"
"errors"
@@ -427,22 +428,27 @@ func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instruct
server.WithInstructions(instructions),
)
// Create a new context for the tools with all relevant information.
clientCtx := toolsdk.WithClient(ctx, client)
// Get the workspace agent token from the environment.
toolOpts := make([]func(*toolsdk.Deps), 0)
var hasAgentClient bool
if agentToken, err := getAgentToken(fs); err == nil && agentToken != "" {
hasAgentClient = true
agentClient := agentsdk.New(client.URL)
agentClient.SetSessionToken(agentToken)
clientCtx = toolsdk.WithAgentClient(clientCtx, agentClient)
toolOpts = append(toolOpts, toolsdk.WithAgentClient(agentClient))
} else {
cliui.Warnf(inv.Stderr, "CODER_AGENT_TOKEN is not set, task reporting will not be available")
}
if appStatusSlug == "" {
cliui.Warnf(inv.Stderr, "CODER_MCP_APP_STATUS_SLUG is not set, task reporting will not be available.")
if appStatusSlug != "" {
toolOpts = append(toolOpts, toolsdk.WithAppStatusSlug(appStatusSlug))
} else {
clientCtx = toolsdk.WithWorkspaceAppStatusSlug(clientCtx, appStatusSlug)
cliui.Warnf(inv.Stderr, "CODER_MCP_APP_STATUS_SLUG is not set, task reporting will not be available.")
}
toolDeps, err := toolsdk.NewDeps(client, toolOpts...)
if err != nil {
return xerrors.Errorf("failed to initialize tool dependencies: %w", err)
}
// Register tools based on the allowlist (if specified)
@@ -455,7 +461,7 @@ func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instruct
if len(allowedTools) == 0 || slices.ContainsFunc(allowedTools, func(t string) bool {
return t == tool.Tool.Name
}) {
mcpSrv.AddTools(mcpFromSDK(tool))
mcpSrv.AddTools(mcpFromSDK(tool, toolDeps))
}
}
@@ -463,7 +469,7 @@ func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instruct
done := make(chan error)
go func() {
defer close(done)
srvErr := srv.Listen(clientCtx, invStdin, invStdout)
srvErr := srv.Listen(ctx, invStdin, invStdout)
done <- srvErr
}()
@@ -726,7 +732,7 @@ func getAgentToken(fs afero.Fs) (string, error) {
// 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 {
func mcpFromSDK(sdkTool toolsdk.GenericTool, tb toolsdk.Deps) server.ServerTool {
// NOTE: some clients will silently refuse to use tools if there is an issue
// with the tool's schema or configuration.
if sdkTool.Schema.Properties == nil {
@@ -743,27 +749,17 @@ func mcpFromSDK(sdkTool toolsdk.Tool[any]) server.ServerTool {
},
},
Handler: func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
result, err := sdkTool.Handler(ctx, request.Params.Arguments)
var buf bytes.Buffer
if err := json.NewEncoder(&buf).Encode(request.Params.Arguments); err != nil {
return nil, xerrors.Errorf("failed to encode request arguments: %w", err)
}
result, err := sdkTool.Handler(ctx, tb, buf.Bytes())
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),
mcp.NewTextContent(string(result)),
},
}, nil
},
+16 -6
View File
@@ -31,12 +31,12 @@ func TestExpMcpServer(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
cmdDone := make(chan struct{})
cancelCtx, cancel := context.WithCancel(ctx)
t.Cleanup(cancel)
// Given: a running coder deployment
client := coderdtest.New(t, nil)
_ = coderdtest.CreateFirstUser(t, client)
owner := 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_get_authenticated_user")
@@ -48,7 +48,6 @@ func TestExpMcpServer(t *testing.T) {
// nolint: gocritic // not the focus of this test
clitest.SetupConfig(t, client, root)
cmdDone := make(chan struct{})
go func() {
defer close(cmdDone)
err := inv.Run()
@@ -61,9 +60,6 @@ func TestExpMcpServer(t *testing.T) {
_ = pty.ReadLine(ctx) // ignore echoed output
output := pty.ReadLine(ctx)
cancel()
<-cmdDone
// Then: we should only see the allowed tools in the response
var toolsResponse struct {
Result struct {
@@ -81,6 +77,20 @@ func TestExpMcpServer(t *testing.T) {
}
slices.Sort(foundTools)
require.Equal(t, []string{"coder_get_authenticated_user"}, foundTools)
// Call the tool and ensure it works.
toolPayload := `{"jsonrpc":"2.0","id":3,"method":"tools/call", "params": {"name": "coder_get_authenticated_user", "arguments": {}}}`
pty.WriteLine(toolPayload)
_ = pty.ReadLine(ctx) // ignore echoed output
output = pty.ReadLine(ctx)
require.NotEmpty(t, output, "should have received a response from the tool")
// Ensure it's valid JSON
_, err = json.Marshal(output)
require.NoError(t, err, "should have received a valid JSON response from the tool")
// Ensure the tool returns the expected user
require.Contains(t, output, owner.UserID.String(), "should have received the expected user ID")
cancel()
<-cmdDone
})
t.Run("OK", func(t *testing.T) {