feat(console): persist terminal session routes for restart recovery

console/cmd/console/main.go

@@ -95,6 +95,12 @@ func main() {
 	registryService.SetHasher(db.Hasher)
 	registryService.SetTaskRetention(time.Duration(cfg.TaskRetentionDays) * 24 * time.Hour)
 	registryService.ConfigureProxy(cfg.ProxyEnabled, cfg.ProxyAllowedWorkerCIDRs, cfg.ProxyAllowedWorkerPorts, cfg.ProxyAllowedDirectDomains)
+	restoreCtx, restoreCancel := context.WithTimeout(context.Background(), 10*time.Second)
+	if err := registryService.RestoreTerminalSessionRoutes(restoreCtx, time.Now()); err != nil {
+		restoreCancel()
+		fatal("failed to restore terminal session routes", "error", err)
+	}
+	restoreCancel()
 	grpcSrv := grpcserver.NewServer(registryService)
 	httpHandler := httpapi.NewWorkerHandler(
 		store,

console/db/migrations/00006_terminal_session_routes.sql

@@ -0,0 +1,20 @@
+-- +goose Up
+CREATE TABLE terminal_session_routes (
+    scoped_session_id TEXT PRIMARY KEY,
+    node_id TEXT NOT NULL,
+    lease_expires_unix_ms INTEGER NOT NULL CHECK (lease_expires_unix_ms > 0),
+    last_used_unix_ms INTEGER NOT NULL,
+    created_at_unix_ms INTEGER NOT NULL,
+    updated_at_unix_ms INTEGER NOT NULL
+);
+
+CREATE INDEX idx_terminal_session_routes_node
+    ON terminal_session_routes(node_id);
+
+CREATE INDEX idx_terminal_session_routes_lease
+    ON terminal_session_routes(lease_expires_unix_ms);
+
+-- +goose Down
+DROP INDEX IF EXISTS idx_terminal_session_routes_lease;
+DROP INDEX IF EXISTS idx_terminal_session_routes_node;
+DROP TABLE IF EXISTS terminal_session_routes;

console/db/queries/terminal_session_routes.sql

@@ -0,0 +1,50 @@
+-- name: UpsertTerminalSessionRoute :execrows
+INSERT INTO terminal_session_routes (
+    scoped_session_id,
+    node_id,
+    lease_expires_unix_ms,
+    last_used_unix_ms,
+    created_at_unix_ms,
+    updated_at_unix_ms
+) VALUES (?, ?, ?, ?, ?, ?)
+ON CONFLICT(scoped_session_id) DO UPDATE SET
+    lease_expires_unix_ms = MAX(terminal_session_routes.lease_expires_unix_ms, excluded.lease_expires_unix_ms),
+    last_used_unix_ms = MAX(terminal_session_routes.last_used_unix_ms, excluded.last_used_unix_ms),
+    updated_at_unix_ms = MAX(terminal_session_routes.updated_at_unix_ms, excluded.updated_at_unix_ms)
+WHERE terminal_session_routes.node_id = excluded.node_id;
+
+-- name: DeleteTerminalSessionRouteBySessionAndNode :execrows
+DELETE FROM terminal_session_routes
+WHERE scoped_session_id = ? AND node_id = ?;
+
+-- name: DeleteTerminalSessionRoutesByNode :execrows
+DELETE FROM terminal_session_routes
+WHERE node_id = ?;
+
+-- name: DeleteExpiredTerminalSessionRoutes :execrows
+DELETE FROM terminal_session_routes
+WHERE lease_expires_unix_ms <= ?;
+
+-- name: ListActiveTerminalSessionRoutes :many
+SELECT
+    scoped_session_id,
+    node_id,
+    lease_expires_unix_ms,
+    last_used_unix_ms,
+    created_at_unix_ms,
+    updated_at_unix_ms
+FROM terminal_session_routes
+WHERE lease_expires_unix_ms > ?
+ORDER BY node_id ASC, scoped_session_id ASC;
+
+-- name: GetTerminalSessionRouteBySession :one
+SELECT
+    scoped_session_id,
+    node_id,
+    lease_expires_unix_ms,
+    last_used_unix_ms,
+    created_at_unix_ms,
+    updated_at_unix_ms
+FROM terminal_session_routes
+WHERE scoped_session_id = ?
+LIMIT 1;

console/internal/grpcserver/connect_service_test.go

@@ -214,8 +214,8 @@ func TestDeleteProvisionedWorkerDisconnectsSessionAndRevokesCredential(t *testin
 		t.Fatalf("connect worker failed: %v", err)
 	}
 
-	if removed := svc.DeleteProvisionedWorker(workerID); !removed {
-		t.Fatalf("expected delete to return true")
+	if removed, err := svc.DeleteProvisionedWorker(workerID); err != nil || !removed {
+		t.Fatalf("expected delete to return true, removed=%t err=%v", removed, err)
 	}
 	if _, ok := svc.GetWorkerSecret(workerID); ok {
 		t.Fatalf("expected credential to be revoked")
@@ -1272,7 +1272,7 @@ func TestDispatchCommandTerminalSessionCapacityDoesNotClearConcurrentProvisional
 
 	session.resolvePending(&registryv1.CommandResult{
 		CommandId:       secondDispatch.GetCommandId(),
-		PayloadJson:     []byte(`{"session_id":"session-shared"}`),
+		PayloadJson:     []byte(`{"session_id":"session-shared","lease_expires_unix_ms":4102444800000}`),
 		CompletedUnixMs: now.UnixMilli(),
 	})
 	second := <-secondDone

console/internal/grpcserver/service.go

@@ -27,6 +27,7 @@ const (
 	defaultCommandDispatchTimeout = 60 * time.Second
 	defaultTerminalRouteTTL       = 30 * time.Minute
 	terminalRoutePruneMinInterval = 1 * time.Minute
+	terminalRouteStoreTimeout     = 5 * time.Second
 	computerUseCapabilityName     = "computeruse"
 	computerUseCapabilityDeclared = "computerUse"
 	readImageCapabilityName       = "readimage"
@@ -69,6 +70,7 @@ type RegistryService struct {
 	terminalNodeToSessionIDIndex map[string]map[string]struct{}
 	terminalRouteReservationSeq  uint64
 	terminalRouteTTL             time.Duration
+	terminalRouteStore           terminalSessionRouteStore
 	lastTerminalRoutePruneUnixMs atomic.Int64
 
 	tasksMu sync.RWMutex
@@ -109,6 +111,7 @@ func NewRegistryService(
 		terminalSessionToNode:        make(map[string]terminalSessionRoute),
 		terminalNodeToSessionIDIndex: make(map[string]map[string]struct{}),
 		terminalRouteTTL:             defaultTerminalRouteTTL,
+		terminalRouteStore:           store,
 		tasks:                        make(map[string]*taskRecord),
 		taskRequestReservations:      make(map[string]struct{}),
 		criticalPersistenceFailureFn: func(error) {},

console/internal/grpcserver/service_connect.go

@@ -71,7 +71,10 @@ func (s *RegistryService) Connect(stream grpc.BidiStreamingServer[registryv1.Con
 	}
 
 	session := newActiveSessionAt(hello.GetNodeId(), sessionID, hello, now)
-	recoveryCandidates := s.beginTerminalSessionRecovery(session.nodeID, now)
+	recoveryCandidates, err := s.beginTerminalSessionRecoveryWithError(session.nodeID, now)
+	if err != nil {
+		return status.Errorf(codes.Internal, "prepare terminal session recovery: %v", err)
+	}
 	session.setRecoveryCandidates(recoveryCandidates)
 	if err := s.configureSessionProxy(session, hello, workerSecret); err != nil {
 		return status.Errorf(codes.InvalidArgument, "invalid proxy endpoint: %v", err)

console/internal/grpcserver/service_credentials_runtime.go

@@ -139,21 +139,24 @@ func normalizeProvisioningWorkerType(workerType string) string {
 	}
 }
 
-func (s *RegistryService) DeleteProvisionedWorker(nodeID string) bool {
+func (s *RegistryService) DeleteProvisionedWorker(nodeID string) (bool, error) {
 	trimmedNodeID := strings.TrimSpace(nodeID)
 	if trimmedNodeID == "" {
-		return false
+		return false, nil
+	}
+	if _, err := s.deleteTerminalSessionRoutesByNode(trimmedNodeID); err != nil {
+		return false, err
 	}
 
 	deletedCredentialInMemory := s.deleteCredential(trimmedNodeID)
 	deletedCredentialInDB := s.store.DeleteCredential(trimmedNodeID)
 	deletedNode := s.store.Delete(trimmedNodeID)
 	if !deletedCredentialInMemory && !deletedCredentialInDB && !deletedNode {
-		return false
+		return false, nil
 	}
 
 	s.disconnectWorker(trimmedNodeID, "worker credential revoked")
-	return true
+	return true, nil
 }
 
 func (s *RegistryService) getCredential(nodeID string) (string, bool) {

console/internal/grpcserver/service_dispatch.go

@@ -236,19 +236,35 @@ func (s *RegistryService) dispatchCommandAttempt(
 			terminalRouteReservationID,
 		)
 	}
-	confirmTerminalRoute := func(resultPayload []byte) {
+	confirmTerminalRouteInMemory := func() {
 		if terminalSessionID != "" {
 			s.confirmTerminalSessionRoute(
 				terminalSessionID,
 				session.nodeID,
 				terminalRouteReservationID,
 				s.nowFn(),
 			)
-			if leaseExpiresUnixMs := terminalSessionLeaseExpiresUnixMs(resultPayload); leaseExpiresUnixMs > 0 {
-				s.updateTerminalSessionRouteLease(terminalSessionID, session.nodeID, leaseExpiresUnixMs, s.nowFn())
-			}
 		}
 	}
+	commitTerminalRoute := func(resultPayload []byte) error {
+		if terminalSessionID == "" {
+			return nil
+		}
+		confirmed, err := s.commitConfirmedTerminalSessionRoute(
+			terminalSessionID,
+			session.nodeID,
+			terminalRouteReservationID,
+			terminalSessionLeaseExpiresUnixMs(resultPayload),
+			s.nowFn(),
+		)
+		if err != nil {
+			return err
+		}
+		if !confirmed {
+			return errors.New("terminal session route changed before persistence")
+		}
+		return nil
+	}
 
 	commandID, err := s.newCommandIDFn()
 	if err != nil {
@@ -297,7 +313,7 @@ func (s *RegistryService) dispatchCommandAttempt(
 	if onDispatched != nil {
 		if err := onDispatched(commandID); err != nil {
 			if terminalRouteReservationID != 0 {
-				confirmTerminalRoute(nil)
+				confirmTerminalRouteInMemory()
 			}
 			return dispatchAttemptResult{}, err
 		}
@@ -308,7 +324,7 @@ func (s *RegistryService) dispatchCommandAttempt(
 		if terminalRouteReservationID != 0 && terminalSessionID != "" {
 			// The dispatch reached the worker stream, so cancellation does not
 			// prove that session creation failed.
-			confirmTerminalRoute(nil)
+			confirmTerminalRouteInMemory()
 		}
 		if errors.Is(commandCtx.Err(), context.DeadlineExceeded) {
 			return dispatchAttemptResult{}, context.DeadlineExceeded
@@ -320,7 +336,13 @@ func (s *RegistryService) dispatchCommandAttempt(
 			return dispatchAttemptResult{}, status.Error(codes.Unavailable, "worker session closed before command result")
 		}
 		if outcome.err == nil && terminalSessionID != "" {
-			confirmTerminalRoute(outcome.payloadJSON)
+			if capability == taskCapabilityTerminalExec {
+				if err := commitTerminalRoute(outcome.payloadJSON); err != nil {
+					return dispatchAttemptResult{}, status.Errorf(codes.Internal, "persist terminal session route: %v", err)
+				}
+			} else {
+				confirmTerminalRouteInMemory()
+			}
 			return dispatchAttemptResult{outcome: outcome}, nil
 		}
 		if outcome.err == nil || terminalSessionID == "" {
@@ -331,8 +353,8 @@ func (s *RegistryService) dispatchCommandAttempt(
 		case isSessionNotFoundCommandError(outcome.err):
 			if terminalRouteReservationID != 0 {
 				rollbackTerminalRouteReservation()
-			} else {
-				s.clearTerminalSessionRoute(terminalSessionID, session.nodeID)
+			} else if err := s.clearTerminalSessionRoute(terminalSessionID, session.nodeID); err != nil {
+				return dispatchAttemptResult{}, status.Errorf(codes.Internal, "delete terminal session route: %v", err)
 			}
 		case isSessionCapacityCommandError(outcome.err):
 			releaseResult := rollbackTerminalRouteReservation()
@@ -344,7 +366,7 @@ func (s *RegistryService) dispatchCommandAttempt(
 			}, nil
 		case terminalRouteReservationID != 0:
 			// Other execution errors do not prove that session creation failed.
-			confirmTerminalRoute(nil)
+			confirmTerminalRouteInMemory()
 		}
 		return dispatchAttemptResult{outcome: outcome}, nil
 	}

console/internal/grpcserver/terminal_capacity_dispatch_test.go

@@ -270,7 +270,7 @@ func TestDispatchCommandRetriesTerminalCapacityOnAnotherWorker(t *testing.T) {
 	}
 	workerB.resolvePending(&registryv1.CommandResult{
 		CommandId:       dispatchB.GetCommandId(),
-		PayloadJson:     []byte(`{"session_id":"session-retry","stdout":"ok"}`),
+		PayloadJson:     []byte(`{"session_id":"session-retry","stdout":"ok","lease_expires_unix_ms":4102444800000}`),
 		CompletedUnixMs: now.UnixMilli(),
 	})
 
@@ -371,7 +371,7 @@ func TestSubmitTaskSkipsReportedFullConnectedWorker(t *testing.T) {
 		Payload: &registryv1.ConnectRequest_CommandResult{
 			CommandResult: &registryv1.CommandResult{
 				CommandId:       dispatchB.GetCommandId(),
-				PayloadJson:     []byte(`{"session_id":"obx:owner-a:session-connected-skip","stdout":"ok"}`),
+				PayloadJson:     []byte(`{"session_id":"obx:owner-a:session-connected-skip","stdout":"ok","lease_expires_unix_ms":4102444800000}`),
 				CompletedUnixMs: now.UnixMilli(),
 			},
 		},
@@ -531,7 +531,7 @@ func TestSubmitTaskRetriesCapacityAcrossConnectedWorkers(t *testing.T) {
 		Payload: &registryv1.ConnectRequest_CommandResult{
 			CommandResult: &registryv1.CommandResult{
 				CommandId:       dispatchB.GetCommandId(),
-				PayloadJson:     []byte(`{"session_id":"obx:owner-a:external-session","stdout":"ok"}`),
+				PayloadJson:     []byte(`{"session_id":"obx:owner-a:external-session","stdout":"ok","lease_expires_unix_ms":4102444800000}`),
 				CompletedUnixMs: now.UnixMilli(),
 			},
 		},
@@ -618,7 +618,7 @@ func TestTerminalCapacityRetryWorksForAllTaskModes(t *testing.T) {
 			workerB.resolvePending(&registryv1.CommandResult{
 				CommandId: dispatchB.GetCommandId(),
 				PayloadJson: []byte(`{"session_id":"obx:owner-a:session-mode-` +
-					string(mode) + `","stdout":"ok"}`),
+					string(mode) + `","stdout":"ok","lease_expires_unix_ms":4102444800000}`),
 				CompletedUnixMs: now.UnixMilli(),
 			})
 
@@ -778,7 +778,7 @@ func TestConcurrentProvisionalCapacityOnlyLastRollbackRetries(t *testing.T) {
 	retryDispatch := receiveCommandDispatch(t, workerB)
 	workerB.resolvePending(&registryv1.CommandResult{
 		CommandId:       retryDispatch.GetCommandId(),
-		PayloadJson:     []byte(`{"session_id":"session-shared-retry","stdout":"ok"}`),
+		PayloadJson:     []byte(`{"session_id":"session-shared-retry","stdout":"ok","lease_expires_unix_ms":4102444800000}`),
 		CompletedUnixMs: now.UnixMilli(),
 	})
 	second := <-secondDone