feat: add mcp write tools with a fan-out guard
This commit is contained in:
@@ -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
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user