From 786b8fbd3539e54272662f6c6557843224b70cae Mon Sep 17 00:00:00 2001 From: jaylonmcshan19-x Date: Tue, 24 Mar 2026 13:26:44 -0500 Subject: [PATCH] add get_sentinel_mock tool --- CHANGELOG.md | 4 + pkg/tools/dynamic_tool.go | 5 + pkg/tools/tfe/get_sentinel_mock.go | 116 ++++++++++++++++++++++++ pkg/tools/tfe/get_sentinel_mock_test.go | 23 +++++ pkg/toolsets/mapping.go | 1 + 5 files changed, 149 insertions(+) create mode 100644 pkg/tools/tfe/get_sentinel_mock.go create mode 100644 pkg/tools/tfe/get_sentinel_mock_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 572d73a4..76895310 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ FIXES * Fix sessions handling in stateless and load balanced environments +FEATURES + +* [New Tool] `get_sentinel_mock` Export and download Sentinel mock bundle data for a Terraform plan + IMPROVEMENTS * Add `--heartbeat-interval` CLI flag and `MCP_HEARTBEAT_INTERVAL` env var for HTTP heartbeat in load-balanced environments diff --git a/pkg/tools/dynamic_tool.go b/pkg/tools/dynamic_tool.go index bcfdb42a..ccc30102 100644 --- a/pkg/tools/dynamic_tool.go +++ b/pkg/tools/dynamic_tool.go @@ -209,6 +209,11 @@ func (r *DynamicToolRegistry) registerTFETools() { r.mcpServer.AddTool(tool.Tool, tool.Handler) } + if toolsets.IsToolEnabled("get_sentinel_mock", r.enabledToolsets) { + tool := r.createDynamicTFETool("get_sentinel_mock", tfeTools.GetSentinelMock) + r.mcpServer.AddTool(tool.Tool, tool.Handler) + } + // Terraform toolset - Variable set tools if toolsets.IsToolEnabled("list_variable_sets", r.enabledToolsets) { tool := r.createDynamicTFETool("list_variable_sets", tfeTools.ListVariableSets) diff --git a/pkg/tools/tfe/get_sentinel_mock.go b/pkg/tools/tfe/get_sentinel_mock.go new file mode 100644 index 00000000..9d97667f --- /dev/null +++ b/pkg/tools/tfe/get_sentinel_mock.go @@ -0,0 +1,116 @@ +// Copyright (c) HashiCorp, Inc. +// SPDX-License-Identifier: MPL-2.0 + +package tools + +import ( + "context" + "encoding/base64" + "fmt" + "time" + + "github.com/hashicorp/go-tfe" + "github.com/hashicorp/terraform-mcp-server/pkg/client" + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + log "github.com/sirupsen/logrus" +) + +// GetSentinelMock creates a tool to export and download Sentinel mock data for a plan. +func GetSentinelMock(logger *log.Logger) server.ServerTool { + return server.ServerTool{ + Tool: mcp.NewTool("get_sentinel_mock", + mcp.WithDescription(`Exports and downloads Sentinel mock bundle data for a Terraform plan. This data can be used to test Sentinel policies against plan output. The export is asynchronous - this tool handles polling until the export is ready.`), + mcp.WithTitleAnnotation("Get Sentinel mock data for a Terraform plan"), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithString("plan_id", + mcp.Required(), + mcp.Description("The ID of the plan to export Sentinel mock data for (e.g., plan-8F5JFydVYAmtTjET)"), + ), + ), + Handler: func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return getSentinelMockHandler(ctx, request, logger) + }, + } +} + +func getSentinelMockHandler(ctx context.Context, request mcp.CallToolRequest, logger *log.Logger) (*mcp.CallToolResult, error) { + planID, err := request.RequireString("plan_id") + if err != nil { + return ToolError(logger, "missing required input: plan_id", err) + } + + tfeClient, err := client.GetTfeClientFromContext(ctx, logger) + if err != nil { + return ToolError(logger, "failed to get Terraform client - ensure TFE_TOKEN and TFE_ADDRESS are configured", err) + } + + // Create the plan export + dataType := tfe.PlanExportSentinelMockBundleV0 + planExport, err := tfeClient.PlanExports.Create(ctx, tfe.PlanExportCreateOptions{ + Plan: &tfe.Plan{ID: planID}, + DataType: &dataType, + }) + if err != nil { + return ToolErrorf(logger, "failed to create plan export for plan %s: %v", planID, err) + } + + logger.WithFields(log.Fields{ + "plan_id": planID, + "plan_export_id": planExport.ID, + }).Debug("Created plan export, polling for completion") + + // Poll until the export is finished + maxAttempts := 30 + pollInterval := 2 * time.Second + + for i := 0; i < maxAttempts; i++ { + planExport, err = tfeClient.PlanExports.Read(ctx, planExport.ID) + if err != nil { + return ToolErrorf(logger, "failed to read plan export status: %v", err) + } + + switch planExport.Status { + case tfe.PlanExportFinished: + // Export is ready, download it + data, err := tfeClient.PlanExports.Download(ctx, planExport.ID) + if err != nil { + return ToolErrorf(logger, "failed to download plan export: %v", err) + } + + // Return as base64-encoded tar.gz + encoded := base64.StdEncoding.EncodeToString(data) + result := fmt.Sprintf(`{"plan_id": "%s", "plan_export_id": "%s", "data_type": "sentinel-mock-bundle-v0", "format": "base64-tar-gz", "data": "%s"}`, planID, planExport.ID, encoded) + return mcp.NewToolResultText(result), nil + + case tfe.PlanExportErrored: + return ToolErrorf(logger, "plan export failed with error status for plan %s", planID) + + case tfe.PlanExportCanceled: + return ToolErrorf(logger, "plan export was canceled for plan %s", planID) + + case tfe.PlanExportExpired: + return ToolErrorf(logger, "plan export expired for plan %s", planID) + + case tfe.PlanExportPending, tfe.PlanExportQueued: + // Still processing, wait and retry + logger.WithFields(log.Fields{ + "status": planExport.Status, + "attempt": i + 1, + }).Debug("Plan export still processing, waiting...") + + select { + case <-ctx.Done(): + return ToolError(logger, "context canceled while waiting for plan export", ctx.Err()) + case <-time.After(pollInterval): + continue + } + + default: + return ToolErrorf(logger, "unexpected plan export status: %s", planExport.Status) + } + } + + return ToolErrorf(logger, "plan export timed out after %d attempts for plan %s", maxAttempts, planID) +} diff --git a/pkg/tools/tfe/get_sentinel_mock_test.go b/pkg/tools/tfe/get_sentinel_mock_test.go new file mode 100644 index 00000000..cd416d2c --- /dev/null +++ b/pkg/tools/tfe/get_sentinel_mock_test.go @@ -0,0 +1,23 @@ +// Copyright (c) HashiCorp, Inc. +// SPDX-License-Identifier: MPL-2.0 + +package tools + +import ( + "testing" + + log "github.com/sirupsen/logrus" + "github.com/stretchr/testify/assert" +) + +func TestGetSentinelMock(t *testing.T) { + logger := log.New() + + t.Run("tool creation", func(t *testing.T) { + tool := GetSentinelMock(logger) + + assert.Equal(t, "get_sentinel_mock", tool.Tool.Name) + assert.Contains(t, tool.Tool.Description, "Sentinel mock") + assert.NotNil(t, tool.Handler) + }) +} diff --git a/pkg/toolsets/mapping.go b/pkg/toolsets/mapping.go index 1536f13f..3e755252 100644 --- a/pkg/toolsets/mapping.go +++ b/pkg/toolsets/mapping.go @@ -34,6 +34,7 @@ var ToolToToolset = map[string]string{ "delete_workspace_safely": Terraform, "list_runs": Terraform, "get_run_details": Terraform, + "get_sentinel_mock": Terraform, "create_run": Terraform, "action_run": Terraform, "list_workspace_variables": Terraform,