fix

apps/api/internal/httpapi/auth_handlers.go

@@ -140,8 +140,9 @@ func (s *Server) beginAuthenticatedSession(ctx context.Context, userID, companyI
 	}
 	s.Sessions.Put(ctx, auth.SessionUserIDKey, userID.String())
 	s.putSessionVersion(ctx, userID)
-	// Fresh login/register clears any prior impersonation chain.
+	// Fresh login/register clears any prior impersonation chain / staff tenant act-as.
 	s.Sessions.Remove(ctx, auth.SessionImpersonatorIDKey)
+	s.Sessions.Remove(ctx, auth.SessionStaffHomeCompanyKey)
 	if companyID == uuid.Nil {
 		s.Sessions.Put(ctx, auth.SessionCompanyIDKey, "")
 		return nil
@@ -302,16 +303,32 @@ func (s *Server) handleMe(w http.ResponseWriter, r *http.Request) {
 		Error(w, http.StatusInternalServerError, "failed to load companies")
 		return
 	}
+	staffTenantSwitch := false
+	access, accessErr := s.Auth.GetStaffAccess(r.Context(), uid)
+	if accessErr == nil && access.FullAdmin && access.Role == auth.StaffRoleAdmin {
+		staffTenantSwitch = true
+		all, listErr := s.Auth.ListCompanies(r.Context(), 500)
+		if listErr != nil {
+			Error(w, http.StatusInternalServerError, "failed to load companies")
+			return
+		}
+		companies = all
+	}
 	cidStr := s.Sessions.GetString(r.Context(), auth.SessionCompanyIDKey)
+	homeStr := strings.TrimSpace(s.Sessions.GetString(r.Context(), auth.SessionStaffHomeCompanyKey))
 	out := map[string]any{
 		"user":              user,
 		"companies":         companies,
 		"active_company_id": cidStr,
 	}
-	if access, err := s.Auth.GetStaffAccess(r.Context(), uid); err == nil && (access.FullAdmin || access.SupportDesk) {
+	if accessErr == nil && (access.FullAdmin || access.SupportDesk) {
 		out["staff_access"] = access
 		out["staff_capabilities"] = auth.StaffCapabilities(access.Role)
 	}
+	if staffTenantSwitch {
+		out["staff_tenant_switch"] = true
+	}
+	staffOverride := false
 	if cid, err := uuid.Parse(cidStr); err == nil {
 		for _, c := range companies {
 			if c.ID == cid {
@@ -319,13 +336,33 @@ func (s *Server) handleMe(w http.ResponseWriter, r *http.Request) {
 				break
 			}
 		}
+		// Staff may have selected a company not in the capped list — resolve by id.
+		if _, ok := out["company"]; !ok && staffTenantSwitch {
+			if c, cerr := s.Auth.CompanyByID(r.Context(), cid); cerr == nil {
+				out["company"] = c
+			}
+		}
 		if credits, err := s.Billing.CreditsOverview(r.Context(), cid, s.Config.LowCreditsThreshold); err == nil {
 			out["credits"] = credits
 		}
 		if m, err := s.Auth.EnsureMembership(r.Context(), uid, cid); err == nil {
 			out["membership"] = map[string]string{"role": m.Role, "status": m.Status}
+		} else if staffTenantSwitch && errors.Is(err, auth.ErrNotCompanyMember) {
+			staffOverride = true
+			out["membership"] = map[string]any{"role": "admin", "status": "active", "staff_override": true}
 		}
 	}
+	if staffTenantSwitch && homeStr != "" {
+		out["staff_home_company_id"] = homeStr
+		if hid, err := uuid.Parse(homeStr); err == nil {
+			if c, cerr := s.Auth.CompanyByID(r.Context(), hid); cerr == nil {
+				out["staff_home_company"] = c
+			}
+		}
+	}
+	if staffOverride {
+		out["staff_tenant_acting"] = true
+	}
 	impersonating := false
 	if impStr := strings.TrimSpace(s.Sessions.GetString(r.Context(), auth.SessionImpersonatorIDKey)); impStr != "" {
 		if impID, err := uuid.Parse(impStr); err == nil && impID != uuid.Nil {
@@ -345,8 +382,7 @@ func (s *Server) handleMe(w http.ResponseWriter, r *http.Request) {
 	if !s.Config.IsProduction() {
 		canSwitch := impersonating
 		if !canSwitch {
-			access, err := s.checkStaffAccess(r.Context(), uid)
-			if err == nil && access.FullAdmin {
+			if accessErr == nil && access.FullAdmin {
 				canSwitch = true
 			} else if isLocalDemoEmail(user.Email) {
 				canSwitch = true
@@ -418,10 +454,48 @@ func (s *Server) handleSelectCompany(w http.ResponseWriter, r *http.Request) {
 		Error(w, http.StatusBadRequest, "invalid company_id")
 		return
 	}
-	if _, err := s.Auth.EnsureMembership(r.Context(), uid, cid); err != nil {
-		Error(w, http.StatusForbidden, "forbidden")
-		return
+	_, memErr := s.Auth.EnsureMembership(r.Context(), uid, cid)
+	staffOK := s.staffMayActAsCompany(r.Context(), uid)
+	if memErr != nil {
+		if !(errors.Is(memErr, auth.ErrNotCompanyMember) && staffOK) {
+			Error(w, http.StatusForbidden, "forbidden")
+			return
+		}
+		if _, cerr := s.Auth.CompanyByID(r.Context(), cid); cerr != nil {
+			if errors.Is(cerr, auth.ErrCompanyNotFound) {
+				Error(w, http.StatusNotFound, "company not found")
+				return
+			}
+			Error(w, http.StatusInternalServerError, "company lookup failed")
+			return
+		}
+		// First hop into a foreign tenant: remember home company for revert.
+		if strings.TrimSpace(s.Sessions.GetString(r.Context(), auth.SessionStaffHomeCompanyKey)) == "" {
+			if home := s.resolveStaffHomeCompany(r.Context(), uid); home != uuid.Nil {
+				s.Sessions.Put(r.Context(), auth.SessionStaffHomeCompanyKey, home.String())
+			}
+		}
+	} else {
+		// Back on a membership company — clear act-as home.
+		s.Sessions.Remove(r.Context(), auth.SessionStaffHomeCompanyKey)
 	}
 	s.Sessions.Put(r.Context(), auth.SessionCompanyIDKey, cid.String())
 	JSON(w, http.StatusOK, map[string]string{"company_id": cid.String()})
 }
+
+// resolveStaffHomeCompany picks the company to restore on revert: current session
+// company when the admin is a member, else their first membership company.
+func (s *Server) resolveStaffHomeCompany(ctx context.Context, userID uuid.UUID) uuid.UUID {
+	if cur := strings.TrimSpace(s.Sessions.GetString(ctx, auth.SessionCompanyIDKey)); cur != "" {
+		if cid, err := uuid.Parse(cur); err == nil {
+			if _, err := s.Auth.EnsureMembership(ctx, userID, cid); err == nil {
+				return cid
+			}
+		}
+	}
+	companies, err := s.Auth.ListUserCompanies(ctx, userID)
+	if err != nil || len(companies) == 0 {
+		return uuid.Nil
+	}
+	return companies[0].ID
+}