Add tool filtering to thv serve API. (#1168)

This change adds support for `"tools"` key in the workload creation payload.

Fixes #1023

Author: blkt

docs/server/docs.go

@@ -217,7 +217,7 @@ func Serve(
 		return err
 	}
 
-	logger.Infof("starting %s server", addrType, address)
+	logger.Infof("starting %s server at %s", addrType, address)
 
 	// Start server.
 	go func() {

docs/server/swagger.json

@@ -265,6 +265,7 @@ func (s *WorkloadRoutes) createWorkload(w http.ResponseWriter, r *http.Request)
 		WithOIDCConfig(req.OIDC.Issuer, req.OIDC.Audience, req.OIDC.JwksURL, req.OIDC.ClientID, req.OIDC.AllowOpaqueTokens,
 										"", "", false). // JWKS auth parameters not exposed through API yet
 		WithTelemetryConfig("", false, "", 0.0, nil, false, nil). // Not exposed through API yet.
+		WithToolsFilter(req.ToolsFilter).
 		Build(ctx, imageMetadata, req.EnvVars, &runner.DetachedEnvVarValidator{})
 	if err != nil {
 		logger.Errorf("Failed to create run config: %v", err)
@@ -485,6 +486,8 @@ type createRequest struct {
 	ProxyMode string `json:"proxy_mode"`
 	// Whether network isolation is turned on. This applies the rules in the permission profile.
 	NetworkIsolation bool `json:"network_isolation"`
+	// Tools filter
+	ToolsFilter []string `json:"tools"`
 }
 
 // oidcOptions represents OIDC configuration options

docs/server/swagger.yaml

@@ -433,18 +433,18 @@ func (d *defaultManager) RunWorkloadDetached(ctx context.Context, runConfig *run
 		detachedArgs = append(detachedArgs, "--enable-audit")
 	}
 
+	if runConfig.ToolsFilter != nil {
+		toolsFilter := strings.Join(runConfig.ToolsFilter, ",")
+		detachedArgs = append(detachedArgs, "--tools", toolsFilter)
+	}
+
 	// Add the image and any arguments
 	detachedArgs = append(detachedArgs, runConfig.Image)
 	if len(runConfig.CmdArgs) > 0 {
 		detachedArgs = append(detachedArgs, "--")
 		detachedArgs = append(detachedArgs, runConfig.CmdArgs...)
 	}
 
-	if runConfig.ToolsFilter != nil {
-		toolsFilter := strings.Join(runConfig.ToolsFilter, ",")
-		detachedArgs = append(detachedArgs, "--tools", toolsFilter)
-	}
-
 	// Create a new command
 	// #nosec G204 - This is safe as execPath is the path to the current binary
 	detachedCmd := exec.Command(execPath, detachedArgs...)

pkg/api/server.go

@@ -66,6 +66,8 @@ type Workload struct {
 	Labels map[string]string `json:"labels,omitempty"`
 	// Group is the name of the group this workload belongs to, if any.
 	Group string `json:"group,omitempty"`
+	// ToolsFilter is the filter on tools applied to the workload.
+	ToolsFilter []string `json:"tools,omitempty"`
 }
 
 // loadGroupFromRunConfig attempts to load group information from the runconfig

pkg/api/v1/workloads.go

@@ -85,6 +85,31 @@ var _ = Describe("FetchMcpServer", func() {
 			})
 		})
 
+		Context("when starting the server from registry with tools filter", func() {
+			It("should start when filters are correct", func() {
+				By("Starting the fetch MCP server")
+				stdout, stderr := e2e.NewTHVCommand(config, "run", "--name", serverName, "fetch", "--tools", "fetch").ExpectSuccess()
+
+				// The command should indicate success
+				Expect(stdout+stderr).To(ContainSubstring("fetch"), "Output should mention the fetch server")
+
+				By("Waiting for the server to be running")
+				err := e2e.WaitForMCPServer(config, serverName, 60*time.Second)
+				Expect(err).ToNot(HaveOccurred(), "Server should be running within 30 seconds")
+
+				By("Verifying the server appears in the list")
+				stdout, _ = e2e.NewTHVCommand(config, "list").ExpectSuccess()
+				Expect(stdout).To(ContainSubstring(serverName), "Server should appear in the list")
+				Expect(stdout).To(ContainSubstring("running"), "Server should be in running state")
+			})
+
+			It("should not start when filters are incorrect", func() {
+				By("Starting the fetch MCP server")
+				_, _, err := e2e.NewTHVCommand(config, "run", "--name", serverName, "fetch", "--tools", "wrong-tool").ExpectFailure()
+				Expect(err).To(HaveOccurred(), "Should fail with non-existent server")
+			})
+		})
+
 		Context("when managing the server lifecycle", func() {
 			BeforeEach(func() {
 				// Start a server for lifecycle tests