Updating sensitive data suppression logic.

Signed-off-by: Eric Evans <[email protected]>

Author: ericevans-nv

src/aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -20,6 +20,8 @@
 import typing
 from abc import ABC
 from abc import abstractmethod
+from collections.abc import Awaitable
+from collections.abc import Callable
 from contextlib import asynccontextmanager
 from functools import partial
 from pathlib import Path
@@ -123,18 +125,8 @@ async def lifespan(starting_app: FastAPI):
         self.set_cors_config(aiq_app)
 
         @aiq_app.middleware("http")
-        async def _suppress_authentication_logs(request: Request, call_next):
-            """
-            Intercepts authentication request and supreses logs that contain sensitive data.
-            """
-            default_log_level = logging.getLogger("uvicorn.access").level
-            if request.url.path in (AuthenticationEndpoint.REDIRECT_URI.value
-                                    or AuthenticationEndpoint.PROMPT_REDIRECT_URI.value):
-                logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
-            response = await call_next(request)
-            logging.getLogger("uvicorn.access").setLevel(default_log_level)
-
-            return response
+        async def authentication_log_filter(request: Request, call_next: Callable[[Request], Awaitable[Response]]):
+            return await self._suppress_authentication_logs(request, call_next)
 
         return aiq_app
 
@@ -170,6 +162,25 @@ def set_cors_config(self, aiq_app: FastAPI) -> None:
             **cors_kwargs,
         )
 
+    async def _suppress_authentication_logs(self, request: Request,
+                                            call_next: Callable[[Request], Awaitable[Response]]) -> Response:
+        """
+        Intercepts authentication request and supreses logs that contain sensitive data.
+        """
+        from aiq.utils.log_utils import LogFilter
+
+        logs_to_suppress: list[str] = [
+            AuthenticationEndpoint.REDIRECT_URI.value, AuthenticationEndpoint.PROMPT_REDIRECT_URI.value
+        ]
+
+        logging.getLogger("uvicorn.access").addFilter(LogFilter(logs_to_suppress))
+        try:
+            response = await call_next(request)
+        finally:
+            logging.getLogger("uvicorn.access").removeFilter(LogFilter(logs_to_suppress))
+
+        return response
+
     @abstractmethod
     async def configure(self, app: FastAPI, builder: WorkflowBuilder):
         pass

src/aiq/runtime/session.py

@@ -13,16 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import httpx
 import asyncio
 import contextvars
 import typing
-
 from collections.abc import Awaitable
 from collections.abc import Callable
 from contextlib import asynccontextmanager
 from contextlib import nullcontext
 
+import httpx
 from fastapi import Request
 
 from aiq.builder.context import AIQContext

src/aiq/utils/log_utils.py

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+
+class LogFilter(logging.Filter):
+    """
+    This class is used to filter log records based on a defined set of criteria.
+    """
+
+    def __init__(self, filter_criteria: list[str]):
+        self._filter_criteria = filter_criteria
+        super().__init__()
+
+    def filter(self, record: logging.LogRecord):
+        """
+        Evaluates whether a log record should be emitted based on the message content.
+
+        Returns:
+            False if the message content contains any of the filter criteria, True otherwise.
+        """
+        if any(match in record.getMessage() for match in self._filter_criteria):
+            return False
+        return True

tests/aiq/authentication/test_authentication.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import webbrowser
 from datetime import datetime
 from datetime import timedelta