diff --git a/server/internal/services/dispatch.go b/server/internal/services/dispatch.go index 2e0d268..332f0b5 100644 --- a/server/internal/services/dispatch.go +++ b/server/internal/services/dispatch.go @@ -62,6 +62,12 @@ func (d *commandDispatcher) dispatch(serverID string, cmd *pb.ServerCommand) err } } +// DispatchRunStep pushes a RunStepCmd to a server's agent. Caller must have +// registered StepResults.Await(commandID) first. +func DispatchRunStep(serverID, commandID string, cmd *pb.RunStepCmd) error { + return Dispatcher.dispatch(serverID, &pb.ServerCommand{CommandId: commandID, RunStep: cmd}) +} + // KeyGenParams carries all options for a generate-key command. type KeyGenParams struct { Label string diff --git a/server/internal/services/workflow_runner.go b/server/internal/services/workflow_runner.go new file mode 100644 index 0000000..de54128 --- /dev/null +++ b/server/internal/services/workflow_runner.go @@ -0,0 +1,391 @@ +package services + +import ( + "fmt" + "strings" + "time" + + "github.com/google/uuid" + "github.com/mrhid6/vantage/server/internal/db" + "github.com/mrhid6/vantage/server/internal/grpc/pb" + "github.com/mrhid6/vantage/server/internal/models" + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo" + "go.mongodb.org/mongo-driver/v2/mongo/options" +) + +const stepDispatchGrace = 15 * time.Second + +// TriggerWorkflow snapshots the workflow, creates a run doc, and starts a +// background goroutine per target server (parallel fan-out). Returns run_id. +func TriggerWorkflow(workflowID, actor string) (string, error) { + wf, err := GetWorkflow(workflowID) + if err != nil { + return "", err + } + if len(wf.TargetServerIDs) == 0 { + return "", fmt.Errorf("workflow has no target servers") + } + if len(wf.Steps) == 0 { + return "", fmt.Errorf("workflow has no steps") + } + + // Reject a concurrent run of the same workflow. + ctx, cancel := wfCtx() + running := db.Col("workflow_runs").FindOne(ctx, bson.M{"workflow_id": workflowID, "status": "running"}) + cancel() + if running.Err() == nil { + return "", fmt.Errorf("workflow already has a run in progress") + } + + resolved, err := resolveSteps(wf) + if err != nil { + return "", err + } + + run := models.WorkflowRun{ + RunID: uuid.New().String(), + WorkflowID: workflowID, + Name: wf.Name, + Steps: resolved, + Status: "running", + TriggeredBy: actor, + StartedAt: time.Now(), + ServerRuns: make([]models.ServerRun, 0, len(wf.TargetServerIDs)), + } + for _, sid := range wf.TargetServerIDs { + hostname := sid + if s, e := GetServer(sid); e == nil { + hostname = s.Hostname + } + sr := models.ServerRun{ServerID: sid, Hostname: hostname, Status: "queued", RunEnv: map[string]string{}} + for _, rs := range resolved { + sr.Steps = append(sr.Steps, models.StepRun{Order: rs.Order, Name: rs.Name, Status: "queued", OutputEnv: map[string]string{}}) + } + run.ServerRuns = append(run.ServerRuns, sr) + } + + ictx, icancel := wfCtx() + defer icancel() + if _, err := db.Col("workflow_runs").InsertOne(ictx, run); err != nil { + return "", err + } + + go executeRun(run.RunID) + return run.RunID, nil +} + +// resolveSteps freezes each workflow step ref into a ResolvedStep by loading the +// library step and applying overrides. +func resolveSteps(wf *models.Workflow) ([]models.ResolvedStep, error) { + ctx, cancel := wfCtx() + defer cancel() + out := make([]models.ResolvedStep, 0, len(wf.Steps)) + for _, ref := range wf.Steps { + lib, err := getStep(ctx, ref.StepID) + if err != nil { + return nil, err + } + rs := models.ResolvedStep{ + Order: ref.Order, + Name: lib.Name, + Interpreter: lib.Interpreter, + Script: lib.Script, + SecretRefs: lib.SecretRefs, + OnFailure: ref.OnFailure, + MaxRetries: ref.MaxRetries, + } + if ref.Overrides != nil { + if ref.Overrides.Script != nil { + rs.Script = *ref.Overrides.Script + } + if ref.Overrides.SecretRefs != nil { + rs.SecretRefs = ref.Overrides.SecretRefs + } + } + if rs.OnFailure == "" { + rs.OnFailure = "stop" + } + out = append(out, rs) + } + return out, nil +} + +// executeRun fans out one goroutine per server run and waits for all to finish. +func executeRun(runID string) { + run, err := GetRun(runID) + if err != nil { + return + } + done := make(chan int, len(run.ServerRuns)) + for i := range run.ServerRuns { + go func(idx int) { + runServer(runID, idx, run.Steps, run.ServerRuns[idx].ServerID) + done <- idx + }(i) + } + for range run.ServerRuns { + <-done + } + + // Aggregate status. + final, _ := GetRun(runID) + status := "success" + for _, sr := range final.ServerRuns { + if sr.Status == "failed" { + status = "failed" + } + } + now := time.Now() + ctx, cancel := wfCtx() + defer cancel() + _, _ = db.Col("workflow_runs").UpdateOne(ctx, bson.M{"run_id": runID}, + bson.M{"$set": bson.M{"status": status, "finished_at": now}}) +} + +// runServer executes the resolved steps sequentially on one server, threading +// output env forward and applying per-step failure policy. +func runServer(runID string, srvIdx int, steps []models.ResolvedStep, serverID string) { + now := time.Now() + setServerRun(runID, srvIdx, bson.M{"server_runs.$.status": "running", "server_runs.$.started_at": now}) + + if !Dispatcher.IsConnected(serverID) { + fin := time.Now() + setServerRun(runID, srvIdx, bson.M{"server_runs.$.status": "skipped", "server_runs.$.finished_at": fin}) + return + } + + runEnv := map[string]string{} + serverFailed := false + + for i, step := range steps { + startStep(runID, serverID, i, "running") + var res *pb.StepResult + attempts := 0 + maxAttempts := 1 + if step.OnFailure == "retry" { + maxAttempts = step.MaxRetries + 1 + } + + // Merge secrets into command env (kept out of persisted logs). + secretVals := resolveSecrets(step.SecretRefs) + cmdEnv := map[string]string{} + for k, v := range runEnv { + cmdEnv[k] = v + } + for k, v := range secretVals { + cmdEnv[k] = v + } + + for attempts < maxAttempts { + attempts++ + res = dispatchAndWait(serverID, &pb.RunStepCmd{ + Interpreter: step.Interpreter, + Script: step.Script, + Env: cmdEnv, + TimeoutSeconds: 0, + }) + if res != nil && res.ExitCode == 0 { + break + } + } + + // Mask secret values before persisting. + stdout, stderr := "", "" + exit := 1 + outEnv := map[string]string{} + if res != nil { + stdout = maskSecrets(res.Stdout, secretVals) + stderr = maskSecrets(res.Stderr, secretVals) + exit = res.ExitCode + for k, v := range res.OutputEnv { + outEnv[k] = v + runEnv[k] = v // implicit: all outputs flow to all later steps + } + } else { + stderr = "[vantage] agent did not return a result" + } + + status := "success" + if exit != 0 { + status = "failed" + } + finishStep(runID, serverID, i, status, attempts, exit, stdout, stderr, outEnv) + + if exit != 0 { + switch step.OnFailure { + case "continue": + // keep going + default: // "stop" or exhausted "retry" + serverFailed = true + } + if serverFailed { + markRemainingSkipped(runID, serverID, i+1) + break + } + } + } + + fin := time.Now() + status := "success" + if serverFailed { + status = "failed" + } + setServerRun(runID, srvIdx, bson.M{ + "server_runs.$.status": status, + "server_runs.$.finished_at": fin, + "server_runs.$.run_env": runEnv, + }) +} + +// dispatchAndWait registers a waiter, dispatches the step, and blocks for the +// result or a timeout. +func dispatchAndWait(serverID string, cmd *pb.RunStepCmd) *pb.StepResult { + commandID := uuid.New().String() + ch := StepResults.Await(commandID) + if err := DispatchRunStep(serverID, commandID, cmd); err != nil { + StepResults.Cancel(commandID) + return &pb.StepResult{ExitCode: 1, Stderr: "[vantage] dispatch failed: " + err.Error()} + } + wait := time.Duration(cmd.TimeoutSeconds)*time.Second + stepDispatchGrace + if cmd.TimeoutSeconds == 0 { + wait = 30*time.Minute + stepDispatchGrace + } + select { + case res := <-ch: + return res + case <-time.After(wait): + StepResults.Cancel(commandID) + return &pb.StepResult{ExitCode: 124, Stderr: "[vantage] timed out waiting for agent result"} + } +} + +func resolveSecrets(refs []string) map[string]string { + out := map[string]string{} + for _, ref := range refs { + // ref format "group/KEY"; resolve via RevealSecret. + parts := strings.SplitN(ref, "/", 2) + if len(parts) != 2 { + continue + } + if v, err := RevealSecret(parts[0], parts[1]); err == nil { + out[parts[1]] = v + } + } + return out +} + +func maskSecrets(s string, secrets map[string]string) string { + for _, v := range secrets { + if v == "" { + continue + } + s = strings.ReplaceAll(s, v, "***") + } + return s +} + +// ---- run doc mutation helpers ---- + +func setServerRun(runID string, srvIdx int, set bson.M) { + ctx, cancel := wfCtx() + defer cancel() + _, _ = db.Col("workflow_runs").UpdateOne(ctx, + bson.M{"run_id": runID, "server_runs.server_id": serverIDAt(runID, srvIdx)}, + bson.M{"$set": set}) +} + +// serverIDAt returns the server_id at an index (positional operator needs a match). +func serverIDAt(runID string, srvIdx int) string { + r, err := GetRun(runID) + if err != nil || srvIdx >= len(r.ServerRuns) { + return "" + } + return r.ServerRuns[srvIdx].ServerID +} + +func startStep(runID, serverID string, order int, status string) { + now := time.Now() + updateStep(runID, serverID, order, bson.M{ + "server_runs.$[s].steps.$[t].status": status, + "server_runs.$[s].steps.$[t].started_at": now, + }) +} + +func finishStep(runID, serverID string, order int, status string, attempts, exit int, stdout, stderr string, outEnv map[string]string) { + now := time.Now() + updateStep(runID, serverID, order, bson.M{ + "server_runs.$[s].steps.$[t].status": status, + "server_runs.$[s].steps.$[t].attempts": attempts, + "server_runs.$[s].steps.$[t].exit_code": exit, + "server_runs.$[s].steps.$[t].stdout": stdout, + "server_runs.$[s].steps.$[t].stderr": stderr, + "server_runs.$[s].steps.$[t].output_env": outEnv, + "server_runs.$[s].steps.$[t].finished_at": now, + }) +} + +func markRemainingSkipped(runID, serverID string, fromOrder int) { + ctx, cancel := wfCtx() + defer cancel() + _, _ = db.Col("workflow_runs").UpdateMany(ctx, + bson.M{"run_id": runID}, + bson.M{"$set": bson.M{"server_runs.$[s].steps.$[t].status": "skipped"}}, + options.UpdateMany().SetArrayFilters([]interface{}{ + bson.M{"s.server_id": serverID}, + bson.M{"t.order": bson.M{"$gte": fromOrder}, "t.status": "queued"}, + }), + ) +} + +func updateStep(runID, serverID string, order int, set bson.M) { + ctx, cancel := wfCtx() + defer cancel() + _, _ = db.Col("workflow_runs").UpdateOne(ctx, + bson.M{"run_id": runID}, + bson.M{"$set": set}, + options.UpdateOne().SetArrayFilters([]interface{}{ + bson.M{"s.server_id": serverID}, + bson.M{"t.order": order}, + }), + ) +} + +// ---- reads ---- + +func GetRun(runID string) (*models.WorkflowRun, error) { + ctx, cancel := wfCtx() + defer cancel() + var r models.WorkflowRun + err := db.Col("workflow_runs").FindOne(ctx, bson.M{"run_id": runID}).Decode(&r) + if err == mongo.ErrNoDocuments { + return nil, fmt.Errorf("run not found") + } + return &r, err +} + +func ListRuns(workflowID string, limit int64) ([]models.WorkflowRun, error) { + ctx, cancel := wfCtx() + defer cancel() + cur, err := db.Col("workflow_runs").Find(ctx, bson.M{"workflow_id": workflowID}, + options.Find().SetSort(bson.D{{Key: "started_at", Value: -1}}).SetLimit(limit)) + if err != nil { + return nil, err + } + defer cur.Close(ctx) + runs := []models.WorkflowRun{} + if err := cur.All(ctx, &runs); err != nil { + return nil, err + } + return runs, nil +} + +func CancelRun(runID string) error { + now := time.Now() + ctx, cancel := wfCtx() + defer cancel() + _, err := db.Col("workflow_runs").UpdateOne(ctx, + bson.M{"run_id": runID, "status": "running"}, + bson.M{"$set": bson.M{"status": "cancelled", "finished_at": now}}) + return err +}