nvidia-nat 1.3.0a20250928__py3-none-any.whl → 1.3.0a20250930__py3-none-any.whl

nat/authentication/oauth2/oauth2_auth_code_flow_provider.py

@@ -13,10 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+from collections.abc import Callable
+from datetime import UTC
 from datetime import datetime
-from datetime import timezone
-from typing import Callable
 
+import httpx
 from authlib.integrations.httpx_client import OAuth2Client as AuthlibOAuth2Client
 from pydantic import SecretStr
 
@@ -28,6 +30,8 @@ from nat.data_models.authentication import AuthFlowType
 from nat.data_models.authentication import AuthResult
 from nat.data_models.authentication import BearerTokenCred
 
+logger = logging.getLogger(__name__)
+
 
 class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConfig]):
 
@@ -41,26 +45,30 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
         if not isinstance(refresh_token, str):
             return None
 
-        with AuthlibOAuth2Client(
-                client_id=self.config.client_id,
-                client_secret=self.config.client_secret,
-        ) as client:
-            try:
+        try:
+            with AuthlibOAuth2Client(
+                    client_id=self.config.client_id,
+                    client_secret=self.config.client_secret,
+            ) as client:
                 new_token_data = client.refresh_token(self.config.token_url, refresh_token=refresh_token)
-            except Exception:
-                # On any failure, we'll fall back to the full auth flow.
-                return None
 
-        expires_at_ts = new_token_data.get("expires_at")
-        new_expires_at = datetime.fromtimestamp(expires_at_ts, tz=timezone.utc) if expires_at_ts else None
+                expires_at_ts = new_token_data.get("expires_at")
+                new_expires_at = datetime.fromtimestamp(expires_at_ts, tz=UTC) if expires_at_ts else None
 
-        new_auth_result = AuthResult(
-            credentials=[BearerTokenCred(token=SecretStr(new_token_data["access_token"]))],
-            token_expires_at=new_expires_at,
-            raw=new_token_data,
-        )
+            new_auth_result = AuthResult(
+                credentials=[BearerTokenCred(token=SecretStr(new_token_data["access_token"]))],
+                token_expires_at=new_expires_at,
+                raw=new_token_data,
+            )
 
-        self._authenticated_tokens[user_id] = new_auth_result
+            self._authenticated_tokens[user_id] = new_auth_result
+        except httpx.HTTPStatusError:
+            return None
+        except httpx.RequestError:
+            return None
+        except Exception:
+            # On any other failure, we'll fall back to the full auth flow.
+            return None
 
         return new_auth_result
 

nat/builder/builder.py

@@ -56,7 +56,7 @@ if typing.TYPE_CHECKING:
     from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
 
 
-class UserManagerHolder():
+class UserManagerHolder:
 
     def __init__(self, context: Context) -> None:
         self._context = context

nat/builder/context.py

@@ -40,12 +40,12 @@ from nat.utils.reactive.subject import Subject
 class Singleton(type):
 
     def __init__(cls, name, bases, dict):
-        super(Singleton, cls).__init__(name, bases, dict)
+        super().__init__(name, bases, dict)
         cls.instance = None
 
     def __call__(cls, *args, **kw):
         if cls.instance is None:
-            cls.instance = super(Singleton, cls).__call__(*args, **kw)
+            cls.instance = super().__call__(*args, **kw)
         return cls.instance
 
 

nat/builder/front_end.py

@@ -37,7 +37,7 @@ class FrontEndBase(typing.Generic[FrontEndConfigT], ABC):
 
         super().__init__()
 
-        self._full_config: "Config" = full_config
+        self._full_config: Config = full_config
         self._front_end_config: FrontEndConfigT = typing.cast(FrontEndConfigT, full_config.general.front_end)
 
     @property

nat/cli/cli_utils/config_override.py

@@ -84,7 +84,7 @@ class LayeredConfig:
                     if lower_value not in ['true', 'false']:
                         raise ValueError(f"Boolean value must be 'true' or 'false', got '{value}'")
                     value = lower_value == 'true'
-                elif isinstance(original_value, (int, float)):
+                elif isinstance(original_value, int | float):
                     value = type(original_value)(value)
                 elif isinstance(original_value, list):
                     value = [v.strip() for v in value.split(',')]

