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.
This commit is contained in:
2026-08-16 20:30:07 -07:00
parent 3827d10e6e
commit 93c160ec51
18 changed files with 775 additions and 72 deletions
+97 -2
View File
@@ -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
}