Refactor: address remaining technical debt items

Fixes #30
Fixes #27
Fixes #21

Author: sangalo20

nexus-bridge/bridge_test.go

@@ -236,15 +236,15 @@ func TestBridge_ContextCancellation(t *testing.T) {
 		errChan <- bridge.MaintainWebSocket(ctx, "conn-123", "ws"+server.URL[4:], handler)
 	}()
 
-	time.Sleep(100 * time.Millisecond)
+	time.Sleep(500 * time.Millisecond)
 	cancel()
 
 	select {
 	case err := <-errChan:
 		if !errors.Is(err, context.Canceled) {
 			t.Errorf("Expected context.Canceled error, but got %v", err)
 		}
-	case <-time.After(1 * time.Second):
+	case <-time.After(5 * time.Second):
 		t.Fatal("Bridge did not exit after context cancellation")
 	}
 	if metrics.connectionStatus.Load() != 0.0 {

nexus-broker/cmd/nexus-broker/main.go

@@ -94,9 +94,13 @@ func main() {
 	store := provider.NewStore(db)
 
 	// Setup handlers
+	redirectPath := os.Getenv("REDIRECT_PATH")
+	if redirectPath == "" {
+		redirectPath = "/auth/callback"
+	}
 	providersHandler := handlers.NewProvidersHandler(store)
-	consentHandler := handlers.NewConsentHandler(db, baseURL, stateKey, cachingClient)
-	callbackHandler := handlers.NewCallbackHandler(db, encryptionKey, stateKey, cachingClient)
+	consentHandler := handlers.NewConsentHandler(db, baseURL, redirectPath, stateKey, cachingClient)
+	callbackHandler := handlers.NewCallbackHandler(db, baseURL, redirectPath, encryptionKey, stateKey, cachingClient)
 
 	// Setup routes
 	router := srv.Router()

nexus-broker/pkg/handlers/callback.go

@@ -9,7 +9,6 @@ import (
 	"net"
 	"net/http"
 	"net/url"
-	"os"
 	"strings"
 	"time"
 
@@ -29,6 +28,8 @@ import (
 // CallbackHandler handles OAuth callback and token exchange
 type CallbackHandler struct {
 	db                    *sqlx.DB
+	baseURL               string
+	redirectPath          string
 	encryptionKey         []byte
 	stateKey              []byte
 	httpClient            *http.Client
@@ -40,7 +41,7 @@ type CallbackHandler struct {
 }
 
 // NewCallbackHandler creates a new callback handler
-func NewCallbackHandler(db *sqlx.DB, encryptionKey, stateKey []byte, httpClient *http.Client) *CallbackHandler {
+func NewCallbackHandler(db *sqlx.DB, baseURL, redirectPath string, encryptionKey, stateKey []byte, httpClient *http.Client) *CallbackHandler {
 	success := prometheus.NewCounter(prometheus.CounterOpts{
 		Name:        "oauth_token_exchanges_total",
 		Help:        "Total OAuth token exchanges",
@@ -76,6 +77,8 @@ func NewCallbackHandler(db *sqlx.DB, encryptionKey, stateKey []byte, httpClient
 
 	return &CallbackHandler{
 		db:                    db,
+		baseURL:               baseURL,
+		redirectPath:          redirectPath,
 		encryptionKey:         encryptionKey,
 		stateKey:              stateKey,
 		httpClient:            httpClient,
@@ -161,11 +164,8 @@ func (h *CallbackHandler) Handle(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Compute redirect_uri to match the auth request
-	redirectPath := os.Getenv("REDIRECT_PATH")
-	if redirectPath == "" {
-		redirectPath = "/auth/callback"
-	}
-	base := strings.TrimSuffix(os.Getenv("BASE_URL"), "/")
+	redirectPath := h.redirectPath
+	base := strings.TrimSuffix(h.baseURL, "/")
 	redirectURI := base + redirectPath
 
 	// Check if provider wants to skip scope on token exchange (e.g., Salesforce rejects it)

nexus-broker/pkg/handlers/callback_test.go

@@ -26,7 +26,7 @@ func TestRefresh_StaticKeyProvider(t *testing.T) {
 	defer db.Close()
 
 	sqlxDB := sqlx.NewDb(db, "sqlmock")
-	handler := NewCallbackHandler(sqlxDB, []byte("test-key"), []byte("test-key"), http.DefaultClient)
+	handler := NewCallbackHandler(sqlxDB, "http://localhost:8080", "/auth/callback", []byte("test-key"), []byte("test-key"), http.DefaultClient)
 
 	// Mock the initial query to find the connection
 
@@ -66,7 +66,7 @@ func TestRefresh_OAuth2Provider(t *testing.T) {
 	}))
 	defer mockProviderServer.Close()
 
-	handler := NewCallbackHandler(sqlxDB, []byte("01234567890123456789012345678901"), []byte("01234567890123456789012345678901"), mockProviderServer.Client())
+	handler := NewCallbackHandler(sqlxDB, "http://localhost:8080", "/auth/callback", []byte("01234567890123456789012345678901"), []byte("01234567890123456789012345678901"), mockProviderServer.Client())
 
 	// Mock the initial query to find the connection
 
@@ -117,7 +117,7 @@ func TestGetCaptureSchema(t *testing.T) {
 	sqlxDB := sqlx.NewDb(db, "sqlmock")
 	// Use a real key for signing/verifying state
 	stateKey := []byte("01234567890123456789012345678901")
-	handler := NewCallbackHandler(sqlxDB, nil, stateKey, http.DefaultClient)
+	handler := NewCallbackHandler(sqlxDB, "http://localhost:8080", "/auth/callback", nil, stateKey, http.DefaultClient)
 
 	providerID := uuid.New()
 	stateData := auth.StateData{
@@ -170,7 +170,7 @@ func TestSaveCredential_ValidState(t *testing.T) {
 	sqlxDB := sqlx.NewDb(db, "sqlmock")
 	stateKey := []byte("01234567890123456789012345678901")
 	encryptionKey := []byte("01234567890123456789012345678901")
-	handler := NewCallbackHandler(sqlxDB, encryptionKey, stateKey, http.DefaultClient)
+	handler := NewCallbackHandler(sqlxDB, "http://localhost:8080", "/auth/callback", encryptionKey, stateKey, http.DefaultClient)
 
 	connectionID := uuid.New()
 	stateData := auth.StateData{
@@ -224,7 +224,7 @@ func TestSaveCredential_ValidState(t *testing.T) {
 }
 
 func TestSaveCredential_InvalidState(t *testing.T) {
-	handler := NewCallbackHandler(nil, nil, []byte("test-key"), http.DefaultClient)
+	handler := NewCallbackHandler(nil, "http://localhost:8080", "/auth/callback", nil, []byte("test-key"), http.DefaultClient)
 
 	creds := map[string]interface{}{"api_key": "test-key"}
 	body := map[string]interface{}{
@@ -245,7 +245,7 @@ func TestSaveCredential_InvalidState(t *testing.T) {
 }
 
 func TestSaveCredential_InvalidJSON(t *testing.T) {
-	handler := NewCallbackHandler(nil, nil, nil, http.DefaultClient)
+	handler := NewCallbackHandler(nil, "http://localhost:8080", "/auth/callback", nil, nil, http.DefaultClient)
 
 	req, err := http.NewRequest("POST", "/auth/capture-credential", bytes.NewBuffer([]byte("not-json")))
 	assert.NoError(t, err)

nexus-broker/pkg/handlers/consent.go

@@ -7,7 +7,6 @@ import (
 	"log"
 	"net/http"
 	"net/url"
-	"os"
 	"strings"
 	"time"
 
@@ -34,14 +33,15 @@ type ConsentSpec struct {
 type ConsentHandler struct {
 	db             *sqlx.DB
 	baseURL        string
+	redirectPath   string
 	stateKey       []byte
 	httpClient     *http.Client
 	consentsMetric prometheus.Counter
 	consentsOpenID prometheus.Counter
 }
 
 // NewConsentHandler creates a new consent handler
-func NewConsentHandler(db *sqlx.DB, baseURL string, stateKey []byte, httpClient *http.Client) *ConsentHandler {
+func NewConsentHandler(db *sqlx.DB, baseURL, redirectPath string, stateKey []byte, httpClient *http.Client) *ConsentHandler {
 	metric := prometheus.NewCounter(prometheus.CounterOpts{
 		Name: "oauth_consents_created_total",
 		Help: "Total OAuth consents created",
@@ -63,6 +63,7 @@ func NewConsentHandler(db *sqlx.DB, baseURL string, stateKey []byte, httpClient
 	return &ConsentHandler{
 		db:             db,
 		baseURL:        baseURL,
+		redirectPath:   redirectPath,
 		stateKey:       stateKey,
 		httpClient:     httpClient,
 		consentsMetric: metric,
@@ -246,10 +247,7 @@ func (h *ConsentHandler) GetSpec(w http.ResponseWriter, r *http.Request) {
 // buildAuthURL constructs the OAuth authorization URL
 func (h *ConsentHandler) buildAuthURL(providerAuthURL, clientID, state, codeChallenge string, scopes []string, providerParams *json.RawMessage) (string, error) {
 	baseURL := strings.TrimSuffix(h.baseURL, "/")
-	redirectPath := os.Getenv("REDIRECT_PATH")
-	if redirectPath == "" {
-		redirectPath = "/auth/callback"
-	}
+	redirectPath := h.redirectPath
 
 	if providerAuthURL == "" {
 		return "", fmt.Errorf("provider auth_url is required for OAuth2")

nexus-broker/pkg/handlers/consent_test.go

@@ -32,14 +32,13 @@ func TestGetSpec_OAuth2(t *testing.T) {
 	defer mockProviderServer.Close()
 
 	// Pass the test server's client to the handler
-	handler := NewConsentHandler(sqlxDB, "http://localhost:8080", []byte("test-key"), mockProviderServer.Client())
+	handler := NewConsentHandler(sqlxDB, "http://localhost:8080", "/auth/callback", []byte("test-key"), mockProviderServer.Client())
 
 	paramsJSON := []byte(`{"access_type": "offline", "prompt": "consent"}`)
 
-
-rows := sqlmock.NewRows([]string{"id", "name", "auth_type", "auth_url", "client_id", "scopes", "params"}).
+	rows := sqlmock.NewRows([]string{"id", "name", "auth_type", "auth_url", "client_id", "scopes", "params"}).
 		AddRow("a0a0a0a0-a0a0-a0a0-a0a0-a0a0a0a0a0a0", "Test OAuth2 Provider", "oauth2", "http://provider.com/auth", "test-client-id", "{openid}", paramsJSON)
-mock.ExpectQuery("SELECT id, name, auth_type, auth_url, client_id, scopes, params FROM provider_profiles WHERE id = \\$1").
+	mock.ExpectQuery("SELECT id, name, auth_type, auth_url, client_id, scopes, params FROM provider_profiles WHERE id = \\$1").
 		WithArgs("a0a0a0a0-a0a0-a0a0-a0a0-a0a0a0a0a0a0").
 		WillReturnRows(rows)
 
@@ -82,12 +81,11 @@ func TestGetSpec_StaticKey(t *testing.T) {
 
 	sqlxDB := sqlx.NewDb(db, "sqlmock")
 	// For static key tests, we can pass a default client as no external calls are made.
-	handler := NewConsentHandler(sqlxDB, "http://localhost:8080", []byte("test-key"), http.DefaultClient)
-
+	handler := NewConsentHandler(sqlxDB, "http://localhost:8080", "/auth/callback", []byte("test-key"), http.DefaultClient)
 
-rows := sqlmock.NewRows([]string{"id", "name", "auth_type", "auth_url", "client_id", "scopes", "params"}).
+	rows := sqlmock.NewRows([]string{"id", "name", "auth_type", "auth_url", "client_id", "scopes", "params"}).
 		AddRow("b1b1b1b1-b1b1-b1b1-b1b1-b1b1b1b1b1b1", "Test API", "api_key", nil, nil, "{}", []byte("{}"))
-mock.ExpectQuery("SELECT id, name, auth_type, auth_url, client_id, scopes, params FROM provider_profiles WHERE id = \\$1").
+	mock.ExpectQuery("SELECT id, name, auth_type, auth_url, client_id, scopes, params FROM provider_profiles WHERE id = \\$1").
 		WithArgs("b1b1b1b1-b1b1-b1b1-b1b1-b1b1b1b1b1b1").
 		WillReturnRows(rows)
 
@@ -134,7 +132,7 @@ func TestGetSpec_MixedOAuth2_Discovery(t *testing.T) {
 				"authorization_endpoint": "http://%s/openid/connect/authorize",
 				"jwks_uri": "http://%s/jwks"
 			}`,
-			r.Host, r.Host, r.Host)
+				r.Host, r.Host, r.Host)
 			w.Write([]byte(oidcConfig))
 			return
 		}
@@ -143,16 +141,16 @@ func TestGetSpec_MixedOAuth2_Discovery(t *testing.T) {
 	defer ts.Close()
 
 	// Handler under test
-	handler := NewConsentHandler(sqlxDB, "http://localhost:8080", []byte("test-key"), ts.Client())
+	handler := NewConsentHandler(sqlxDB, "http://localhost:8080", "/auth/callback", []byte("test-key"), ts.Client())
 
 	// Define the configured (legacy) auth URL
 	configuredAuthURL := ts.URL + "/oauth/v2/authorize"
 
 	// 1. Mock DB Provider Query
 
-rows := sqlmock.NewRows([]string{"id", "name", "auth_type", "auth_url", "client_id", "scopes", "params"}).
+	rows := sqlmock.NewRows([]string{"id", "name", "auth_type", "auth_url", "client_id", "scopes", "params"}).
 		AddRow("00000000-0000-0000-0000-000000000000", "Slack", "oauth2", configuredAuthURL, "slack-client", "{chat:write}", []byte("{}"))
-	
+
 	// Use regex to avoid strict string matching issues with sqlmock
 	mock.ExpectQuery("SELECT .* FROM provider_profiles WHERE id = .*").
 		WithArgs("00000000-0000-0000-0000-000000000000").
@@ -192,4 +190,4 @@ rows := sqlmock.NewRows([]string{"id", "name", "auth_type", "auth_url", "client_
 	if !strings.HasPrefix(response.AuthURL, configuredAuthURL) {
 		t.Errorf("Expected AuthURL to start with configured URL %s, but got %s", configuredAuthURL, response.AuthURL)
 	}
-}
+}

nexus-gateway/pkg/grpc/server_grpc.go

@@ -31,6 +31,21 @@ func NewService(handler *usecase.Handler) *Service {
 	return &Service{usecaseHandler: handler}
 }
 
+func mapUsecaseError(err error, msg string) error {
+	switch {
+	case errors.Is(err, usecase.ErrProviderNotFound):
+		return status.Errorf(codes.NotFound, "%s: %v", msg, err)
+	case errors.Is(err, usecase.ErrInvalidState):
+		return status.Errorf(codes.InvalidArgument, "%s: %v", msg, err)
+	case errors.Is(err, usecase.ErrProviderAmbiguous):
+		return status.Errorf(codes.FailedPrecondition, "%s: %v", msg, err)
+	case errors.Is(err, usecase.ErrBrokerUnavailable):
+		return status.Errorf(codes.Unavailable, "%s: %v", msg, err)
+	default:
+		return status.Errorf(codes.Internal, "%s: %v", msg, err)
+	}
+}
+
 // RequestConnection implements NexusServiceServer.RequestConnection.
 func (s *Service) RequestConnection(ctx context.Context, req *nexuspb.RequestConnectionRequest) (*nexuspb.RequestConnectionResponse, error) {
 	if req == nil {
@@ -45,7 +60,7 @@ func (s *Service) RequestConnection(ctx context.Context, req *nexuspb.RequestCon
 		Action:       req.GetAction(),
 	})
 	if err != nil {
-		return nil, status.Errorf(codes.Internal, "request connection failed: %v", err)
+		return nil, mapUsecaseError(err, "request connection failed")
 	}
 	return &nexuspb.RequestConnectionResponse{
 		AuthUrl:      out.AuthURL,
@@ -64,7 +79,7 @@ func (s *Service) CheckConnection(ctx context.Context, req *nexuspb.CheckConnect
 	}
 	statusStr, err := s.usecaseHandler.CheckConnectionCore(ctx, req.GetConnectionId())
 	if err != nil {
-		return nil, status.Errorf(codes.Internal, "check connection failed: %v", err)
+		return nil, mapUsecaseError(err, "check connection failed")
 	}
 	return &nexuspb.CheckConnectionResponse{Status: statusStr}, nil
 }
@@ -77,7 +92,7 @@ func (s *Service) GetToken(ctx context.Context, req *nexuspb.GetTokenRequest) (*
 	data, code, err := s.usecaseHandler.GetTokenCore(ctx, req.GetConnectionId())
 	if err != nil {
 		_ = code // keep the HTTP status for potential mapping if needed later
-		return nil, status.Errorf(codes.Internal, "get token failed: %v", err)
+		return nil, mapUsecaseError(err, "get token failed")
 	}
 	st, err := structpb.NewStruct(data)
 	if err != nil {
@@ -94,7 +109,7 @@ func (s *Service) RefreshConnection(ctx context.Context, req *nexuspb.RefreshCon
 	data, code, err := s.usecaseHandler.RefreshConnectionCore(ctx, req.GetConnectionId())
 	if err != nil {
 		_ = code // unused
-		return nil, status.Errorf(codes.Internal, "refresh connection failed: %v", err)
+		return nil, mapUsecaseError(err, "refresh connection failed")
 	}
 	st, err := structpb.NewStruct(data)
 	if err != nil {

nexus-sdk/client.go

@@ -149,18 +149,7 @@ func (c *Client) GetToken(ctx context.Context, connectionID string) (*TokenRespo
     defer resp.Body.Close()
     var raw map[string]any
     if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil { return nil, err }
-    tr := &TokenResponse{Raw: raw}
-    if v, ok := raw["access_token"].(string); ok { tr.AccessToken = v }
-    if v, ok := raw["token_type"].(string); ok { tr.TokenType = &v }
-    if v, ok := raw["expires_in"].(float64); ok { vv := int64(v); tr.ExpiresIn = &vv }
-    if v, ok := raw["expires_at"]; ok { tr.ExpiresAt = v }
-    if v, ok := raw["scope"].(string); ok { tr.Scope = &v }
-    if v, ok := raw["id_token"].(string); ok { tr.IDToken = &v }
-    if v, ok := raw["refresh_token"].(string); ok { tr.RefreshToken = &v }
-    if v, ok := raw["provider"].(string); ok { tr.Provider = &v }
-    if v, ok := raw["strategy"].(map[string]interface{}); ok { tr.Strategy = v }
-    if v, ok := raw["credentials"].(map[string]interface{}); ok { tr.Credentials = v }
-    return tr, nil
+    return parseTokenResponse(raw), nil
 }
 
 // RefreshConnection calls the Gateway to force a token refresh.
@@ -171,6 +160,10 @@ func (c *Client) RefreshConnection(ctx context.Context, connectionID string) (*T
     defer resp.Body.Close()
     var raw map[string]any
     if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil { return nil, err }
+    return parseTokenResponse(raw), nil
+}
+
+func parseTokenResponse(raw map[string]any) *TokenResponse {
     tr := &TokenResponse{Raw: raw}
     if v, ok := raw["access_token"].(string); ok { tr.AccessToken = v }
     if v, ok := raw["token_type"].(string); ok { tr.TokenType = &v }
@@ -182,7 +175,7 @@ func (c *Client) RefreshConnection(ctx context.Context, connectionID string) (*T
     if v, ok := raw["provider"].(string); ok { tr.Provider = &v }
     if v, ok := raw["strategy"].(map[string]interface{}); ok { tr.Strategy = v }
     if v, ok := raw["credentials"].(map[string]interface{}); ok { tr.Credentials = v }
-    return tr, nil
+    return tr
 }
 
 // RefreshViaBroker calls RefreshConnection (Gateway Proxy).