chore(codersdk/toolsdk): add panic recovery tool middleware

johnstcn · johnstcn · commit 968abbb5d133 · 2025-04-24T16:19:13.000+01:00
diff --git a/codersdk/toolsdk/toolsdk.go b/codersdk/toolsdk/toolsdk.go
@@ -150,10 +150,31 @@ type UploadTarFileArgs struct {
 	Files map[string]string `json:"files"`
 }
 
+// WithRecover wraps a HandlerFunc to recover from panics and return an error.
+func WithRecover[Arg, Ret any](h HandlerFunc[Arg, Ret]) HandlerFunc[Arg, Ret] {
+	return func(tb Toolbox, args Arg) (ret Ret, err error) {
+		defer func() {
+			if r := recover(); r != nil {
+				err = xerrors.Errorf("tool handler panic: %v", r)
+			}
+		}()
+		return h(tb, args)
+	}
+}
+
+// wrapAll wraps all provided tools with the given middleware function.
+func wrapAll(mw func(HandlerFunc[any, any]) HandlerFunc[any, any], tools ...Tool[any, any]) []Tool[any, any] {
+	for i, t := range tools {
+		t.Handler = mw(t.Handler)
+		tools[i] = t
+	}
+	return tools
+}
+
 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 = []Tool[any, any]{
+	All = wrapAll(WithRecover,
 		CreateTemplate.Generic(),
 		CreateTemplateVersion.Generic(),
 		CreateWorkspace.Generic(),
@@ -170,7 +191,7 @@ var (
 		ReportTask.Generic(),
 		UploadTarFile.Generic(),
 		UpdateTemplateActiveVersion.Generic(),
-	}
+	)
 
 	ReportTask = Tool[ReportTaskArgs, string]{
 		Tool: aisdk.Tool{
diff --git a/codersdk/toolsdk/toolsdk_test.go b/codersdk/toolsdk/toolsdk_test.go
@@ -9,6 +9,8 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/kylecarbs/aisdk-go"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/v2/coderd/coderdtest"
@@ -379,3 +381,61 @@ func TestMain(m *testing.M) {
 
 	os.Exit(code)
 }
+
+func TestWithRecovery(t *testing.T) {
+	t.Parallel()
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+		fakeTool := toolsdk.Tool[string, string]{
+			Tool: aisdk.Tool{
+				Name:        "fake_tool",
+				Description: "Returns a string for testing.",
+			},
+			Handler: func(tb toolsdk.Toolbox, args string) (string, error) {
+				require.Equal(t, "test", args)
+				return "ok", nil
+			},
+		}
+
+		wrapped := toolsdk.WithRecover(fakeTool.Handler)
+		v, err := wrapped(nil, "test")
+		require.NoError(t, err)
+		require.Equal(t, "ok", v)
+	})
+
+	t.Run("Error", func(t *testing.T) {
+		t.Parallel()
+		fakeTool := toolsdk.Tool[string, string]{
+			Tool: aisdk.Tool{
+				Name:        "fake_tool",
+				Description: "Returns an error for testing.",
+			},
+			Handler: func(tb toolsdk.Toolbox, args string) (string, error) {
+				require.Equal(t, "test", args)
+				return "", assert.AnError
+			},
+		}
+		wrapped := toolsdk.WithRecover(fakeTool.Handler)
+		v, err := wrapped(nil, "test")
+		require.Empty(t, v)
+		require.ErrorIs(t, err, assert.AnError)
+	})
+
+	t.Run("Panic", func(t *testing.T) {
+		t.Parallel()
+		panicTool := toolsdk.Tool[string, string]{
+			Tool: aisdk.Tool{
+				Name:        "panic_tool",
+				Description: "Panics for testing.",
+			},
+			Handler: func(tb toolsdk.Toolbox, args string) (string, error) {
+				panic("you can't sweat this fever out")
+			},
+		}
+
+		wrapped := toolsdk.WithRecover(panicTool.Handler)
+		v, err := wrapped(nil, "disco")
+		require.Empty(t, v)
+		require.ErrorContains(t, err, "you can't sweat this fever out")
+	})
+}
