Add agent restart lifecycle command

Extends the existing CheckIn RPC with a one-shot AgentCommand (restart only -- stop/uninstall need real per-platform OS service-manager integration and stay deliberately out of scope), delivered at-most-once: cleared the instant it's handed to the agent in a response, since a restarting agent's process is gone before it could ever confirm receipt. On restart, the agent flushes whatever's buffered, aborts its source task, and exits cleanly, relying entirely on the host's own service manager to bring it back up. Issuing a command is gated at RoleAdmin (stricter than config editing's RoleEditor) and logged into the same audit_log table Phase 7's AI interactions use, via a new agent_command event type. A real bug was found and fixed during live verification: the first implementation tried to atomically read-and-clear pending_command in a single INSERT...ON CONFLICT statement using a sibling CTE referenced only from RETURNING, on the assumption that Postgres evaluates every part of a WITH query against one pre-statement snapshot. That's wrong specifically for FOR UPDATE, which always reads the latest row version including one written earlier in the same statement -- confirmed empirically (a restart command was always coming back empty even when genuinely pending, so the agent never received it). Fixed by splitting into two real, ordered statements inside one explicit transaction. See /docs/agent-management-design.md's "Lifecycle commands" section.
2026-08-16 20:30:07 -07:00
parent 3827d10e6e
commit 93c160ec51
18 changed files with 775 additions and 72 deletions
@@ -18,28 +18,58 @@ type store interface {
 	Get(ctx context.Context, tenantID, host string) (*Agent, error)
 	SetOverride(ctx context.Context, tenantID, host string, override ConfigOverride, updatedBy string) (*Agent, error)
 	ClearOverride(ctx context.Context, tenantID, host string) error
+	IssueCommand(ctx context.Context, tenantID, host, command, issuedBy string) (*Agent, error)
+}
+
+// CommandLogger records an issued lifecycle command into the Phase 4
+// audit trail -- same nil-by-default, fail-open shape as
+// aiapi.InteractionLogger and queryapi.AuditLogger: a single-tenant
+// deployment with no enterprise/ configured just doesn't log these.
+// enterprise/internal/audit supplies the real implementation
+// (event_type = 'agent_command', see
+// metadata/migrations/0039_add_agent_command_event_type.sql) -- this is
+// the one entry point in this package genuinely worth logging even
+// without enterprise/ wired, given "strict RBAC, full audit trail" was
+// the explicit precondition for building lifecycle commands at all (see
+// /docs/agent-management-design.md); it degrades gracefully rather than
+// being required, matching every other optional audit hook in this
+// codebase, but a real deployment should wire it.
+type CommandLogger interface {
+	LogCommand(ctx context.Context, entry CommandLogEntry) error
+}
+
+type CommandLogEntry struct {
+	Host     string
+	Command  string
+	IssuedBy string
 }

 type Handler struct {
 	logger     *slog.Logger
 	store      store
 	authorizer authz.Authorizer
+	commands   CommandLogger
 }

-func NewHandler(logger *slog.Logger, store store, authorizer authz.Authorizer) *Handler {
-	return &Handler{logger: logger, store: store, authorizer: authorizer}
+// commands may be nil -- see CommandLogger's doc comment.
+func NewHandler(logger *slog.Logger, store store, authorizer authz.Authorizer, commands CommandLogger) *Handler {
+	return &Handler{logger: logger, store: store, authorizer: authorizer, commands: commands}
 }

 // RegisterRoutes: viewing inventory is RoleViewer (same bar as viewing
 // a dashboard); editing an agent's remote config is RoleEditor -- an
 // operational-tuning action, not an admin-only one, matching the RBAC
 // matrix's treatment of alert rules/notification targets rather than
-// user/role management.
+// user/role management. Issuing a lifecycle command is RoleAdmin --
+// stricter than config editing, matching the matrix's treatment of
+// similarly consequential actions (e.g. deleting a notification
+// target) rather than day-to-day tuning.
 func (h *Handler) RegisterRoutes(mux *http.ServeMux) {
 	mux.HandleFunc("GET /agents", authz.RequireRole(h.authorizer, authz.RoleViewer, h.handleList))
 	mux.HandleFunc("GET /agents/{host}", authz.RequireRole(h.authorizer, authz.RoleViewer, h.handleGet))
 	mux.HandleFunc("PUT /agents/{host}/config", authz.RequireRole(h.authorizer, authz.RoleEditor, h.handleSetConfig))
 	mux.HandleFunc("DELETE /agents/{host}/config", authz.RequireRole(h.authorizer, authz.RoleEditor, h.handleClearConfig))
+	mux.HandleFunc("PUT /agents/{host}/command", authz.RequireRole(h.authorizer, authz.RoleAdmin, h.handleIssueCommand))
 }

 // tenantID mirrors dashboards.Handler.tenantID exactly -- resolved from
@@ -106,6 +136,43 @@ func (h *Handler) handleSetConfig(w http.ResponseWriter, r *http.Request) {
 	writeJSON(w, http.StatusOK, a)
 }

+type issueCommandRequest struct {
+	Command string `json:"command"`
+}
+
+// handleIssueCommand queues a one-shot lifecycle command -- see
+// Store.IssueCommand's doc comment for the delivery/clearing semantics.
+// Logs to CommandLogger fail-open (a write failure is logged server-
+// side and otherwise ignored, same posture as aiapi's interaction
+// logging): audit-trail completeness matters, but it shouldn't be able
+// to turn a legitimate restart request into a 500.
+func (h *Handler) handleIssueCommand(w http.ResponseWriter, r *http.Request) {
+	var req issueCommandRequest
+	if !decodeJSON(w, r, &req) {
+		return
+	}
+	if !validCommand(req.Command) {
+		writeError(w, http.StatusBadRequest, `command must be "restart"`)
+		return
+	}
+
+	host := r.PathValue("host")
+	issuedBy := h.updatedBy(r)
+	a, err := h.store.IssueCommand(r.Context(), h.tenantID(r), host, req.Command, issuedBy)
+	if err != nil {
+		h.writeStoreErr(w, err, "issuing agent command")
+		return
+	}
+
+	if h.commands != nil {
+		if err := h.commands.LogCommand(r.Context(), CommandLogEntry{Host: host, Command: req.Command, IssuedBy: issuedBy}); err != nil {
+			h.logger.Error("logging agent command to audit trail", "host", host, "command", req.Command, "error", err)
+		}
+	}
+
+	writeJSON(w, http.StatusOK, a)
+}
+
 func (h *Handler) handleClearConfig(w http.ResponseWriter, r *http.Request) {
 	if err := h.store.ClearOverride(r.Context(), h.tenantID(r), r.PathValue("host")); err != nil {
 		h.writeStoreErr(w, err, "clearing agent config")
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"io"
 	"log/slog"
 	"net/http"
@@ -78,8 +79,33 @@ func (f *fakeStore) ClearOverride(_ context.Context, tenantID, host string) erro
 	return nil
 }

+func (f *fakeStore) IssueCommand(_ context.Context, tenantID, host, command, issuedBy string) (*Agent, error) {
+	a, ok := f.agents[tenantID+"/"+host]
+	if !ok {
+		return nil, ErrNotFound
+	}
+	a.PendingCommand = command
+	a.CommandIssuedBy = issuedBy
+	cp := *a
+	return &cp, nil
+}
+
+// fakeCommandLogger records LogCommand calls for assertions; nil-safe
+// callers should use a nil *fakeCommandLogger the same way production
+// code treats a nil CommandLogger, but tests that want to assert
+// logging happened construct a real one.
+type fakeCommandLogger struct {
+	entries []CommandLogEntry
+	err     error
+}
+
+func (f *fakeCommandLogger) LogCommand(_ context.Context, entry CommandLogEntry) error {
+	f.entries = append(f.entries, entry)
+	return f.err
+}
+
 func newTestHandler(s *fakeStore) *Handler {
-	return NewHandler(discardLogger(), s, nil)
+	return NewHandler(discardLogger(), s, nil, nil)
 }

 func doRequest(t *testing.T, h *Handler, method, path string, body any) *httptest.ResponseRecorder {
@@ -200,7 +226,7 @@ func TestRequireEditorRoleForConfigWrites(t *testing.T) {
 	s := newFakeStore()
 	s.put(Agent{TenantID: "default", Host: "web-01"})
 	authorizer := fakeAuthorizer{role: authz.RoleViewer}
-	h := NewHandler(discardLogger(), s, authorizer)
+	h := NewHandler(discardLogger(), s, authorizer, nil)

 	interval := int64(30000)
 	rec := doRequest(t, h, "PUT", "/agents/web-01/config", ConfigOverride{HeartbeatIntervalMS: &interval})
@@ -209,6 +235,75 @@ func TestRequireEditorRoleForConfigWrites(t *testing.T) {
 	}
 }

+func TestHandleIssueCommandRoundTrips(t *testing.T) {
+	s := newFakeStore()
+	s.put(Agent{TenantID: "default", Host: "web-01"})
+	logger := &fakeCommandLogger{}
+	h := NewHandler(discardLogger(), s, nil, logger)
+
+	rec := doRequest(t, h, "PUT", "/agents/web-01/command", map[string]string{"command": "restart"})
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
+	}
+	var got Agent
+	if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
+		t.Fatalf("decoding response: %v", err)
+	}
+	if got.PendingCommand != "restart" {
+		t.Fatalf("PendingCommand = %q, want restart", got.PendingCommand)
+	}
+	if len(logger.entries) != 1 || logger.entries[0].Command != "restart" || logger.entries[0].Host != "web-01" {
+		t.Fatalf("unexpected audit log entries: %+v", logger.entries)
+	}
+}
+
+func TestHandleIssueCommandRejectsUnknownCommand(t *testing.T) {
+	s := newFakeStore()
+	s.put(Agent{TenantID: "default", Host: "web-01"})
+	h := newTestHandler(s)
+
+	rec := doRequest(t, h, "PUT", "/agents/web-01/command", map[string]string{"command": "uninstall"})
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want 400 (uninstall is not a supported command yet)", rec.Code)
+	}
+}
+
+func TestHandleIssueCommandUnknownHostIsNotFound(t *testing.T) {
+	h := newTestHandler(newFakeStore())
+	rec := doRequest(t, h, "PUT", "/agents/nope/command", map[string]string{"command": "restart"})
+	if rec.Code != http.StatusNotFound {
+		t.Fatalf("status = %d, want 404", rec.Code)
+	}
+}
+
+// TestHandleIssueCommandFailOpenOnLoggerError is the regression test
+// for CommandLogger's documented fail-open posture: an audit-log write
+// failure must not turn a legitimate command issuance into an error
+// response.
+func TestHandleIssueCommandFailOpenOnLoggerError(t *testing.T) {
+	s := newFakeStore()
+	s.put(Agent{TenantID: "default", Host: "web-01"})
+	logger := &fakeCommandLogger{err: errors.New("audit db unreachable")}
+	h := NewHandler(discardLogger(), s, nil, logger)
+
+	rec := doRequest(t, h, "PUT", "/agents/web-01/command", map[string]string{"command": "restart"})
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want 200 even though the audit logger failed", rec.Code)
+	}
+}
+
+func TestRequireAdminRoleForCommands(t *testing.T) {
+	s := newFakeStore()
+	s.put(Agent{TenantID: "default", Host: "web-01"})
+	authorizer := fakeAuthorizer{role: authz.RoleEditor}
+	h := NewHandler(discardLogger(), s, authorizer, nil)
+
+	rec := doRequest(t, h, "PUT", "/agents/web-01/command", map[string]string{"command": "restart"})
+	if rec.Code != http.StatusForbidden {
+		t.Fatalf("status = %d, want 403 (Editor must not be able to issue lifecycle commands, only Admin+)", rec.Code)
+	}
+}
+
 type fakeAuthorizer struct {
 	role authz.Role
 }
