Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 12 additions & 28 deletions go/internal/store/agent_activity.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ package store
import (
"context"
"fmt"

"github.com/RigelBuild/compass/go/internal/store/db"
)

// AgentActivity is an agent's durable activity: the free-text string it last
Expand All @@ -24,14 +26,11 @@ type AgentActivity struct {
// presence/activity split (DL-074) — presence lives in memory, but the activity
// survives a restart because it lands here.
func (s *Store) SetActivity(ctx context.Context, agentAccountID AccountID, activity string, atUnixMs int64) error {
if _, err := s.pool.Exec(ctx,
`INSERT INTO agent_activity (agent_account_id, activity, activity_at_unix_ms)
VALUES ($1, $2, $3)
ON CONFLICT (agent_account_id)
DO UPDATE SET activity = EXCLUDED.activity,
activity_at_unix_ms = EXCLUDED.activity_at_unix_ms`,
string(agentAccountID), activity, atUnixMs,
); err != nil {
if err := s.q.SetActivity(ctx, db.SetActivityParams{
AgentAccountID: string(agentAccountID),
Activity: activity,
ActivityAtUnixMs: atUnixMs,
}); err != nil {
return fmt.Errorf("store: set agent activity: %w", err)
}
return nil
Expand All @@ -54,30 +53,15 @@ func (s *Store) ActivityFor(ctx context.Context, accountIDs []AccountID) (map[Ac
ids[i] = string(id)
}

rows, err := s.pool.Query(ctx,
`SELECT agent_account_id, activity, activity_at_unix_ms
FROM agent_activity
WHERE agent_account_id = ANY($1)`,
ids,
)
rows, err := s.q.ActivityFor(ctx, ids)
if err != nil {
return nil, fmt.Errorf("store: read agent activity: %w", err)
}
defer rows.Close()

for rows.Next() {
var (
id string
activity string
atMs int64
)
if err := rows.Scan(&id, &activity, &atMs); err != nil {
return nil, fmt.Errorf("store: scan agent activity: %w", err)
for _, row := range rows {
out[AccountID(row.AgentAccountID)] = AgentActivity{
Activity: row.Activity,
ActivityAtUnixMs: row.ActivityAtUnixMs,
}
out[AccountID(id)] = AgentActivity{Activity: activity, ActivityAtUnixMs: atMs}
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("store: iterate agent activity: %w", err)
}
return out, nil
}
23 changes: 9 additions & 14 deletions go/internal/store/agent_config.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import (
"sort"
"strings"

"github.com/RigelBuild/compass/go/internal/store/db"
yaml "go.yaml.in/yaml/v3"
)

Expand Down Expand Up @@ -136,13 +137,10 @@ func (s *Store) PutAgentConfig(ctx context.Context, actor AccountID, bundle []by
if err != nil {
return "", err
}
if _, err := s.pool.Exec(ctx,
`INSERT INTO agent_config_bundle (singleton, version, bundle)
VALUES (TRUE, $1, $2)
ON CONFLICT (singleton)
DO UPDATE SET version = EXCLUDED.version, bundle = EXCLUDED.bundle, updated_at = now()`,
version, bundle,
); err != nil {
if err := s.q.PutAgentConfig(ctx, db.PutAgentConfigParams{
Version: version,
Bundle: bundle,
}); err != nil {
return "", fmt.Errorf("store: put agent config: %w", err)
}
return version, nil
Expand All @@ -163,15 +161,14 @@ func ValidateConfigBundle(bundle []byte) (version string, err error) {
// downstream (the fetch path then materializes an empty config dir), but the
// store still reports the absence; the caller decides empty-is-ok.
func (s *Store) CurrentAgentConfig(ctx context.Context) (version string, bundle []byte, err error) {
if err := s.pool.QueryRow(ctx,
`SELECT version, bundle FROM agent_config_bundle WHERE singleton = TRUE`,
).Scan(&version, &bundle); err != nil {
row, err := s.q.CurrentAgentConfig(ctx)
if err != nil {
if noRows(err) {
return "", nil, fmt.Errorf("%w: no agent config bundle declared", ErrNotFound)
}
return "", nil, fmt.Errorf("store: read agent config: %w", err)
}
return version, bundle, nil
return row.Version, row.Bundle, nil
}

// DeleteAgentConfig clears the fleet config bundle, returning the store to the
Expand All @@ -183,9 +180,7 @@ func (s *Store) CurrentAgentConfig(ctx context.Context) (version string, bundle
// return-to-unconfigured path (RIG-1625 T2), chosen over blessing an
// empty-tarball push.
func (s *Store) DeleteAgentConfig(ctx context.Context) error {
if _, err := s.pool.Exec(ctx,
`DELETE FROM agent_config_bundle WHERE singleton = TRUE`,
); err != nil {
if err := s.q.DeleteAgentConfig(ctx); err != nil {
return fmt.Errorf("store: delete agent config: %w", err)
}
return nil
Expand Down
65 changes: 21 additions & 44 deletions go/internal/store/agent_placements.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ package store
import (
"context"
"fmt"

"github.com/RigelBuild/compass/go/internal/store/db"
)

// Agent placement: the durable record of WHERE each agent runs — which Runner,
Expand Down Expand Up @@ -62,15 +64,11 @@ func (s *Store) RecordAgentPlacement(ctx context.Context, agentAccountID Account
if containerName == "" {
return fmt.Errorf("%w: container name is required", ErrInvalidArgument)
}
if _, err := s.pool.Exec(ctx,
`INSERT INTO agent_placements (agent_account_id, runner_id, container_name)
VALUES ($1, $2, $3)
ON CONFLICT (agent_account_id) DO UPDATE
SET runner_id = EXCLUDED.runner_id,
container_name = EXCLUDED.container_name,
updated_at = now()`,
string(agentAccountID), runnerID, containerName,
); err != nil {
if err := s.q.RecordAgentPlacement(ctx, db.RecordAgentPlacementParams{
AgentAccountID: string(agentAccountID),
RunnerID: runnerID,
ContainerName: containerName,
}); err != nil {
if pgErrIs(err, pgForeignKeyViolation) {
return fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, agentAccountID)
}
Expand All @@ -95,11 +93,8 @@ func (s *Store) AgentForContainer(ctx context.Context, containerName string) (Ac
if containerName == "" {
return "", fmt.Errorf("%w: container name is required", ErrInvalidArgument)
}
var accountID string
if err := s.pool.QueryRow(ctx,
`SELECT agent_account_id FROM agent_placements WHERE container_name = $1`,
containerName,
).Scan(&accountID); err != nil {
accountID, err := s.q.AgentForContainer(ctx, containerName)
if err != nil {
if noRows(err) {
return "", fmt.Errorf("%w: container %q is not placed", ErrNotFound, containerName)
}
Expand All @@ -119,30 +114,17 @@ func (s *Store) ListAgentPlacementsForRunner(ctx context.Context, runnerID strin
if runnerID == "" {
return nil, fmt.Errorf("%w: runner id is required", ErrInvalidArgument)
}
rows, err := s.pool.Query(ctx,
`SELECT agent_account_id, runner_id, container_name
FROM agent_placements
WHERE runner_id = $1
ORDER BY agent_account_id`,
runnerID,
)
rows, err := s.q.ListAgentPlacementsForRunner(ctx, runnerID)
if err != nil {
return nil, fmt.Errorf("store: list agent placements: %w", err)
}
defer rows.Close()

placements := []AgentPlacement{}
for rows.Next() {
var p AgentPlacement
var accountID string
if err := rows.Scan(&accountID, &p.RunnerID, &p.ContainerName); err != nil {
return nil, fmt.Errorf("store: scan agent placement: %w", err)
}
p.AgentAccountID = AccountID(accountID)
placements = append(placements, p)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("store: iterate agent placements: %w", err)
placements := make([]AgentPlacement, 0, len(rows))
for _, row := range rows {
placements = append(placements, AgentPlacement{
AgentAccountID: AccountID(row.AgentAccountID),
RunnerID: row.RunnerID,
ContainerName: row.ContainerName,
})
}
return placements, nil
}
Expand All @@ -158,10 +140,7 @@ func (s *Store) DeleteAgentPlacement(ctx context.Context, containerName string)
if containerName == "" {
return fmt.Errorf("%w: container name is required", ErrInvalidArgument)
}
if _, err := s.pool.Exec(ctx,
`DELETE FROM agent_placements WHERE container_name = $1`,
containerName,
); err != nil {
if err := s.q.DeleteAgentPlacement(ctx, containerName); err != nil {
return fmt.Errorf("store: delete agent placement: %w", err)
}
return nil
Expand All @@ -177,14 +156,12 @@ func (s *Store) PlacementForAgent(ctx context.Context, agentAccountID AccountID)
if agentAccountID == "" {
return "", "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument)
}
if err := s.pool.QueryRow(ctx,
`SELECT runner_id, container_name FROM agent_placements WHERE agent_account_id = $1`,
string(agentAccountID),
).Scan(&runnerID, &containerName); err != nil {
row, err := s.q.PlacementForAgent(ctx, string(agentAccountID))
if err != nil {
if noRows(err) {
return "", "", fmt.Errorf("%w: agent %q is not placed", ErrNotFound, agentAccountID)
}
return "", "", fmt.Errorf("store: resolve placement for agent: %w", err)
}
return runnerID, containerName, nil
return row.RunnerID, row.ContainerName, nil
}
37 changes: 14 additions & 23 deletions go/internal/store/agent_sessions.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ import (
"context"
"fmt"
"time"

"github.com/RigelBuild/compass/go/internal/store/db"
)

// The durable session-ownership chain: the persistent
Expand Down Expand Up @@ -41,10 +43,11 @@ func (s *Store) RecordAgentSession(ctx context.Context, sessionID string, agentA
if agentAccountID == "" {
return fmt.Errorf("%w: agent account id is required", ErrInvalidArgument)
}
if _, err := s.pool.Exec(ctx,
`INSERT INTO agent_sessions (session_id, agent_account_id, recorded_at_unix_ms) VALUES ($1, $2, $3)`,
sessionID, string(agentAccountID), time.Now().UnixMilli(),
); err != nil {
if err := s.q.InsertAgentSession(ctx, db.InsertAgentSessionParams{
SessionID: sessionID,
AgentAccountID: string(agentAccountID),
RecordedAtUnixMs: time.Now().UnixMilli(),
}); err != nil {
if pgErrIs(err, pgUniqueViolation) {
return fmt.Errorf("%w: session %q already recorded", ErrConflict, sessionID)
}
Expand All @@ -69,14 +72,8 @@ func (s *Store) LatestSessionForAccount(ctx context.Context, agent AccountID) (s
if agent == "" {
return "", false, fmt.Errorf("%w: agent account id is required", ErrInvalidArgument)
}
if err := s.pool.QueryRow(ctx,
`SELECT session_id
FROM agent_sessions
WHERE agent_account_id = $1
ORDER BY recorded_at_unix_ms DESC, session_id DESC
LIMIT 1`,
string(agent),
).Scan(&sessionID); err != nil {
sessionID, err = s.q.LatestSessionForAccount(ctx, string(agent))
if err != nil {
if noRows(err) {
return "", false, nil
}
Expand Down Expand Up @@ -105,17 +102,11 @@ func (s *Store) RequireAgentSessionSubscriber(ctx context.Context, caller Accoun
if sessionID == "" {
return fmt.Errorf("%w: session id is required", ErrInvalidArgument)
}
var authorized bool
if err := s.pool.QueryRow(ctx,
`SELECT EXISTS (
SELECT 1
FROM agent_sessions se
JOIN agent_accounts ag ON ag.account_id = se.agent_account_id
JOIN channel_members cm ON cm.channel_id = ag.home_channel_id
AND cm.account_id = $2
WHERE se.session_id = $1)`,
sessionID, string(caller),
).Scan(&authorized); err != nil {
authorized, err := s.q.RequireAgentSessionSubscriber(ctx, db.RequireAgentSessionSubscriberParams{
SessionID: sessionID,
AccountID: string(caller),
})
if err != nil {
return fmt.Errorf("store: authorize agent session subscriber: %w", err)
}
if !authorized {
Expand Down
Loading
Loading