From 59d38bf599c79da82966ac1237700dce399b5319 Mon Sep 17 00:00:00 2001
From: Joseph Lombrozo <joe@djeebus.net>
Date: Tue, 9 Sep 2025 16:44:35 +0000
Subject: [PATCH 1/6] add the forcetypeassert linter

---
 .golangci.yml                                 |  1 +
 packages/api/internal/auth/constants.go       | 82 ++++++++++++++++++-
 packages/api/internal/auth/middleware.go      | 24 +++---
 packages/api/internal/handlers/accesstoken.go |  5 +-
 packages/api/internal/handlers/apikey.go      | 11 +--
 packages/api/internal/handlers/auth.go        | 16 ++--
 .../api/internal/handlers/sandbox_create.go   |  5 +-
 packages/api/internal/handlers/sandbox_get.go |  4 +-
 .../api/internal/handlers/sandbox_kill.go     |  3 +-
 .../api/internal/handlers/sandbox_logs.go     |  9 +-
 .../api/internal/handlers/sandbox_metrics.go  |  3 +-
 .../api/internal/handlers/sandbox_pause.go    |  4 +-
 .../api/internal/handlers/sandbox_resume.go   |  3 +-
 .../api/internal/handlers/sandboxes_list.go   | 12 ++-
 .../handlers/sandboxes_list_metrics.go        |  3 +-
 packages/api/internal/handlers/store.go       |  3 +-
 .../api/internal/handlers/team_metrics.go     |  3 +-
 .../api/internal/handlers/team_metrics_max.go |  3 +-
 packages/api/internal/handlers/teams.go       |  3 +-
 .../handlers/template_request_build_v2.go     |  7 +-
 .../api/internal/handlers/templates_list.go   |  2 +-
 .../api/internal/middleware/launchdarkly.go   | 10 +--
 .../api/internal/orchestrator/orchestrator.go |  2 +-
 .../placement/placement_benchmark_test.go     |  8 +-
 .../placement/placement_best_of_K.go          |  2 +-
 .../placement/placement_best_of_K_test.go     | 16 ++--
 .../orchestrator/placement/placement_test.go  |  7 +-
 packages/api/main.go                          |  7 +-
 packages/clickhouse/pkg/batcher/timer.go      |  4 +-
 packages/client-proxy/internal/edge/http.go   | 10 ++-
 packages/client-proxy/internal/proxy/proxy.go |  7 +-
 packages/envd/internal/api/download.go        |  4 +-
 packages/envd/internal/api/envs.go            |  2 +-
 packages/envd/internal/api/init.go            |  2 +-
 packages/envd/internal/api/upload.go          |  6 +-
 packages/envd/internal/logs/interceptor.go    | 70 ++++++++++------
 .../internal/services/filesystem/watch.go     |  2 +-
 .../services/filesystem/watch_sync.go         |  3 +-
 .../envd/internal/services/process/input.go   |  2 +-
 .../envd/internal/services/process/start.go   |  8 +-
 packages/envd/internal/utils/map.go           | 32 ++++++--
 .../internal/sandbox/block/cache.go           |  3 +-
 .../internal/sandbox/uffd/handler.go          |  5 +-
 packages/shared/pkg/logger/grpc.go            | 17 ++--
 packages/shared/pkg/logs/loki.go              |  7 +-
 packages/shared/pkg/proxy/proxy_test.go       | 17 ++--
 .../shared/pkg/storage/gcp_multipart_test.go  | 21 +++--
 packages/shared/pkg/telemetry/tracing.go      |  9 +-
 .../shared/scripts/seed/postgres/seed-db.go   | 13 ++-
 49 files changed, 326 insertions(+), 176 deletions(-)

diff --git a/.golangci.yml b/.golangci.yml
index 9c472a7c7a..7f7f21a06a 100644
--- a/.golangci.yml
+++ b/.golangci.yml
@@ -23,6 +23,7 @@ linters:
     - durationcheck
     - errname
     - errorlint
+    - forcetypeassert
     - govet
     - staticcheck
     - testifylint
diff --git a/packages/api/internal/auth/constants.go b/packages/api/internal/auth/constants.go
index 832b42b20f..c2e51ac66f 100644
--- a/packages/api/internal/auth/constants.go
+++ b/packages/api/internal/auth/constants.go
@@ -1,6 +1,84 @@
 package auth
 
+import (
+	"errors"
+	"fmt"
+
+	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
+	"go.uber.org/zap"
+
+	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
+)
+
 const (
-	TeamContextKey   string = "team"
-	UserIDContextKey string = "user_id"
+	teamContextKey   string = "team"
+	userIDContextKey string = "user_id"
+)
+
+var (
+	ErrNotFoundInContext = errors.New("not found in context")
+	ErrInvalidType       = errors.New("unexpected type")
 )
+
+type ginContextValueHelper[T any] struct {
+	contextKey string
+}
+
+func (g *ginContextValueHelper[T]) set(c *gin.Context, val T) {
+	c.Set(g.contextKey, val)
+}
+
+func (g *ginContextValueHelper[T]) get(c *gin.Context) (T, error) {
+	var t T
+
+	v := c.Value(teamContextKey)
+	if v == nil {
+		return t, ErrNotFoundInContext
+	}
+
+	t, ok := v.(T)
+	if !ok {
+		return t, fmt.Errorf("%w: wanted %T, got %T",
+			ErrInvalidType, t, v)
+	}
+
+	return t, nil
+}
+
+func (g *ginContextValueHelper[T]) safeGet(c *gin.Context) T {
+	v, err := g.get(c)
+	if err != nil {
+		zap.L().Warn("failed to "+g.contextKey, zap.Error(err))
+	}
+	return v
+}
+
+var (
+	teamInfoHelper = ginContextValueHelper[authcache.AuthTeamInfo]{"team"}
+	userIDHelper   = ginContextValueHelper[uuid.UUID]{"user_id"}
+)
+
+func setTeamInfo(c *gin.Context, teamInfo authcache.AuthTeamInfo) {
+	teamInfoHelper.set(c, teamInfo)
+}
+
+func GetTeamInfo(c *gin.Context) (authcache.AuthTeamInfo, error) {
+	return teamInfoHelper.get(c)
+}
+
+func SafeGetTeamInfo(c *gin.Context) authcache.AuthTeamInfo {
+	return teamInfoHelper.safeGet(c)
+}
+
+func setUserID(c *gin.Context, userID uuid.UUID) {
+	userIDHelper.set(c, userID)
+}
+
+func GetUserID(c *gin.Context) (uuid.UUID, error) {
+	return userIDHelper.get(c)
+}
+
+func SafeGetUserID(c *gin.Context) uuid.UUID {
+	return userIDHelper.safeGet(c)
+}
diff --git a/packages/api/internal/auth/middleware.go b/packages/api/internal/auth/middleware.go
index 6027af9e40..b0a16907e7 100644
--- a/packages/api/internal/auth/middleware.go
+++ b/packages/api/internal/auth/middleware.go
@@ -44,7 +44,7 @@ type commonAuthenticator[T any] struct {
 	securitySchemeName string
 	headerKey          headerKey
 	validationFunction func(context.Context, string) (T, *api.APIError)
-	contextKey         string
+	setContext         func(ginContext *gin.Context, result T)
 	errorMessage       string
 }
 
