fixup! add WithCleanContext middleware func

johnstcn · johnstcn · commit 6c97ac48db0e · 2025-04-25T12:00:47.000+01:00
diff --git a/cli/exp_mcp.go b/cli/exp_mcp.go
@@ -713,8 +713,8 @@ func mcpFromSDK(sdkTool toolsdk.Tool[any, any], tb toolsdk.Deps) server.ServerTo
 				Required:   sdkTool.Schema.Required,
 			},
 		},
-		Handler: func(_ context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-			result, err := sdkTool.Handler(tb, request.Params.Arguments)
+		Handler: func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			result, err := sdkTool.Handler(ctx, tb, request.Params.Arguments)
 			if err != nil {
 				return nil, err
 			}
diff --git a/codersdk/toolsdk/toolsdk.go b/codersdk/toolsdk/toolsdk.go
@@ -21,7 +21,7 @@ type Deps struct {
 }
 
 // HandlerFunc is a function that handles a tool call.
-type HandlerFunc[Arg, Ret any] func(tb Deps, args Arg) (Ret, error)
+type HandlerFunc[Arg, Ret any] func(context.Context, Deps, Arg) (Ret, error)
 
 type Tool[Arg, Ret any] struct {
 	aisdk.Tool
@@ -32,12 +32,12 @@ type Tool[Arg, Ret any] struct {
 func (t Tool[Arg, Ret]) Generic() Tool[any, any] {
 	return Tool[any, any]{
 		Tool: t.Tool,
-		Handler: func(tb Deps, args any) (any, error) {
+		Handler: func(ctx context.Context, tb Deps, args any) (any, error) {
 			typedArg, ok := args.(Arg)
 			if !ok {
 				return nil, xerrors.Errorf("developer error: invalid argument type for tool %s", t.Tool.Name)
 			}
-			return t.Handler(tb, typedArg)
+			return t.Handler(ctx, tb, typedArg)
 		},
 	}
 }
@@ -165,7 +165,7 @@ func wrapAll(mw func(HandlerFunc[any, any]) HandlerFunc[any, any], tools ...Tool
 var (
 	// All is a list of all tools that can be used in the Coder CLI.
 	// When you add a new tool, be sure to include it here!
-	All = wrapAll(WithRecover,
+	All = wrapAll(WithCleanContext, wrapAll(WithRecover,
 		CreateTemplate.Generic(),
 		CreateTemplateVersion.Generic(),
 		CreateWorkspace.Generic(),
@@ -182,7 +182,7 @@ var (
 		ReportTask.Generic(),
 		UploadTarFile.Generic(),
 		UpdateTemplateActiveVersion.Generic(),
-	)
+	)...)
 
 	ReportTask = Tool[ReportTaskArgs, string]{
 		Tool: aisdk.Tool{
