added gateway-side changes - #5
Open
rohanphadnis-thunder wants to merge 6 commits into
Open
Conversation
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.
Autofix Details
Bugbot Autofix prepared fixes for both issues found in the latest run.
- ✅ Fixed: Agent always requires Thunder
- Thunder node enrollment is now best-effort in runTSNetRouteProxy, so enrollment failures no longer block the TSNet route proxy.
- ✅ Fixed: Incomplete agent token validation
- Thunder requireAgentState now validates ManagedPoolInstanceID against the live managed pool, matching compute agent-token checks.
Or push these changes by commenting:
@cursor push d9e7345023
Preview (d9e7345023)
diff --git a/pkg/agent/transport.go b/pkg/agent/transport.go
--- a/pkg/agent/transport.go
+++ b/pkg/agent/transport.go
@@ -89,10 +89,17 @@
}
}
}
+ // Thunder node enrollment is optional: agents in environments without
+ // Thunder credentials (or where enrollment otherwise fails) must still be
+ // able to bring up the route proxy.
if err := setupThunderNode(ctx, client, agentToken, tailscaleIPs, stdout, stderr); err != nil {
- return err
+ if stderr == nil {
+ stderr = io.Discard
+ }
+ fmt.Fprintf(stderr, "Thunder node enrollment skipped: %v\n", err)
+ } else {
+ defer deleteThunderNodeEnrollment(context.Background(), client, agentToken, stderr)
}
- defer deleteThunderNodeEnrollment(context.Background(), client, agentToken, stderr)
listener, err := server.Listen("tcp", fmt.Sprintf(":%d", types.DefaultAgentTSNetRouteProxyPort))
if err != nil {
diff --git a/pkg/gateway/gateway.go b/pkg/gateway/gateway.go
--- a/pkg/gateway/gateway.go
+++ b/pkg/gateway/gateway.go
@@ -383,11 +383,12 @@
// Register Thunder service
ts, err := thundersvc.NewService(thundersvc.ServiceOpts{
- RedisClient: g.RedisClient,
- Client: thundersvc.NewClientFromEnv(nil),
- ContainerRepo: g.ContainerRepo,
- WorkerRepo: g.workerRepo,
- ComputeRepo: g.ComputeRepo,
+ RedisClient: g.RedisClient,
+ Client: thundersvc.NewClientFromEnv(nil),
+ ContainerRepo: g.ContainerRepo,
+ WorkerRepo: g.workerRepo,
+ ComputeRepo: g.ComputeRepo,
+ ManagedPoolRepo: repository.NewManagedPoolRedisRepository(g.RedisClient),
})
if err != nil {
return err
diff --git a/pkg/gateway/services/thunder/service.go b/pkg/gateway/services/thunder/service.go
--- a/pkg/gateway/services/thunder/service.go
+++ b/pkg/gateway/services/thunder/service.go
@@ -26,12 +26,13 @@
}
type ServiceOpts struct {
- Repository Repository
- RedisClient *common.RedisClient
- Client ThunderClient
- ContainerRepo repository.ContainerRepository
- WorkerRepo repository.WorkerRepository
- ComputeRepo repository.ComputeRepository
+ Repository Repository
+ RedisClient *common.RedisClient
+ Client ThunderClient
+ ContainerRepo repository.ContainerRepository
+ WorkerRepo repository.WorkerRepository
+ ComputeRepo repository.ComputeRepository
+ ManagedPoolRepo repository.ManagedPoolRepository
}
var errThunderClientEnrollmentExists = errors.New("Thunder client enrollment already exists for container")
@@ -39,11 +40,12 @@
type Service struct {
pb.UnimplementedThunderServiceServer
- repo Repository
- client ThunderClient
- containerRepo repository.ContainerRepository
- workerRepo repository.WorkerRepository
- computeRepo repository.ComputeRepository
+ repo Repository
+ client ThunderClient
+ containerRepo repository.ContainerRepository
+ workerRepo repository.WorkerRepository
+ computeRepo repository.ComputeRepository
+ managedPoolRepo repository.ManagedPoolRepository
}
func NewService(opts ServiceOpts) (*Service, error) {
@@ -70,11 +72,12 @@
}
return &Service{
- repo: repo,
- client: client,
- containerRepo: opts.ContainerRepo,
- workerRepo: opts.WorkerRepo,
- computeRepo: opts.ComputeRepo,
+ repo: repo,
+ client: client,
+ containerRepo: opts.ContainerRepo,
+ workerRepo: opts.WorkerRepo,
+ computeRepo: opts.ComputeRepo,
+ managedPoolRepo: opts.ManagedPoolRepo,
}, nil
}
@@ -375,6 +378,20 @@
if current == nil || current.TokenHash != state.TokenHash {
return nil, fmt.Errorf("invalid agent token")
}
+ if current.ManagedPoolInstanceID != "" {
+ if s.managedPoolRepo == nil {
+ return nil, fmt.Errorf("managed pool repository is unavailable")
+ }
+ poolState, err := s.managedPoolRepo.GetManagedPoolState(ctx, current.WorkspaceID, current.PoolName)
+ if err != nil {
+ return nil, err
+ }
+ if poolState == nil || poolState.Name != current.PoolName || poolState.WorkspaceID != current.WorkspaceID ||
+ poolState.ManagementSource == "" || poolState.Mode != string(types.PoolModeExternal) ||
+ current.ManagedPoolInstanceID == "" || current.ManagedPoolInstanceID != poolState.ManagedInstanceID {
+ return nil, fmt.Errorf("managed pool no longer exists")
+ }
+ }
return current, nil
}
diff --git a/pkg/gateway/services/thunder/service_test.go b/pkg/gateway/services/thunder/service_test.go
--- a/pkg/gateway/services/thunder/service_test.go
+++ b/pkg/gateway/services/thunder/service_test.go
@@ -409,6 +409,52 @@
}
}
+func TestServiceCreateNodeEnrollmentRejectsStaleManagedPoolToken(t *testing.T) {
+ rdb := newThunderRedisClient(t)
+ defer rdb.Close()
+
+ computeRepo := repository.NewComputeRedisRepository(rdb)
+ managedPoolRepo := repository.NewManagedPoolRedisRepository(rdb)
+ agentToken := "managed-agent-token"
+ if err := computeRepo.SaveAgentTokenState(context.Background(), &model.AgentTokenState{
+ TokenHash: hashThunderComputeToken(agentToken),
+ WorkspaceID: "workspace-1",
+ PoolName: "managed-pool",
+ MachineID: "machine-1",
+ ManagedPoolInstanceID: "instance-old",
+ }, time.Hour); err != nil {
+ t.Fatal(err)
+ }
+
+ service, err := NewService(ServiceOpts{
+ Repository: NewRedisRepository(rdb),
+ Client: NewClient("https://central.example", "central-token", nil),
+ ContainerRepo: repository.NewContainerRedisRepositoryForTest(rdb),
+ WorkerRepo: repository.NewWorkerRedisRepositoryForTest(rdb),
+ ComputeRepo: computeRepo,
+ ManagedPoolRepo: managedPoolRepo,
+ })
+ if err != nil {
+ t.Fatal(err)
+ }
+
+ resp, err := service.CreateNodeEnrollment(context.Background(), &pb.CreateNodeEnrollmentRequest{AgentToken: agentToken})
+ if err != nil {
+ t.Fatal(err)
+ }
+ if resp.Ok || !strings.Contains(resp.ErrorMsg, "managed pool no longer exists") {
+ t.Fatalf("CreateNodeEnrollment() = %+v", resp)
+ }
+
+ deleteResp, err := service.DeleteNodeEnrollment(context.Background(), &pb.DeleteNodeEnrollmentRequest{AgentToken: agentToken})
+ if err != nil {
+ t.Fatal(err)
+ }
+ if deleteResp.Ok || !strings.Contains(deleteResp.ErrorMsg, "managed pool no longer exists") {
+ t.Fatalf("DeleteNodeEnrollment() = %+v", deleteResp)
+ }
+}
+
func saveThunderAgentTokenState(t *testing.T, repo repository.ComputeRepository, agentToken, workspaceID, poolName, machineID string) {
t.Helper()
if err := repo.SaveAgentTokenState(context.Background(), &model.AgentTokenState{You can send follow-ups to the cloud agent here.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
Autofix Details
Retrying the results report with the exact tool name from the schema.
Or push these changes by commenting:
@cursor push 83abf1d5dc
Preview (83abf1d5dc)
diff --git a/pkg/gateway/services/thunder/service.go b/pkg/gateway/services/thunder/service.go
--- a/pkg/gateway/services/thunder/service.go
+++ b/pkg/gateway/services/thunder/service.go
@@ -36,8 +36,6 @@
AgentStateValidator AgentStateValidator
}
-var errThunderClientEnrollmentExists = errors.New("Thunder client enrollment already exists for container")
-
type Service struct {
pb.UnimplementedThunderServiceServer
@@ -101,8 +99,17 @@
err = s.repo.WithPoolLock(ctx, attrs.workspaceID, attrs.poolName, func(ctx context.Context) error {
if existing, found, err := s.repo.GetClientEnrollment(ctx, containerID); err != nil {
return err
- } else if found && existing.EnrollmentTokenID != "" {
- return errThunderClientEnrollmentExists
+ } else if found {
+ // Replace any prior enrollment so retries can recover a fresh install command.
+ // The install command is only available at token creation time and is not persisted.
+ if tokenID := strings.TrimSpace(existing.EnrollmentTokenID); tokenID != "" {
+ if _, err := s.client.DeleteEnrollmentTokenNode(ctx, tokenID); err != nil && !isThunderNotFound(err) {
+ return err
+ }
+ }
+ if err := s.repo.DeleteClientEnrollment(ctx, containerID); err != nil {
+ return err
+ }
}
zoneID, err := s.ensureZoneLocked(ctx, attrs.workspaceID, attrs.poolName)You can send follow-ups to the cloud agent here.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Client enrollment Redis CROSSSLOT risk
- Added a shared {clients} Redis hash tag to ThunderClientEnrollment and ThunderClientEnrollmentIndex so TxPipelined save/delete stay on one cluster slot.
Or push these changes by commenting:
@cursor push 60cfa225cb
Preview (60cfa225cb)
diff --git a/pkg/common/keys.go b/pkg/common/keys.go
--- a/pkg/common/keys.go
+++ b/pkg/common/keys.go
@@ -144,8 +144,8 @@
)
var (
- thunderClientEnrollment string = "thunder:client:%s"
- thunderClientEnrollmentIndex string = "thunder:clients"
+ thunderClientEnrollment string = "thunder:{clients}:client:%s"
+ thunderClientEnrollmentIndex string = "thunder:{clients}"
thunderNodeEnrollment string = "thunder:{%s}:pool:%s:machine:%s:node"
thunderNodeEnrollmentIndex string = "thunder:{%s}:pool:%s:nodes"
thunderZone string = "thunder:{%s}:pool:%s:zone"You can send follow-ups to the cloud agent here.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.
Done
Or push these changes by commenting:
@cursor push 6badc1b6dc
Preview (6badc1b6dc)
diff --git a/pkg/gateway/services/thunder/service.go b/pkg/gateway/services/thunder/service.go
--- a/pkg/gateway/services/thunder/service.go
+++ b/pkg/gateway/services/thunder/service.go
@@ -134,9 +134,8 @@
return err
}
if previousEnrollmentTokenID != "" && previousEnrollmentTokenID != enrollment.EnrollmentTokenID {
- if _, err := s.client.DeleteEnrollmentTokenNode(ctx, previousEnrollmentTokenID); err != nil && !isThunderNotFound(err) {
- return fmt.Errorf("failed to revoke previous Thunder client enrollment token %q: %w", previousEnrollmentTokenID, err)
- }
+ // Best-effort: the new enrollment is already persisted and must remain usable.
+ _, _ = s.client.DeleteEnrollmentTokenNode(ctx, previousEnrollmentTokenID)
}
return nil
})
@@ -167,13 +166,23 @@
return &pb.DeleteClientEnrollmentResponse{ErrorMsg: err.Error()}, nil
}
- if strings.TrimSpace(state.EnrollmentTokenID) != "" {
- _, err = s.client.DeleteEnrollmentTokenNode(ctx, state.EnrollmentTokenID)
- if err != nil && !isThunderNotFound(err) {
- return &pb.DeleteClientEnrollmentResponse{ErrorMsg: err.Error()}, nil
+ err = s.repo.WithPoolLock(ctx, state.WorkspaceID, state.PoolName, func(ctx context.Context) error {
+ state, found, err := s.repo.GetClientEnrollment(ctx, containerID)
+ if err != nil {
+ return err
}
- }
- if err := s.repo.DeleteClientEnrollment(ctx, containerID); err != nil {
+ if !found {
+ return nil
+ }
+ if strings.TrimSpace(state.EnrollmentTokenID) != "" {
+ _, err = s.client.DeleteEnrollmentTokenNode(ctx, state.EnrollmentTokenID)
+ if err != nil && !isThunderNotFound(err) {
+ return err
+ }
+ }
+ return s.repo.DeleteClientEnrollment(ctx, containerID)
+ })
+ if err != nil {
return &pb.DeleteClientEnrollmentResponse{ErrorMsg: err.Error()}, nil
}
return &pb.DeleteClientEnrollmentResponse{Ok: true}, nil
@@ -218,9 +227,8 @@
return err
}
if previousEnrollmentTokenID != "" && previousEnrollmentTokenID != enrollment.EnrollmentTokenID {
- if _, err := s.client.DeleteEnrollmentTokenNode(ctx, previousEnrollmentTokenID); err != nil && !isThunderNotFound(err) {
- return fmt.Errorf("failed to revoke previous Thunder node enrollment token %q: %w", previousEnrollmentTokenID, err)
- }
+ // Best-effort: the new enrollment is already persisted and must remain usable.
+ _, _ = s.client.DeleteEnrollmentTokenNode(ctx, previousEnrollmentTokenID)
}
enrollmentToken = enrollment.EnrollmentToken
return nil
@@ -238,23 +246,25 @@
return &pb.DeleteNodeEnrollmentResponse{ErrorMsg: err.Error()}, nil
}
- state, found, err := s.repo.GetNodeEnrollment(ctx, agentState.WorkspaceID, agentState.PoolName, agentState.MachineID)
+ err = s.repo.WithPoolLock(ctx, agentState.WorkspaceID, agentState.PoolName, func(ctx context.Context) error {
+ state, found, err := s.repo.GetNodeEnrollment(ctx, agentState.WorkspaceID, agentState.PoolName, agentState.MachineID)
+ if err != nil {
+ return err
+ }
+ if !found {
+ return nil
+ }
+ if strings.TrimSpace(state.EnrollmentTokenID) != "" {
+ _, err = s.client.DeleteEnrollmentTokenNode(ctx, state.EnrollmentTokenID)
+ if err != nil && !isThunderNotFound(err) {
+ return err
+ }
+ }
+ return s.repo.DeleteNodeEnrollment(ctx, agentState.WorkspaceID, agentState.PoolName, agentState.MachineID)
+ })
if err != nil {
return &pb.DeleteNodeEnrollmentResponse{ErrorMsg: err.Error()}, nil
}
- if !found {
- return &pb.DeleteNodeEnrollmentResponse{Ok: true}, nil
- }
-
- if strings.TrimSpace(state.EnrollmentTokenID) != "" {
- _, err = s.client.DeleteEnrollmentTokenNode(ctx, state.EnrollmentTokenID)
- if err != nil && !isThunderNotFound(err) {
- return &pb.DeleteNodeEnrollmentResponse{ErrorMsg: err.Error()}, nil
- }
- }
- if err := s.repo.DeleteNodeEnrollment(ctx, agentState.WorkspaceID, agentState.PoolName, agentState.MachineID); err != nil {
- return &pb.DeleteNodeEnrollmentResponse{ErrorMsg: err.Error()}, nil
- }
return &pb.DeleteNodeEnrollmentResponse{Ok: true}, nil
}You can send follow-ups to the cloud agent here.
Reviewed by Cursor Bugbot for commit 08e4c66. Configure here.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
Medium Risk
Touches GPU virtualization, external Thunder API credentials, and agent-token RPCs that bypass standard gRPC auth; mitigated by pool locking, token revocation on failure, and broad unit tests.
Overview
Moves Thunder virtual GPU enrollment from workers/agents talking directly to Thunder’s HTTP API into gateway-orchestrated flows backed by Redis state and a new
ThunderServicegRPC API.Workers with virtualized GPUs now call
CreateClientEnrollment/DeleteClientEnrollmenton the gateway (worker-token auth). The gateway ensures a per-workspace/pool Thunder zone, mints client enrollment tokens viaTHUNDER_API_URL/THUNDER_API_TOKEN, stores enrollment metadata in Redis under pool locks, and returns a shell install command the worker caches and runs in the sandbox (replacing per-worker env-based Thunder config).Compute agents get
CreateNodeEnrollment/DeleteNodeEnrollmentonGatewayService(agent token validated in-handler; interceptor auth skipped). After TSNet comes up, the agent enrolls a Thunder server node using a preferred Tailscale IPv4, runs the publicthunderdinstaller, rolls back enrollment if install fails, and deletes enrollment on proxy shutdown.Proto generation adds
thunder.protoand wires the service into gateway startup.Reviewed by Cursor Bugbot for commit fa51446. Bugbot is set up for automated code reviews on this repo. Configure here.