diff --git a/go/internal/runnerhub/commands.go b/go/internal/runnerhub/commands.go index a0ea9e09..44f39614 100644 --- a/go/internal/runnerhub/commands.go +++ b/go/internal/runnerhub/commands.go @@ -25,9 +25,15 @@ import ( ) // Provision relays a ProvisionAgentWorkspace command to the owning Runner and -// returns the container name it created. requestID is the client_request_id -// idempotency key: a timeout-retry with the same id returns the same container -// (no duplicate); empty mints a fresh id (no dedup). +// returns the container name it created plus the id of the Runner that served +// the call. The Runner id is returned rather than looked up afterwards so the +// caller records the agent's durable placement against the Runner that ACTUALLY +// ran the provision — re-reading the registry after the round trip could name a +// different Runner if one re-enrolled in the meantime, and a placement pointing +// at the wrong Runner is worse than none (reattach would re-drive the wrong +// set). requestID is the client_request_id idempotency key: a timeout-retry with +// the same id returns the same container (no duplicate); empty mints a fresh id +// (no dedup). // // The client_request_id alone must NOT be the dedup key: it is a client-chosen // string, and two provisions that reuse one value for DIFFERENT workspaces @@ -39,26 +45,28 @@ import ( // This mirrors the comms store's (author_account_id, client_request_id) // idempotency scoping (store/migrations/0001_init.sql), strengthened to the full // provision identity since a provision creates a real isolated container. -func (h *Hub) Provision(ctx context.Context, requestID string, req *compassv1.ProvisionAgentWorkspaceRequest) (*compassv1.ProvisionAgentWorkspaceResponse, error) { - result, err := h.relay(ctx, "", &compassv1internal.SessionsResponse{ +func (h *Hub) Provision(ctx context.Context, requestID string, req *compassv1.ProvisionAgentWorkspaceRequest) (*compassv1.ProvisionAgentWorkspaceResponse, string, error) { + result, runnerID, err := h.relay(ctx, "", &compassv1internal.SessionsResponse{ RequestId: provisionDedupID(requestID, req), Command: &compassv1internal.SessionsResponse_Provision{Provision: req}, }) if err != nil { - return nil, err + return nil, "", err } resp := result.GetProvision() // Record which agent account this container was provisioned for, so a later // Start can promote it to a session binding RelayCommsCall resolves against // (comms-tools design T2). The Runner never asserts this account; it is the - // Server's own record, keyed by the container name the Runner returned. + // Server's own record, keyed by the container name the Runner returned. This + // binding is the LIVE comms binding only, cleared on re-enroll; the DURABLE + // container/Runner placement is the caller's store write. h.bindContainer(resp.GetContainerName(), store.AccountID(req.GetAgentAccountId())) - return resp, nil + return resp, runnerID, nil } // Start relays a StartAgentSession command to the owning Runner. func (h *Hub) Start(ctx context.Context, requestID string, req *compassv1.StartAgentSessionRequest) (*compassv1.StartAgentSessionResponse, error) { - result, err := h.relay(ctx, req.GetContainerName(), &compassv1internal.SessionsResponse{ + result, _, err := h.relay(ctx, req.GetContainerName(), &compassv1internal.SessionsResponse{ RequestId: orNewRequestID(requestID), Command: &compassv1internal.SessionsResponse_Start{Start: req}, }) @@ -76,7 +84,7 @@ func (h *Hub) Start(ctx context.Context, requestID string, req *compassv1.StartA // Stop relays a StopAgentSession command to the owning Runner. func (h *Hub) Stop(ctx context.Context, requestID string, req *compassv1.StopAgentSessionRequest) (*compassv1.StopAgentSessionResponse, error) { - result, err := h.relay(ctx, req.GetSessionId(), &compassv1internal.SessionsResponse{ + result, _, err := h.relay(ctx, req.GetSessionId(), &compassv1internal.SessionsResponse{ RequestId: orNewRequestID(requestID), Command: &compassv1internal.SessionsResponse_Stop{Stop: req}, }) @@ -92,7 +100,7 @@ func (h *Hub) Stop(ctx context.Context, requestID string, req *compassv1.StopAge // Reload relays a ReloadAgentSession command to the owning Runner. func (h *Hub) Reload(ctx context.Context, requestID string, req *compassv1.ReloadAgentSessionRequest) (*compassv1.ReloadAgentSessionResponse, error) { - result, err := h.relay(ctx, req.GetSessionId(), &compassv1internal.SessionsResponse{ + result, _, err := h.relay(ctx, req.GetSessionId(), &compassv1internal.SessionsResponse{ RequestId: orNewRequestID(requestID), Command: &compassv1internal.SessionsResponse_Reload{Reload: req}, }) @@ -105,7 +113,7 @@ func (h *Hub) Reload(ctx context.Context, requestID string, req *compassv1.Reloa // Status relays a GetAgentStatus command to the owning Runner — the Runner is // authoritative for live session truth, so the Server reconciles to its answer. func (h *Hub) Status(ctx context.Context, requestID string, req *compassv1.GetAgentStatusRequest) (*compassv1.GetAgentStatusResponse, error) { - result, err := h.relay(ctx, req.GetSessionId(), &compassv1internal.SessionsResponse{ + result, _, err := h.relay(ctx, req.GetSessionId(), &compassv1internal.SessionsResponse{ RequestId: orNewRequestID(requestID), Command: &compassv1internal.SessionsResponse_Status{Status: req}, }) @@ -119,20 +127,22 @@ func (h *Hub) Status(ctx context.Context, requestID string, req *compassv1.GetAg // the outcome to a Connect status: a RunnerError result becomes the mapped // Connect code; a transport failure (no Runner, stream drop) becomes // Unavailable. sessionKey selects the owning Runner (single-Runner MVP: any -// non-empty key resolves the one Runner). -func (h *Hub) relay(ctx context.Context, sessionKey string, cmd *compassv1internal.SessionsResponse) (*compassv1internal.SessionsRequest, error) { - router, err := h.routerFor(sessionKey) +// non-empty key resolves the one Runner). The served Runner's id is returned +// alongside the result for the one caller that must attribute the command to a +// Runner (Provision, recording a durable placement); the rest discard it. +func (h *Hub) relay(ctx context.Context, sessionKey string, cmd *compassv1internal.SessionsResponse) (*compassv1internal.SessionsRequest, string, error) { + router, runnerID, err := h.routerFor(sessionKey) if err != nil { - return nil, connect.NewError(connect.CodeUnavailable, err) + return nil, "", connect.NewError(connect.CodeUnavailable, err) } result, err := router.dispatch(ctx, cmd) if err != nil { - return nil, connect.NewError(connect.CodeUnavailable, err) + return nil, "", connect.NewError(connect.CodeUnavailable, err) } if runnerErr := result.GetError(); runnerErr != nil { - return nil, runnerErrorToConnect(runnerErr) + return nil, "", runnerErrorToConnect(runnerErr) } - return result, nil + return result, runnerID, nil } // runnerErrorToConnect maps a RunnerError to the Connect status the client sees. diff --git a/go/internal/runnerhub/commands_test.go b/go/internal/runnerhub/commands_test.go index c854d215..3aaa72db 100644 --- a/go/internal/runnerhub/commands_test.go +++ b/go/internal/runnerhub/commands_test.go @@ -57,7 +57,7 @@ func TestCommandsNoRunnerIsUnavailable(t *testing.T) { call func() error }{ {"provision", func() error { - _, err := hub.Provision(ctx, "r1", &compassv1.ProvisionAgentWorkspaceRequest{}) + _, _, err := hub.Provision(ctx, "r1", &compassv1.ProvisionAgentWorkspaceRequest{}) return err }}, {"start", func() error { @@ -99,7 +99,7 @@ func TestStartRelaySurfacesAlreadyRunningAsAlreadyExists(t *testing.T) { // Enroll a Runner and bind a send that answers every command with an // ALREADY_RUNNING error result correlated by the pushed request id. hub.enroll("runner-1", store.Subject{Kind: store.SubjectRunner, ID: "runner-1"}) - router, err := hub.routerFor("any") + router, _, err := hub.routerFor("any") if err != nil { t.Fatalf("routerFor after enroll = %v, want a router", err) } @@ -130,7 +130,7 @@ func TestStartRelaySurfacesAlreadyRunningAsAlreadyExists(t *testing.T) { func TestStartRelayReturnsSessionIdOnSuccess(t *testing.T) { hub := newHubOnly() hub.enroll("runner-1", store.Subject{Kind: store.SubjectRunner, ID: "runner-1"}) - router, _ := hub.routerFor("any") + router, _, _ := hub.routerFor("any") router.attach(func(cmd *compassv1internal.SessionsResponse) error { go router.complete(&compassv1internal.SessionsRequest{ RequestId: cmd.GetRequestId(), diff --git a/go/internal/runnerhub/concurrent_dispatch_test.go b/go/internal/runnerhub/concurrent_dispatch_test.go index 6f5c9026..f1c2cd26 100644 --- a/go/internal/runnerhub/concurrent_dispatch_test.go +++ b/go/internal/runnerhub/concurrent_dispatch_test.go @@ -59,7 +59,7 @@ func TestConcurrentDispatchOverRealStreamNoDataRace(t *testing.T) { // before router.attach ran would (correctly) find no live send. waitRouterAttached(t, hub) - router, err := hub.routerFor("runner-1") + router, _, err := hub.routerFor("runner-1") if err != nil { t.Fatalf("routerFor after enroll+attach = %v, want the live router", err) } diff --git a/go/internal/runnerhub/handler.go b/go/internal/runnerhub/handler.go index 30d6f4d2..ca2979fa 100644 --- a/go/internal/runnerhub/handler.go +++ b/go/internal/runnerhub/handler.go @@ -70,7 +70,7 @@ func (h *Handler) Sessions(ctx context.Context, stream *connect.BidiStream[compa if !ok { return errUnauthenticated } - router, err := h.hub.routerFor(subj.ID) + router, _, err := h.hub.routerFor(subj.ID) if err != nil { // A Sessions stream with no enrolled Runner — the Runner must Enroll // before opening Sessions. diff --git a/go/internal/runnerhub/helpers_test.go b/go/internal/runnerhub/helpers_test.go index 41d08df3..2931c7bf 100644 --- a/go/internal/runnerhub/helpers_test.go +++ b/go/internal/runnerhub/helpers_test.go @@ -51,7 +51,7 @@ func waitRouterAttached(t *testing.T, hub *Hub) { t.Fatal("runner command router never attached a live Sessions send") default: } - if router, err := hub.routerFor("gate"); err == nil { + if router, _, err := hub.routerFor("gate"); err == nil { router.mu.Lock() attached := router.send != nil router.mu.Unlock() diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index 27fceea8..205b813a 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -262,13 +262,17 @@ func (h *Hub) enroll(id string, subject store.Subject) (reattached bool) { return reattached } -// router returns the attached Runner's command router, or an error when no -// Runner is enrolled (a session command with no Runner to serve it). -func (h *Hub) routerFor(sessionID string) (*commandRouter, error) { +// routerFor returns the attached Runner's command router and its id, or an error +// when no Runner is enrolled (a session command with no Runner to serve it). The +// id travels out with the router so a caller that must attribute the call to a +// Runner (Provision, recording an agent's durable placement) names the Runner +// that actually served it, rather than re-reading the registry afterwards and +// racing a re-enroll onto the wrong id. +func (h *Hub) routerFor(sessionID string) (*commandRouter, string, error) { h.mu.Lock() defer h.mu.Unlock() if h.runner == nil { - return nil, fmt.Errorf("no runner enrolled to serve session %q", sessionID) + return nil, "", fmt.Errorf("no runner enrolled to serve session %q", sessionID) } - return h.runner.router, nil + return h.runner.router, h.runner.id, nil } diff --git a/go/internal/runnerhub/hub_test.go b/go/internal/runnerhub/hub_test.go index 2957ff7a..6ee88b44 100644 --- a/go/internal/runnerhub/hub_test.go +++ b/go/internal/runnerhub/hub_test.go @@ -269,7 +269,7 @@ func TestEnrollDuplicateReattaches(t *testing.T) { } // A router is resolvable after enrollment (a session command has a Runner to // serve it). - if _, err := hub.routerFor("any"); err != nil { + if _, _, err := hub.routerFor("any"); err != nil { t.Fatalf("routerFor after enroll = %v, want a live router", err) } } diff --git a/go/internal/runnerhub/integration_pgtest_test.go b/go/internal/runnerhub/integration_pgtest_test.go index 79ae5ec6..6f2090f3 100644 --- a/go/internal/runnerhub/integration_pgtest_test.go +++ b/go/internal/runnerhub/integration_pgtest_test.go @@ -212,7 +212,7 @@ func provisionWhenSeamLive(t *testing.T, ctx context.Context, hub *runnerhub.Hub t.Helper() deadline := time.After(integrationTimeout) for { - resp, err := hub.Provision(ctx, "prov-1", &compassv1.ProvisionAgentWorkspaceRequest{ + resp, _, err := hub.Provision(ctx, "prov-1", &compassv1.ProvisionAgentWorkspaceRequest{ AgentAccountId: string(agentID), Repo: &compassv1.ProvisionAgentWorkspaceRequest_LocalPath{LocalPath: "/mirror/repo.git"}, }) diff --git a/go/internal/runnerhub/provision_dedup_test.go b/go/internal/runnerhub/provision_dedup_test.go index ab4d351b..b21a6a47 100644 --- a/go/internal/runnerhub/provision_dedup_test.go +++ b/go/internal/runnerhub/provision_dedup_test.go @@ -53,7 +53,7 @@ type provisionOutcome struct { func enrollAttached(t *testing.T, hub *Hub, send *recordingSend) *commandRouter { t.Helper() hub.enroll("runner-1", store.Subject{Kind: store.SubjectRunner, ID: "runner-1"}) - router, err := hub.routerFor("any") + router, _, err := hub.routerFor("any") if err != nil { t.Fatalf("routerFor after enroll = %v, want the live router", err) } @@ -83,7 +83,7 @@ func TestProvisionSameClientRequestIdDedups(t *testing.T) { } outcomes := make(chan provisionOutcome, 2) call := func() { - resp, err := hub.Provision(context.Background(), id, req) + resp, _, err := hub.Provision(context.Background(), id, req) outcomes <- provisionOutcome{resp, err} } @@ -140,7 +140,7 @@ func TestProvisionEmptyClientRequestIdDoesNotDedup(t *testing.T) { req := &compassv1.ProvisionAgentWorkspaceRequest{} outcomes := make(chan provisionOutcome, 2) call := func() { - resp, err := hub.Provision(context.Background(), "", req) + resp, _, err := hub.Provision(context.Background(), "", req) outcomes <- provisionOutcome{resp, err} } @@ -211,7 +211,7 @@ func TestProvisionSameIdDifferentWorkspaceDoesNotDedup(t *testing.T) { } outcomes := make(chan provisionOutcome, 2) call := func(req *compassv1.ProvisionAgentWorkspaceRequest) { - resp, err := hub.Provision(context.Background(), id, req) + resp, _, err := hub.Provision(context.Background(), id, req) outcomes <- provisionOutcome{resp, err} } diff --git a/go/internal/runnerhub/relay_comms_test.go b/go/internal/runnerhub/relay_comms_test.go index f778f0b9..e3f1d9fd 100644 --- a/go/internal/runnerhub/relay_comms_test.go +++ b/go/internal/runnerhub/relay_comms_test.go @@ -331,7 +331,7 @@ func TestRelayCommsCallStoppedSessionFailsClosedNotFound(t *testing.T) { func TestProvisionThenStartBindsSessionToProvisionedAccount(t *testing.T) { hub := newHubOnly() hub.enroll("runner-1", store.Subject{Kind: store.SubjectRunner, ID: "runner-1"}) - router, err := hub.routerFor("any") + router, _, err := hub.routerFor("any") if err != nil { t.Fatalf("routerFor after enroll = %v, want a router", err) } @@ -356,7 +356,7 @@ func TestProvisionThenStartBindsSessionToProvisionedAccount(t *testing.T) { }) ctx := context.Background() - if _, err := hub.Provision(ctx, "req-prov", &compassv1.ProvisionAgentWorkspaceRequest{AgentAccountId: "acct-agent"}); err != nil { + if _, _, err := hub.Provision(ctx, "req-prov", &compassv1.ProvisionAgentWorkspaceRequest{AgentAccountId: "acct-agent"}); err != nil { t.Fatalf("Provision = %v, want success", err) } if _, err := hub.Start(ctx, "req-start", &compassv1.StartAgentSessionRequest{ContainerName: "cont-1"}); err != nil { @@ -378,7 +378,7 @@ func TestProvisionThenStartBindsSessionToProvisionedAccount(t *testing.T) { func TestProvisionWithEmptyAccountLeavesNoBindingAndFailsClosed(t *testing.T) { hub, comms := newHubWithComms() hub.enroll("runner-1", store.Subject{Kind: store.SubjectRunner, ID: "runner-1"}) - router, err := hub.routerFor("any") + router, _, err := hub.routerFor("any") if err != nil { t.Fatalf("routerFor after enroll = %v, want a router", err) } @@ -401,7 +401,7 @@ func TestProvisionWithEmptyAccountLeavesNoBindingAndFailsClosed(t *testing.T) { }) ctx := context.Background() - if _, err := hub.Provision(ctx, "req-prov", &compassv1.ProvisionAgentWorkspaceRequest{AgentAccountId: ""}); err != nil { + if _, _, err := hub.Provision(ctx, "req-prov", &compassv1.ProvisionAgentWorkspaceRequest{AgentAccountId: ""}); err != nil { t.Fatalf("Provision (empty account) = %v, want success", err) } if _, err := hub.Start(ctx, "req-start", &compassv1.StartAgentSessionRequest{ContainerName: "cont-1"}); err != nil { diff --git a/go/internal/store/agent_placements.go b/go/internal/store/agent_placements.go new file mode 100644 index 00000000..500097b0 --- /dev/null +++ b/go/internal/store/agent_placements.go @@ -0,0 +1,148 @@ +package store + +import ( + "context" + "fmt" +) + +// Agent placement: the durable record of WHERE each agent runs — which Runner, +// under which container name — written at ProvisionAgentWorkspace, the one hop +// where every fact is in hand (agent_account_id is the Server's own request +// field, container_name the Runner's response, runner_id the Runner it relayed +// to). Before 0004 the Server learned that triple and kept it only in RAM (the +// RunnerHub's in-memory container->account binding), so a Server restart or a +// Runner re-enroll lost it. +// +// Placement is NOT authorization. SubscribeAgentSession authorizes through +// agent_sessions -> agent_accounts -> channel_members and never reads this +// table; keeping the two apart is what stops the container hop 0003 introduced +// from growing back into the security boundary. What placement is for is the +// two reads below: StartAgentSession resolving the account that owns an incoming +// container_name, and reattach recovery (SEA-1516) naming every agent stranded +// by a Runner restart. + +// AgentPlacement is one agent's placement: the Runner it runs on and the +// container name it runs under. Returned by ListAgentPlacementsForRunner, which +// is the reattach read — re-driving Provision needs both fields. +type AgentPlacement struct { + AgentAccountID AccountID + RunnerID string + ContainerName string +} + +// RecordAgentPlacement persists where an agent is placed, at +// ProvisionAgentWorkspace. It is an UPSERT keyed on the agent, because an agent +// is on at most one Runner under one container name: a re-provision REPLACES the +// placement, updating runner_id and container_name TOGETHER so a row can never +// pair a fresh Runner with the name from a previous one. That also makes the +// call idempotent under the client_request_id provision-retry contract — a retry +// rewrites the same values rather than conflicting. +// +// An unknown agent_account_id is ErrInvalidArgument (the FK). +// +// A container_name already placed for a DIFFERENT agent is ErrConflict (the +// unique index) — but that is a GUARD, not a live safety property today, and +// the distinction matters for anyone changing either side. Placements are +// PERMANENT for now: nothing anywhere deletes a row, so a name is never +// released. And container names are DERIVED, not allocated: BuildSpec computes +// NamePrefix + accountID (internal/runner/spec.go:73), so two distinct agents +// cannot produce the same name and the conflict is unreachable in production. +// The unique index is therefore leaning on that derivation. If name derivation +// ever changes — a Runner-scoped prefix, a random suffix, an operator-supplied +// name — the conflict becomes reachable, and a re-provision that yields a new +// name for the same agent would leave the OLD name owned forever with no +// release path: that change needs a delete path added alongside it. +func (s *Store) RecordAgentPlacement(ctx context.Context, agentAccountID AccountID, runnerID, containerName string) error { + if agentAccountID == "" { + return fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) + } + if runnerID == "" { + return fmt.Errorf("%w: runner id is required", ErrInvalidArgument) + } + 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 pgErrIs(err, pgForeignKeyViolation) { + return fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, agentAccountID) + } + if pgErrIs(err, pgUniqueViolation) { + return fmt.Errorf("%w: container %q is already placed for another agent", ErrConflict, containerName) + } + return fmt.Errorf("store: record agent placement: %w", err) + } + return nil +} + +// AgentForContainer resolves the agent account a placed container belongs to. +// StartAgentSession is the caller: its request carries only container_name (the +// frozen StartAgentSessionRequest field), so this is how it learns whose session +// it is about to record. Reading it from the durable placement rather than an +// in-memory binding is what makes the ownership record survive a Server restart +// or a Runner re-enroll between Provision and Start. +// +// An unplaced container is ErrNotFound — the Server never provisioned it (or +// the placement was superseded), so no session may bind to it. +func (s *Store) AgentForContainer(ctx context.Context, containerName string) (AccountID, error) { + 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 { + if noRows(err) { + return "", fmt.Errorf("%w: container %q is not placed", ErrNotFound, containerName) + } + return "", fmt.Errorf("store: resolve agent for container: %w", err) + } + return AccountID(accountID), nil +} + +// ListAgentPlacementsForRunner returns every agent placed on runnerID. This is +// the reattach read (SEA-1516): after a Runner restart its surviving containers +// are orphaned, and the Server re-drives Provision for exactly this set — which +// is why each row carries the container name and not just the account. Ordered +// by agent account id so a recovery pass is deterministic and its logs diffable. +// A Runner with no placements yields an empty slice, not an error: nothing to +// reattach is a normal outcome, not a failure. +func (s *Store) ListAgentPlacementsForRunner(ctx context.Context, runnerID string) ([]AgentPlacement, error) { + 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, + ) + 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) + } + return placements, nil +} diff --git a/go/internal/store/agent_placements_test.go b/go/internal/store/agent_placements_test.go new file mode 100644 index 00000000..a2c9ca55 --- /dev/null +++ b/go/internal/store/agent_placements_test.go @@ -0,0 +1,209 @@ +//go:build pgtest + +package store + +// Agent placement: the durable record of which Runner an agent is on and the +// container name it runs under. Two things must hold for the SEA-1516 reattach +// work that reads it — a placement is SINGULAR per agent (a re-provision +// replaces it, never accumulates), and the container -> agent mapping it +// exposes is exclusive. Both are database invariants (the PK and the unique +// index), so both are pgtest-backed; a mock would only re-assert the Go code. + +import ( + "testing" +) + +// TestRecordAgentPlacementRoundTrips pins the base contract both reads depend +// on: what Provision wrote is what Start and reattach read back — the account +// resolvable from the container name, and the placement listed under its Runner +// carrying the container name re-driving Provision needs. +func TestRecordAgentPlacementRoundTrips(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + if err := s.RecordAgentPlacement(ctx, agent.ID, "runner-1", "compass-agent-"+string(agent.ID)); err != nil { + t.Fatalf("RecordAgentPlacement: %v", err) + } + + got, err := s.AgentForContainer(ctx, "compass-agent-"+string(agent.ID)) + if err != nil { + t.Fatalf("AgentForContainer: %v", err) + } + if got != agent.ID { + t.Fatalf("AgentForContainer = %q, want the placed agent %q", got, agent.ID) + } + + placements, err := s.ListAgentPlacementsForRunner(ctx, "runner-1") + if err != nil { + t.Fatalf("ListAgentPlacementsForRunner: %v", err) + } + if len(placements) != 1 { + t.Fatalf("runner-1 placements = %d, want 1", len(placements)) + } + want := AgentPlacement{AgentAccountID: agent.ID, RunnerID: "runner-1", ContainerName: "compass-agent-" + string(agent.ID)} + if placements[0] != want { + t.Fatalf("placement = %+v, want %+v", placements[0], want) + } +} + +// TestRecordAgentPlacementReplacesOnRePlacement is the invariant reattach hangs +// on: an agent is on AT MOST ONE Runner. Re-provisioning it onto a different +// Runner must REPLACE its row — updating runner_id and container_name TOGETHER — +// not add a second. If placements accumulated, reattach on the OLD Runner would +// re-drive Provision for an agent that has already moved; if the two columns +// updated independently, a row would pair the new Runner with the old container +// name and reattach would try to recover a container that is not there. +func TestRecordAgentPlacementReplacesOnRePlacement(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + if err := s.RecordAgentPlacement(ctx, agent.ID, "runner-1", "cont-old"); err != nil { + t.Fatalf("first RecordAgentPlacement: %v", err) + } + if err := s.RecordAgentPlacement(ctx, agent.ID, "runner-2", "cont-new"); err != nil { + t.Fatalf("re-placement onto runner-2: %v", err) + } + + // The old Runner must no longer claim it — otherwise a reattach pass on + // runner-1 re-drives Provision for an agent that has moved away. + old, err := s.ListAgentPlacementsForRunner(ctx, "runner-1") + if err != nil { + t.Fatalf("ListAgentPlacementsForRunner(runner-1): %v", err) + } + if len(old) != 0 { + t.Fatalf("runner-1 still holds %+v after the agent moved, want no placements", old) + } + + current, err := s.ListAgentPlacementsForRunner(ctx, "runner-2") + if err != nil { + t.Fatalf("ListAgentPlacementsForRunner(runner-2): %v", err) + } + want := AgentPlacement{AgentAccountID: agent.ID, RunnerID: "runner-2", ContainerName: "cont-new"} + if len(current) != 1 || current[0] != want { + t.Fatalf("runner-2 placements = %+v, want exactly [%+v]", current, want) + } + + // The superseded container name resolves to nothing: a Start arriving late + // for the old container cannot record a session against a placement that no + // longer exists. + if _, err := s.AgentForContainer(ctx, "cont-old"); err == nil { + t.Fatal("AgentForContainer(cont-old) = nil error, want ErrNotFound (the old placement was replaced)") + } else { + sentinelIs(t, err, ErrNotFound, "superseded container lookup") + } +} + +// TestRecordAgentPlacementRejectsAContainerClaimedByAnotherAgent pins the +// exclusivity of the container -> agent mapping. StartAgentSession resolves the +// session's OWNER through this mapping, and that owner picks the home channel +// the authz JOIN authorizes against — so a second agent claiming a live +// container name would let a session be recorded under the wrong owner and +// hand a foreign channel's members the right to watch it. The unique index +// refuses it as ErrConflict rather than letting the write land. +func TestRecordAgentPlacementRejectsAContainerClaimedByAnotherAgent(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + first := mustAgent(t, s, owner.ID, "first") + second := mustAgent(t, s, owner.ID, "second") + + if err := s.RecordAgentPlacement(ctx, first.ID, "runner-1", "cont-shared"); err != nil { + t.Fatalf("first placement: %v", err) + } + + err := s.RecordAgentPlacement(ctx, second.ID, "runner-1", "cont-shared") + sentinelIs(t, err, ErrConflict, "second agent claiming a live container name") + + // The first agent keeps the container: the refused write changed nothing. + got, err := s.AgentForContainer(ctx, "cont-shared") + if err != nil { + t.Fatalf("AgentForContainer after the refused claim: %v", err) + } + if got != first.ID { + t.Fatalf("cont-shared resolves to %q, want the original owner %q", got, first.ID) + } +} + +// TestRecordAgentPlacementUnknownAgentIsInvalidArgument pins the FK: a placement +// for an account that is not an agent cannot land. Without it a row could name a +// nonexistent agent and reattach would re-drive Provision for nobody. +func TestRecordAgentPlacementUnknownAgentIsInvalidArgument(t *testing.T) { + s := newTestStore(t) + + err := s.RecordAgentPlacement(t.Context(), "no-such-agent", "runner-1", "cont-1") + sentinelIs(t, err, ErrInvalidArgument, "placement for an unknown agent") +} + +// TestAgentForContainerUnplacedIsNotFound pins Start's fail-closed path: a +// container this Server never placed resolves no owner, so no session row can be +// recorded against it. Distinct from a conflict — nothing is wrong with the +// request, the container is simply not ours. +func TestAgentForContainerUnplacedIsNotFound(t *testing.T) { + s := newTestStore(t) + + _, err := s.AgentForContainer(t.Context(), "never-provisioned") + sentinelIs(t, err, ErrNotFound, "unplaced container lookup") +} + +// TestListAgentPlacementsForRunnerScopesToItsRunner is the reattach read itself: +// a recovery pass must get EXACTLY the agents on the restarted Runner. Too few +// strands an agent; too many re-drives Provision for agents on a healthy Runner, +// disrupting live work. Ordering is asserted because a recovery pass should be +// deterministic and its logs diffable across runs. +func TestListAgentPlacementsForRunnerScopesToItsRunner(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + a := mustAgent(t, s, owner.ID, "agent-a") + b := mustAgent(t, s, owner.ID, "agent-b") + elsewhere := mustAgent(t, s, owner.ID, "agent-elsewhere") + + for _, p := range []AgentPlacement{ + {a.ID, "runner-1", "cont-a"}, + {b.ID, "runner-1", "cont-b"}, + {elsewhere.ID, "runner-2", "cont-elsewhere"}, + } { + if err := s.RecordAgentPlacement(ctx, p.AgentAccountID, p.RunnerID, p.ContainerName); err != nil { + t.Fatalf("RecordAgentPlacement(%+v): %v", p, err) + } + } + + got, err := s.ListAgentPlacementsForRunner(ctx, "runner-1") + if err != nil { + t.Fatalf("ListAgentPlacementsForRunner: %v", err) + } + if len(got) != 2 { + t.Fatalf("runner-1 placements = %+v, want exactly the 2 agents on it", got) + } + for _, p := range got { + if p.AgentAccountID == elsewhere.ID { + t.Fatalf("runner-1 listing included %q, which is placed on runner-2", elsewhere.ID) + } + if p.ContainerName == "" { + t.Fatalf("placement %+v carries no container name; reattach cannot re-drive Provision without it", p) + } + } + if got[0].AgentAccountID > got[1].AgentAccountID { + t.Fatalf("placements = %+v, want ascending agent order (a recovery pass must be deterministic)", got) + } +} + +// TestListAgentPlacementsForRunnerEmptyForUnknownRunner pins that "nothing to +// reattach" is a normal outcome, not an error: a Runner enrolling for the first +// time holds no placements, and a recovery pass over it must proceed quietly +// rather than treat the empty set as a failure. +func TestListAgentPlacementsForRunnerEmptyForUnknownRunner(t *testing.T) { + s := newTestStore(t) + + got, err := s.ListAgentPlacementsForRunner(t.Context(), "runner-never-seen") + if err != nil { + t.Fatalf("ListAgentPlacementsForRunner(unknown) = %v, want nil (no placements is not a failure)", err) + } + if len(got) != 0 { + t.Fatalf("unknown runner placements = %+v, want empty", got) + } +} diff --git a/go/internal/store/agent_sessions.go b/go/internal/store/agent_sessions.go index 4e7555c1..d6be9f38 100644 --- a/go/internal/store/agent_sessions.go +++ b/go/internal/store/agent_sessions.go @@ -6,64 +6,49 @@ import ( ) // The durable session-ownership chain: the persistent -// session_id -> container_name -> agent_account_id -> home_channel_id mapping -// that SubscribeAgentSession resolves to authorize a subscriber. It is rooted -// non-spoofably — agent_account_id is a request field, but container_name and -// session_id are server-minted response values recorded only after the Runner -// call succeeds (agent_sessions.sql / 0003_agent_ownership.sql). The rows -// survive a Server restart, so the authz boundary does not depend on the -// in-memory RunnerHub enrollment. +// session_id -> agent_account_id -> home_channel_id mapping that +// SubscribeAgentSession resolves to authorize a subscriber. It is rooted +// non-spoofably — session_id is a server-minted response value recorded only +// after the Runner call succeeds (0004_agent_placement.sql). The rows survive a +// Server restart, so the authz boundary does not depend on the in-memory +// RunnerHub enrollment. +// +// The chain used to hop through a container_name row (0003_agent_ownership.sql). +// That hop was removed in 0004 on a SCHEMA fact, not a naming one: in +// agent_containers, container_name was the PRIMARY KEY and agent_account_id a +// NOT NULL FK to the account, so the hop was a provable 1:1 pass-through — it +// could resolve exactly one account for a name, and never fewer. Removing it +// therefore authorizes the identical set of (session, caller) pairs; only the +// table count differs. That argument holds whatever container names look like, +// which is why it is the one the collapse rests on. +// +// Secondarily, on where a container name comes from at all: it is derived from +// the agent account (internal/runner/spec.go BuildSpec, NamePrefix + accountID), +// so nothing is lost by not storing it — where one is genuinely needed it is +// recomputed. But that is a Runner-side convention the Server never enforces, +// so it is a remark, not the justification. -// RecordAgentContainer persists the container_name -> agent_account_id mapping -// at ProvisionAgentWorkspace, where the agent identity is known (the request -// field) and container_name is the server-minted response. It is idempotent: -// one container belongs to exactly one agent and never changes owner, so a -// client_request_id retry that returns the same container_name re-records the -// same row as a no-op (ON CONFLICT DO NOTHING) rather than an ErrConflict. An -// unknown agent_account_id is ErrInvalidArgument (the FK). -func (s *Store) RecordAgentContainer(ctx context.Context, containerName string, agentAccountID AccountID) error { - if containerName == "" { - return fmt.Errorf("%w: container name is required", ErrInvalidArgument) - } - if agentAccountID == "" { - return fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) - } - if _, err := s.pool.Exec(ctx, - `INSERT INTO agent_containers (container_name, agent_account_id) - VALUES ($1, $2) - ON CONFLICT (container_name) DO NOTHING`, - containerName, string(agentAccountID), - ); err != nil { - if pgErrIs(err, pgForeignKeyViolation) { - return fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, agentAccountID) - } - return fmt.Errorf("store: record agent container: %w", err) - } - return nil -} - -// RecordAgentSession persists the session_id -> container_name mapping at -// StartAgentSession, where container_name is the request handle and session_id -// is the server-minted response. The FK to agent_containers means a session -// cannot bind to a container the Server never provisioned — the chain is -// complete or it does not exist. An unknown container_name is -// ErrInvalidArgument (the FK); a re-used session_id is ErrConflict. -func (s *Store) RecordAgentSession(ctx context.Context, sessionID, containerName string) error { +// RecordAgentSession persists the session_id -> agent_account_id mapping at +// StartAgentSession, where session_id is the server-minted response and the +// agent account is the one the started session belongs to. An unknown +// agent_account_id is ErrInvalidArgument (the FK); a re-used session_id is +// ErrConflict. +func (s *Store) RecordAgentSession(ctx context.Context, sessionID string, agentAccountID AccountID) error { if sessionID == "" { return fmt.Errorf("%w: session id is required", ErrInvalidArgument) } - if containerName == "" { - return fmt.Errorf("%w: container name is required", ErrInvalidArgument) + if agentAccountID == "" { + return fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) } if _, err := s.pool.Exec(ctx, - `INSERT INTO agent_sessions (session_id, container_name) VALUES ($1, $2)`, - sessionID, containerName, + `INSERT INTO agent_sessions (session_id, agent_account_id) VALUES ($1, $2)`, + sessionID, string(agentAccountID), ); err != nil { if pgErrIs(err, pgUniqueViolation) { return fmt.Errorf("%w: session %q already recorded", ErrConflict, sessionID) } if pgErrIs(err, pgForeignKeyViolation) { - return fmt.Errorf("%w: container %q was not provisioned", ErrInvalidArgument, containerName) + return fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, agentAccountID) } return fmt.Errorf("store: record agent session: %w", err) } @@ -72,21 +57,20 @@ func (s *Store) RecordAgentSession(ctx context.Context, sessionID, containerName // RequireAgentSessionSubscriber is the read-path authorization primitive for // SubscribeAgentSession — the streaming sibling of requireChannelMember. In ONE -// query it resolves the ownership chain (session_id -> container_name -> -// agent_account_id -> home_channel_id) and checks the caller's membership on -// that home channel, returning ErrNotFound for BOTH an unknown session_id AND a -// non-member caller. Merging the two into one indistinguishable error is -// load-bearing: it must not leak session existence to a caller who holds a -// foreign session_id — neither via the error class NOR via timing skew. A -// two-step shape (resolve, then separately check membership) would take one -// round-trip for an unknown session but two for a known-but-foreign one, so the -// latency itself would distinguish "does not exist" from "exists but -// forbidden". The single EXISTS query below is constant-shape: it returns true -// only when the session exists AND the caller is a member, and every other -// outcome is the same false -> the same ErrNotFound in the same one round-trip. -// The caller (SubscribeAgentSession) maps ErrNotFound to its stream rejection, -// so a foreign or unknown session is refused identically and enumerates nothing -// (the not-found/forbidden merge, D9). +// query it resolves the ownership chain (session_id -> agent_account_id -> +// home_channel_id) and checks the caller's membership on that home channel, +// returning ErrNotFound for BOTH an unknown session_id AND a non-member caller. +// Merging the two into one indistinguishable error is load-bearing: it must not +// leak session existence to a caller who holds a foreign session_id — neither +// via the error class NOR via timing skew. A two-step shape (resolve, then +// separately check membership) would take one round-trip for an unknown session +// but two for a known-but-foreign one, so the latency itself would distinguish +// "does not exist" from "exists but forbidden". The single EXISTS query below is +// constant-shape: it returns true only when the session exists AND the caller is +// a member, and every other outcome is the same false -> the same ErrNotFound in +// the same one round-trip. The caller (SubscribeAgentSession) maps ErrNotFound to +// its stream rejection, so a foreign or unknown session is refused identically +// and enumerates nothing (the not-found/forbidden merge, D9). func (s *Store) RequireAgentSessionSubscriber(ctx context.Context, caller AccountID, sessionID string) error { if sessionID == "" { return fmt.Errorf("%w: session id is required", ErrInvalidArgument) @@ -96,8 +80,7 @@ func (s *Store) RequireAgentSessionSubscriber(ctx context.Context, caller Accoun `SELECT EXISTS ( SELECT 1 FROM agent_sessions se - JOIN agent_containers c ON c.container_name = se.container_name - JOIN agent_accounts ag ON ag.account_id = c.agent_account_id + 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)`, diff --git a/go/internal/store/agent_sessions_test.go b/go/internal/store/agent_sessions_test.go index 87159a39..9f8f7d15 100644 --- a/go/internal/store/agent_sessions_test.go +++ b/go/internal/store/agent_sessions_test.go @@ -3,50 +3,46 @@ package store // Read-path authorization: RequireAgentSessionSubscriber resolves -// the ownership chain (session_id -> container_name -> agent_account_id -> -// home_channel_id) and authorizes a caller iff it is a member of that home -// channel. The load-bearing contract is the not-found/forbidden MERGE: an -// unknown session and a known-but-foreign session are refused with the SAME -// ErrNotFound, so a caller holding a foreign session_id cannot probe whether it -// exists. Pgtest-backed — the authz lives in the JOIN, provable only against a -// real database. +// the ownership chain (session_id -> agent_account_id -> home_channel_id) and +// authorizes a caller iff it is a member of that home channel. The load-bearing +// contract is the not-found/forbidden MERGE: an unknown session and a +// known-but-foreign session are refused with the SAME ErrNotFound, so a caller +// holding a foreign session_id cannot probe whether it exists. Pgtest-backed — +// the authz lives in the JOIN, provable only against a real database. import ( "context" + "fmt" "testing" ) -// recordChain seeds a full ownership chain for agent and returns the session_id -// a subscriber would present. container_name and session_id are opaque handles -// here (the server mints them in production). -func recordChain(t *testing.T, s *Store, agent Account, container, sessionID string) { +// recordSession seeds the ownership chain for agent and returns the session_id a +// subscriber would present. session_id is an opaque handle here (the server +// mints it in production). The chain used to hop through a container_name row; +// 0004 collapsed that hop, so seeding is now the single session write. +func recordSession(t *testing.T, s *Store, agent Account, sessionID string) { t.Helper() - ctx := context.Background() - if err := s.RecordAgentContainer(ctx, container, agent.ID); err != nil { - t.Fatalf("RecordAgentContainer(%q): %v", container, err) - } - if err := s.RecordAgentSession(ctx, sessionID, container); err != nil { + if err := s.RecordAgentSession(t.Context(), sessionID, agent.ID); err != nil { t.Fatalf("RecordAgentSession(%q): %v", sessionID, err) } } // TestRequireAgentSessionSubscriberAuthorizesHomeChannelMember pins the happy -// path AND, by construction, that the resolve walks all three hops: the caller -// is authorized only because the full session->container->agent->home_channel -// chain resolves to a channel it belongs to. The owner is seeded into the -// agent's home channel at CreateAgent. A bug in ANY of the three joins would -// drop the row and turn this legitimate subscribe into an ErrNotFound, reddening -// this test. +// path AND, by construction, that the resolve walks every hop: the caller is +// authorized only because the full session->agent->home_channel chain resolves +// to a channel it belongs to. The owner is seeded into the agent's home channel +// at CreateAgent. A bug in EITHER join would drop the row and turn this +// legitimate subscribe into an ErrNotFound, reddening this test. func TestRequireAgentSessionSubscriberAuthorizesHomeChannelMember(t *testing.T) { ctx := context.Background() s := newTestStore(t) owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - recordChain(t, s, agent, "cont-1", "sess-1") + recordSession(t, s, agent, "sess-1") // owner is a member of the agent's home channel (seeded at CreateAgent). if err := s.RequireAgentSessionSubscriber(ctx, owner.ID, "sess-1"); err != nil { - t.Fatalf("home-channel member authorize = %v, want nil (full chain resolves + member)", err) + t.Fatalf("home-channel member authorize = %v, want nil (chain resolves + member)", err) } } @@ -76,7 +72,7 @@ func TestRequireAgentSessionSubscriberNonMemberSameAsUnknown(t *testing.T) { owner := mustUser(t, s, "owner") outsider := mustUser(t, s, "outsider") agent := mustAgent(t, s, owner.ID, "agent") - recordChain(t, s, agent, "cont-1", "real-session") + recordSession(t, s, agent, "real-session") // outsider even owns an unrelated channel — so it IS a member of something, // proving membership is checked against the RESOLVED home channel, not any @@ -111,3 +107,74 @@ func TestRequireAgentSessionSubscriberEmptySessionIDIsInvalidArgument(t *testing err := s.RequireAgentSessionSubscriber(ctx, owner.ID, "") sentinelIs(t, err, ErrInvalidArgument, "empty session id subscribe") } + +// TestRequireAgentSessionSubscriberCollapseIsUniformAcrossRefusalCauses is the +// 0004 regression guard for the visibility collapse. Collapsing the chain from +// session -> container -> agent -> home_channel down to +// session -> agent -> home_channel removed a JOIN from a SECURITY query, and +// the risk of that edit is not that a refusal becomes an authorization — it is +// that the refusals stop being IDENTICAL, so the error a caller gets back +// starts discriminating between causes it must not distinguish. +// +// The three tests above each pin one refusal in isolation. This one pins them +// against EACH OTHER: every way of failing to be authorized — the session was +// never recorded, the session exists but belongs to another owner's agent, the +// session exists and the caller is a total stranger to the system — must produce +// one indistinguishable ErrNotFound. If a future rewrite of the JOIN surfaced, +// say, a distinct error once the agent row resolves but membership does not, a +// caller could subtract the two answers and enumerate live session ids. Every +// refusal below is checked against the same sentinel AND against the same +// message, so a divergence in either reddens this test. +func TestRequireAgentSessionSubscriberCollapseIsUniformAcrossRefusalCauses(t *testing.T) { + ctx := context.Background() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + recordSession(t, s, agent, "live-session") + + // Three unauthorized callers, each failing for a structurally DIFFERENT + // reason, all presenting the same real session id. + rival := mustUser(t, s, "rival") + rivalAgent := mustAgent(t, s, rival.ID, "rival-agent") // owns a DIFFERENT agent + home channel + recordSession(t, s, rivalAgent, "rival-session") + stranger := mustUser(t, s, "stranger") // member of nothing at all + + refusals := []struct { + name string + sessionID string + err error + }{ + // The session does not exist: no chain resolves. + {"unknown session", "no-such-session", s.RequireAgentSessionSubscriber(ctx, owner.ID, "no-such-session")}, + // The session exists; the caller is a member of a home channel, just not + // THIS one. The chain resolves fully and membership is what fails — the + // case a two-step resolve-then-check would answer differently. + {"foreign session, caller owns another agent", "live-session", s.RequireAgentSessionSubscriber(ctx, rival.ID, "live-session")}, + // The session exists; the caller belongs to nothing. + {"foreign session, caller belongs to nothing", "live-session", s.RequireAgentSessionSubscriber(ctx, stranger.ID, "live-session")}, + // Symmetric: the owner probing the RIVAL's session. Proves the refusal is + // not a property of one privileged caller. + {"owner probing a rival's session", "rival-session", s.RequireAgentSessionSubscriber(ctx, owner.ID, "rival-session")}, + } + + for _, r := range refusals { + sentinelIs(t, r.err, ErrNotFound, r.name) + } + + // Uniformity, the actual contract: every refusal reads identically apart from + // the session id the caller itself supplied. A message that varied by cause + // would leak the cause even while the sentinel matched. + for _, r := range refusals { + want := fmt.Sprintf("%v: session %q", ErrNotFound, r.sessionID) + if got := r.err.Error(); got != want { + t.Fatalf("%s refusal = %q, want %q (every refusal must read identically, echoing only the presented session id)", r.name, got, want) + } + } + + // And the collapse must not have cost the happy path: the legitimate owner + // still resolves through the shortened chain. Without this, a query that + // refused EVERYTHING would satisfy the uniformity check above. + if err := s.RequireAgentSessionSubscriber(ctx, owner.ID, "live-session"); err != nil { + t.Fatalf("owner on its own agent's session = %v, want nil (uniform refusal must not mean universal refusal)", err) + } +} diff --git a/go/internal/store/migrations/0004_agent_placement.sql b/go/internal/store/migrations/0004_agent_placement.sql new file mode 100644 index 00000000..e883f22c --- /dev/null +++ b/go/internal/store/migrations/0004_agent_placement.sql @@ -0,0 +1,151 @@ +-- 0004_agent_placement: collapse the container hop out of the session-ownership +-- chain, and record which Runner each agent is placed on (SEA-1516 reattach). +-- +-- Two changes, one motivation: agent_containers sat in the middle of a security +-- boundary to carry a fact that is not authorization, while the fact reattach +-- actually needs was stored nowhere. +-- +-- * agent_containers held container_name -> agent_account_id, and its ONLY +-- read anywhere was the SubscribeAgentSession authz JOIN — where it could +-- never do more than pass the account through (container_name PK, NOT NULL +-- FK to the account: a 1:1 hop). So agent_sessions now points at the agent +-- account DIRECTLY and the chain shortens from +-- session_id -> container_name -> agent_account_id -> home_channel_id to +-- session_id -> agent_account_id -> home_channel_id. The name/account +-- mapping itself is not lost — it moves to agent_placements below, out of +-- the authz path and into the operational state where it belongs. +-- +-- * NOTHING recorded which Runner an agent runs on. After a Runner restart +-- its surviving containers are orphaned, and the Server must re-drive +-- Provision for exactly the agents that Runner held — a set it could not +-- name. agent_placements records it, alongside the container name that +-- re-drive needs. +-- +-- The authz semantics are UNCHANGED. Shortening the chain removes a hop that +-- could only ever be a 1:1 pass-through (container_name PK, NOT NULL FK to the +-- account), so the set of (session, caller) pairs the JOIN authorizes is +-- identical; only the table count differs. The rooting argument from 0003 also +-- survives: session_id is still the SERVER-MINTED StartAgentSession response, +-- written only after the Runner call succeeds, so a row still never claims a +-- session the Runner failed to create. +-- +-- Convention (0001_init.sql:7-11, 0003_agent_ownership.sql:15-16): text ids, FK +-- ON DELETE RESTRICT so a referenced agent cannot be orphaned out from under a +-- mapping. + +-- ── agent_sessions: session_id -> agent_account_id ────────────────────────── +-- Added nullable, backfilled through the old chain so no existing session row +-- loses its ownership, then tightened to NOT NULL. Doing it in that order (not +-- a bare NOT NULL add) is what makes the migration non-destructive on a +-- database that already has sessions. +ALTER TABLE agent_sessions + ADD COLUMN agent_account_id TEXT REFERENCES agent_accounts (account_id) ON DELETE RESTRICT; + +UPDATE agent_sessions se + SET agent_account_id = c.agent_account_id + FROM agent_containers c + WHERE c.container_name = se.container_name; + +ALTER TABLE agent_sessions + ALTER COLUMN agent_account_id SET NOT NULL; + +-- Dropping the column takes its FK and agent_sessions_container_idx with it. +ALTER TABLE agent_sessions DROP COLUMN container_name; + +-- The equivalent of the old agent_sessions_container_idx: look a session up by +-- the party that owns it. Also the "sessions of this agent" direction reattach +-- reads once it knows which agents a Runner held. +CREATE INDEX agent_sessions_agent_idx ON agent_sessions (agent_account_id); + +-- ── agent_placements: where each agent runs, and under what name ──────────── +-- Operational placement state, written at ProvisionAgentWorkspace — the one hop +-- where every fact is in hand: agent_account_id is the Server's own request +-- field, container_name is the Runner's response, and runner_id is the Runner +-- the Server relayed to. Before this the Server learned that triple and +-- remembered it only in RAM (runnerhub's in-memory container->account binding), +-- so a Server restart or a Runner re-enroll lost it. +-- +-- Deliberately NOT part of the authz chain. SubscribeAgentSession authorizes +-- through agent_sessions -> agent_accounts -> channel_members and never reads +-- this table; placement is where an agent runs, not who may watch it. Keeping +-- the two apart is what stops the container hop 0003 introduced from growing +-- back into the security boundary. +-- +-- PK on the agent, not a surrogate: an agent is on AT MOST ONE Runner under one +-- container name, so a re-provision REPLACES the row (upsert on +-- agent_account_id, updating runner_id AND container_name together) rather than +-- accumulating a second placement that would make "which Runner owns this +-- agent, under which name" ambiguous — the exact ambiguity reattach must not +-- face, and the reason a stale container_name alongside a fresh runner_id would +-- be worse than no row at all. +-- +-- runner_id is deliberately NOT a FK: Runners are enrolled in memory under a +-- token subject (store.SubjectRunner), with no runners table to reference, and +-- a placement must OUTLIVE the Runner's attachment — that it survives the +-- Runner going away is the whole point of recording it. +-- +-- This table is created BEFORE agent_containers is dropped, because the +-- backfill below reads it. Order is load-bearing, and the whole file is one +-- transaction, so the create/backfill/drop is atomic. +CREATE TABLE agent_placements ( + agent_account_id TEXT PRIMARY KEY REFERENCES agent_accounts (account_id) ON DELETE RESTRICT, + runner_id TEXT NOT NULL, + container_name TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +-- The reattach query is "every agent on this Runner", so runner_id is the read +-- direction that needs the index (the agent direction is already the PK). +CREATE INDEX agent_placements_runner_idx ON agent_placements (runner_id); + +-- StartAgentSession arrives holding only container_name (the frozen +-- StartAgentSessionRequest field), so it resolves the owning account by this +-- index before recording the session. UNIQUE because the mapping is 1:1 — a +-- container belongs to exactly one agent — and enforcing it here means a +-- second agent can never claim a live container name, which would let Start +-- record a session under the wrong owner and hand the authz JOIN the wrong +-- home channel. +CREATE UNIQUE INDEX agent_placements_container_key ON agent_placements (container_name); + +-- ── Backfill: every pre-upgrade container keeps working ───────────────────── +-- Without this, the upgrade BREAKS every already-provisioned container: +-- StartAgentSession now resolves its owner through AgentForContainer +-- (agent_placements.go), so a container with no placement row is permanently +-- un-Startable — and the mapping it would need is right here in +-- agent_containers, about to be dropped. +-- +-- runner_id = '' is the deliberate UNKNOWN-RUNNER SENTINEL, not a placeholder +-- we tolerate. agent_containers recorded no Runner, and we genuinely do not +-- know which one held the container, so the row must be readable in the one +-- direction where the answer is known and invisible in the one where it is not: +-- +-- * AgentForContainer resolves it (it reads only container_name), so Start +-- keeps working for every pre-upgrade container. This is the break we fix. +-- * ListAgentPlacementsForRunner CANNOT return it: that read rejects an empty +-- runner id outright (agent_placements.go ErrInvalidArgument guard), and no +-- enrolled Runner has an empty id, so a backfilled row is never re-driven +-- against a Runner we only guessed at. Exactly right — a wrong reattach is +-- worse than none. +-- +-- Self-healing: the next ProvisionAgentWorkspace upserts on agent_account_id +-- and overwrites '' with the Runner that actually served the call, so the +-- sentinel is transient per agent. runner_id therefore stays NOT NULL ('' being +-- a value that satisfies it) rather than becoming nullable — nullable would +-- widen the column's contract permanently to encode a one-time migration state. +-- +-- ON CONFLICT DO NOTHING: agent_containers keyed on container_name, so one +-- agent could in principle hold several containers while a placement is one per +-- agent (PK). Keeping the first rather than failing the migration is right — +-- any of them restores Start, and the next provision replaces it with truth. +INSERT INTO agent_placements (agent_account_id, runner_id, container_name) +SELECT c.agent_account_id, '', c.container_name + FROM agent_containers c +ON CONFLICT DO NOTHING; + +-- Superseded, not merely redundant: agent_placements above records the same +-- container_name -> agent_account_id fact and MORE (which Runner, and when), +-- keyed on the agent rather than the name — and the backfill has just carried +-- every existing mapping across. Dropping the table drops +-- agent_containers_agent_idx with it. +DROP TABLE agent_containers; diff --git a/go/server/service.go b/go/server/service.go index a6528664..f009cce7 100644 --- a/go/server/service.go +++ b/go/server/service.go @@ -18,6 +18,7 @@ import ( "context" "errors" "fmt" + "log/slog" "time" "connectrpc.com/connect" @@ -39,6 +40,14 @@ const apiVersion = "compass.v1" // sinceSeq = 0, so it carries no meaningful position. const resyncSeq uint64 = 0 +// rollbackStopTimeout bounds the best-effort Stop in abandonStartedSession. The +// runnerhub dispatch path (internal/runnerhub) has NO deadline of its own — a +// wedged-but-connected Runner that accepts Stop but never answers its result +// would otherwise hang StartAgentSession (and stall graceful-shutdown drain) +// forever. A package var rather than a const only so a test can shorten it; it +// is never reassigned in production. +var rollbackStopTimeout = 30 * time.Second + // errNoRunnerHub is returned by the container-lifecycle RPCs when no Runner door // is mounted (the socket-only path) — the RPC is Unavailable, never a panic. var errNoRunnerHub = errors.New("no runner hub configured on this server") @@ -48,6 +57,11 @@ var errNoRunnerHub = errors.New("no runner hub configured on this server") // with Unauthenticated rather than stream an unauthorized session. var errNoCaller = errors.New("no caller identity in request context") +// errNoSessionID is the abandon-path cause when a Runner returned success for a +// Start but no session id: there is nothing to Stop, and nothing the caller +// could ever address, so the roll-back reduces to logging it loudly. +var errNoSessionID = errors.New("runner returned no session id to stop") + // busPayload is the bus's event type. Go's generated oneof interface is // unexported, so the bus carries the whole response message with only its // Payload oneof set at publish time — the bus stamps Seq/AtUnixMs/InstanceEpoch @@ -97,17 +111,27 @@ func (s *service) ProvisionAgentWorkspace( if s.hub == nil { return nil, connect.NewError(connect.CodeUnavailable, errNoRunnerHub) } - resp, err := s.hub.Provision(ctx, req.Msg.GetClientRequestId(), req.Msg) + resp, runnerID, err := s.hub.Provision(ctx, req.Msg.GetClientRequestId(), req.Msg) if err != nil { return nil, err } - // Record the container_name -> agent_account_id link of the durable - // ownership chain, only now that the Runner has created the - // container. agent_account_id is the request field; container_name is the - // server-minted response — so the row is rooted on a value the client cannot - // forge. Idempotent, matching the client_request_id provision-retry contract. - if err := s.store.RecordAgentContainer(ctx, resp.GetContainerName(), store.AccountID(req.Msg.GetAgentAccountId())); err != nil { - return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("recording agent container ownership: %w", err)) + // Record the agent's durable PLACEMENT — which Runner it is on and the + // container name it runs under — only now that the Runner has created the + // container. Every field is trustworthy at exactly this point: + // agent_account_id is the Server's own request field, container_name is the + // Runner's response, and runnerID names the Runner that actually served this + // call, so the row is rooted on values the client cannot forge. Idempotent + // (upsert on the agent), matching the client_request_id provision-retry + // contract. + // + // This is the write that makes the Server's knowledge survive itself: before + // it, the container -> account mapping lived only in the RunnerHub's + // in-memory binding, so a Server restart or Runner re-enroll between + // Provision and Start left StartAgentSession unable to say whose session it + // was recording. It is also what SEA-1516 reattach recovery reads to name + // every agent stranded by a Runner restart. + if err := s.store.RecordAgentPlacement(ctx, store.AccountID(req.Msg.GetAgentAccountId()), runnerID, resp.GetContainerName()); err != nil { + return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("recording agent placement: %w", err)) } return connect.NewResponse(resp), nil } @@ -131,12 +155,34 @@ func (s *service) StartAgentSession( if err != nil { return nil, err } - // Bind the session_id -> container_name link of the durable ownership chain - // now that the Runner has started the session. container_name - // is the request handle; session_id is the server-minted response. Completes - // the chain SubscribeAgentSession resolves to authorize a subscriber. - if err := s.store.RecordAgentSession(ctx, resp.GetSessionId(), req.Msg.GetContainerName()); err != nil { - return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("recording agent session ownership: %w", err)) + // Complete the durable ownership chain now that the Runner has started the + // session: session_id (the server-minted response) -> agent_account_id, the + // chain SubscribeAgentSession resolves to authorize a subscriber. The + // request carries only container_name, so the owning account comes from the + // placement recorded at Provision — a durable read, not an in-memory + // binding, so the ownership record is written correctly even across a Server + // restart or a Runner re-enroll since the container was provisioned. + // + // Either store step can fail AFTER the Runner has irreversibly started an + // agent, and unlike Provision this does NOT self-heal: Provision's write is + // an idempotent upsert and its client_request_id dedups a retry at the + // router, whereas StartAgentSessionRequest carries no client_request_id (the + // relay id is minted per call above), so a client retry issues a genuinely + // NEW Start rather than rejoining this one. Returning the error alone would + // therefore discard the only handle to a LIVE session — the response, and + // with it the session id, never reaches the caller, so it can never Stop, + // Reload or Subscribe it, while the hub's promoteSession binding inside + // hub.Start already considers it live. So we tear the session back down and + // keep the invariant: either the session exists AND is recorded, or it does + // not exist. + agentAccountID, err := s.store.AgentForContainer(ctx, req.Msg.GetContainerName()) + if err != nil { + return nil, s.abandonStartedSession(ctx, req.Msg.GetContainerName(), resp.GetSessionId(), + fmt.Errorf("resolving agent for container: %w", err)) + } + if err := s.store.RecordAgentSession(ctx, resp.GetSessionId(), agentAccountID); err != nil { + return nil, s.abandonStartedSession(ctx, req.Msg.GetContainerName(), resp.GetSessionId(), + fmt.Errorf("recording agent session ownership: %w", err)) } return connect.NewResponse(resp), nil } @@ -268,9 +314,9 @@ func (s *service) SubscribeEvents( // SubscribeAgentSession streams one agent session's typed observation trace to a // caller authorized to see it. It first resolves-and-authorizes in one step — // RequireAgentSessionSubscriber walks the durable ownership chain (session_id -> -// container_name -> agent_account_id -> home_channel_id) and checks the caller's -// membership on that home channel, returning the SAME not-found for an unknown -// session and a non-member so neither can probe session existence. Only past +// agent_account_id -> home_channel_id) and checks the caller's membership on +// that home channel, returning the SAME not-found for an unknown session and a +// non-member so neither can probe session existence. Only past // that gate does it subscribe to the session's live tail and // forward frames until the client disconnects, the session ends, or the // subscriber lags past its buffer. No snapshot replay: the observation pane is a @@ -320,6 +366,37 @@ func (s *service) SubscribeAgentSession( } } +// abandonStartedSession tears down a session StartAgentSession already started +// on the Runner but could not record, and returns the Internal error to fail +// with. Best-effort: the Stop's own failure must not mask the original cause, +// which is why the returned error NAMES the session id — a session the store +// never recorded and Stop could not kill is only reapable by an operator who +// can see its id. +// +// The teardown runs on context.WithoutCancel(ctx): the caller's context may +// already be cancelled (a client that gave up is one plausible reason the store +// write failed at all), and the session is live regardless — the same reasoning +// AgentRuntime.Launch uses to remove a half-provisioned container +// (internal/runtime/agent.go). The dispatch path has no independent deadline, so +// the Stop is bounded here by rollbackStopTimeout — StartAgentSession returns +// within that bound even against a Runner that accepts Stop but never answers. +func (s *service) abandonStartedSession(ctx context.Context, containerName, sessionID string, cause error) error { + stopErr := errNoSessionID + if sessionID != "" { + stopCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackStopTimeout) + defer cancel() + _, stopErr = s.hub.Stop(stopCtx, "", &compassv1.StopAgentSessionRequest{SessionId: sessionID}) + } + if stopErr != nil { + slog.ErrorContext(ctx, "started agent session could not be recorded or stopped; it is running unreapable", + "session_id", sessionID, "container_name", containerName, "cause", cause, "stop_error", stopErr) + } else { + slog.ErrorContext(ctx, "started agent session could not be recorded; stopped it to avoid stranding", + "session_id", sessionID, "container_name", containerName, "cause", cause) + } + return connect.NewError(connect.CodeInternal, fmt.Errorf("%w (started session %q was rolled back)", cause, sessionID)) +} + // forward drains the replay snapshot (oldest first), then forwards the live tail // until the client disconnects, the server shuts down, or the subscriber lags // past the ring. The lag case emits a final ResyncRequired. Both phases select diff --git a/go/server/service_agentsession_pgtest_test.go b/go/server/service_agentsession_pgtest_test.go index 8cfccea2..74f485fd 100644 --- a/go/server/service_agentsession_pgtest_test.go +++ b/go/server/service_agentsession_pgtest_test.go @@ -75,11 +75,8 @@ func newAgentSessionFixture(t *testing.T) agentSessionFixture { if err != nil { t.Fatalf("CreateAgent: %v", err) } - if err := st.RecordAgentContainer(ctx, "cont-1", agent.ID); err != nil { - t.Fatalf("RecordAgentContainer: %v", err) - } const sessionID = "sess-1" - if err := st.RecordAgentSession(ctx, sessionID, "cont-1"); err != nil { + if err := st.RecordAgentSession(ctx, sessionID, agent.ID); err != nil { t.Fatalf("RecordAgentSession: %v", err) } diff --git a/go/server/service_placement_pgtest_test.go b/go/server/service_placement_pgtest_test.go new file mode 100644 index 00000000..757ba1cf --- /dev/null +++ b/go/server/service_placement_pgtest_test.go @@ -0,0 +1,607 @@ +//go:build pgtest && unix + +package server + +// The two placement-dependent handler seams, against a real Postgres AND a real +// Runner door: ProvisionAgentWorkspace's placement write, and StartAgentSession's +// placement read plus its post-relay roll-back. Both defects these cover live +// exactly here — the store primitives are individually correct under either bug, +// so only a test that drives the handler with a relay behind it can observe them. +// +// The Runner is a fake Sessions loop dialed into the REAL mounted RunnerService +// door (the runnerhub seam's own test shape, runnerhub/seam_test.go): it records +// every command the Server pushes and answers Provision/Start/Stop, so "the +// handler stopped the session it could not record" is an observed wire command, +// not a mock expectation. +// +// Two facts are seeded/read with pgx directly rather than through the Store: the +// backfilled placement (runner_id = '' is exactly what RecordAgentPlacement +// REFUSES, by design — only the migration writes it) and the raw agent_sessions +// row (the store exposes only the authz predicate over it). Both go through the +// pgtest DSN, so they land in this test's isolated schema. + +import ( + "context" + "crypto/sha256" + "errors" + "io" + "log/slog" + "net" + "net/http" + "net/http/httptest" + "runtime" + "strings" + "sync" + "testing" + "time" + + "connectrpc.com/connect" + "github.com/jackc/pgx/v5" + + "github.com/sealedsecurity/compass/go/events" + compassv1 "github.com/sealedsecurity/compass/go/gen/compass/v1" + "github.com/sealedsecurity/compass/go/gen/compass/v1/compassv1connect" + "github.com/sealedsecurity/compass/go/internal/auth" + "github.com/sealedsecurity/compass/go/internal/board" + compassv1internal "github.com/sealedsecurity/compass/go/internal/gen/compass/v1" + "github.com/sealedsecurity/compass/go/internal/gen/compass/v1/compassv1internalconnect" + "github.com/sealedsecurity/compass/go/internal/pgtest" + "github.com/sealedsecurity/compass/go/internal/runnerhub" + "github.com/sealedsecurity/compass/go/internal/store" +) + +// The fixed identities the fake Runner answers with, so assertions name them +// directly rather than threading values through the harness. +const ( + fakeRunnerID = "runner-1" + fakeRunnerToken = "cnVubmVyLXRva2Vu" //nolint:gosec // a test bearer, hashed into this test's own schema + fakeContainer = "compass-agent-c1" + fakeSessionID = "sess-relayed" +) + +// placementFixture is one service wired to a real store and a real Runner door, +// with a fake Runner attached and its command router proven live. +type placementFixture struct { + dsn string + store *store.Store + client compassv1connect.CompassServiceClient + runner *recordingRunner + // agentID is a real agent account (with its home channel) that every + // placement and session in these tests belongs to. + agentID store.AccountID +} + +// newPlacementFixture stands up store + hub + service + the mounted Runner door, +// dials the door with a fake Runner, and returns only once that Runner's command +// router is attached — so a dispatched command reaches the fake rather than +// racing the stream open. +func newPlacementFixture(t *testing.T) placementFixture { + return newPlacementFixtureWith(t, false) +} + +// newPlacementFixtureWith is newPlacementFixture with a fake Runner that either +// answers Stop (withholdStop=false) or accepts but never answers it +// (withholdStop=true) — the wedged-Runner shape the rollback-bound test drives. +func newPlacementFixtureWith(t *testing.T, withholdStop bool) placementFixture { + ctx := context.Background() // the test root context + dsn := pgtest.RequireDSN(t) + st, err := store.Open(ctx, dsn) + if err != nil { + t.Fatalf("store Open: %v", err) + } + t.Cleanup(st.Close) + + admin, err := st.BootstrapAdmin(ctx, store.NewUser{Handle: "admin", DisplayName: "admin"}) + if err != nil { + t.Fatalf("BootstrapAdmin: %v", err) + } + agent, err := st.CreateAgent(ctx, admin.ID, store.NewAgent{Handle: "atlas", DisplayName: "Atlas"}) + if err != nil { + t.Fatalf("CreateAgent: %v", err) + } + + // The Runner's bearer, resolvable through the production resolver: the door + // authenticates the fake Runner exactly as it authenticates a real one. + if err := st.PutTokenHash(ctx, sha256.Sum256([]byte(fakeRunnerToken)), + store.Subject{Kind: store.SubjectRunner, ID: fakeRunnerID}); err != nil { + t.Fatalf("PutTokenHash(runner): %v", err) + } + + bus := events.NewBus[busPayload]() + t.Cleanup(bus.Close) + brd := board.NewProjection(bus) + tail := newSessionTail() + hub := newRunnerHub(brd, tail, nil, slog.New(slog.DiscardHandler)) + svc := newService("test", bus, st, hub, brd, tail) + + return placementFixture{ + dsn: dsn, + store: st, + client: newH2CClient(t, newH2CTestServer(t, svc)), + runner: attachFakeRunner(t, st, hub, withholdStop), + agentID: agent.ID, + } +} + +// TestStartAgentSessionWithBackfilledPlacementRecordsSession is the upgrade-path +// assertion at the handler seam: a placement whose runner_id is the EMPTY +// SENTINEL — byte-for-byte what 0004's backfill writes for a container +// provisioned before the upgrade — still lets StartAgentSession resolve its +// owner and record the session. Without that backfill such a container has no +// placement row at all, AgentForContainer is ErrNotFound, and the container is +// permanently un-Startable; this test is what that regression reddens. +// +// It also pins the other half of why ” is the RIGHT sentinel rather than a +// tolerated placeholder: the same row Start resolves is invisible to reattach, +// so it is never re-driven against a Runner we only guessed at. +func TestStartAgentSessionWithBackfilledPlacementRecordsSession(t *testing.T) { + f := newPlacementFixture(t) + ctx := context.Background() // the test root context + + // Exactly the row 0004's backfill INSERT produces: real agent, real + // container name, empty runner id. Seeded with SQL because + // RecordAgentPlacement refuses an empty runner id — only the migration + // writes this shape. + execSQL(t, ctx, f.dsn, + `INSERT INTO agent_placements (agent_account_id, runner_id, container_name) VALUES ($1, '', $2)`, + string(f.agentID), fakeContainer) + + resp, err := f.client.StartAgentSession(ctx, connect.NewRequest(&compassv1.StartAgentSessionRequest{ + ContainerName: fakeContainer, + })) + if err != nil { + t.Fatalf("StartAgentSession on a backfilled placement = %v, want success (a pre-upgrade container must stay Startable)", err) + } + if got := resp.Msg.GetSessionId(); got != fakeSessionID { + t.Fatalf("session id = %q, want %q (the Runner's answer)", got, fakeSessionID) + } + + // The ownership row exists and names the placement's agent: keeping the + // container Startable is only worth anything if its session is owned. + if got := sessionOwner(t, ctx, f.dsn, fakeSessionID); got != string(f.agentID) { + t.Fatalf("session %q is owned by %q, want %q", fakeSessionID, got, f.agentID) + } + + // Reattach cannot see the sentinel row: the guard rejects an empty runner id + // outright, and no enrolled Runner has an empty id (it is the token + // subject), so a backfilled placement is never re-driven at a guess. + if _, err := f.store.ListAgentPlacementsForRunner(ctx, ""); !errors.Is(err, store.ErrInvalidArgument) { + t.Fatalf(`ListAgentPlacementsForRunner("") = %v, want ErrInvalidArgument (the sentinel must be invisible to reattach)`, err) + } + placements, err := f.store.ListAgentPlacementsForRunner(ctx, fakeRunnerID) + if err != nil { + t.Fatalf("ListAgentPlacementsForRunner(%q): %v", fakeRunnerID, err) + } + if len(placements) != 0 { + t.Fatalf("runner %q sees %d placements, want 0 (a backfilled row belongs to no known Runner)", fakeRunnerID, len(placements)) + } +} + +// TestStartAgentSessionStopsTheSessionItCannotRecord pins the anti-stranding +// invariant: when a post-relay store step fails, the Runner has ALREADY started +// an agent, so the handler must tear it back down rather than return an error +// that discards the only handle to it. Driven by the real failure — an UNPLACED +// container, so AgentForContainer is ErrNotFound — and asserted on the wire: the +// fake Runner must have received a Stop naming the session Start just minted, +// and the returned error must name that session id so an operator can reap it +// even if the Stop itself had failed. +// +// Under the pre-fix handler the Runner sees a Start and no Stop and the error +// carries no session id: the session runs on, unaddressable forever. +func TestStartAgentSessionStopsTheSessionItCannotRecord(t *testing.T) { + f := newPlacementFixture(t) + ctx := context.Background() // the test root context + + // No placement is seeded, so the post-relay AgentForContainer read fails + // after the Runner has started the agent. + _, err := f.client.StartAgentSession(ctx, connect.NewRequest(&compassv1.StartAgentSessionRequest{ + ContainerName: fakeContainer, + })) + if err == nil { + t.Fatal("StartAgentSession on an unplaced container = success, want an error") + } + if got := connect.CodeOf(err); got != connect.CodeInternal { + t.Fatalf("error code = %v, want Internal", got) + } + if !strings.Contains(err.Error(), fakeSessionID) { + t.Fatalf("error %q does not name the started session %q; an operator could not reap it", err, fakeSessionID) + } + + // The Runner was told to stop the session it just started — the invariant. + if !f.runner.sawStop(fakeSessionID) { + t.Fatalf("runner never received a Stop for session %q; the started session is stranded (commands seen: %v)", + fakeSessionID, f.runner.commands()) + } + // And nothing was recorded, so the store agrees the session does not exist: + // either it exists AND is recorded, or it does not exist. + if got := sessionRowCount(t, ctx, f.dsn, fakeSessionID); got != 0 { + t.Fatalf("agent_sessions has %d rows for the rolled-back session %q, want 0", got, fakeSessionID) + } +} + +// TestStartAgentSessionRollbackStopIsBounded pins that the anti-stranding Stop +// cannot itself hang StartAgentSession: against a Runner that ACCEPTS the Stop +// but never answers its result (the wedged-but-connected shape), the handler +// must still return within rollbackStopTimeout rather than blocking forever on +// a dispatch path that has no deadline of its own. Without the bound this test +// hangs until the go-test binary's own timeout kills it; with it the handler +// returns the same Internal error naming the session promptly. The timeout is +// shortened here so the test is fast and deterministic — no real long sleep. +func TestStartAgentSessionRollbackStopIsBounded(t *testing.T) { + prev := rollbackStopTimeout + rollbackStopTimeout = 200 * time.Millisecond + t.Cleanup(func() { rollbackStopTimeout = prev }) + + f := newPlacementFixtureWith(t, true) // Runner withholds the Stop result + ctx := context.Background() // the test root context + + // A generous ceiling relative to the 200ms bound: comfortably above it, far + // below any real 30s hang, so a regression to unbounded reddens as a + // timeout here rather than stalling the whole suite. + done := make(chan error, 1) + go func() { + _, err := f.client.StartAgentSession(ctx, connect.NewRequest(&compassv1.StartAgentSessionRequest{ + ContainerName: fakeContainer, + })) + done <- err + }() + + select { + case err := <-done: + if err == nil { + t.Fatal("StartAgentSession on an unplaced container = success, want an error") + } + if got := connect.CodeOf(err); got != connect.CodeInternal { + t.Fatalf("error code = %v, want Internal", got) + } + if !strings.Contains(err.Error(), fakeSessionID) { + t.Fatalf("error %q does not name the started session %q; an operator could not reap it", err, fakeSessionID) + } + case <-time.After(10 * time.Second): + t.Fatalf("StartAgentSession did not return within 10s against a Runner that never answers Stop; the rollback Stop is unbounded") + } + + // The Stop was still pushed — the invariant holds; it is merely bounded. + if !f.runner.sawStop(fakeSessionID) { + t.Fatalf("runner never received a Stop for session %q (commands seen: %v)", fakeSessionID, f.runner.commands()) + } +} + +// TestProvisionAgentWorkspaceRecordsPlacementNamingServingRunner pins that the +// placement Provision writes names the Runner that ACTUALLY served the relay +// (the id the hub returns from the round trip), not a re-read of the registry — +// a placement pointing at the wrong Runner would make reattach re-drive the +// wrong set. Asserted through the reattach read itself, runner_id's only +// consumer, then carried one step further: the provisioned container is now +// Startable, because Provision's write is precisely what Start's read needs. +func TestProvisionAgentWorkspaceRecordsPlacementNamingServingRunner(t *testing.T) { + f := newPlacementFixture(t) + ctx := context.Background() // the test root context + + resp, err := f.client.ProvisionAgentWorkspace(ctx, connect.NewRequest(&compassv1.ProvisionAgentWorkspaceRequest{ + AgentAccountId: string(f.agentID), + Repo: &compassv1.ProvisionAgentWorkspaceRequest_LocalPath{LocalPath: "/mirror/repo.git"}, + ClientRequestId: "prov-1", + })) + if err != nil { + t.Fatalf("ProvisionAgentWorkspace = %v, want success", err) + } + if got := resp.Msg.GetContainerName(); got != fakeContainer { + t.Fatalf("container name = %q, want %q (the Runner's answer)", got, fakeContainer) + } + + placements, err := f.store.ListAgentPlacementsForRunner(ctx, fakeRunnerID) + if err != nil { + t.Fatalf("ListAgentPlacementsForRunner(%q): %v", fakeRunnerID, err) + } + if len(placements) != 1 { + t.Fatalf("runner %q has %d placements after one provision, want 1", fakeRunnerID, len(placements)) + } + got := placements[0] + if got.AgentAccountID != f.agentID { + t.Fatalf("placement agent = %q, want %q", got.AgentAccountID, f.agentID) + } + if got.RunnerID != fakeRunnerID { + t.Fatalf("placement runner = %q, want %q (the Runner that served the relay)", got.RunnerID, fakeRunnerID) + } + if got.ContainerName != fakeContainer { + t.Fatalf("placement container = %q, want %q", got.ContainerName, fakeContainer) + } + + startResp, err := f.client.StartAgentSession(ctx, connect.NewRequest(&compassv1.StartAgentSessionRequest{ + ContainerName: fakeContainer, + })) + if err != nil { + t.Fatalf("StartAgentSession after Provision = %v, want success", err) + } + if owner := sessionOwner(t, ctx, f.dsn, startResp.Msg.GetSessionId()); owner != string(f.agentID) { + t.Fatalf("session owner = %q, want %q", owner, f.agentID) + } +} + +// --- direct SQL (the two facts the Store deliberately does not expose) -------- + +// execSQL runs one statement against the test's isolated schema. Used only to +// seed the backfilled placement, whose empty runner_id RecordAgentPlacement +// refuses by design — the migration is its only production writer. +func execSQL(t *testing.T, ctx context.Context, dsn, sql string, args ...any) { + t.Helper() + conn := connectPG(t, ctx, dsn) + if _, err := conn.Exec(ctx, sql, args...); err != nil { + t.Fatalf("exec %q: %v", sql, err) + } +} + +// sessionOwner reads the agent_sessions ownership row for sessionID, failing if +// there is none. The Store exposes only the authz predicate over this table, so +// asserting "the row exists and names this agent" reads it directly. +func sessionOwner(t *testing.T, ctx context.Context, dsn, sessionID string) string { + t.Helper() + conn := connectPG(t, ctx, dsn) + var owner string + if err := conn.QueryRow(ctx, + `SELECT agent_account_id FROM agent_sessions WHERE session_id = $1`, sessionID, + ).Scan(&owner); err != nil { + t.Fatalf("session %q has no recorded owner: %v", sessionID, err) + } + return owner +} + +// sessionRowCount counts ownership rows for sessionID (0 or 1). +func sessionRowCount(t *testing.T, ctx context.Context, dsn, sessionID string) int { + t.Helper() + conn := connectPG(t, ctx, dsn) + var n int + if err := conn.QueryRow(ctx, + `SELECT count(*) FROM agent_sessions WHERE session_id = $1`, sessionID, + ).Scan(&n); err != nil { + t.Fatalf("count agent_sessions: %v", err) + } + return n +} + +// connectPG opens a pgx connection on the pgtest DSN (so it lands in this test's +// isolated schema), closed on cleanup. +func connectPG(t *testing.T, ctx context.Context, dsn string) *pgx.Conn { + t.Helper() + conn, err := pgx.Connect(ctx, dsn) + if err != nil { + t.Fatalf("pgx.Connect: %v", err) + } + t.Cleanup(func() { _ = conn.Close(ctx) }) // cleanup close; nothing actionable remains + return conn +} + +// --- the fake Runner --------------------------------------------------------- + +// attachFakeRunner mounts the real RunnerService door over hub, dials it with a +// fake Runner, runs its Sessions loop, and returns once the router is live. +func attachFakeRunner(t *testing.T, st *store.Store, hub *runnerhub.Hub, withholdStop bool) *recordingRunner { + t.Helper() + path, handler := runnerhub.NewMountedHandler(hub, + func(ctx context.Context, presented string, want store.SubjectKind) (store.Subject, error) { + return auth.ResolveToken(ctx, st, presented, want) + }) + mux := http.NewServeMux() + mux.Handle(path, handler) + srv := httptest.NewUnstartedServer(mux) + srv.Config.Protocols = cleartextHTTP2() + srv.Start() + t.Cleanup(srv.Close) + + // The loop's own context: cancelled at cleanup to end the Sessions stream. + ctx, cancel := context.WithCancel(context.Background()) // the test root context + t.Cleanup(cancel) + + tr := h2cTransport(func(ctx context.Context, network, addr string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, network, addr) + }) + t.Cleanup(tr.CloseIdleConnections) + rc := compassv1internalconnect.NewRunnerServiceClient( + &http.Client{Transport: tr}, srv.URL, + connect.WithInterceptors(runnerBearer(fakeRunnerToken)), + ) + if _, err := rc.Enroll(ctx, connect.NewRequest(&compassv1internal.EnrollRequest{RunnerId: fakeRunnerID})); err != nil { + t.Fatalf("Enroll = %v, want success", err) + } + + rec := &recordingRunner{attached: make(chan struct{}), withholdStop: withholdStop} + stream := rc.Sessions(ctx) + loopDone := make(chan error, 1) + go rec.serve(stream, loopDone) + t.Cleanup(func() { + // Close the request half: the loop's Receive then sees a clean EOF and + // the server's Sessions handler detaches the router — the seam test's + // proven teardown. (Cancelling ctx alone does not unblock a client-side + // Receive already parked on the HTTP/2 body.) + if err := stream.CloseRequest(); err != nil { + t.Errorf("CloseRequest = %v", err) + } + if err := <-loopDone; err != nil { + t.Errorf("fake runner sessions loop ended with %v, want a clean EOF", err) + } + cancel() + }) + + // Gate on the SERVER-side router being live, not merely the client having + // sent. connect-go initiates the request on the loop's bootstrap Send, but + // the handler's router.attach runs asynchronously once that reaches the + // server; a command dispatched into that window gets a retriable Unavailable + // ("no live runner sessions stream"). Probing with a real round trip — a + // read-only Status the fake answers — is the observed attach, the same gate + // shape runnerhub's integration test uses. Yield between probes; the + // deadline turns a genuinely wedged seam into a fast failure, never a sleep. + select { + case <-rec.attached: + case <-timeAfter(): + t.Fatal("fake runner never opened its Sessions stream") + } + deadline := timeAfter() + for { + _, err := hub.Status(ctx, "attach-probe", &compassv1.GetAgentStatusRequest{SessionId: "attach-probe"}) + if err == nil { + break + } + if connect.CodeOf(err) != connect.CodeUnavailable { + t.Fatalf("attach probe = %v, want success or a transient Unavailable", err) + } + select { + case <-deadline: + t.Fatalf("runner command router never attached a live Sessions stream: %v", err) + default: + } + runtime.Gosched() + } + // The probe is bookkeeping, not a command under test. + rec.forget() + return rec +} + +// recordingRunner is a minimal Runner-side Sessions loop that records every +// command the Server pushes and answers the three this suite drives. Recording +// on the wire is what makes "the handler stopped the session" an observed fact +// rather than a mock expectation. +type recordingRunner struct { + // attached closes once the bootstrap Send has flushed the stream open. + attached chan struct{} + + // withholdStop makes the loop record a Stop but never answer its result — + // the wedged-but-connected Runner that hangs an unbounded rollback Stop. + withholdStop bool + + mu sync.Mutex + seen []*compassv1internal.SessionsResponse +} + +// serve runs the dispatch loop. Like the seam test's loop it opens with one +// bootstrap Send, because connect-go does not initiate the request (and so the +// server's router.attach never runs) until the client's first Send; the empty +// result carries no request id, so router.complete treats it as an unknown-id +// no-op. It ends on the clean EOF the caller's CloseRequest produces. +func (r *recordingRunner) serve( + stream *connect.BidiStreamForClient[compassv1internal.SessionsRequest, compassv1internal.SessionsResponse], + done chan<- error, +) { + if err := stream.Send(&compassv1internal.SessionsRequest{}); err != nil { + done <- err + return + } + close(r.attached) + for { + cmd, err := stream.Receive() + if err != nil { + if errors.Is(err, io.EOF) { + done <- nil + return + } + done <- err + return + } + r.record(cmd) + if r.withholdStop && cmd.GetStop() != nil { + continue // record it, but never answer: the wedged-Runner shape + } + if err := stream.Send(answer(cmd)); err != nil { + done <- err + return + } + } +} + +func (r *recordingRunner) record(cmd *compassv1internal.SessionsResponse) { + r.mu.Lock() + defer r.mu.Unlock() + r.seen = append(r.seen, cmd) +} + +// forget drops every recorded command — used once, to discard the attach probe +// so the assertions see only the commands their handler drove. +func (r *recordingRunner) forget() { + r.mu.Lock() + defer r.mu.Unlock() + r.seen = nil +} + +// commands summarizes every command seen, for failure output. +func (r *recordingRunner) commands() []string { + r.mu.Lock() + defer r.mu.Unlock() + out := make([]string, 0, len(r.seen)) + for _, c := range r.seen { + switch v := c.GetCommand().(type) { + case *compassv1internal.SessionsResponse_Provision: + out = append(out, "provision "+v.Provision.GetAgentAccountId()) + case *compassv1internal.SessionsResponse_Start: + out = append(out, "start "+v.Start.GetContainerName()) + case *compassv1internal.SessionsResponse_Stop: + out = append(out, "stop "+v.Stop.GetSessionId()) + default: + out = append(out, "other") + } + } + return out +} + +// sawStop reports whether the Server pushed a Stop for sessionID. +func (r *recordingRunner) sawStop(sessionID string) bool { + r.mu.Lock() + defer r.mu.Unlock() + for _, c := range r.seen { + if stop := c.GetStop(); stop != nil && stop.GetSessionId() == sessionID { + return true + } + } + return false +} + +// answer builds the correlated result for one command: a fixed container name +// for Provision, a fixed session id for Start, success for Stop, and an explicit +// error for anything this suite does not drive, so an unexpected command is a +// loud failure rather than a hang. +func answer(cmd *compassv1internal.SessionsResponse) *compassv1internal.SessionsRequest { + out := &compassv1internal.SessionsRequest{RequestId: cmd.GetRequestId()} + switch cmd.GetCommand().(type) { + case *compassv1internal.SessionsResponse_Provision: + out.Result = &compassv1internal.SessionsRequest_Provision{ + Provision: &compassv1.ProvisionAgentWorkspaceResponse{ContainerName: fakeContainer}, + } + case *compassv1internal.SessionsResponse_Start: + out.Result = &compassv1internal.SessionsRequest_Start{ + Start: &compassv1.StartAgentSessionResponse{SessionId: fakeSessionID}, + } + case *compassv1internal.SessionsResponse_Status: + out.Result = &compassv1internal.SessionsRequest_Status{Status: &compassv1.GetAgentStatusResponse{}} + case *compassv1internal.SessionsResponse_Stop: + out.Result = &compassv1internal.SessionsRequest_Stop{Stop: &compassv1.StopAgentSessionResponse{}} + default: + out.Result = &compassv1internal.SessionsRequest_Error{Error: &compassv1internal.RunnerError{ + Code: compassv1internal.RunnerErrorCode_RUNNER_ERROR_CODE_INTERNAL, + Message: "unexpected command", + }} + } + return out +} + +// runnerBearer stamps the Runner's bearer on every outbound RPC (unary + +// streaming), mirroring the Runner-side interceptor so the door authenticates. +type runnerBearer string + +func (b runnerBearer) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc { + return func(ctx context.Context, req connect.AnyRequest) (connect.AnyResponse, error) { + req.Header().Set("Authorization", "Bearer "+string(b)) + return next(ctx, req) + } +} + +func (b runnerBearer) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc { + return func(ctx context.Context, spec connect.Spec) connect.StreamingClientConn { + conn := next(ctx, spec) + conn.RequestHeader().Set("Authorization", "Bearer "+string(b)) + return conn + } +} + +func (b runnerBearer) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc { + return next +}