From b3651ab58c3fcd8b56eb9c562e0d63c6a19177ce Mon Sep 17 00:00:00 2001 From: mrhid6 Date: Wed, 9 Sep 2026 07:28:48 +0000 Subject: [PATCH] feat: add mcp write tools with a fan-out guard --- server/internal/mcp/tools_write.go | 298 ++++++++++++++++++++++++ server/internal/mcp/tools_write_test.go | 44 ++++ server/internal/mcp/transport.go | 6 +- 3 files changed, 347 insertions(+), 1 deletion(-) create mode 100644 server/internal/mcp/tools_write.go create mode 100644 server/internal/mcp/tools_write_test.go diff --git a/server/internal/mcp/tools_write.go b/server/internal/mcp/tools_write.go new file mode 100644 index 0000000..abfb2fd --- /dev/null +++ b/server/internal/mcp/tools_write.go @@ -0,0 +1,298 @@ +package mcp + +import ( + "context" + "fmt" + + "gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/services" +) + +// stringSliceArg reads a JSON array-of-strings argument, ignoring any element +// that is not a string. Missing or wrongly-typed input decodes to nil, which +// every caller here treats as "no targets named this way". +func stringSliceArg(args map[string]any, key string) []string { + raw, ok := args[key].([]any) + if !ok { + return nil + } + out := make([]string, 0, len(raw)) + for _, v := range raw { + if s, ok := v.(string); ok { + out = append(out, s) + } + } + return out +} + +func mustLookup(name string) Tool { + t, ok := All().Lookup(name) + if !ok { + panic("mcp: unknown tool " + name) + } + return t +} + +type runStartedResult struct { + RunID string `json:"run_id"` + Note string `json:"note"` +} + +// run_workflow. The real REST run route (internal/api/workflows.go's +// runWorkflow) does not take an ad-hoc target list at all: it calls +// services.TriggerWorkflow(instanceID, workflowID, actor), which resolves the +// workflow's own configured target_server_ids/target_tags via +// services.ResolveTargets (unscoped) and runs against exactly that set. There +// is no per-call server_ids/tags override to plumb through, so this tool takes +// only workflow_id. To keep the token's scope meaningful — TriggerWorkflow +// itself does not consult it — this handler first loads the workflow and +// resolves its configured targets through ResolveTargetsScoped with the +// caller's TokenScope, and refuses the run outright if that scoped view does +// not cover every server the unscoped resolution would touch. That is the +// fan-out and scope check; the actual dispatch is the same single call the +// REST route makes, so there is exactly one path that starts a run. +func init() { + All().Register(Tool{ + Name: "run_workflow", + Write: true, + Scope: "workflows:write", + Description: "Run a workflow against the servers it is already configured to target " + + "(its saved server list and tags — this call does not let you pick different " + + "targets). This EXECUTES COMMANDS on real machines and cannot be undone from " + + "here. Returns a run ID immediately; poll get_run for progress and get_run_logs " + + "for output. Refused if the workflow's targets reach outside this token's own " + + "server scope, or if it would affect more than the fan-out limit without confirm:true.", + Handler: func(ctx context.Context, c Caller, args map[string]any) (any, error) { + workflowID := stringArg(args, "workflow_id") + if workflowID == "" { + return nil, fmt.Errorf("workflow_id is required") + } + + wf, err := services.GetWorkflow(c.InstanceID, workflowID) + if err != nil { + return nil, fmt.Errorf("workflow not found") + } + + allTargets, err := services.ResolveTargets(c.InstanceID, wf.TargetServerIDs, wf.TargetTags) + if err != nil { + return nil, fmt.Errorf("this workflow matches no servers") + } + scopedTargets, err := services.ResolveTargetsScoped(c.InstanceID, wf.TargetServerIDs, wf.TargetTags, c.TokenScope) + if err != nil || len(scopedTargets) != len(allTargets) { + return nil, fmt.Errorf("no servers visible to this token matched the request") + } + + if err := CheckFanOut(len(scopedTargets), args); err != nil { + return nil, err + } + + runID, err := services.TriggerWorkflow(c.InstanceID, workflowID, c.TokenName) + if err != nil { + return nil, fmt.Errorf("could not start the run: %w", err) + } + + LogCall(c, mustLookup("run_workflow"), args, len(scopedTargets)) + + return runStartedResult{ + RunID: runID, + Note: "The run is in progress. Poll get_run with this run_id; do not assume it succeeded.", + }, nil + }, + }) +} + +type cancelledResult struct { + Cancelled bool `json:"cancelled"` +} + +// cancel_run. The REST cancel route (workflows.go's cancelRun) calls +// services.CancelRun(instanceID, runID) directly; that call is already scoped +// to the caller's instance by instanceID, which is what "verifies the run +// belongs to the caller's instance" reduces to here — there is no separate +// per-server scope to check, since cancelling touches the run record, not a +// server. +func init() { + All().Register(Tool{ + Name: "cancel_run", + Write: true, + Scope: "workflows:write", + Description: "Cancel an in-progress workflow run. This stops further steps from " + + "being dispatched to real machines but cannot undo steps that already ran, and " + + "cannot be undone from here.", + Handler: func(ctx context.Context, c Caller, args map[string]any) (any, error) { + runID := stringArg(args, "run_id") + if runID == "" { + return nil, fmt.Errorf("run_id is required") + } + + if err := services.CancelRun(c.InstanceID, runID); err != nil { + return nil, fmt.Errorf("could not cancel the run: %w", err) + } + + LogCall(c, mustLookup("cancel_run"), args, 0) + + return cancelledResult{Cancelled: true}, nil + }, + }) +} + +type updateBatchResult struct { + Servers int `json:"servers"` + Succeeded []string `json:"succeeded"` + Failed map[string]string `json:"failed,omitempty"` +} + +// apply_updates. The REST route (internal/api/handlers.go's applyUpdates) is +// per-server: POST /servers/:id/apply-updates resolves one server with +// services.GetServerScoped and calls services.DispatchApplyUpdates(serverID). +// There is no fleet-wide variant of that service call to invoke once, so this +// tool resolves the requested targets through ResolveTargetsScoped exactly as +// the brief describes, then calls the same DispatchApplyUpdates the REST route +// calls, once per resolved server — the identical dispatch, just looped +// instead of hardcoded to one server_id from the URL. +func init() { + All().Register(Tool{ + Name: "apply_updates", + Write: true, + Scope: "servers:write", + Description: "Apply pending OS package updates on real servers, selected by " + + "server_ids and/or tags. This installs packages on real machines right now and " + + "cannot be undone from here. A server may need a reboot afterward, which this " + + "tool does not do.", + Handler: func(ctx context.Context, c Caller, args map[string]any) (any, error) { + ids := stringSliceArg(args, "server_ids") + targets, err := services.ResolveTargetsScoped(c.InstanceID, ids, tagArg(args), c.TokenScope) + if err != nil { + return nil, fmt.Errorf("no servers visible to this token matched the request") + } + if err := CheckFanOut(len(targets), args); err != nil { + return nil, err + } + + result := updateBatchResult{Servers: len(targets), Failed: map[string]string{}} + for _, srv := range targets { + if err := services.DispatchApplyUpdates(srv.ServerID); err != nil { + result.Failed[srv.ServerID] = err.Error() + continue + } + result.Succeeded = append(result.Succeeded, srv.ServerID) + } + if len(result.Failed) == 0 { + result.Failed = nil + } + + LogCall(c, mustLookup("apply_updates"), args, len(targets)) + + return result, nil + }, + }) +} + +type agentUpdateResult struct { + Servers int `json:"servers"` + Succeeded map[string]string `json:"succeeded,omitempty"` + Failed map[string]string `json:"failed,omitempty"` +} + +// update_agent. Same shape as apply_updates: the REST route +// (handlers.go's updateAgent) resolves one server and calls +// services.DispatchUpdateAgent(serverID), so this tool loops the same call +// over the resolved, scoped target set. +func init() { + All().Register(Tool{ + Name: "update_agent", + Write: true, + Scope: "servers:write", + Description: "Trigger the Vantage agent on real servers, selected by server_ids " + + "and/or tags, to download and replace itself with the latest version. This " + + "restarts the agent process on real machines and cannot be undone from here.", + Handler: func(ctx context.Context, c Caller, args map[string]any) (any, error) { + ids := stringSliceArg(args, "server_ids") + targets, err := services.ResolveTargetsScoped(c.InstanceID, ids, tagArg(args), c.TokenScope) + if err != nil { + return nil, fmt.Errorf("no servers visible to this token matched the request") + } + if err := CheckFanOut(len(targets), args); err != nil { + return nil, err + } + + result := agentUpdateResult{Servers: len(targets), Succeeded: map[string]string{}, Failed: map[string]string{}} + for _, srv := range targets { + version, err := services.DispatchUpdateAgent(srv.ServerID) + if err != nil { + result.Failed[srv.ServerID] = err.Error() + continue + } + result.Succeeded[srv.ServerID] = version + } + if len(result.Succeeded) == 0 { + result.Succeeded = nil + } + if len(result.Failed) == 0 { + result.Failed = nil + } + + LogCall(c, mustLookup("update_agent"), args, len(targets)) + + return result, nil + }, + }) +} + +type assignKeyResult struct { + Servers int `json:"servers"` + Succeeded []string `json:"succeeded"` + Failed map[string]string `json:"failed,omitempty"` +} + +// assign_key. The REST route (handlers.go's assignKey) takes one server_id in +// the body and calls services.AssignKey(instanceID, keyID, serverID) directly +// — AssignKey itself resolves the server with the unscoped services.GetServer, +// not GetServerScoped, so the REST route carries no token-scope check of its +// own (session auth has no server-scope restriction; only API tokens do). For +// the MCP surface, this tool resolves every named target through +// ResolveTargetsScoped first — the same chokepoint every other target- +// resolving write tool goes through — so a token whose scope excludes a server +// cannot reach it here even though the REST handler's own server lookup would +// not have stopped it. Then it calls the identical AssignKey once per resolved +// server. +func init() { + All().Register(Tool{ + Name: "assign_key", + Write: true, + Scope: "keys:write", + Description: "Assign an SSH key to real servers, selected by server_ids and/or " + + "tags. The agent rewrites /root/.ssh/authorized_keys on each targeted machine " + + "and this cannot be undone from here — use revoke to remove it afterward.", + Handler: func(ctx context.Context, c Caller, args map[string]any) (any, error) { + keyID := stringArg(args, "key_id") + if keyID == "" { + return nil, fmt.Errorf("key_id is required") + } + + ids := stringSliceArg(args, "server_ids") + targets, err := services.ResolveTargetsScoped(c.InstanceID, ids, tagArg(args), c.TokenScope) + if err != nil { + return nil, fmt.Errorf("no servers visible to this token matched the request") + } + if err := CheckFanOut(len(targets), args); err != nil { + return nil, err + } + + result := assignKeyResult{Servers: len(targets), Failed: map[string]string{}} + for _, srv := range targets { + if _, err := services.AssignKey(c.InstanceID, keyID, srv.ServerID); err != nil { + result.Failed[srv.ServerID] = err.Error() + continue + } + result.Succeeded = append(result.Succeeded, srv.ServerID) + } + if len(result.Failed) == 0 { + result.Failed = nil + } + + LogCall(c, mustLookup("assign_key"), args, len(targets)) + + return result, nil + }, + }) +} diff --git a/server/internal/mcp/tools_write_test.go b/server/internal/mcp/tools_write_test.go new file mode 100644 index 0000000..18282da --- /dev/null +++ b/server/internal/mcp/tools_write_test.go @@ -0,0 +1,44 @@ +package mcp + +import "testing" + +func TestWriteToolsAreMarkedAsWrites(t *testing.T) { + want := []string{"run_workflow", "cancel_run", "apply_updates", "update_agent", "assign_key"} + for _, name := range want { + tool, ok := All().Lookup(name) + if !ok { + t.Errorf("tool %q is not registered", name) + continue + } + if !tool.Write { + t.Errorf("tool %q is not marked as a write, so it would be listed to a read-only agent", name) + } + } +} + +// A tool description is prompt text. A model choosing between tools must be +// told which ones touch real machines. +func TestWriteToolDescriptionsStateBlastRadius(t *testing.T) { + for _, tool := range All().Tools() { + if !tool.Write { + continue + } + if len(tool.Description) < 40 { + t.Errorf("tool %q has a %d-char description; write tools must state what they affect", + tool.Name, len(tool.Description)) + } + } +} + +// A caller holding every resource scope but not mcp:write must still see no +// write tools at all. +func TestWriteToolsHiddenWithoutMCPWrite(t *testing.T) { + c := Caller{Scopes: []string{ + "mcp:read", "servers:write", "workflows:write", "keys:write", + }} + for _, tool := range All().Visible(c) { + if tool.Write { + t.Errorf("write tool %q visible without mcp:write", tool.Name) + } + } +} diff --git a/server/internal/mcp/transport.go b/server/internal/mcp/transport.go index 575025b..5a5cdf7 100644 --- a/server/internal/mcp/transport.go +++ b/server/internal/mcp/transport.go @@ -80,7 +80,11 @@ func registerSDKTool(srv *sdk.Server, tool Tool, caller Caller) { if err != nil { return nil, nil, err } - LogCall(caller, tool, args, 0) + if !tool.Write { + // Write tools log their own call with a resolved server count, + // which this layer cannot know. + LogCall(caller, tool, args, 0) + } return nil, out, nil }) }