coffey-labs/cairnobs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add agent restart lifecycle command

Browse Source 
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:
jcoffey-dev
2026-08-16 20:30:07 -07:00
parent 3827d10e6e
commit 93c160ec51
18 changed files with 775 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ingest/internal/agentregistry/agentregistry.go
+52 -8
View File
@@ -10,9 +10,11 @@ package agentregistry
import (
"context"
"encoding/json"
"errors"
"fmt"
"github.com/google/uuid"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/sentry/sentry/ingest/internal/grpcserver"
@@ -49,16 +51,47 @@ func New(pool *pgxpool.Pool) *Registry {
var _ grpcserver.AgentRegistry = (*Registry)(nil)
func (r *Registry) CheckIn(ctx context.Context, tenantID string, info grpcserver.AgentCheckIn) (grpcserver.AgentOverride, error) {
func (r *Registry) CheckIn(ctx context.Context, tenantID string, info grpcserver.AgentCheckIn) (grpcserver.CheckInResult, error) {
if tenantID == "" {
tenantID = defaultTenantID
}
tx, err := r.pool.Begin(ctx)
if err != nil {
return grpcserver.CheckInResult{}, fmt.Errorf("agentregistry: beginning transaction: %w", err)
}
defer tx.Rollback(ctx)
// Read (and lock) whatever pending_command exists BEFORE the
// upsert below clears it -- as two real, ordered statements in one
// transaction, not one clever statement. An earlier version tried
// to do this with a single INSERT...ON CONFLICT plus a sibling CTE
// referenced only from RETURNING, on the assumption that Postgres
// evaluates every part of a WITH query against the same pre-
// statement snapshot; that assumption is wrong specifically for
// FOR UPDATE, which always locks (and therefore reads) the latest
// row version to do its job, including versions written earlier in
// the SAME statement -- confirmed empirically against a live
// Postgres (the CTE's FOR UPDATE was reading its own sibling
// UPDATE's just-cleared NULL, always reporting "no command" even
// when one was genuinely pending). Two statements in one
// transaction has no such ambiguity: the SELECT strictly
// happens-before the UPDATE, full stop. FOR UPDATE against zero
// rows (a brand-new agent's first-ever check-in) is a harmless
// no-op -- there's nothing to lock or have a pending command yet.
var pendingCommand *string
err = tx.QueryRow(ctx, `SELECT pending_command FROM agents WHERE tenant_id = $1 AND host = $2 FOR UPDATE`,
tenantID, info.Host,
).Scan(&pendingCommand)
if err != nil && !errors.Is(err, pgx.ErrNoRows) {
return grpcserver.CheckInResult{}, fmt.Errorf("agentregistry: reading pending command: %w", err)
}
var (
desiredOverride []byte
desiredVersion *string
)
err := r.pool.QueryRow(ctx, `
err = tx.QueryRow(ctx, `
INSERT INTO agents (
id, tenant_id, host, service,
reported_agent_version, reported_source_kind, reported_source_detail,
@@ -76,7 +109,8 @@ func (r *Registry) CheckIn(ctx context.Context, tenantID string, info grpcserver
reported_heartbeat_on = EXCLUDED.reported_heartbeat_on,
reported_heartbeat_ms = EXCLUDED.reported_heartbeat_ms,
last_seen_at = now(),
applied_override_version = EXCLUDED.applied_override_version
applied_override_version = EXCLUDED.applied_override_version,
pending_command = NULL
RETURNING desired_override, desired_override_version`,
uuid.NewString(), tenantID, info.Host, info.Service,
info.AgentVersion, info.SourceKind, info.SourceDetail,
@@ -85,18 +119,27 @@ func (r *Registry) CheckIn(ctx context.Context, tenantID string, info grpcserver
info.AppliedOverrideVersion,
).Scan(&desiredOverride, &desiredVersion)
if err != nil {
return grpcserver.AgentOverride{}, fmt.Errorf("agentregistry: upserting check-in: %w", err)
return grpcserver.CheckInResult{}, fmt.Errorf("agentregistry: upserting check-in: %w", err)
}
if err := tx.Commit(ctx); err != nil {
return grpcserver.CheckInResult{}, fmt.Errorf("agentregistry: committing check-in: %w", err)
}
result := grpcserver.CheckInResult{}
if pendingCommand != nil {
result.Command = *pendingCommand
}
if desiredVersion == nil || len(desiredOverride) == 0 {
return grpcserver.AgentOverride{HasOverride: false}, nil
return result, nil
}
var fields overrideFields
if err := json.Unmarshal(desiredOverride, &fields); err != nil {
return grpcserver.AgentOverride{}, fmt.Errorf("agentregistry: parsing stored override: %w", err)
return grpcserver.CheckInResult{}, fmt.Errorf("agentregistry: parsing stored override: %w", err)
}
return grpcserver.AgentOverride{
result.Override = grpcserver.AgentOverride{
HasOverride: true,
BatchMaxSize: fields.BatchMaxSize,
BatchFlushIntervalMS: fields.BatchFlushIntervalMS,
@@ -104,5 +147,6 @@ func (r *Registry) CheckIn(ctx context.Context, tenantID string, info grpcserver
HeartbeatIntervalMS: fields.HeartbeatIntervalMS,
JournaldUnit: fields.JournaldUnit,
Version: *desiredVersion,
}, nil
}
return result, nil
}
ingest/internal/grpcserver/server.go
+45 -15
View File
@@ -91,9 +91,32 @@ type TenantResolver interface {
// simply doesn't get agent inventory/management, exactly like one
// without ENTERPRISE_AUTH_URL doesn't get tenant-tagged records.
type AgentRegistry interface {
CheckIn(ctx context.Context, tenantID string, info AgentCheckIn) (AgentOverride, error)
CheckIn(ctx context.Context, tenantID string, info AgentCheckIn) (CheckInResult, error)
}
// CheckInResult bundles the two independent things a CheckIn can hand
// back to an agent -- a persistent config override to converge to, and
// a one-shot command to act on immediately. Kept as two separate
// concepts (not folded into one "override" shape) since their delivery
// semantics differ: Override is re-offered every CheckIn until the
// agent's applied_override_version matches; Command is cleared the
// instant it's handed out (see agent_control.proto's CheckInResponse
// comment).
type CheckInResult struct {
Override AgentOverride
// Command is AgentCommandRestart or "" (nothing pending). A plain
// string, not the generated proto enum type, so AgentRegistry
// implementations don't need to import the gRPC-facing package --
// same reasoning as AgentCheckIn/AgentOverride below.
Command string
}
// AgentCommandRestart is the one supported value for
// CheckInResult.Command / the agents table's pending_command column --
// see agent_control.proto's AgentCommand enum comment for why STOP/
// UNINSTALL aren't here yet.
const AgentCommandRestart = "restart"
// AgentCheckIn is what an agent reports about itself on each CheckIn --
// a plain-Go mirror of agentv1.ReportedConfig plus the identity/
// tenant fields, kept separate from the proto type so AgentRegistry
@@ -250,7 +273,7 @@ func (s *Server) CheckIn(ctx context.Context, req *agentv1.CheckInRequest) (*age
}
cfg := req.GetCurrentConfig()
override, err := s.agents.CheckIn(ctx, tenantID, AgentCheckIn{
result, err := s.agents.CheckIn(ctx, tenantID, AgentCheckIn{
Host: req.GetHost(),
Service: req.GetService(),
AgentVersion: cfg.GetAgentVersion(),
@@ -267,20 +290,27 @@ func (s *Server) CheckIn(ctx context.Context, req *agentv1.CheckInRequest) (*age
return nil, status.Errorf(codes.Internal, "recording check-in: %v", err)
}
if !override.HasOverride {
return &agentv1.CheckInResponse{HasOverride: false}, nil
resp := &agentv1.CheckInResponse{HasOverride: result.Override.HasOverride}
if result.Override.HasOverride {
resp.Override = &agentv1.DesiredOverride{
BatchMaxSize: result.Override.BatchMaxSize,
BatchFlushIntervalMs: result.Override.BatchFlushIntervalMS,
HeartbeatEnabled: result.Override.HeartbeatEnabled,
HeartbeatIntervalMs: result.Override.HeartbeatIntervalMS,
JournaldUnit: result.Override.JournaldUnit,
Version: result.Override.Version,
}
}
return &agentv1.CheckInResponse{
HasOverride: true,
Override: &agentv1.DesiredOverride{
BatchMaxSize: override.BatchMaxSize,
BatchFlushIntervalMs: override.BatchFlushIntervalMS,
HeartbeatEnabled: override.HeartbeatEnabled,
HeartbeatIntervalMs: override.HeartbeatIntervalMS,
JournaldUnit: override.JournaldUnit,
Version: override.Version,
},
}, nil
switch result.Command {
case AgentCommandRestart:
resp.PendingCommand = agentv1.AgentCommand_AGENT_COMMAND_RESTART
s.logger.Info("delivering restart command to agent", "host", req.GetHost())
case "":
// nothing pending
default:
s.logger.Error("agent registry returned an unknown command, ignoring", "host", req.GetHost(), "command", result.Command)
}
return resp, nil
}
// bearerTokenFromContext reads the same "authorization: Bearer <token>"
ingest/internal/grpcserver/server_test.go
+34 -5
View File
@@ -59,20 +59,23 @@ type fakeAgentRegistry struct {
mu sync.Mutex
checkIns []AgentCheckIn
overrides map[string]AgentOverride
commands map[string]string
err error
}
func (f *fakeAgentRegistry) CheckIn(_ context.Context, tenantID string, info AgentCheckIn) (AgentOverride, error) {
func (f *fakeAgentRegistry) CheckIn(_ context.Context, tenantID string, info AgentCheckIn) (CheckInResult, error) {
f.mu.Lock()
defer f.mu.Unlock()
if f.err != nil {
return AgentOverride{}, f.err
return CheckInResult{}, f.err
}
f.checkIns = append(f.checkIns, info)
if f.overrides == nil {
return AgentOverride{HasOverride: false}, nil
key := tenantID + "/" + info.Host
result := CheckInResult{Command: f.commands[key]}
if f.overrides != nil {
result.Override = f.overrides[key]
}
return f.overrides[tenantID+"/"+info.Host], nil
return result, nil
}
func newTestServer(p batchProducer) *Server {
@@ -374,3 +377,29 @@ func TestCheckInWithResolverRejectsMissingToken(t *testing.T) {
t.Fatalf("CheckIn() error = %v, want Unauthenticated", err)
}
}
func TestCheckInDeliversPendingRestartCommand(t *testing.T) {
reg := &fakeAgentRegistry{commands: map[string]string{"/web-01": AgentCommandRestart}}
s := newTestServerWithAgents(reg)
resp, err := s.CheckIn(context.Background(), &agentv1.CheckInRequest{Host: "web-01", CurrentConfig: &agentv1.ReportedConfig{}})
if err != nil {
t.Fatalf("CheckIn() error = %v", err)
}
if resp.GetPendingCommand() != agentv1.AgentCommand_AGENT_COMMAND_RESTART {
t.Fatalf("PendingCommand = %v, want AGENT_COMMAND_RESTART", resp.GetPendingCommand())
}
}
func TestCheckInNoCommandReportsUnspecified(t *testing.T) {
reg := &fakeAgentRegistry{}
s := newTestServerWithAgents(reg)
resp, err := s.CheckIn(context.Background(), &agentv1.CheckInRequest{Host: "web-01", CurrentConfig: &agentv1.ReportedConfig{}})
if err != nil {
t.Fatalf("CheckIn() error = %v", err)
}
if resp.GetPendingCommand() != agentv1.AgentCommand_AGENT_COMMAND_UNSPECIFIED {
t.Fatalf("PendingCommand = %v, want AGENT_COMMAND_UNSPECIFIED", resp.GetPendingCommand())
}
}