Enhance company context handling in customer middleware and routing

- Introduced CompanyContextMiddleware to resolve the active company context for customer requests, ensuring that tender recommendations and company-scoped APIs remain in sync with the database.
- Updated public routes to utilize the new CompanyContextMiddleware alongside the existing AuthMiddleware, improving the handling of company-specific requests.
- Added unit tests for the pickActiveCompanyID function to validate the logic for selecting the appropriate company context based on customer assignments and requested company IDs.

This update enhances the accuracy and reliability of company context management in the application, improving user experience and data consistency.

cmd/web/router/routes.go

@@ -281,7 +281,7 @@ func RegisterPublicRoutes(e *echo.Echo, customerHandler *customer.Handler, tende
 
 	// Public Tenders Routes
 	tendersGP := v1.Group("/tenders")
-	tendersGP.Use(customerHandler.AuthMiddleware())
+	tendersGP.Use(customerHandler.AuthMiddleware(), customerHandler.CompanyContextMiddleware())
 	{
 		tendersGP.GET("", tenderHandler.GetPublicTenders)
 		tendersGP.GET("/recommend", tenderHandler.RecommendTenders)
@@ -308,7 +308,7 @@ func RegisterPublicRoutes(e *echo.Echo, customerHandler *customer.Handler, tende
 	// Public Tender Approvals Routes
 	tenderApprovalGP := v1.Group("/tender-approvals")
 	{
-		tenderApprovalGP.Use(customerHandler.AuthMiddleware())
+		tenderApprovalGP.Use(customerHandler.AuthMiddleware(), customerHandler.CompanyContextMiddleware())
 		tenderApprovalGP.POST("", tenderApprovalHandler.ToggleTenderApproval)
 		tenderApprovalGP.GET("", tenderApprovalHandler.PublicGetTenderApprovals)
 		tenderApprovalGP.GET("/:id", tenderApprovalHandler.GetPublicTenderApproval)
@@ -319,13 +319,13 @@ func RegisterPublicRoutes(e *echo.Echo, customerHandler *customer.Handler, tende
 	// Public Company Routes
 	companiesGP := v1.Group("/companies")
 	{
-		companiesGP.Use(customerHandler.AuthMiddleware())
+		companiesGP.Use(customerHandler.AuthMiddleware(), customerHandler.CompanyContextMiddleware())
 		companiesGP.GET("", companyHandler.MyCompany)
 	}
 
 	// AI tender recommendation routes
 	recommendationGP := v1.Group("")
-	recommendationGP.Use(customerHandler.AuthMiddleware())
+	recommendationGP.Use(customerHandler.AuthMiddleware(), customerHandler.CompanyContextMiddleware())
 	{
 		recommendationGP.POST("/onboarding", companyHandler.StartOnboarding)
 		recommendationGP.POST("/recommend", companyHandler.RecommendTenders)

internal/customer/company_context.go

@@ -0,0 +1,49 @@
+package customer
+
+import (
+	"context"
+	"errors"
+	"strings"
+)
+
+var errCompanyNotAssigned = errors.New("company is not assigned to customer")
+
+// pickActiveCompanyID chooses the company context for a customer request.
+// When requestedCompanyID is set it must be in assigned; otherwise a still-valid
+// token company is kept; if the token company was removed, the first assignment is used.
+func pickActiveCompanyID(assigned []string, tokenCompanyID, requestedCompanyID string) (string, error) {
+	if len(assigned) == 0 {
+		return "", nil
+	}
+
+	requested := strings.TrimSpace(requestedCompanyID)
+	if requested != "" {
+		for _, id := range assigned {
+			if id == requested {
+				return requested, nil
+			}
+		}
+		return "", errCompanyNotAssigned
+	}
+
+	token := strings.TrimSpace(tokenCompanyID)
+	if token != "" {
+		for _, id := range assigned {
+			if id == token {
+				return token, nil
+			}
+		}
+	}
+
+	return assigned[0], nil
+}
+
+// ResolveActiveCompanyID resolves the active company for an authenticated customer.
+func (s *customerService) ResolveActiveCompanyID(ctx context.Context, customerID, tokenCompanyID, requestedCompanyID string) (string, error) {
+	customer, err := s.repository.GetByID(ctx, customerID)
+	if err != nil {
+		return "", errors.New("customer not found")
+	}
+
+	return pickActiveCompanyID(customer.Companies, tokenCompanyID, requestedCompanyID)
+}

internal/customer/company_context_test.go

@@ -0,0 +1,77 @@
+package customer
+
+import "testing"
+
+func TestPickActiveCompanyID(t *testing.T) {
+	assigned := []string{"company-a", "company-b"}
+
+	tests := []struct {
+		name        string
+		assigned    []string
+		token       string
+		requested   string
+		want        string
+		wantErr     bool
+		wantErrType error
+	}{
+		{
+			name:      "uses requested company when assigned",
+			assigned:  assigned,
+			token:     "company-a",
+			requested: "company-b",
+			want:      "company-b",
+		},
+		{
+			name:        "rejects requested company that is not assigned",
+			assigned:    assigned,
+			token:       "company-a",
+			requested:   "company-c",
+			wantErr:     true,
+			wantErrType: errCompanyNotAssigned,
+		},
+		{
+			name:     "keeps valid token company",
+			assigned: assigned,
+			token:    "company-b",
+			want:     "company-b",
+		},
+		{
+			name:     "falls back to first assignment when token company was removed",
+			assigned: []string{"company-b"},
+			token:    "company-a",
+			want:     "company-b",
+		},
+		{
+			name:     "returns empty when customer has no companies",
+			assigned: nil,
+			token:    "company-a",
+			want:     "",
+		},
+		{
+			name:     "uses first assignment when token is empty",
+			assigned: assigned,
+			want:     "company-a",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := pickActiveCompanyID(tt.assigned, tt.token, tt.requested)
+			if tt.wantErr {
+				if err == nil {
+					t.Fatalf("expected error, got nil")
+				}
+				if tt.wantErrType != nil && err != tt.wantErrType {
+					t.Fatalf("expected error %v, got %v", tt.wantErrType, err)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+			if got != tt.want {
+				t.Fatalf("pickActiveCompanyID() = %q, want %q", got, tt.want)
+			}
+		})
+	}
+}

internal/customer/middleware.go

@@ -1,6 +1,7 @@
 package customer
 
 import (
+	"errors"
 	"net/http"
 	"strings"
 	"tm/pkg/audit"
@@ -61,6 +62,39 @@ func (h *Handler) AuthMiddleware() echo.MiddlewareFunc {
 	}
 }
 
+// CompanyContextMiddleware resolves company_id from current customer assignments.
+// JWT company_id can be stale after admin reassigns companies; this keeps tender
+// recommendations and other company-scoped APIs in sync with the database.
+func (h *Handler) CompanyContextMiddleware() echo.MiddlewareFunc {
+	return func(next echo.HandlerFunc) echo.HandlerFunc {
+		return func(c echo.Context) error {
+			customerID, err := GetCustomerIDFromContext(c)
+			if err != nil {
+				return response.Unauthorized(c, "User not authenticated")
+			}
+
+			tokenCompanyID, _ := c.Get("company_id").(string)
+			requestedCompanyID := strings.TrimSpace(c.QueryParam("company_id"))
+
+			companyID, err := h.service.ResolveActiveCompanyID(
+				c.Request().Context(),
+				customerID,
+				tokenCompanyID,
+				requestedCompanyID,
+			)
+			if err != nil {
+				if errors.Is(err, errCompanyNotAssigned) {
+					return response.Forbidden(c, "Company is not assigned to customer")
+				}
+				return response.Unauthorized(c, "User not authenticated")
+			}
+
+			c.Set("company_id", companyID)
+			return next(c)
+		}
+	}
+}
+
 // GetCustomerIDFromContext extracts customer ID from Echo context
 func GetCustomerIDFromContext(c echo.Context) (string, error) {
 	customerID, ok := c.Get("customer_id").(string)

internal/customer/service.go

@@ -53,6 +53,9 @@ type Service interface {
 	// Role assignment operations
 	AssignRole(ctx context.Context, customerID string, form *AssignRoleForm) (*AssignRoleResponse, error)
 	AssignCompanies(ctx context.Context, customerID string, Companies []string) error
+
+	// ResolveActiveCompanyID returns the company context for an authenticated customer.
+	ResolveActiveCompanyID(ctx context.Context, customerID, tokenCompanyID, requestedCompanyID string) (string, error)
 }
 
 // customerService implements the Service interface
@@ -797,10 +800,9 @@ func (s *customerService) Login(ctx context.Context, form *LoginForm) (*AuthResp
 		return nil, errors.New("invalid credentials")
 	}
 
-	// Generate JWT tokens using authorization service
+	companyID, err := pickActiveCompanyID(customer.Companies, "", "")
-	companyID := ""
+	if err != nil {
-	if len(customer.Companies) > 0 {
+		return nil, errors.New("failed to resolve customer company")
-		companyID = customer.Companies[0] // Use first company ID for JWT token
 	}
 
 	tokenPair, err := s.authService.GenerateTokenPair(
@@ -863,32 +865,20 @@ func (s *customerService) RefreshToken(ctx context.Context, refreshToken string)
 		"refresh_token": refreshToken[:10] + "...", // Log only first 10 chars for security
 	})
 
-	// Validate refresh token and generate new access token
+	validationResult, err := s.authService.ValidateRefreshToken(refreshToken)
-	newAccessToken, err := s.authService.RefreshAccessToken(refreshToken)
 	if err != nil {
-		s.logger.Warn("Failed to refresh access token", map[string]interface{}{
+		s.logger.Warn("Failed to validate refresh token", map[string]interface{}{
 			"error": err.Error(),
 		})
 		return nil, errors.New("invalid or expired refresh token")
 	}
-
-	// Parse the new access token to get customer information
-	validationResult, err := s.authService.ValidateAccessToken(newAccessToken)
-	if err != nil {
-		s.logger.Error("Failed to validate new access token", map[string]interface{}{
-			"error": err.Error(),
-		})
-		return nil, errors.New("failed to generate valid access token")
-	}
-
 	if !validationResult.Valid {
-		s.logger.Error("New access token validation failed", map[string]interface{}{
+		s.logger.Warn("Invalid refresh token provided", map[string]interface{}{
 			"error": validationResult.Error,
 		})
-		return nil, errors.New("failed to generate valid access token")
+		return nil, errors.New("invalid or expired refresh token")
 	}
 
-	// Get customer information
 	customerID, err := bson.ObjectIDFromHex(validationResult.Claims.UserID)
 	if err != nil {
 		s.logger.Error("Failed to parse customer ID from token", map[string]interface{}{
@@ -906,11 +896,29 @@ func (s *customerService) RefreshToken(ctx context.Context, refreshToken string)
 		return nil, errors.New("customer not found")
 	}
 
-	// Check if customer is still active
 	if customer.Status != CustomerStatusActive {
 		return nil, errors.New("customer account is inactive")
 	}
 
+	companyID, err := pickActiveCompanyID(customer.Companies, validationResult.Claims.CompanyID, "")
+	if err != nil {
+		return nil, errors.New("failed to resolve customer company")
+	}
+
+	tokenPair, err := s.authService.GenerateTokenPair(
+		customer.ID.Hex(),
+		customer.Email,
+		string(customer.Role),
+		companyID,
+	)
+	if err != nil {
+		s.logger.Error("Failed to generate JWT tokens during refresh", map[string]interface{}{
+			"error":       err.Error(),
+			"customer_id": customer.ID,
+		})
+		return nil, errors.New("failed to generate authentication tokens")
+	}
+
 	companies, err := s.loadCompaniesForCustomer(ctx, customer)
 	if err != nil {
 		s.logger.Error("Failed to load companies for customer", map[string]interface{}{
@@ -922,13 +930,14 @@ func (s *customerService) RefreshToken(ctx context.Context, refreshToken string)
 	s.logger.Info("Access token refreshed successfully", map[string]interface{}{
 		"customer_id":   customer.ID,
 		"customer_type": string(customer.Type),
+		"company_id":    companyID,
 	})
 
 	return &AuthResponse{
 		Customer:     customer.ToResponse(companies),
-		AccessToken:  newAccessToken,
+		AccessToken:  tokenPair.AccessToken,
-		RefreshToken: refreshToken,                            // Return the same refresh token
+		RefreshToken: tokenPair.RefreshToken,
-		ExpiresAt:    time.Now().Add(15 * time.Minute).Unix(), // 15 minutes from now
+		ExpiresAt:    time.Now().Add(15 * time.Minute).Unix(),
 	}, nil
 }