@@ -17,6 +17,16 @@ import (

 var ErrNotFound = errors.New("not found")

+// CommandRestart is the one supported lifecycle command -- see
+// ingest/internal/grpcserver.AgentCommandRestart and
+// agent_control.proto's AgentCommand enum comment for why STOP/
+// UNINSTALL aren't here yet.
+const CommandRestart = "restart"
+
+func validCommand(c string) bool {
+	return c == CommandRestart
+}
+
 // ConfigOverride is the remotely-editable subset of an agent's config --
 // a plain-Go mirror of ingest/internal/agentregistry's overrideFields
 // and agent_control.proto's DesiredOverride. Deliberately duplicated
@@ -57,6 +67,17 @@ type Agent struct {
 	// recomputing the same string comparison itself.
 	Pending   bool   `json:"pending"`
 	UpdatedBy string `json:"updated_by,omitempty"`
+	// PendingCommand is "" when nothing is queued, or CommandRestart
+	// while a restart hasn't yet been delivered to the agent. Unlike
+	// Pending (config), there's no way to observe "delivered" from this
+	// table alone -- ingest clears pending_command the instant it hands
+	// the command out (see ingest/internal/agentregistry.Registry.
+	// CheckIn), so PendingCommand flipping back to "" just as plausibly
+	// means "delivered a moment ago" as "never issued." CommandIssuedAt
+	// is what the web UI shows instead, as a last-issued record.
+	PendingCommand  string     `json:"pending_command,omitempty"`
+	CommandIssuedAt *time.Time `json:"command_issued_at,omitempty"`
+	CommandIssuedBy string     `json:"command_issued_by,omitempty"`
 }

 type Store struct {
@@ -73,7 +94,8 @@ const selectColumns = `
 	reported_batch_max_size, reported_batch_flush_ms,
 	reported_heartbeat_on, reported_heartbeat_ms,
 	first_seen_at, last_seen_at,
-	desired_override, desired_override_version, applied_override_version, updated_by`
+	desired_override, desired_override_version, applied_override_version, updated_by,
+	pending_command, command_issued_at, command_issued_by`

 func (s *Store) List(ctx context.Context, tenantID string) ([]Agent, error) {
 	rows, err := s.pool.Query(ctx, `
@@ -145,6 +167,27 @@ func (s *Store) SetOverride(ctx context.Context, tenantID, host string, override
 	return s.Get(ctx, tenantID, host)
 }

+// IssueCommand queues a one-shot lifecycle command for host, delivered
+// on its next CheckIn and cleared atomically by ingest the instant
+// that happens (see ingest/internal/agentregistry.Registry.CheckIn) --
+// unlike SetOverride, there's no "applied" confirmation to wait for,
+// since a restarting agent's process is gone before it could send one.
+// command_issued_at/by are overwritten on every call, forming a
+// last-issued record even after pending_command itself clears.
+func (s *Store) IssueCommand(ctx context.Context, tenantID, host, command, issuedBy string) (*Agent, error) {
+	tag, err := s.pool.Exec(ctx, `
+		UPDATE agents SET pending_command = $1, command_issued_at = now(), command_issued_by = $2
+		WHERE tenant_id = $3 AND host = $4`,
+		command, issuedBy, tenantID, host)
+	if err != nil {
+		return nil, err
+	}
+	if tag.RowsAffected() == 0 {
+		return nil, ErrNotFound
+	}
+	return s.Get(ctx, tenantID, host)
+}
+
 // ClearOverride reverts an agent to running its local agent.toml
 // untouched -- the next CheckIn gets has_override=false.
 func (s *Store) ClearOverride(ctx context.Context, tenantID, host string) error {
@@ -167,7 +210,7 @@ type rowScanner interface {
 func scanAgent(row rowScanner) (Agent, error) {
 	var a Agent
 	var desiredOverride []byte
-	var desiredVersion, updatedBy *string
+	var desiredVersion, updatedBy, pendingCommand, commandIssuedBy *string
 	if err := row.Scan(
 		&a.ID, &a.TenantID, &a.Host, &a.Service,
 		&a.AgentVersion, &a.SourceKind, &a.SourceDetail,
@@ -175,6 +218,7 @@ func scanAgent(row rowScanner) (Agent, error) {
 		&a.HeartbeatEnabled, &a.HeartbeatIntervalMS,
 		&a.FirstSeenAt, &a.LastSeenAt,
 		&desiredOverride, &desiredVersion, &a.AppliedOverrideVersion, &updatedBy,
+		&pendingCommand, &a.CommandIssuedAt, &commandIssuedBy,
 	); err != nil {
 		if errors.Is(err, pgx.ErrNoRows) {
 			return Agent{}, ErrNotFound
@@ -184,6 +228,12 @@ func scanAgent(row rowScanner) (Agent, error) {
 	if updatedBy != nil {
 		a.UpdatedBy = *updatedBy
 	}
+	if pendingCommand != nil {
+		a.PendingCommand = *pendingCommand
+	}
+	if commandIssuedBy != nil {
+		a.CommandIssuedBy = *commandIssuedBy
+	}
 	if desiredVersion != nil {
 		a.DesiredOverrideVersion = *desiredVersion
 		a.Pending = *desiredVersion != a.AppliedOverrideVersion