@@ -107,8 +107,8 @@ func (a *commonAuthenticator[T]) Authenticate(ctx context.Context, input *openap
 	telemetry.ReportEvent(ctx, "api key validated")
 
 	// Set the property on the gin context
-	if a.contextKey != "" {
-		middleware.GetGinContext(ctx).Set(a.contextKey, result)
+	if a.setContext != nil {
+		a.setContext(middleware.GetGinContext(ctx), result)
 	}
 
 	return nil
@@ -146,7 +146,7 @@ func CreateAuthenticationFunc(
 				removePrefix: "",
 			},
 			validationFunction: teamValidationFunction,
-			contextKey:         TeamContextKey,
+			setContext:         setTeamInfo,
 			errorMessage:       "Invalid API key, please visit https://e2b.dev/docs/api-key for more information.",
 		},
 		&commonAuthenticator[uuid.UUID]{
@@ -157,7 +157,7 @@ func CreateAuthenticationFunc(
 				removePrefix: "Bearer ",
 			},
 			validationFunction: userValidationFunction,
-			contextKey:         UserIDContextKey,
+			setContext:         setUserID,
 			errorMessage:       "Invalid Access token, try to login again by running `e2b auth login`.",
 		},
 		&commonAuthenticator[uuid.UUID]{
@@ -168,7 +168,7 @@ func CreateAuthenticationFunc(
 				removePrefix: "",
 			},
 			validationFunction: supabaseTokenValidationFunction,
-			contextKey:         UserIDContextKey,
+			setContext:         setUserID,
 			errorMessage:       "Invalid Supabase token.",
 		},
 		&commonAuthenticator[authcache.AuthTeamInfo]{
@@ -179,7 +179,7 @@ func CreateAuthenticationFunc(
 				removePrefix: "",
 			},
 			validationFunction: supabaseTeamValidationFunction,
-			contextKey:         TeamContextKey,
+			setContext:         setTeamInfo,
 			errorMessage:       "Invalid Supabase token teamID.",
 		},
 		&commonAuthenticator[struct{}]{
@@ -190,15 +190,19 @@ func CreateAuthenticationFunc(
 				removePrefix: "",
 			},
 			validationFunction: adminValidationFunction,
-			contextKey:         "",
+			setContext:         nil,
 			errorMessage:       "Invalid Access token.",
 		},
 	}
 
 	return func(ctx context.Context, input *openapi3filter.AuthenticationInput) error {
-		ginContext := ctx.Value(middleware.GinContextKey).(*gin.Context)
-		requestContext := ginContext.Request.Context()
+		value := ctx.Value(middleware.GinContextKey)
+		ginContext, ok := value.(*gin.Context)
+		if !ok {
+			return fmt.Errorf("%w: received %T", ErrInvalidType, value)
+		}
 
+		requestContext := ginContext.Request.Context()
 		_, span := tracer.Start(requestContext, "authenticate")
 		defer span.End()
 
diff --git a/packages/api/internal/handlers/accesstoken.go b/packages/api/internal/handlers/accesstoken.go
index f8c5496fd0..a0c9c3faf6 100644
--- a/packages/api/internal/handlers/accesstoken.go
+++ b/packages/api/internal/handlers/accesstoken.go
@@ -8,6 +8,7 @@ import (
 	"github.com/google/uuid"
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	"github.com/e2b-dev/infra/packages/shared/pkg/keys"
 	"github.com/e2b-dev/infra/packages/shared/pkg/models/accesstoken"
@@ -17,7 +18,7 @@ import (
 func (a *APIStore) PostAccessTokens(c *gin.Context) {
 	ctx := c.Request.Context()
 
-	userID := a.GetUserID(c)
+	userID := auth.SafeGetUserID(c)
 
 	body, err := utils.ParseBody[api.NewAccessToken](ctx, c)
 	if err != nil {
@@ -74,7 +75,7 @@ func (a *APIStore) PostAccessTokens(c *gin.Context) {
 func (a *APIStore) DeleteAccessTokensAccessTokenID(c *gin.Context, accessTokenID string) {
 	ctx := c.Request.Context()
 
-	userID := a.GetUserID(c)
+	userID := auth.SafeGetUserID(c)
 
 	accessTokenIDParsed, err := uuid.Parse(accessTokenID)
 	if err != nil {
diff --git a/packages/api/internal/handlers/apikey.go b/packages/api/internal/handlers/apikey.go
index 61f2ce292b..e9c2a53335 100644
--- a/packages/api/internal/handlers/apikey.go
+++ b/packages/api/internal/handlers/apikey.go
@@ -12,6 +12,7 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/team"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	"github.com/e2b-dev/infra/packages/db/queries"
@@ -39,7 +40,7 @@ func (a *APIStore) PatchApiKeysApiKeyID(c *gin.Context, apiKeyID string) {
 		return
 	}
 
-	teamID := a.GetTeamInfo(c).Team.ID
+	teamID := auth.SafeGetTeamInfo(c).Team.ID
 
 	now := time.Now()
 	_, err = a.sqlcDB.UpdateTeamApiKey(ctx, queries.UpdateTeamApiKeyParams{
@@ -64,7 +65,7 @@ func (a *APIStore) PatchApiKeysApiKeyID(c *gin.Context, apiKeyID string) {
 func (a *APIStore) GetApiKeys(c *gin.Context) {
 	ctx := c.Request.Context()
 
-	teamID := a.GetTeamInfo(c).Team.ID
+	teamID := auth.SafeGetTeamInfo(c).Team.ID
 
 	apiKeysDB, err := a.db.Client.TeamAPIKey.
 		Query().
@@ -116,7 +117,7 @@ func (a *APIStore) DeleteApiKeysApiKeyID(c *gin.Context, apiKeyID string) {
 		return
 	}
 
-	teamID := a.GetTeamInfo(c).Team.ID
+	teamID := auth.SafeGetTeamInfo(c).Team.ID
 
 	err = a.db.Client.TeamAPIKey.DeleteOneID(apiKeyIDParsed).Where(teamapikey.TeamID(teamID)).Exec(ctx)
 	if models.IsNotFound(err) {
@@ -135,8 +136,8 @@ func (a *APIStore) DeleteApiKeysApiKeyID(c *gin.Context, apiKeyID string) {
 func (a *APIStore) PostApiKeys(c *gin.Context) {
 	ctx := c.Request.Context()
 
-	userID := a.GetUserID(c)
-	teamID := a.GetTeamInfo(c).Team.ID
+	userID := auth.SafeGetUserID(c)
+	teamID := auth.SafeGetTeamInfo(c).Team.ID
 
 	body, err := utils.ParseBody[api.NewTeamAPIKey](ctx, c)
 	if err != nil {
diff --git a/packages/api/internal/handlers/auth.go b/packages/api/internal/handlers/auth.go
index 896fdbd5f1..efd11fa264 100644
--- a/packages/api/internal/handlers/auth.go
+++ b/packages/api/internal/handlers/auth.go
@@ -1,22 +1,24 @@
 package handlers
 
 import (
+	"errors"
 	"fmt"
 
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/db/queries"
 )
 
-func (a *APIStore) GetUserID(c *gin.Context) uuid.UUID {
-	return c.Value(auth.UserIDContextKey).(uuid.UUID)
-}
+var ErrNoUserIDInContext = errors.New("no user id in context")
 
 func (a *APIStore) GetUserAndTeams(c *gin.Context) (*uuid.UUID, []queries.GetTeamsWithUsersTeamsWithTierRow, error) {
-	userID := a.GetUserID(c)
+	userID, err := auth.GetUserID(c)
+	if err != nil {
+		return nil, nil, ErrNoUserIDInContext
+	}
+
 	ctx := c.Request.Context()
 
 	teams, err := a.sqlcDB.GetTeamsWithUsersTeamsWithTier(ctx, userID)
@@ -26,7 +28,3 @@ func (a *APIStore) GetUserAndTeams(c *gin.Context) (*uuid.UUID, []queries.GetTea
 
 	return &userID, teams, err
 }
-
-func (a *APIStore) GetTeamInfo(c *gin.Context) authcache.AuthTeamInfo {
-	return c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo)
-}
diff --git a/packages/api/internal/handlers/sandbox_create.go b/packages/api/internal/handlers/sandbox_create.go
index f067191ba0..b5ec5ac172 100644
--- a/packages/api/internal/handlers/sandbox_create.go
+++ b/packages/api/internal/handlers/sandbox_create.go
@@ -13,7 +13,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/cache/instance"
 	"github.com/e2b-dev/infra/packages/api/internal/middleware/otel/metrics"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
@@ -42,8 +41,8 @@ var mostUsedTemplates = map[string]struct{}{
 func (a *APIStore) PostSandboxes(c *gin.Context) {
 	ctx := c.Request.Context()
 
-	// Get team from context, use TeamContextKey
-	teamInfo := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo)
+	// Get team from context
+	teamInfo := auth.SafeGetTeamInfo(c)
 
 	c.Set("teamID", teamInfo.Team.ID.String())
 
diff --git a/packages/api/internal/handlers/sandbox_get.go b/packages/api/internal/handlers/sandbox_get.go
index 94a4b398f4..b5b0089a5a 100644
--- a/packages/api/internal/handlers/sandbox_get.go
+++ b/packages/api/internal/handlers/sandbox_get.go
@@ -9,7 +9,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/cache/instance"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	"github.com/e2b-dev/infra/packages/db/queries"
@@ -21,8 +20,7 @@ import (
 func (a *APIStore) GetSandboxesSandboxID(c *gin.Context, id string) {
 	ctx := c.Request.Context()
 
-	teamInfo := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo)
-	team := teamInfo.Team
+	team := auth.SafeGetTeamInfo(c).Team
 
 	telemetry.ReportEvent(ctx, "get sandbox")
 
diff --git a/packages/api/internal/handlers/sandbox_kill.go b/packages/api/internal/handlers/sandbox_kill.go
index 914fe49e9f..f92f8264e0 100644
--- a/packages/api/internal/handlers/sandbox_kill.go
+++ b/packages/api/internal/handlers/sandbox_kill.go
@@ -11,7 +11,6 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/cache/instance"
 	template_manager "github.com/e2b-dev/infra/packages/api/internal/template-manager"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
@@ -72,7 +71,7 @@ func (a *APIStore) DeleteSandboxesSandboxID(
 	ctx := c.Request.Context()
 	sandboxID = utils.ShortID(sandboxID)
 
-	team := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo).Team
+	team := auth.SafeGetTeamInfo(c).Team
 	teamID := team.ID
 
 	telemetry.SetAttributes(ctx,
diff --git a/packages/api/internal/handlers/sandbox_logs.go b/packages/api/internal/handlers/sandbox_logs.go
index beae1012dc..626ef971da 100644
--- a/packages/api/internal/handlers/sandbox_logs.go
+++ b/packages/api/internal/handlers/sandbox_logs.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	apiedge "github.com/e2b-dev/infra/packages/shared/pkg/http/edge"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
@@ -21,7 +20,7 @@ func (a *APIStore) GetSandboxesSandboxIDLogs(c *gin.Context, sandboxID string, p
 	ctx := c.Request.Context()
 	sandboxID = utils.ShortID(sandboxID)
 
-	team := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo).Team
+	team := auth.SafeGetTeamInfo(c).Team
 
 	telemetry.SetAttributes(ctx,
 		attribute.String("instance.id", sandboxID),
@@ -29,9 +28,9 @@ func (a *APIStore) GetSandboxesSandboxIDLogs(c *gin.Context, sandboxID string, p
 	)
 
 	/// Sandboxes living in a cluster
-	sbxLogs, err := a.getClusterSandboxLogs(ctx, sandboxID, team.ID.String(), utils.WithClusterFallback(team.ClusterID), params.Limit, params.Start)
-	if err != nil {
-		a.sendAPIStoreError(c, int(err.Code), err.Message)
+	sbxLogs, apiErr := a.getClusterSandboxLogs(ctx, sandboxID, team.ID.String(), utils.WithClusterFallback(team.ClusterID), params.Limit, params.Start)
+	if apiErr != nil {
+		a.sendAPIStoreError(c, int(apiErr.Code), apiErr.Message)
 		return
 	}
 
diff --git a/packages/api/internal/handlers/sandbox_metrics.go b/packages/api/internal/handlers/sandbox_metrics.go
index b76afa2334..1ee81211a0 100644
--- a/packages/api/internal/handlers/sandbox_metrics.go
+++ b/packages/api/internal/handlers/sandbox_metrics.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	clickhouse "github.com/e2b-dev/infra/packages/clickhouse/pkg"
 	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
@@ -25,7 +24,7 @@ func (a *APIStore) GetSandboxesSandboxIDMetrics(c *gin.Context, sandboxID string
 	defer span.End()
 	sandboxID = utils.ShortID(sandboxID)
 
-	team := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo).Team
+	team := auth.SafeGetTeamInfo(c).Team
 
 	metricsReadFlag, err := a.featureFlags.BoolFlag(ctx, featureflags.MetricsReadFlagName,
 		featureflags.SandboxContext(sandboxID))
diff --git a/packages/api/internal/handlers/sandbox_pause.go b/packages/api/internal/handlers/sandbox_pause.go
index 1a91c14566..b05d617d4d 100644
--- a/packages/api/internal/handlers/sandbox_pause.go
+++ b/packages/api/internal/handlers/sandbox_pause.go
@@ -11,6 +11,7 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/cache/instance"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	"github.com/e2b-dev/infra/packages/db/queries"
@@ -22,8 +23,7 @@ func (a *APIStore) PostSandboxesSandboxIDPause(c *gin.Context, sandboxID api.San
 	ctx := c.Request.Context()
 	// Get team from context, use TeamContextKey
 
-	teamID := a.GetTeamInfo(c).Team.ID
-
+	teamID := auth.SafeGetTeamInfo(c).Team.ID
 	sandboxID = utils.ShortID(sandboxID)
 
 	span := trace.SpanFromContext(ctx)
diff --git a/packages/api/internal/handlers/sandbox_resume.go b/packages/api/internal/handlers/sandbox_resume.go
index 1c734c45b4..df276bcce7 100644
--- a/packages/api/internal/handlers/sandbox_resume.go
+++ b/packages/api/internal/handlers/sandbox_resume.go
@@ -13,7 +13,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/cache/instance"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	"github.com/e2b-dev/infra/packages/db/queries"
@@ -26,7 +25,7 @@ func (a *APIStore) PostSandboxesSandboxIDResume(c *gin.Context, sandboxID api.Sa
 	ctx := c.Request.Context()
 
 	// Get team from context, use TeamContextKey
-	teamInfo := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo)
+	teamInfo := auth.SafeGetTeamInfo(c)
 
 	span := trace.SpanFromContext(ctx)
 	traceID := span.SpanContext().TraceID().String()
diff --git a/packages/api/internal/handlers/sandboxes_list.go b/packages/api/internal/handlers/sandboxes_list.go
index 3999f2463a..20711f3131 100644
--- a/packages/api/internal/handlers/sandboxes_list.go
+++ b/packages/api/internal/handlers/sandboxes_list.go
@@ -16,7 +16,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/cache/instance"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	"github.com/e2b-dev/infra/packages/db/queries"
@@ -71,7 +70,13 @@ func (a *APIStore) GetSandboxes(c *gin.Context, params api.GetSandboxesParams) {
 	ctx := c.Request.Context()
 	telemetry.ReportEvent(ctx, "list sandboxes")
 
-	teamInfo := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo)
+	teamInfo, err := auth.GetTeamInfo(c)
+	if err != nil {
+		telemetry.ReportError(ctx, "failed to get team info", err)
+		a.sendAPIStoreError(c, http.StatusBadRequest, err.Error())
+		return
+	}
+
 	team := teamInfo.Team
 
 	a.posthog.IdentifyAnalyticsTeam(team.ID.String(), team.Name)
@@ -102,8 +107,7 @@ func (a *APIStore) GetV2Sandboxes(c *gin.Context, params api.GetV2SandboxesParam
 	ctx := c.Request.Context()
 	telemetry.ReportEvent(ctx, "list sandboxes")
 
-	teamInfo := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo)
-	team := teamInfo.Team
+	team := auth.SafeGetTeamInfo(c).Team
 
 	a.posthog.IdentifyAnalyticsTeam(team.ID.String(), team.Name)
 	properties := a.posthog.GetPackageToPosthogProperties(&c.Request.Header)
diff --git a/packages/api/internal/handlers/sandboxes_list_metrics.go b/packages/api/internal/handlers/sandboxes_list_metrics.go
index 25fcca6ccd..7db1b9cbe1 100644
--- a/packages/api/internal/handlers/sandboxes_list_metrics.go
+++ b/packages/api/internal/handlers/sandboxes_list_metrics.go
@@ -12,7 +12,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
@@ -82,7 +81,7 @@ func (a *APIStore) GetSandboxesMetrics(c *gin.Context, params api.GetSandboxesMe
 	ctx := c.Request.Context()
 	telemetry.ReportEvent(ctx, "list running instances with metrics")
 
-	team := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo).Team
+	team := auth.SafeGetTeamInfo(c).Team
 
 	if len(params.SandboxIds) > maxSandboxMetricsCount {
 		zap.L().Error("Too many sandboxes requested", zap.Int("requested_count", len(params.SandboxIds)), zap.Int("max_count", maxSandboxMetricsCount), logger.WithTeamID(team.ID.String()))
diff --git a/packages/api/internal/handlers/store.go b/packages/api/internal/handlers/store.go
index fe965fb73d..454a7be582 100644
--- a/packages/api/internal/handlers/store.go
+++ b/packages/api/internal/handlers/store.go
@@ -20,6 +20,7 @@ import (
 
 	analyticscollector "github.com/e2b-dev/infra/packages/api/internal/analytics_collector"
 	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/auth"
 	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	templatecache "github.com/e2b-dev/infra/packages/api/internal/cache/templates"
 	dbapi "github.com/e2b-dev/infra/packages/api/internal/db"
@@ -403,7 +404,7 @@ func (a *APIStore) GetUserIDFromSupabaseToken(ctx context.Context, supabaseToken
 }
 
 func (a *APIStore) GetTeamFromSupabaseToken(ctx context.Context, teamID string) (authcache.AuthTeamInfo, *api.APIError) {
-	userID := a.GetUserID(middleware.GetGinContext(ctx))
+	userID := auth.SafeGetUserID(middleware.GetGinContext(ctx))
 
 	cacheKey := fmt.Sprintf("%s-%s", userID.String(), teamID)
 	team, tier, err := a.authCache.GetOrSet(ctx, cacheKey, func(ctx context.Context, key string) (*queries.Team, *queries.Tier, error) {
diff --git a/packages/api/internal/handlers/team_metrics.go b/packages/api/internal/handlers/team_metrics.go
index acfe1f8013..6b4e184126 100644
--- a/packages/api/internal/handlers/team_metrics.go
+++ b/packages/api/internal/handlers/team_metrics.go
@@ -10,7 +10,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
@@ -23,7 +22,7 @@ func (a *APIStore) GetTeamsTeamIDMetrics(c *gin.Context, teamID string, params a
 	ctx, span := a.Tracer.Start(ctx, "team-metrics")
 	defer span.End()
 
-	team := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo).Team
+	team := auth.SafeGetTeamInfo(c).Team
 
 	if teamID != team.ID.String() {
 		telemetry.ReportError(ctx, "team ids mismatch", fmt.Errorf("you (%s) are not authorized to access this team's (%s) metrics", team.ID, teamID), telemetry.WithTeamID(team.ID.String()))
diff --git a/packages/api/internal/handlers/team_metrics_max.go b/packages/api/internal/handlers/team_metrics_max.go
index 7f5b333adf..3e0dabf1d3 100644
--- a/packages/api/internal/handlers/team_metrics_max.go
+++ b/packages/api/internal/handlers/team_metrics_max.go
@@ -10,7 +10,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/metrics"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	clickhouse "github.com/e2b-dev/infra/packages/clickhouse/pkg"
@@ -23,7 +22,7 @@ func (a *APIStore) GetTeamsTeamIDMetricsMax(c *gin.Context, teamID string, param
 	ctx, span := a.Tracer.Start(ctx, "team-metrics-max")
 	defer span.End()
 
-	team := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo).Team
+	team := auth.SafeGetTeamInfo(c).Team
 
 	if teamID != team.ID.String() {
 		telemetry.ReportError(ctx, "team ids mismatch", fmt.Errorf("you (%s) are not authorized to access this team's (%s) metrics", team.ID, teamID), telemetry.WithTeamID(team.ID.String()))
diff --git a/packages/api/internal/handlers/teams.go b/packages/api/internal/handlers/teams.go
index f537b22e33..5a08d95f28 100644
--- a/packages/api/internal/handlers/teams.go
+++ b/packages/api/internal/handlers/teams.go
@@ -6,6 +6,7 @@ import (
 	"github.com/gin-gonic/gin"
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
+	"github.com/e2b-dev/infra/packages/api/internal/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/team"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
@@ -13,7 +14,7 @@ import (
 func (a *APIStore) GetTeams(c *gin.Context) {
 	ctx := c.Request.Context()
 
-	userID := a.GetUserID(c)
+	userID := auth.SafeGetUserID(c)
 
 	results, err := a.sqlcDB.GetTeamsWithUsersTeams(ctx, userID)
 	if err != nil {
diff --git a/packages/api/internal/handlers/template_request_build_v2.go b/packages/api/internal/handlers/template_request_build_v2.go
index 0ad53c7649..f67c813cc5 100644
--- a/packages/api/internal/handlers/template_request_build_v2.go
+++ b/packages/api/internal/handlers/template_request_build_v2.go
@@ -9,7 +9,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	apiutils "github.com/e2b-dev/infra/packages/api/internal/utils"
 	"github.com/e2b-dev/infra/packages/db/queries"
 	"github.com/e2b-dev/infra/packages/shared/pkg/db"
@@ -112,11 +111,9 @@ func (a *APIStore) GetTeamAndTier(
 	_, span := a.Tracer.Start(c.Request.Context(), "get-team-and-tier")
 	defer span.End()
 
-	if c.Value(auth.TeamContextKey) != nil {
-		teamInfo := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo)
-
+	if teamInfo, err := auth.GetTeamInfo(c); err != nil {
 		return teamInfo.Team, teamInfo.Tier, nil
-	} else if c.Value(auth.UserIDContextKey) != nil {
+	} else if _, err := auth.GetUserID(c); err != nil {
 		_, teams, err := a.GetUserAndTeams(c)
 		if err != nil {
 			return nil, nil, &api.APIError{
diff --git a/packages/api/internal/handlers/templates_list.go b/packages/api/internal/handlers/templates_list.go
index 1bc27a188d..d6d2efe048 100644
--- a/packages/api/internal/handlers/templates_list.go
+++ b/packages/api/internal/handlers/templates_list.go
@@ -19,7 +19,7 @@ import (
 func (a *APIStore) GetTemplates(c *gin.Context, params api.GetTemplatesParams) {
 	ctx := c.Request.Context()
 
-	userID := c.Value(auth.UserIDContextKey).(uuid.UUID)
+	userID := auth.SafeGetUserID(c)
 
 	var team *queries.Team
 	teams, err := a.sqlcDB.GetTeamsWithUsersTeams(ctx, userID)
diff --git a/packages/api/internal/middleware/launchdarkly.go b/packages/api/internal/middleware/launchdarkly.go
index fafb690e49..f49d0428b5 100644
--- a/packages/api/internal/middleware/launchdarkly.go
+++ b/packages/api/internal/middleware/launchdarkly.go
@@ -2,11 +2,9 @@ package middleware
 
 import (
 	"github.com/gin-gonic/gin"
-	"github.com/google/uuid"
 	"github.com/launchdarkly/go-sdk-common/v3/ldcontext"
 
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 )
 
@@ -30,8 +28,8 @@ func InitLaunchDarklyContext(c *gin.Context) {
 }
 
 func createLaunchDarklyUserContext(c *gin.Context) (ldcontext.Context, bool) {
-	userID, ok := c.Value(auth.UserIDContextKey).(uuid.UUID)
-	if !ok {
+	userID, err := auth.GetUserID(c)
+	if err != nil {
 		return ldcontext.Context{}, false
 	}
 
@@ -39,8 +37,8 @@ func createLaunchDarklyUserContext(c *gin.Context) (ldcontext.Context, bool) {
 }
 
 func createLaunchDarklyTeamContext(c *gin.Context) (ldcontext.Context, bool) {
-	authTeamInfo, ok := c.Value(auth.TeamContextKey).(authcache.AuthTeamInfo)
-	if !ok {
+	authTeamInfo, err := auth.GetTeamInfo(c)
+	if err != nil {
 		return ldcontext.Context{}, false
 	}
 
diff --git a/packages/api/internal/orchestrator/orchestrator.go b/packages/api/internal/orchestrator/orchestrator.go
index 68540086c4..ce14e2af48 100644
--- a/packages/api/internal/orchestrator/orchestrator.go
+++ b/packages/api/internal/orchestrator/orchestrator.go
@@ -102,7 +102,7 @@ func New(
 
 	// Initialize both placement algorithms
 	leastBusyAlgorithm := &placement.LeastBusyAlgorithm{}
-	bestOfKAlgorithm := placement.NewBestOfK(getBestOfKConfig(ctx, featureFlags)).(*placement.BestOfK)
+	bestOfKAlgorithm := placement.NewBestOfK(getBestOfKConfig(ctx, featureFlags))
 
 	o := Orchestrator{
 		httpClient:         httpClient,
diff --git a/packages/api/internal/orchestrator/placement/placement_benchmark_test.go b/packages/api/internal/orchestrator/placement/placement_benchmark_test.go
index 67b4fb1bbb..2d5f4e4501 100644
--- a/packages/api/internal/orchestrator/placement/placement_benchmark_test.go
+++ b/packages/api/internal/orchestrator/placement/placement_benchmark_test.go
@@ -11,6 +11,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"go.opentelemetry.io/otel/trace/noop"
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
@@ -193,7 +194,9 @@ func (n *SimulatedNode) getUtilization() (cpuUtil, memUtil float64) {
 }
 
 // runBenchmark runs a comprehensive placement benchmark with lifecycle tracking
-func runBenchmark(_ *testing.B, algorithm Algorithm, config BenchmarkConfig) *BenchmarkMetrics {
+func runBenchmark(b *testing.B, algorithm Algorithm, config BenchmarkConfig) *BenchmarkMetrics {
+	b.Helper()
+
 	ctx, cancel := context.WithTimeout(context.Background(), config.BenchmarkDuration)
 	defer cancel()
 
@@ -237,7 +240,8 @@ func runBenchmark(_ *testing.B, algorithm Algorithm, config BenchmarkConfig) *Be
 				now := time.Now()
 				// Check and remove expired sandboxes
 				activeSandboxes.Range(func(key, value interface{}) bool {
-					sandbox := value.(*LiveSandbox)
+					sandbox, ok := value.(*LiveSandbox)
+					assert.True(b, ok)
 					if now.Sub(sandbox.StartTime) > sandbox.PlannedDuration {
 						// Remove from node
 						if node, exists := nodeMap[sandbox.NodeID]; exists {
diff --git a/packages/api/internal/orchestrator/placement/placement_best_of_K.go b/packages/api/internal/orchestrator/placement/placement_best_of_K.go
index 233db6731f..97a0362619 100644
--- a/packages/api/internal/orchestrator/placement/placement_best_of_K.go
+++ b/packages/api/internal/orchestrator/placement/placement_best_of_K.go
@@ -84,7 +84,7 @@ type BestOfK struct {
 var _ Algorithm = &BestOfK{}
 
 // NewBestOfK creates a new placement algorithm with the given config
-func NewBestOfK(config BestOfKConfig) Algorithm {
+func NewBestOfK(config BestOfKConfig) *BestOfK {
 	return &BestOfK{
 		config: config,
 	}
diff --git a/packages/api/internal/orchestrator/placement/placement_best_of_K_test.go b/packages/api/internal/orchestrator/placement/placement_best_of_K_test.go
index 3793efd104..1778136719 100644
--- a/packages/api/internal/orchestrator/placement/placement_best_of_K_test.go
+++ b/packages/api/internal/orchestrator/placement/placement_best_of_K_test.go
@@ -13,7 +13,7 @@ import (
 
 func TestBestOfK_Score(t *testing.T) {
 	config := DefaultBestOfKConfig()
-	algo := NewBestOfK(config).(*BestOfK)
+	algo := NewBestOfK(config)
 
 	// Create a test node with known metrics
 	// The node's CpuUsage is set to 50 via the constructor
@@ -39,7 +39,7 @@ func TestBestOfK_Score(t *testing.T) {
 
 func TestBestOfK_Score_PreferBiggerNode(t *testing.T) {
 	config := DefaultBestOfKConfig()
-	algo := NewBestOfK(config).(*BestOfK)
+	algo := NewBestOfK(config)
 
 	// Create a test node with known metrics
 	// The node's CpuUsage is set to 50 via the constructor
@@ -65,7 +65,7 @@ func TestBestOfK_Score_PreferBiggerNode(t *testing.T) {
 
 func TestBestOfK_CanFit(t *testing.T) {
 	config := DefaultBestOfKConfig()
-	algo := NewBestOfK(config).(*BestOfK)
+	algo := NewBestOfK(config)
 
 	// Create a test node with moderate CPU usage
 	node := nodemanager.NewTestNode("test-node", api.NodeStatusReady, 5, 4)
@@ -96,7 +96,7 @@ func TestBestOfK_ChooseNode(t *testing.T) {
 		Alpha: 0.5,
 		K:     3, // Sample all nodes
 	}
-	algo := NewBestOfK(config).(*BestOfK)
+	algo := NewBestOfK(config)
 
 	// Create test nodes with different loads
 	node1 := nodemanager.NewTestNode("node1", api.NodeStatusReady, 8, 4)
@@ -124,7 +124,7 @@ func TestBestOfK_ChooseNode_WithExclusions(t *testing.T) {
 		Alpha: 0.5,
 		K:     3,
 	}
-	algo := NewBestOfK(config).(*BestOfK)
+	algo := NewBestOfK(config)
 
 	// Create test nodes
 	node1 := nodemanager.NewTestNode("node1", api.NodeStatusReady, 8, 4)
@@ -153,7 +153,7 @@ func TestBestOfK_ChooseNode_WithExclusions(t *testing.T) {
 func TestBestOfK_ChooseNode_NoAvailableNodes(t *testing.T) {
 	ctx := context.Background()
 	config := DefaultBestOfKConfig()
-	algo := NewBestOfK(config).(*BestOfK)
+	algo := NewBestOfK(config)
 
 	// Create unhealthy nodes
 	node1 := nodemanager.NewTestNode("node1", api.NodeStatusUnhealthy, 8, 4)
@@ -174,7 +174,7 @@ func TestBestOfK_ChooseNode_NoAvailableNodes(t *testing.T) {
 
 func TestBestOfK_Sample(t *testing.T) {
 	config := DefaultBestOfKConfig()
-	algo := NewBestOfK(config).(*BestOfK)
+	algo := NewBestOfK(config)
 
 	// Create many test nodes
 	var nodes []*nodemanager.Node
@@ -218,7 +218,7 @@ func TestBestOfK_PowerOfKChoices(t *testing.T) {
 		Alpha: 0.5,
 		K:     3,
 	}
-	algo := NewBestOfK(config).(*BestOfK)
+	algo := NewBestOfK(config)
 
 	// Create many nodes with varying loads
 	var nodes []*nodemanager.Node
diff --git a/packages/api/internal/orchestrator/placement/placement_test.go b/packages/api/internal/orchestrator/placement/placement_test.go
index 4c7e3e4212..4d802a56f0 100644
--- a/packages/api/internal/orchestrator/placement/placement_test.go
+++ b/packages/api/internal/orchestrator/placement/placement_test.go
@@ -3,6 +3,7 @@ package placement
 import (
 	"context"
 	"errors"
+	"fmt"
 	"testing"
 	"time"
 
@@ -30,7 +31,11 @@ func (m *mockAlgorithm) chooseNode(ctx context.Context, nodes []*nodemanager.Nod
 	if args.Get(0) == nil {
 		return nil, args.Error(1)
 	}
-	return args.Get(0).(*nodemanager.Node), args.Error(1)
+	node, ok := args.Get(0).(*nodemanager.Node)
+	if !ok {
+		return nil, fmt.Errorf("expected %T, got %T", node, args.Get(0))
+	}
+	return node, args.Error(1)
 }
 
 func TestPlaceSandbox_SuccessfulPlacement(t *testing.T) {
diff --git a/packages/api/main.go b/packages/api/main.go
index 7e10966a6a..00dc454e0f 100644
--- a/packages/api/main.go
+++ b/packages/api/main.go
@@ -28,7 +28,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/auth"
-	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 	"github.com/e2b-dev/infra/packages/api/internal/handlers"
 	customMiddleware "github.com/e2b-dev/infra/packages/api/internal/middleware"
 	metricsMiddleware "github.com/e2b-dev/infra/packages/api/internal/middleware/otel/metrics"
@@ -151,9 +150,9 @@ func NewGinServer(ctx context.Context, tel *telemetry.Client, logger *zap.Logger
 				teamID := ""
 
 				// Get team from context, use TeamContextKey
-				teamInfo := c.Value(auth.TeamContextKey)
-				if teamInfo != nil {
-					teamID = teamInfo.(authcache.AuthTeamInfo).Team.ID.String()
+				teamInfo, err := auth.GetTeamInfo(c)
+				if err == nil {
+					teamID = teamInfo.Team.ID.String()
 				}
 
 				reqLogger := logger
diff --git a/packages/clickhouse/pkg/batcher/timer.go b/packages/clickhouse/pkg/batcher/timer.go
index 0779a74196..d1c6f44acc 100644
--- a/packages/clickhouse/pkg/batcher/timer.go
+++ b/packages/clickhouse/pkg/batcher/timer.go
@@ -9,12 +9,12 @@ import (
 var timerPool sync.Pool
 
 func acquireTimer(timeout time.Duration) *time.Timer {
-	tv := timerPool.Get()
+	tv, _ := timerPool.Get().(*time.Timer)
 	if tv == nil {
 		return time.NewTimer(timeout)
 	}
 
-	t := tv.(*time.Timer)
+	t := tv
 	if t.Reset(timeout) {
 		log.Printf("Active timer trapped into AcquireTimer() with timeout %s", timeout)
 	}
diff --git a/packages/client-proxy/internal/edge/http.go b/packages/client-proxy/internal/edge/http.go
index 30fcc3b0dd..c82f78205a 100644
--- a/packages/client-proxy/internal/edge/http.go
+++ b/packages/client-proxy/internal/edge/http.go
@@ -87,11 +87,13 @@ func NewGinServer(logger *zap.Logger, store *handlers.APIStore, swagger *openapi
 
 func ginBuildAuthenticationHandler(tracer trace.Tracer, auth authorization.AuthorizationService) func(ctx context.Context, input *openapi3filter.AuthenticationInput) error {
 	return func(ctx context.Context, input *openapi3filter.AuthenticationInput) error {
-		ginContext := ctx.Value(middleware.GinContextKey).(*gin.Context)
-		requestContext := ginContext.Request.Context()
+		ginContext, ok := ctx.Value(middleware.GinContextKey).(*gin.Context)
+		if ok && ginContext != nil {
+			requestContext := ginContext.Request.Context()
 
-		_, span := tracer.Start(requestContext, "authenticate")
-		defer span.End()
+			_, span := tracer.Start(requestContext, "authenticate")
+			defer span.End()
+		}
 
 		if input.SecuritySchemeName != securitySchemaName {
 			return fmt.Errorf("invalid security scheme name '%s'", input.SecuritySchemeName)
diff --git a/packages/client-proxy/internal/proxy/proxy.go b/packages/client-proxy/internal/proxy/proxy.go
index 3c5182b4f3..a32ca46f21 100644
--- a/packages/client-proxy/internal/proxy/proxy.go
+++ b/packages/client-proxy/internal/proxy/proxy.go
@@ -64,7 +64,12 @@ func dnsResolution(sandboxId string, logger *zap.Logger) (string, error) {
 			continue
 		}
 
-		node = resp.Answer[0].(*dns.A).A.String()
+		answer, ok := resp.Answer[0].(*dns.A)
+		if !ok {
+			logger.Warn(fmt.Sprintf("received invalid answer type: %T", answer))
+			continue
+		}
+		node := answer.A.String()
 
 		// the sandbox was not found, we want to return this information to the user
 		if node == "127.0.0.1" {
diff --git a/packages/envd/internal/api/download.go b/packages/envd/internal/api/download.go
index cc3d9e1cbb..b38818eece 100644
--- a/packages/envd/internal/api/download.go
+++ b/packages/envd/internal/api/download.go
@@ -28,7 +28,7 @@ func (a *API) GetFiles(w http.ResponseWriter, r *http.Request, params GetFilesPa
 	// signing authorization if needed
 	err := a.validateSigning(r, params.Signature, params.SignatureExpiration, params.Username, path, SigningReadOperation)
 	if err != nil {
-		a.logger.Error().Err(err).Str(string(logs.OperationIDKey), operationID).Msg("error during auth validation")
+		a.logger.Error().Err(err).Str("operation_id", operationID).Msg("error during auth validation")
 		jsonError(w, http.StatusUnauthorized, err)
 		return
 	}
@@ -37,7 +37,7 @@ func (a *API) GetFiles(w http.ResponseWriter, r *http.Request, params GetFilesPa
 		l := a.logger.
 			Err(errMsg).
 			Str("method", r.Method+" "+r.URL.Path).
-			Str(string(logs.OperationIDKey), operationID).
+			Str("operation_id", operationID).
 			Str("path", path).
 			Str("username", params.Username)
 
diff --git a/packages/envd/internal/api/envs.go b/packages/envd/internal/api/envs.go
index 4a571b08e5..465b600776 100644
--- a/packages/envd/internal/api/envs.go
+++ b/packages/envd/internal/api/envs.go
@@ -10,7 +10,7 @@ import (
 func (a *API) GetEnvs(w http.ResponseWriter, _ *http.Request) {
 	operationID := logs.AssignOperationID()
 
-	a.logger.Debug().Str(string(logs.OperationIDKey), operationID).Msg("Getting env vars")
+	a.logger.Debug().Str("operation_id", operationID).Msg("Getting env vars")
 
 	envs := make(EnvVars)
 	a.envVars.Range(func(key, value string) bool {
diff --git a/packages/envd/internal/api/init.go b/packages/envd/internal/api/init.go
index b50e998839..9d195d7921 100644
--- a/packages/envd/internal/api/init.go
+++ b/packages/envd/internal/api/init.go
@@ -16,7 +16,7 @@ func (a *API) PostInit(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
 
 	operationID := logs.AssignOperationID()
-	logger := a.logger.With().Str(string(logs.OperationIDKey), operationID).Logger()
+	logger := a.logger.With().Str("operation_id", operationID).Logger()
 
 	if r.Body != nil {
 		var initRequest PostInitJSONBody
diff --git a/packages/envd/internal/api/upload.go b/packages/envd/internal/api/upload.go
index bcf79d5d58..7640272130 100644
--- a/packages/envd/internal/api/upload.go
+++ b/packages/envd/internal/api/upload.go
@@ -165,7 +165,7 @@ func (a *API) PostFiles(w http.ResponseWriter, r *http.Request, params PostFiles
 	// signing authorization if needed
 	err := a.validateSigning(r, params.Signature, params.SignatureExpiration, params.Username, path, SigningWriteOperation)
 	if err != nil {
-		a.logger.Error().Err(err).Str(string(logs.OperationIDKey), operationID).Msg("error during auth validation")
+		a.logger.Error().Err(err).Str("operation_id", operationID).Msg("error during auth validation")
 		jsonError(w, http.StatusUnauthorized, err)
 		return
 	}
@@ -174,7 +174,7 @@ func (a *API) PostFiles(w http.ResponseWriter, r *http.Request, params PostFiles
 		l := a.logger.
 			Err(errMsg).
 			Str("method", r.Method+" "+r.URL.Path).
-			Str(string(logs.OperationIDKey), operationID).
+			Str("operation_id", operationID).
 			Str("path", path).
 			Str("username", params.Username)
 
@@ -230,7 +230,7 @@ func (a *API) PostFiles(w http.ResponseWriter, r *http.Request, params PostFiles
 				return
 			}
 
-			status, processErr := processFile(r, filePath, part, u, a.logger.With().Str(string(logs.OperationIDKey), operationID).Str("event_type", "file_processing").Logger())
+			status, processErr := processFile(r, filePath, part, u, a.logger.With().Str("operation_id", operationID).Str("event_type", "file_processing").Logger())
 			if processErr != nil {
 				errorCode = status
 				errMsg = processErr
diff --git a/packages/envd/internal/logs/interceptor.go b/packages/envd/internal/logs/interceptor.go
index 5a2c6f7f2d..1aabad7c04 100644
--- a/packages/envd/internal/logs/interceptor.go
+++ b/packages/envd/internal/logs/interceptor.go
@@ -2,6 +2,7 @@ package logs
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strconv"
 	"strings"
@@ -11,11 +12,8 @@ import (
 	"github.com/rs/zerolog"
 )
 
-type OperationID string
-
 const (
-	OperationIDKey    OperationID = "operation_id"
-	DefaultHTTPMethod string      = "POST"
+	DefaultHTTPMethod string = "POST"
 )
 
 var operationID = atomic.Int32{}
@@ -26,8 +24,36 @@ func AssignOperationID() string {
 	return strconv.Itoa(int(id))
 }
 
-func AddRequestIDToContext(ctx context.Context) context.Context {
-	return context.WithValue(ctx, OperationIDKey, AssignOperationID())
+type operationIDKey struct{}
+
+func AddRequestIDToContext(ctx context.Context) (context.Context, string) {
+	operationID := AssignOperationID()
+	return context.WithValue(ctx, operationIDKey{}, operationID), operationID
+}
+
+var (
+	ErrOperationIDNotInContext = errors.New("operation id not in context")
+	ErrUnexpectedType          = errors.New("unexpected type")
+)
+
+func GetRequestID(ctx context.Context) (string, error) {
+	value := ctx.Value(operationIDKey{})
+	if value == nil {
+		return "", ErrOperationIDNotInContext
+	}
+
+	requestID, ok := value.(string)
+	if !ok {
+		return "", fmt.Errorf("%w: expected %T, received %T",
+			ErrUnexpectedType, requestID, value)
+	}
+
+	return requestID, nil
+}
+
+func SafeGetRequestID(ctx context.Context) string {
+	requestID, _ := GetRequestID(ctx)
+	return requestID
 }
 
 func formatMethod(method string) string {
@@ -52,18 +78,19 @@ func formatMethod(method string) string {
 
 func NewUnaryLogInterceptor(logger *zerolog.Logger) connect.UnaryInterceptorFunc {
 	interceptor := func(next connect.UnaryFunc) connect.UnaryFunc {
-		return connect.UnaryFunc(func(
+		return func(
 			ctx context.Context,
 			req connect.AnyRequest,
 		) (connect.AnyResponse, error) {
-			ctx = AddRequestIDToContext(ctx)
+			ctx, operationID := AddRequestIDToContext(ctx)
 
 			res, err := next(ctx, req)
 
 			l := logger.
 				Err(err).
-				Str("method", DefaultHTTPMethod+" "+req.Spec().Procedure).
-				Str(string(OperationIDKey), ctx.Value(OperationIDKey).(string))
+				Str("method", DefaultHTTPMethod+" "+req.Spec().Procedure)
+
+			l = l.Str("operation_id", string(operationID))
 
 			if err != nil {
 				l = l.Int("error_code", int(connect.CodeOf(err)))
@@ -84,10 +111,10 @@ func NewUnaryLogInterceptor(logger *zerolog.Logger) connect.UnaryInterceptorFunc
 			l.Msg(formatMethod(req.Spec().Procedure))
 
 			return res, err
-		})
+		}
 	}
 
-	return connect.UnaryInterceptorFunc(interceptor)
+	return interceptor
 }
 
 func LogServerStreamWithoutEvents[T any, R any](
@@ -97,23 +124,19 @@ func LogServerStreamWithoutEvents[T any, R any](
 	stream *connect.ServerStream[T],
 	handler func(ctx context.Context, req *connect.Request[R], stream *connect.ServerStream[T]) error,
 ) error {
-	ctx = AddRequestIDToContext(ctx)
+	ctx, operationID := AddRequestIDToContext(ctx)
 
 	l := logger.Debug().
 		Str("method", DefaultHTTPMethod+" "+req.Spec().Procedure).
-		Str(string(OperationIDKey), ctx.Value(OperationIDKey).(string))
-
-	if req != nil {
-		l = l.Interface("request", req.Any())
-	}
+		Str("operation_id", operationID).
+		Interface("request", req.Any())
 
 	l.Msg(fmt.Sprintf("%s (server stream start)", formatMethod(req.Spec().Procedure)))
 
 	err := handler(ctx, req, stream)
 
 	logEvent := getErrDebugLogEvent(logger, err).
-		Str("method", DefaultHTTPMethod+" "+req.Spec().Procedure).
-		Str(string(OperationIDKey), ctx.Value(OperationIDKey).(string))
+		Str("operation_id", operationID)
 
 	if err != nil {
 		logEvent = logEvent.Int("error_code", int(connect.CodeOf(err)))
@@ -132,18 +155,17 @@ func LogClientStreamWithoutEvents[T any, R any](
 	stream *connect.ClientStream[T],
 	handler func(ctx context.Context, stream *connect.ClientStream[T]) (*connect.Response[R], error),
 ) (*connect.Response[R], error) {
-	ctx = AddRequestIDToContext(ctx)
-
+	ctx, operationID := AddRequestIDToContext(ctx)
 	logger.Debug().
 		Str("method", DefaultHTTPMethod+" "+stream.Spec().Procedure).
-		Str(string(OperationIDKey), ctx.Value(OperationIDKey).(string)).
+		Str("operation_id", operationID).
 		Msg(fmt.Sprintf("%s (client stream start)", formatMethod(stream.Spec().Procedure)))
 
 	res, err := handler(ctx, stream)
 
 	logEvent := getErrDebugLogEvent(logger, err).
 		Str("method", DefaultHTTPMethod+" "+stream.Spec().Procedure).
-		Str(string(OperationIDKey), ctx.Value(OperationIDKey).(string))
+		Str("operation_id", operationID)
 
 	if err != nil {
 		logEvent = logEvent.Int("error_code", int(connect.CodeOf(err)))
diff --git a/packages/envd/internal/services/filesystem/watch.go b/packages/envd/internal/services/filesystem/watch.go
index a578706f97..925555b2a9 100644
--- a/packages/envd/internal/services/filesystem/watch.go
+++ b/packages/envd/internal/services/filesystem/watch.go
@@ -135,7 +135,7 @@ func (s Service) watchHandler(ctx context.Context, req *connect.Request[rpc.Watc
 				s.logger.
 					Debug().
 					Str("event_type", "filesystem_event").
-					Str(string(logs.OperationIDKey), ctx.Value(logs.OperationIDKey).(string)).
+					Str("operation_id", logs.SafeGetRequestID(ctx)).
 					Interface("filesystem_event", event).
 					Msg("Streaming filesystem event")
 
diff --git a/packages/envd/internal/services/filesystem/watch_sync.go b/packages/envd/internal/services/filesystem/watch_sync.go
index 2c7334f45e..e8bf112f44 100644
--- a/packages/envd/internal/services/filesystem/watch_sync.go
+++ b/packages/envd/internal/services/filesystem/watch_sync.go
@@ -11,7 +11,6 @@ import (
 	"github.com/e2b-dev/fsnotify"
 	"github.com/rs/zerolog"
 
-	"github.com/e2b-dev/infra/packages/envd/internal/logs"
 	"github.com/e2b-dev/infra/packages/envd/internal/permissions"
 	rpc "github.com/e2b-dev/infra/packages/envd/internal/services/spec/filesystem"
 	"github.com/e2b-dev/infra/packages/envd/internal/utils"
@@ -119,7 +118,7 @@ func CreateFileWatcher(ctx context.Context, watchPath string, recursive bool, op
 					logger.
 						Debug().
 						Str("event_type", "filesystem_event").
-						Str(string(logs.OperationIDKey), operationID).
+						Str("operation_id", operationID).
 						Interface("filesystem_event", event).
 						Msg("Streaming filesystem event")
 				}
diff --git a/packages/envd/internal/services/process/input.go b/packages/envd/internal/services/process/input.go
index e2a08268ee..ff67e2c14f 100644
--- a/packages/envd/internal/services/process/input.go
+++ b/packages/envd/internal/services/process/input.go
@@ -29,7 +29,7 @@ func handleInput(ctx context.Context, process *handler.Handler, in *rpc.ProcessI
 		logger.Debug().
 			Str("event_type", "stdin").
 			Interface("stdin", in.GetStdin()).
-			Str(string(logs.OperationIDKey), ctx.Value(logs.OperationIDKey).(string)).
+			Str("operation_id", logs.SafeGetRequestID(ctx)).
 			Msg("Streaming input to process")
 
 	default:
diff --git a/packages/envd/internal/services/process/start.go b/packages/envd/internal/services/process/start.go
index ce5e2c26f7..46d73203a7 100644
--- a/packages/envd/internal/services/process/start.go
+++ b/packages/envd/internal/services/process/start.go
@@ -20,15 +20,15 @@ import (
 func (s *Service) InitializeStartProcess(ctx context.Context, user *user.User, req *rpc.StartRequest) error {
 	var err error
 
-	ctx = logs.AddRequestIDToContext(ctx)
+	ctx, operationID := logs.AddRequestIDToContext(ctx)
 
 	defer s.logger.
 		Err(err).
 		Interface("request", req).
-		Str(string(logs.OperationIDKey), ctx.Value(logs.OperationIDKey).(string)).
+		Str("operation_id", operationID).
 		Msg("Initialized startCmd")
 
-	handlerL := s.logger.With().Str(string(logs.OperationIDKey), ctx.Value(logs.OperationIDKey).(string)).Logger()
+	handlerL := s.logger.With().Str("operation_id", operationID).Logger()
 
 	startProcCtx, startProcCancel := context.WithCancel(ctx)
 	proc, err := handler.New(startProcCtx, user, req, &handlerL, nil, startProcCancel)
@@ -60,7 +60,7 @@ func (s *Service) handleStart(ctx context.Context, req *connect.Request[rpc.Star
 	ctx, cancel := context.WithCancelCause(ctx)
 	defer cancel(nil)
 
-	handlerL := s.logger.With().Str(string(logs.OperationIDKey), ctx.Value(logs.OperationIDKey).(string)).Logger()
+	handlerL := s.logger.With().Str("operation_id", logs.SafeGetRequestID(ctx)).Logger()
 
 	u, err := permissions.GetAuthUser(ctx)
 	if err != nil {
diff --git a/packages/envd/internal/utils/map.go b/packages/envd/internal/utils/map.go
index 982019c409..8be4e3f404 100644
--- a/packages/envd/internal/utils/map.go
+++ b/packages/envd/internal/utils/map.go
@@ -1,6 +1,9 @@
 package utils
 
-import "sync"
+import (
+	"fmt"
+	"sync"
+)
 
 type Map[K comparable, V any] struct {
 	m sync.Map
@@ -18,28 +21,39 @@ func (m *Map[K, V]) Delete(key K) {
 
 func (m *Map[K, V]) Load(key K) (value V, ok bool) {
 	v, ok := m.m.Load(key)
-	if !ok {
-		return value, ok
+	if ok {
+		value, ok = v.(V)
 	}
-	return v.(V), ok
+	return
 }
 
 func (m *Map[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
 	v, loaded := m.m.LoadAndDelete(key)
-	if !loaded {
-		return value, loaded
+	if loaded {
+		value, loaded = v.(V)
 	}
-	return v.(V), loaded
+	return
 }
 
 func (m *Map[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
 	a, loaded := m.m.LoadOrStore(key, value)
-	return a.(V), loaded
+	if loaded {
+		actual, loaded = a.(V)
+	}
+	return
 }
 
 func (m *Map[K, V]) Range(f func(key K, value V) bool) {
 	m.m.Range(func(key, value any) bool {
-		return f(key.(K), value.(V))
+		k, ok := key.(K)
+		if !ok {
+			panic(fmt.Sprintf("key: expected %T got %T", k, key))
+		}
+		v, ok := value.(V)
+		if !ok {
+			panic(fmt.Sprintf("value: expected %T got %T", v, value))
+		}
+		return f(k, v)
 	})
 }
 
diff --git a/packages/orchestrator/internal/sandbox/block/cache.go b/packages/orchestrator/internal/sandbox/block/cache.go
index 14679ff44c..911e9a4314 100644
--- a/packages/orchestrator/internal/sandbox/block/cache.go
+++ b/packages/orchestrator/internal/sandbox/block/cache.go
@@ -236,7 +236,8 @@ func (m *Cache) WriteAtWithoutLock(b []byte, off int64) (int, error) {
 func (m *Cache) dirtySortedKeys() []int64 {
 	var keys []int64
 	m.dirty.Range(func(key, _ any) bool {
-		keys = append(keys, key.(int64))
+		n, _ := key.(int64)
+		keys = append(keys, n)
 		return true
 	})
 	sort.Slice(keys, func(i, j int) bool {
diff --git a/packages/orchestrator/internal/sandbox/uffd/handler.go b/packages/orchestrator/internal/sandbox/uffd/handler.go
index 6425c182b9..b5b6cb9fce 100644
--- a/packages/orchestrator/internal/sandbox/uffd/handler.go
+++ b/packages/orchestrator/internal/sandbox/uffd/handler.go
@@ -95,7 +95,10 @@ func (u *Uffd) handle(sandboxId string) error {
 		return fmt.Errorf("failed accepting firecracker connection: %w", err)
 	}
 
-	unixConn := conn.(*net.UnixConn)
+	unixConn, ok := conn.(*net.UnixConn)
+	if !ok {
+		return fmt.Errorf("expected *net.UnixConn, received %T", conn)
+	}
 
 	mappingsBuf := make([]byte, mappingsSize)
 	uffdBuf := make([]byte, syscall.CmsgSpace(fdSize))
diff --git a/packages/shared/pkg/logger/grpc.go b/packages/shared/pkg/logger/grpc.go
index 950509939e..dd6ace6c5f 100644
--- a/packages/shared/pkg/logger/grpc.go
+++ b/packages/shared/pkg/logger/grpc.go
@@ -24,23 +24,26 @@ func GRPCLogger(l *zap.Logger) logging.Logger {
 		}
 
 		for i := 0; i < len(fields)-1; i += 2 {
-			key := fields[i]
+			key, ok := fields[i].(string)
+			if !ok {
+				panic(fmt.Sprintf("key must be a string, not %T", fields[i]))
+			}
 			value := fields[i+1]
 
 			switch v := value.(type) {
 			case string:
-				f = append(f, zap.String(key.(string), v))
+				f = append(f, zap.String(key, v))
 
-				_, ok := methodFullNameMap[key.(string)]
+				_, ok := methodFullNameMap[key]
 				if ok {
-					methodFullNameMap[key.(string)] = v
+					methodFullNameMap[key] = v
 				}
 			case int:
-				f = append(f, zap.Int(key.(string), v))
+				f = append(f, zap.Int(key, v))
 			case bool:
-				f = append(f, zap.Bool(key.(string), v))
+				f = append(f, zap.Bool(key, v))
 			default:
-				f = append(f, zap.Any(key.(string), v))
+				f = append(f, zap.Any(key, v))
 			}
 		}
 
diff --git a/packages/shared/pkg/logs/loki.go b/packages/shared/pkg/logs/loki.go
index d083e960ab..5be7f1dc8a 100644
--- a/packages/shared/pkg/logs/loki.go
+++ b/packages/shared/pkg/logs/loki.go
@@ -18,7 +18,12 @@ func LokiResponseMapper(res *loghttp.QueryResponse, offset int32, level *LogLeve
 		return nil, fmt.Errorf("unexpected value type received from loki query fetch: %s", res.Data.Result.Type())
 	}
 
-	for _, stream := range res.Data.Result.(loghttp.Streams) {
+	streams, ok := res.Data.Result.(loghttp.Streams)
+	if !ok {
+		return nil, fmt.Errorf("expected loghttp.Streams, received %T", res.Data.Result)
+	}
+
+	for _, stream := range streams {
 		for _, entry := range stream.Entries {
 			fields, err := lokiFlatJsonLineParser(entry.Line)
 			if err != nil {
diff --git a/packages/shared/pkg/proxy/proxy_test.go b/packages/shared/pkg/proxy/proxy_test.go
index 59ac6a6787..03e6f0701e 100644
--- a/packages/shared/pkg/proxy/proxy_test.go
+++ b/packages/shared/pkg/proxy/proxy_test.go
@@ -13,6 +13,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
 	"gotest.tools/assert"
 
@@ -107,13 +108,17 @@ func assertBackendOutput(t *testing.T, backend *testBackend, resp *http.Response
 }
 
 // newTestProxy creates a new proxy server for testing
-func newTestProxy(getDestination func(r *http.Request) (*pool.Destination, error)) (*Proxy, uint, error) {
+func newTestProxy(t *testing.T, getDestination func(r *http.Request) (*pool.Destination, error)) (*Proxy, uint, error) {
+	t.Helper()
+
 	// Find a free port for the proxy
 	l, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		return nil, 0, fmt.Errorf("failed to get free port: %w", err)
 	}
-	port := l.Addr().(*net.TCPAddr).Port
+	addr, ok := l.Addr().(*net.TCPAddr)
+	require.True(t, ok)
+	port := addr.Port
 
 	// Set up the proxy server
 	proxy := New(
@@ -152,7 +157,7 @@ func TestProxyRoutesToTargetServer(t *testing.T) {
 		}, nil
 	}
 
-	proxy, port, err := newTestProxy(getDestination)
+	proxy, port, err := newTestProxy(t, getDestination)
 	if err != nil {
 		t.Fatalf("failed to create proxy: %v", err)
 	}
@@ -197,7 +202,7 @@ func TestProxyReusesConnections(t *testing.T) {
 		}, nil
 	}
 
-	proxy, port, err := newTestProxy(getDestination)
+	proxy, port, err := newTestProxy(t, getDestination)
 	if err != nil {
 		t.Fatalf("failed to create proxy: %v", err)
 	}
@@ -270,7 +275,7 @@ func TestProxyReuseConnectionsWhenBackendChangesFails(t *testing.T) {
 	}
 
 	// Create proxy with the initial routing function
-	proxy, port, err := newTestProxy(getDestination)
+	proxy, port, err := newTestProxy(t, getDestination)
 	if err != nil {
 		t.Fatalf("failed to create proxy: %v", err)
 	}
@@ -355,7 +360,7 @@ func TestProxyDoesNotReuseConnectionsWhenBackendChanges(t *testing.T) {
 	}
 
 	// Create proxy with the initial routing function
-	proxy, port, err := newTestProxy(getDestination)
+	proxy, port, err := newTestProxy(t, getDestination)
 	if err != nil {
 		t.Fatalf("failed to create proxy: %v", err)
 	}
diff --git a/packages/shared/pkg/storage/gcp_multipart_test.go b/packages/shared/pkg/storage/gcp_multipart_test.go
index da833fab6d..7bb19ff8fc 100644
--- a/packages/shared/pkg/storage/gcp_multipart_test.go
+++ b/packages/shared/pkg/storage/gcp_multipart_test.go
@@ -321,8 +321,11 @@ func TestMultipartUploader_HighConcurrency_StressTest(t *testing.T) {
 	// Verify content integrity
 	var reconstructed strings.Builder
 	for i := 1; i <= int(atomic.LoadInt32(&partCalls)); i++ {
-		if part, ok := receivedParts.Load(i); ok {
-			reconstructed.WriteString(part.(string))
+		if val, ok := receivedParts.Load(i); ok {
+			part, ok := val.(string)
+			if ok {
+				reconstructed.WriteString(part)
+			}
 		}
 	}
 	require.Equal(t, testContent, reconstructed.String())
@@ -415,7 +418,8 @@ func TestMultipartUploader_PartialFailures_Recovery(t *testing.T) {
 
 			// Track attempts per part
 			val, _ := partAttempts.LoadOrStore(partNumStr, new(int32))
-			attempts := val.(*int32)
+			attempts, ok := val.(*int32)
+			assert.True(t, ok)
 			currentAttempts := atomic.AddInt32(attempts, 1)
 
 			// Fail first few attempts for each part, then succeed
@@ -446,7 +450,9 @@ func TestMultipartUploader_PartialFailures_Recovery(t *testing.T) {
 
 	// Verify that all parts eventually succeeded after retries
 	partAttempts.Range(func(key, value interface{}) bool {
-		attempts := atomic.LoadInt32(value.(*int32))
+		n, ok := value.(*int32)
+		require.True(t, ok)
+		attempts := atomic.LoadInt32(n)
 		require.Equal(t, int32(maxAttempts-1), attempts, "Part %s should have exactly %d attempts", key, maxAttempts-1)
 		return true
 	})
@@ -677,7 +683,8 @@ func TestMultipartUploader_ConcurrentRetries_RaceCondition(t *testing.T) {
 
 			// Track retry attempts per part with race-safe operations
 			val, _ := retryAttempts.LoadOrStore(partNumStr, new(int32))
-			attempts := val.(*int32)
+			attempts, ok := val.(*int32)
+			assert.True(t, ok)
 			currentAttempt := atomic.AddInt32(attempts, 1)
 
 			// Fail first 2 attempts to force retries under high concurrency
@@ -711,7 +718,9 @@ func TestMultipartUploader_ConcurrentRetries_RaceCondition(t *testing.T) {
 
 	// Verify that retries happened correctly under concurrent conditions
 	retryAttempts.Range(func(key, value interface{}) bool {
-		attempts := atomic.LoadInt32(value.(*int32))
+		n, ok := value.(*int32)
+		require.True(t, ok)
+		attempts := atomic.LoadInt32(n)
 		require.GreaterOrEqual(t, attempts, int32(3), "Part %s should have at least 3 attempts", key)
 		return true
 	})
diff --git a/packages/shared/pkg/telemetry/tracing.go b/packages/shared/pkg/telemetry/tracing.go
index 04c609a995..1bf806a937 100644
--- a/packages/shared/pkg/telemetry/tracing.go
+++ b/packages/shared/pkg/telemetry/tracing.go
@@ -16,13 +16,8 @@ var OTELTracingPrint = os.Getenv("OTEL_TRACING_PRINT") != "false"
 const DebugID = "debug_id"
 
 func getDebugID(ctx context.Context) *string {
-	if ctx.Value(DebugID) == nil {
-		return nil
-	}
-
-	value := ctx.Value(DebugID).(string)
-
-	return &value
+	debugID, _ := ctx.Value(DebugID).(string)
+	return &debugID
 }
 
 func debugFormat(debugID *string, msg string) string {
diff --git a/packages/shared/scripts/seed/postgres/seed-db.go b/packages/shared/scripts/seed/postgres/seed-db.go
index f0e8ed2145..a61d65b674 100644
--- a/packages/shared/scripts/seed/postgres/seed-db.go
+++ b/packages/shared/scripts/seed/postgres/seed-db.go
@@ -56,10 +56,10 @@ func main() {
 		panic(err)
 	}
 
-	email := config["email"].(string)
-	teamID := config["teamId"].(string)
-	accessToken := config["accessToken"].(string)
-	teamAPIKey := config["teamApiKey"].(string)
+	email := safeString(config["email"])
+	teamID := safeString(config["teamId"])
+	accessToken := safeString(config["accessToken"])
+	teamAPIKey := safeString(config["teamApiKey"])
 	teamUUID := uuid.MustParse(teamID)
 
 	// Open .e2b/config.json
@@ -159,3 +159,8 @@ func main() {
 
 	fmt.Printf("Database seeded.\n")
 }
+
+func safeString(a any) string {
+	s, _ := a.(string)
+	return s
+}

From 41484000316e53e459b193ca3f05ec6494e16057 Mon Sep 17 00:00:00 2001
From: Joseph Lombrozo <joe@djeebus.net>
Date: Tue, 9 Sep 2025 17:45:19 +0000
Subject: [PATCH 2/6] cursor pr review changes

---
 packages/api/internal/auth/constants.go                   | 8 +-------
 .../api/internal/handlers/template_request_build_v2.go    | 4 ++--
 packages/client-proxy/internal/proxy/proxy.go             | 2 +-
 3 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/packages/api/internal/auth/constants.go b/packages/api/internal/auth/constants.go
index c2e51ac66f..2394a5bef4 100644
--- a/packages/api/internal/auth/constants.go
+++ b/packages/api/internal/auth/constants.go
@@ -11,11 +11,6 @@ import (
 	authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
 )
 
-const (
-	teamContextKey   string = "team"
-	userIDContextKey string = "user_id"
-)
-
 var (
 	ErrNotFoundInContext = errors.New("not found in context")
 	ErrInvalidType       = errors.New("unexpected type")
@@ -32,7 +27,7 @@ func (g *ginContextValueHelper[T]) set(c *gin.Context, val T) {
 func (g *ginContextValueHelper[T]) get(c *gin.Context) (T, error) {
 	var t T
 
-	v := c.Value(teamContextKey)
+	v := c.Value(g.contextKey)
 	if v == nil {
 		return t, ErrNotFoundInContext
 	}
@@ -70,7 +65,6 @@ func GetTeamInfo(c *gin.Context) (authcache.AuthTeamInfo, error) {
 func SafeGetTeamInfo(c *gin.Context) authcache.AuthTeamInfo {
 	return teamInfoHelper.safeGet(c)
 }
-
 func setUserID(c *gin.Context, userID uuid.UUID) {
 	userIDHelper.set(c, userID)
 }
diff --git a/packages/api/internal/handlers/template_request_build_v2.go b/packages/api/internal/handlers/template_request_build_v2.go
index f67c813cc5..0eb2163d8c 100644
--- a/packages/api/internal/handlers/template_request_build_v2.go
+++ b/packages/api/internal/handlers/template_request_build_v2.go
@@ -111,9 +111,9 @@ func (a *APIStore) GetTeamAndTier(
 	_, span := a.Tracer.Start(c.Request.Context(), "get-team-and-tier")
 	defer span.End()
 
-	if teamInfo, err := auth.GetTeamInfo(c); err != nil {
+	if teamInfo, err := auth.GetTeamInfo(c); err == nil {
 		return teamInfo.Team, teamInfo.Tier, nil
-	} else if _, err := auth.GetUserID(c); err != nil {
+	} else if _, err := auth.GetUserID(c); err == nil {
 		_, teams, err := a.GetUserAndTeams(c)
 		if err != nil {
 			return nil, nil, &api.APIError{
diff --git a/packages/client-proxy/internal/proxy/proxy.go b/packages/client-proxy/internal/proxy/proxy.go
index a32ca46f21..3cc82a783d 100644
--- a/packages/client-proxy/internal/proxy/proxy.go
+++ b/packages/client-proxy/internal/proxy/proxy.go
@@ -69,7 +69,7 @@ func dnsResolution(sandboxId string, logger *zap.Logger) (string, error) {
 			logger.Warn(fmt.Sprintf("received invalid answer type: %T", answer))
 			continue
 		}
-		node := answer.A.String()
+		node = answer.A.String()
 
 		// the sandbox was not found, we want to return this information to the user
 		if node == "127.0.0.1" {

From e0c39854a374c81f3929b4d9a5fcfacba6aa5446 Mon Sep 17 00:00:00 2001
From: Joseph Lombrozo <joe@djeebus.net>
Date: Tue, 9 Sep 2025 17:54:26 +0000
Subject: [PATCH 3/6] linting

---
 packages/api/internal/auth/constants.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/packages/api/internal/auth/constants.go b/packages/api/internal/auth/constants.go
index 2394a5bef4..ed1c2f86ee 100644
--- a/packages/api/internal/auth/constants.go
+++ b/packages/api/internal/auth/constants.go
@@ -65,6 +65,7 @@ func GetTeamInfo(c *gin.Context) (authcache.AuthTeamInfo, error) {
 func SafeGetTeamInfo(c *gin.Context) authcache.AuthTeamInfo {
 	return teamInfoHelper.safeGet(c)
 }
+
 func setUserID(c *gin.Context, userID uuid.UUID) {
 	userIDHelper.set(c, userID)
 }

From 38de2ed13682b135f66aa20f132af3e51bb0cf5d Mon Sep 17 00:00:00 2001
From: Joseph Lombrozo <joe@djeebus.net>
Date: Tue, 9 Sep 2025 17:59:36 +0000
Subject: [PATCH 4/6] fix debug id usage

---
 packages/shared/pkg/telemetry/tracing.go | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/packages/shared/pkg/telemetry/tracing.go b/packages/shared/pkg/telemetry/tracing.go
index 1bf806a937..80ae372d8c 100644
--- a/packages/shared/pkg/telemetry/tracing.go
+++ b/packages/shared/pkg/telemetry/tracing.go
@@ -16,7 +16,11 @@ var OTELTracingPrint = os.Getenv("OTEL_TRACING_PRINT") != "false"
 const DebugID = "debug_id"
 
 func getDebugID(ctx context.Context) *string {
-	debugID, _ := ctx.Value(DebugID).(string)
+	debugID, ok := ctx.Value(DebugID).(string)
+	if !ok {
+		return nil
+	}
+
 	return &debugID
 }
 

From b2b7397615dccb025c1e796f2682e3c08e0884b3 Mon Sep 17 00:00:00 2001
From: Joe Lombrozo <joe.lombrozo@e2b.dev>
Date: Tue, 16 Sep 2025 10:05:12 -0700
Subject: [PATCH 5/6] interesting ...

---
 .../internal/orchestrator/placement/placement_benchmark_test.go | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/packages/api/internal/orchestrator/placement/placement_benchmark_test.go b/packages/api/internal/orchestrator/placement/placement_benchmark_test.go
index 376c2ad55d..d7f10d50ec 100644
--- a/packages/api/internal/orchestrator/placement/placement_benchmark_test.go
+++ b/packages/api/internal/orchestrator/placement/placement_benchmark_test.go
@@ -11,6 +11,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/cache/instance"
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"

From 8452fdcc1f92cd9650df2046c645537fe5dab39d Mon Sep 17 00:00:00 2001
From: Joe Lombrozo <joe.lombrozo@e2b.dev>
Date: Tue, 16 Sep 2025 12:47:46 -0700
Subject: [PATCH 6/6] technically cursor is correct ...

---
 packages/envd/internal/utils/map.go           | 20 +++++++++++++++----
 .../internal/sandbox/block/cache.go           |  5 ++++-
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/packages/envd/internal/utils/map.go b/packages/envd/internal/utils/map.go
index 8be4e3f404..0f2bf27af1 100644
--- a/packages/envd/internal/utils/map.go
+++ b/packages/envd/internal/utils/map.go
@@ -21,17 +21,29 @@ func (m *Map[K, V]) Delete(key K) {
 
 func (m *Map[K, V]) Load(key K) (value V, ok bool) {
 	v, ok := m.m.Load(key)
-	if ok {
-		value, ok = v.(V)
+	if !ok {
+		return
 	}
+
+	value, ok = v.(V)
+	if !ok {
+		panic(fmt.Sprintf("invalid value type: %T", v))
+	}
+
 	return
 }
 
 func (m *Map[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
 	v, loaded := m.m.LoadAndDelete(key)
-	if loaded {
-		value, loaded = v.(V)
+	if !loaded {
+		return
 	}
+
+	value, cast := v.(V)
+	if !cast {
+		panic(fmt.Sprintf("invalid value type: %T", v))
+	}
+
 	return
 }
 
diff --git a/packages/orchestrator/internal/sandbox/block/cache.go b/packages/orchestrator/internal/sandbox/block/cache.go
index 911e9a4314..414c760fac 100644
--- a/packages/orchestrator/internal/sandbox/block/cache.go
+++ b/packages/orchestrator/internal/sandbox/block/cache.go
@@ -236,7 +236,10 @@ func (m *Cache) WriteAtWithoutLock(b []byte, off int64) (int, error) {
 func (m *Cache) dirtySortedKeys() []int64 {
 	var keys []int64
 	m.dirty.Range(func(key, _ any) bool {
-		n, _ := key.(int64)
+		n, ok := key.(int64)
+		if !ok {
+			panic(fmt.Sprintf("invalid map key type: %T", key))
+		}
 		keys = append(keys, n)
 		return true
 	})