nat/cli/commands/mcp/mcp.py

@@ -297,7 +297,7 @@ async def list_tools_via_function_group(
             if fn is not None:
                 tools.append(to_tool_entry(full, fn))
         else:
-            for full, fn in fns.items():
+            for full, fn in (await fns).items():
                 tools.append(to_tool_entry(full, fn))
 
     return tools
@@ -443,7 +443,7 @@ async def ping_mcp_server(url: str,
         # Apply timeout to the entire ping operation
         return await asyncio.wait_for(_ping_operation(), timeout=timeout)
 
-    except asyncio.TimeoutError:
+    except TimeoutError:
         return MCPPingResult(url=url,
                              status="unhealthy",
                              response_time_ms=None,

nat/cli/commands/start.py

@@ -111,7 +111,7 @@ class StartCommandGroup(click.Group):
             elif (issubclass(decomposed_type.root, Path)):
                 param_type = click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path)
 
-            elif (issubclass(decomposed_type.root, (list, tuple, set))):
+            elif (issubclass(decomposed_type.root, list | tuple | set)):
                 if (len(decomposed_type.args) == 1):
                     inner = DecomposedType(decomposed_type.args[0])
                     # Support containers of Literal values -> multiple Choice

nat/cli/type_registry.py

@@ -992,7 +992,7 @@ class TypeRegistry:
             if (short_names[key.local_name] == 1):
                 type_list.append((key.local_name, key.config_type))
 
-        return typing.Union[tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
+        return typing.Union[*tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
 
     def compute_annotation(self, cls: type[TypedBaseModelT]):
 

nat/control_flow/router_agent/register.py

@@ -81,7 +81,7 @@ async def router_agent_workflow(config: RouterAgentWorkflowConfig, builder: Buil
             logger.exception("%s Router Agent failed with exception: %s", AGENT_LOG_PREFIX, ex)
             if config.verbose:
                 return str(ex)
-            return "Router agent failed with exception: %s" % ex
+            return f"Router agent failed with exception: {ex}"
 
     try:
         yield FunctionInfo.from_fn(_response_fn, description=config.description)

nat/data_models/api_server.py

@@ -273,7 +273,7 @@ class ChatResponse(ResponseBaseModelOutput):
         if model is None:
             model = ""
         if created is None:
-            created = datetime.datetime.now(datetime.timezone.utc)
+            created = datetime.datetime.now(datetime.UTC)
 
         return ChatResponse(id=id_,
                             object=object_,
@@ -317,7 +317,7 @@ class ChatResponseChunk(ResponseBaseModelOutput):
         if id_ is None:
             id_ = str(uuid.uuid4())
         if created is None:
-            created = datetime.datetime.now(datetime.timezone.utc)
+            created = datetime.datetime.now(datetime.UTC)
         if model is None:
             model = ""
         if object_ is None:
@@ -343,7 +343,7 @@ class ChatResponseChunk(ResponseBaseModelOutput):
         if id_ is None:
             id_ = str(uuid.uuid4())
         if created is None:
-            created = datetime.datetime.now(datetime.timezone.utc)
+            created = datetime.datetime.now(datetime.UTC)
         if model is None:
             model = ""
 
@@ -485,7 +485,7 @@ class WebSocketUserMessage(BaseModel):
     security: Security = Security()
     error: Error = Error()
     schema_version: str = "1.0.0"
-    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
 
 class WebSocketUserInteractionResponseMessage(BaseModel):
@@ -501,7 +501,7 @@ class WebSocketUserInteractionResponseMessage(BaseModel):
     security: Security = Security()
     error: Error = Error()
     schema_version: str = "1.0.0"
-    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
 
 class SystemIntermediateStepContent(BaseModel):
@@ -527,7 +527,7 @@ class WebSocketSystemIntermediateStepMessage(BaseModel):
     conversation_id: str | None = None
     content: SystemIntermediateStepContent
     status: WebSocketMessageStatus
-    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
 
 class SystemResponseContent(BaseModel):
@@ -551,7 +551,7 @@ class WebSocketSystemResponseTokenMessage(BaseModel):
     conversation_id: str | None = None
     content: SystemResponseContent | Error | GenerateResponse
     status: WebSocketMessageStatus
-    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
     @field_validator("content")
     @classmethod
@@ -560,7 +560,7 @@ class WebSocketSystemResponseTokenMessage(BaseModel):
             raise ValueError(f"Field: content must be 'Error' when type is {WebSocketMessageType.ERROR_MESSAGE}")
 
         if info.data.get("type") == WebSocketMessageType.RESPONSE_MESSAGE and not isinstance(
-                value, (SystemResponseContent, GenerateResponse)):
+                value, SystemResponseContent | GenerateResponse):
             raise ValueError(
                 f"Field: content must be 'SystemResponseContent' when type is {WebSocketMessageType.RESPONSE_MESSAGE}")
         return value
@@ -582,7 +582,7 @@ class WebSocketSystemInteractionMessage(BaseModel):
     conversation_id: str | None = None
     content: HumanPrompt
     status: WebSocketMessageStatus
-    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
 
 # ======== GenerateResponse Converters ========

nat/data_models/authentication.py

@@ -14,8 +14,8 @@
 # limitations under the License.
 
 import typing
+from datetime import UTC
 from datetime import datetime
-from datetime import timezone
 from enum import Enum
 
 import httpx
@@ -166,13 +166,7 @@ class BearerTokenCred(_CredBase):
 
 
 Credential = typing.Annotated[
-    typing.Union[
-        HeaderCred,
-        QueryCred,
-        CookieCred,
-        BasicAuthCred,
-        BearerTokenCred,
-    ],
+    HeaderCred | QueryCred | CookieCred | BasicAuthCred | BearerTokenCred,
     Field(discriminator="kind"),
 ]
 
@@ -213,7 +207,7 @@ class AuthResult(BaseModel):
         """
         Checks if the authentication token has expired.
         """
-        return bool(self.token_expires_at and datetime.now(timezone.utc) >= self.token_expires_at)
+        return bool(self.token_expires_at and datetime.now(UTC) >= self.token_expires_at)
 
     def as_requests_kwargs(self) -> dict[str, typing.Any]:
         """

nat/data_models/dataset_handler.py

@@ -80,7 +80,7 @@ class EvalDatasetJsonConfig(EvalDatasetBaseConfig, name="json"):
 
 
 def read_jsonl(file_path: FilePath):
-    with open(file_path, 'r', encoding='utf-8') as f:
+    with open(file_path, encoding='utf-8') as f:
         data = [json.loads(line) for line in f]
     return pd.DataFrame(data)
 

nat/eval/evaluator/base_evaluator.py

@@ -71,7 +71,7 @@ class BaseEvaluator(ABC):
             TqdmPositionRegistry.release(tqdm_position)
 
         # Compute average if possible
-        numeric_scores = [item.score for item in output_items if isinstance(item.score, (int, float))]
+        numeric_scores = [item.score for item in output_items if isinstance(item.score, int | float)]
         avg_score = round(sum(numeric_scores) / len(numeric_scores), 2) if numeric_scores else None
 
         return EvalOutput(average_score=avg_score, eval_output_items=output_items)

nat/eval/swe_bench_evaluator/evaluate.py

@@ -204,7 +204,7 @@ class SweBenchEvaluator:
         # if report file is not present, return empty EvalOutput
         avg_score = 0.0
         if report_file.exists():
-            with open(report_file, "r", encoding="utf-8") as f:
+            with open(report_file, encoding="utf-8") as f:
                 report = json.load(f)
                 resolved_instances = report.get("resolved_instances", 0)
                 total_instances = report.get("total_instances", 0)

nat/eval/tunable_rag_evaluator/evaluate.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import logging
-from typing import Callable
+from collections.abc import Callable
 
 from langchain.output_parsers import ResponseSchema
 from langchain.output_parsers import StructuredOutputParser

nat/experimental/decorators/experimental_warning_decorator.py

@@ -137,8 +137,7 @@ def experimental(func: Any = None, *, feature_name: str | None = None, metadata:
         @functools.wraps(func)
         def sync_gen_wrapper(*args, **kwargs) -> Generator[Any, Any, Any]:
             issue_experimental_warning(function_name, feature_name, metadata)
-            for item in func(*args, **kwargs):
-                yield item  # yield the original item
+            yield from func(*args, **kwargs)  # yield the original item
 
         return sync_gen_wrapper
 

nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py

@@ -71,7 +71,7 @@ class LLMBasedOutputMergingSelector(StrategyBase):
             raise ImportError("langchain-core is not installed. Please install it to use SingleShotMultiPlanPlanner.\n"
                               "This error can be resolved by installing nvidia-nat-langchain.")
 
-        from typing import Callable
+        from collections.abc import Callable
 
         from pydantic import BaseModel
 

nat/front_ends/console/authentication_flow_handler.py

@@ -14,13 +14,16 @@
 # limitations under the License.
 
 import asyncio
+import logging
 import secrets
 import webbrowser
 from dataclasses import dataclass
 from dataclasses import field
 
 import click
+import httpx
 import pkce
+from authlib.common.errors import AuthlibBaseError as OAuthError
 from authlib.integrations.httpx_client import AsyncOAuth2Client
 from fastapi import FastAPI
 from fastapi import Request
@@ -32,6 +35,8 @@ from nat.data_models.authentication import AuthFlowType
 from nat.data_models.authentication import AuthProviderBaseConfig
 from nat.front_ends.fastapi.fastapi_front_end_controller import _FastApiFrontEndController
 
+logger = logging.getLogger(__name__)
+
 
 # --------------------------------------------------------------------------- #
 # Helpers                                                                     #
@@ -87,17 +92,53 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
         """
         Separated for easy overriding in tests (to inject ASGITransport).
         """
-        client = AsyncOAuth2Client(
-            client_id=cfg.client_id,
-            client_secret=cfg.client_secret,
-            redirect_uri=cfg.redirect_uri,
-            scope=" ".join(cfg.scopes) if cfg.scopes else None,
-            token_endpoint=cfg.token_url,
-            token_endpoint_auth_method=cfg.token_endpoint_auth_method,
-            code_challenge_method="S256" if cfg.use_pkce else None,
-        )
-        self._oauth_client = client
-        return client
+        try:
+            client = AsyncOAuth2Client(
+                client_id=cfg.client_id,
+                client_secret=cfg.client_secret,
+                redirect_uri=cfg.redirect_uri,
+                scope=" ".join(cfg.scopes) if cfg.scopes else None,
+                token_endpoint=cfg.token_url,
+                token_endpoint_auth_method=cfg.token_endpoint_auth_method,
+                code_challenge_method="S256" if cfg.use_pkce else None,
+            )
+            self._oauth_client = client
+            return client
+        except (OAuthError, ValueError, TypeError) as e:
+            raise RuntimeError(f"Invalid OAuth2 configuration: {e}") from e
+        except Exception as e:
+            raise RuntimeError(f"Failed to create OAuth2 client: {e}") from e
+
+    def _create_authorization_url(self,
+                                  client: AsyncOAuth2Client,
+                                  config: OAuth2AuthCodeFlowProviderConfig,
+                                  state: str,
+                                  verifier: str | None = None,
+                                  challenge: str | None = None) -> str:
+        """
+        Create OAuth authorization URL with proper error handling.
+
+        Args:
+            client: The OAuth2 client instance
+            config: OAuth2 configuration
+            state: OAuth state parameter
+            verifier: PKCE verifier (if using PKCE)
+            challenge: PKCE challenge (if using PKCE)
+
+        Returns:
+            The authorization URL
+        """
+        try:
+            auth_url, _ = client.create_authorization_url(
+                config.authorization_url,
+                state=state,
+                code_verifier=verifier if config.use_pkce else None,
+                code_challenge=challenge if config.use_pkce else None,
+                **(config.authorization_kwargs or {})
+            )
+            return auth_url
+        except (OAuthError, ValueError, TypeError) as e:
+            raise RuntimeError(f"Error creating OAuth authorization URL: {e}") from e
 
     # --------------------------- HTTP Basic ------------------------------ #
     @staticmethod
@@ -131,13 +172,12 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
             flow_state.verifier = verifier
             flow_state.challenge = challenge
 
-        auth_url, _ = client.create_authorization_url(
-            cfg.authorization_url,
-            state=state,
-            code_verifier=flow_state.verifier if cfg.use_pkce else None,
-            code_challenge=flow_state.challenge if cfg.use_pkce else None,
-            **(cfg.authorization_kwargs or {})
-        )
+        # Create authorization URL using helper function
+        auth_url = self._create_authorization_url(client=client,
+                                                  config=cfg,
+                                                  state=state,
+                                                  verifier=flow_state.verifier,
+                                                  challenge=flow_state.challenge)
 
         # Register flow + maybe spin up redirect handler
         async with self._server_lock:
@@ -149,14 +189,18 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
             self._flows[state] = flow_state
             self._active_flows += 1
 
-        click.echo("Your browser has been opened for authentication.")
-        webbrowser.open(auth_url)
+        try:
+            webbrowser.open(auth_url)
+            click.echo("Your browser has been opened for authentication.")
+        except Exception as e:
+            logger.error("Browser open failed: %s", e)
+            raise RuntimeError(f"Browser open failed: {e}") from e
 
         # Wait for the redirect to land
         try:
             token = await asyncio.wait_for(flow_state.future, timeout=300)
-        except asyncio.TimeoutError:
-            raise RuntimeError("Authentication timed out (5 min).")
+        except TimeoutError as exc:
+            raise RuntimeError("Authentication timed out (5 min).") from exc
         finally:
             async with self._server_lock:
                 self._flows.pop(state, None)
@@ -175,9 +219,9 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
     # --------------- redirect server / in‑process app -------------------- #
     async def _build_redirect_app(self) -> FastAPI:
         """
-        * If cfg.run_redirect_local_server == True → start a uvicorn server (old behaviour).
-        * Else → only build the FastAPI app and save it to `self._redirect_app`
-                 for in‑process testing with ASGITransport.
+        * If cfg.run_redirect_local_server == True → start a local server.
+        * Else → only build the redirect app and save it to `self._redirect_app`
+                 for in‑process testing.
         """
         app = FastAPI()
 
@@ -195,8 +239,16 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
                     state=state,
                 )
                 flow_state.future.set_result(token)
-            except Exception as exc:  # noqa: BLE001
-                flow_state.future.set_exception(exc)
+            except OAuthError as e:
+                flow_state.future.set_exception(
+                    RuntimeError(f"Authorization server rejected request: {e.error} ({e.description})"))
+                return "Authentication failed: Authorization server rejected the request. You may close this tab."
+            except httpx.HTTPError as e:
+                flow_state.future.set_exception(RuntimeError(f"Network error during token fetch: {e}"))
+                return "Authentication failed: Network error occurred. You may close this tab."
+            except Exception as e:
+                flow_state.future.set_exception(RuntimeError(f"Authentication failed: {e}"))
+                return "Authentication failed: An unexpected error occurred. You may close this tab."
             return "Authentication successful – you may close this tab."
 
         return app
@@ -213,7 +265,7 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
 
             asyncio.create_task(self._server_controller.start_server(host="localhost", port=8000))
 
-            # Give uvicorn a moment to bind sockets before we return
+            # Give the server a moment to bind sockets before we return
             await asyncio.sleep(0.3)
         except Exception as exc:  # noqa: BLE001
             raise RuntimeError(f"Failed to start redirect server: {exc}") from exc
@@ -227,7 +279,7 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
     @property
     def redirect_app(self) -> FastAPI | None:
         """
-        In “test‑mode” (run_redirect_local_server=False) the in‑memory FastAPI
-        app is exposed so you can mount it on `httpx.ASGITransport`.
+        In test mode (run_redirect_local_server=False) the in‑memory redirect
+        app is exposed for testing purposes.
         """
         return self._redirect_app

nat/front_ends/console/console_front_end_plugin.py

@@ -88,7 +88,7 @@ class ConsoleFrontEndPlugin(SimpleFrontEndPluginBase[ConsoleFrontEndConfig]):
         elif (self.front_end_config.input_file):
 
             # Run the workflow
-            with open(self.front_end_config.input_file, "r", encoding="utf-8") as f:
+            with open(self.front_end_config.input_file, encoding="utf-8") as f:
 
                 async with session_manager.workflow.run(f) as runner:
                     runner_outputs = await runner.result(to_type=str)

nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py

@@ -22,6 +22,7 @@ from dataclasses import dataclass
 from dataclasses import field
 
 import pkce
+from authlib.common.errors import AuthlibBaseError as OAuthError
 from authlib.integrations.httpx_client import AsyncOAuth2Client
 
 from nat.authentication.interfaces import FlowHandlerBase
@@ -61,14 +62,50 @@ class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
 
         raise NotImplementedError(f"Authentication method '{method}' is not supported by the websocket frontend.")
 
-    def create_oauth_client(self, config: OAuth2AuthCodeFlowProviderConfig):
-        return AsyncOAuth2Client(client_id=config.client_id,
-                                 client_secret=config.client_secret,
-                                 redirect_uri=config.redirect_uri,
-                                 scope=" ".join(config.scopes) if config.scopes else None,
-                                 token_endpoint=config.token_url,
-                                 code_challenge_method='S256' if config.use_pkce else None,
-                                 token_endpoint_auth_method=config.token_endpoint_auth_method)
+    def create_oauth_client(self, config: OAuth2AuthCodeFlowProviderConfig) -> AsyncOAuth2Client:
+        try:
+            return AsyncOAuth2Client(client_id=config.client_id,
+                                     client_secret=config.client_secret,
+                                     redirect_uri=config.redirect_uri,
+                                     scope=" ".join(config.scopes) if config.scopes else None,
+                                     token_endpoint=config.token_url,
+                                     code_challenge_method='S256' if config.use_pkce else None,
+                                     token_endpoint_auth_method=config.token_endpoint_auth_method)
+        except (OAuthError, ValueError, TypeError) as e:
+            raise RuntimeError(f"Invalid OAuth2 configuration: {e}") from e
+        except Exception as e:
+            raise RuntimeError(f"Failed to create OAuth2 client: {e}") from e
+
+    def _create_authorization_url(self,
+                                  client: AsyncOAuth2Client,
+                                  config: OAuth2AuthCodeFlowProviderConfig,
+                                  state: str,
+                                  verifier: str = None,
+                                  challenge: str = None) -> str:
+        """
+        Create OAuth authorization URL with proper error handling.
+
+        Args:
+            client: The OAuth2 client instance
+            config: OAuth2 configuration
+            state: OAuth state parameter
+            verifier: PKCE verifier (if using PKCE)
+            challenge: PKCE challenge (if using PKCE)
+
+        Returns:
+            The authorization URL
+        """
+        try:
+            authorization_url, _ = client.create_authorization_url(
+                config.authorization_url,
+                state=state,
+                code_verifier=verifier if config.use_pkce else None,
+                code_challenge=challenge if config.use_pkce else None,
+                **(config.authorization_kwargs or {})
+            )
+            return authorization_url
+        except (OAuthError, ValueError, TypeError) as e:
+            raise RuntimeError(f"Error creating OAuth authorization URL: {e}") from e
 
     async def _handle_oauth2_auth_code_flow(self, config: OAuth2AuthCodeFlowProviderConfig) -> AuthenticatedContext:
 
@@ -82,21 +119,19 @@ class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
             flow_state.verifier = verifier
             flow_state.challenge = challenge
 
-        authorization_url, _ = flow_state.client.create_authorization_url(
-            config.authorization_url,
-            state=state,
-            code_verifier=flow_state.verifier if config.use_pkce else None,
-            code_challenge=flow_state.challenge if config.use_pkce else None,
-            **(config.authorization_kwargs or {})
-        )
+        authorization_url = self._create_authorization_url(client=flow_state.client,
+                                                           config=config,
+                                                           state=state,
+                                                           verifier=flow_state.verifier,
+                                                           challenge=flow_state.challenge)
 
         await self._add_flow_cb(state, flow_state)
         await self._web_socket_message_handler.create_websocket_message(_HumanPromptOAuthConsent(text=authorization_url)
                                                                         )
         try:
             token = await asyncio.wait_for(flow_state.future, timeout=300)
-        except asyncio.TimeoutError:
-            raise RuntimeError("Authentication flow timed out after 5 minutes.")
+        except TimeoutError as exc:
+            raise RuntimeError("Authentication flow timed out after 5 minutes.") from exc
         finally:
 
             await self._remove_flow_cb(state)