Log custom headers (#5966)

* add custom headers

* fix: send on startup and defaults

* pylint

* unittests

5 files changed, 57 insertions, 8 deletions

diff --git a/openbb_platform/core/openbb_core/api/router/commands.py b/openbb_platform/core/openbb_core/api/router/commands.py
index 99a1f28f0b7..bd627630542 100644
--- a/openbb_platform/core/openbb_core/api/router/commands.py
+++ b/openbb_platform/core/openbb_core/api/router/commands.py
@@ -1,9 +1,9 @@
 import inspect
 from functools import partial, wraps
 from inspect import Parameter, Signature, signature
-from typing import Any, Callable, Dict, Tuple, TypeVar
+from typing import Any, Callable, Dict, Optional, Tuple, TypeVar
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, Header
 from fastapi.routing import APIRoute
 from openbb_core.app.charting_service import ChartingService
 from openbb_core.app.command_runner import CommandRunner
@@ -70,6 +70,17 @@ def build_new_signature(path: str, func: Callable) -> Signature:
             )
         )
 
+    if custom_headers := SystemService().system_settings.api_settings.custom_headers:
+        for name, default in custom_headers.items():
+            new_parameter_list.append(
+                Parameter(
+                    name.replace("-", "_"),
+                    kind=Parameter.POSITIONAL_OR_KEYWORD,
+                    default=default,
+                    annotation=Annotated[Optional[str], Header()],
+                )
+            )
+
     if Env().API_AUTH:
         new_parameter_list.append(
             Parameter(
diff --git a/openbb_platform/core/openbb_core/app/command_runner.py b/openbb_platform/core/openbb_core/app/command_runner.py
index ae95313b5aa..8e41623f49a 100644
--- a/openbb_platform/core/openbb_core/app/command_runner.py
+++ b/openbb_platform/core/openbb_core/app/command_runner.py
@@ -305,6 +305,15 @@ class StaticCommandRunner:
         # If we're on the api we need to remove "chart" here because the parameter is added on
         # commands.py and the function signature does not expect "chart"
         kwargs.pop("chart", None)
+        # We also pop custom headers
+
+        model_headers = (
+            SystemService().system_settings.api_settings.custom_headers or {}
+        )
+        custom_headers = {
+            name: kwargs.pop(name.replace("-", "_"), default)
+            for name, default in model_headers.items() or {}
+        } or None
 
         try:
             obbject = await cls._command(
@@ -334,6 +343,7 @@ class StaticCommandRunner:
                 func=func,
                 kwargs=kwargs,
                 exec_info=exc_info(),
+                custom_headers=custom_headers,
             )
 
         return obbject
@@ -414,6 +424,7 @@ class CommandRunner:
     def user_settings(self, user_settings: UserSettings) -> None:
         self._user_settings = user_settings
 
+    # pylint: disable=W1113
     async def run(
         self,
         route: str,
@@ -434,6 +445,7 @@ class CommandRunner:
 
         return await StaticCommandRunner.run(execution_context, *args, **kwargs)
 
+    # pylint: disable=W1113
     def sync_run(
         self,
         route: str,
diff --git a/openbb_platform/core/openbb_core/app/logs/logging_service.py b/openbb_platform/core/openbb_core/app/logs/logging_service.py
index 854c7f55aac..569411081e1 100644
--- a/openbb_platform/core/openbb_core/app/logs/logging_service.py
+++ b/openbb_platform/core/openbb_core/app/logs/logging_service.py
@@ -135,7 +135,11 @@ class LoggingService(metaclass=SingletonMeta):
 
         return handlers_manager
 
-    def _log_startup(self, route: Optional[str] = None) -> None:
+    def _log_startup(
+        self,
+        route: Optional[str] = None,
+        custom_headers: Optional[Dict[str, Any]] = None,
+    ) -> None:
         """
         Log startup information.
         """
@@ -165,6 +169,7 @@ class LoggingService(metaclass=SingletonMeta):
                         else {}
                     ),
                     "SYSTEM": self._system_settings,
+                    "custom_headers": custom_headers,
                 },
                 default=to_jsonable_python,
             ),
@@ -180,6 +185,7 @@ class LoggingService(metaclass=SingletonMeta):
         exec_info: Optional[
             Tuple[Type[BaseException], BaseException, Optional[TracebackType]]
         ] = None,
+        custom_headers: Optional[Dict[str, Any]] = None,
     ):
         """
         Log command output and relevant information.
@@ -208,7 +214,7 @@ class LoggingService(metaclass=SingletonMeta):
         self._handlers_manager.update_handlers(self._logging_settings)
 
         if "login" in route:
-            self._log_startup(route)
+            self._log_startup(route, custom_headers)
         else:
             logger = logging.getLogger(__name__)
 
@@ -230,6 +236,7 @@ class LoggingService(metaclass=SingletonMeta):
                     "route": route,
                     "input": kwargs,
                     "error": str(openbb_error.original) if openbb_error else None,
+                    "custom_headers": custom_headers,
                 },
                 default=to_jsonable_python,
             )
diff --git a/openbb_platform/core/openbb_core/app/model/fast_api_settings.py b/openbb_platform/core/openbb_core/app/model/fast_api_settings.py
index 565118dfc57..2795ef310ad 100644
--- a/openbb_platform/core/openbb_core/app/model/fast_api_settings.py
+++ b/openbb_platform/core/openbb_core/app/model/fast_api_settings.py
@@ -1,5 +1,5 @@
 """FastAPI configuration settings model."""
-from typing import List
+from typing import Dict, List, Optional
 
 from pydantic import BaseModel, ConfigDict, Field, computed_field
 
@@ -41,6 +41,9 @@ class FastAPISettings(BaseModel):
     )
     servers: List[Servers] = Field(default_factory=lambda: [Servers()])
     cors: Cors = Field(default_factory=Cors)
+    custom_headers: Optional[Dict[str, str]] = Field(
+        default=None, description="Custom headers and respective default value."
+    )
 
     @computed_field  # type: ignore[misc]
     @property
diff --git a/openbb_platform/core/tests/app/logs/test_logging_service.py b/openbb_platform/core/tests/app/logs/test_logging_service.py
index 67d90d8ce71..bd7460bc3c2 100644
--- a/openbb_platform/core/tests/app/logs/test_logging_service.py
+++ b/openbb_platform/core/tests/app/logs/test_logging_service.py
@@ -86,13 +86,16 @@ def test_log_startup(logging_service):
         )
         logging_service._system_settings = "your_system_settings"
 
-        logging_service._log_startup(route="test_route")
+        logging_service._log_startup(
+            route="test_route", custom_headers={"X-OpenBB-Test": "test"}
+        )
 
         expected_log_data = {
             "route": "test_route",
             "PREFERENCES": "your_preferences",
             "KEYS": {"username": "defined", "password": "defined"},
             "SYSTEM": "your_system_settings",
+            "custom_headers": {"X-OpenBB-Test": "test"},
         }
         mock_info.assert_called_once_with(
             "STARTUP: %s ",
@@ -102,7 +105,7 @@ def test_log_startup(logging_service):
 
 
 @pytest.mark.parametrize(
-    "user_settings, system_settings, route, func, kwargs, exec_info",
+    "user_settings, system_settings, route, func, kwargs, exec_info, custom_headers",
     [
         (
             "mock_settings",
@@ -111,6 +114,7 @@ def test_log_startup(logging_service):
             "mock_func",
             {},
             None,
+            None,
         ),
         (
             "mock_settings",
@@ -119,6 +123,7 @@ def test_log_startup(logging_service):
             "mock_func",
             {},
             (OpenBBError, OpenBBError("mock_error")),
+            {"X-OpenBB-Test": "test"},
         ),
         (
             "mock_settings",
@@ -127,11 +132,19 @@ def test_log_startup(logging_service):
             "mock_func",
             {},
             None,
+            {"X-OpenBB-Test1": "test1", "X-OpenBB-Test2": "test2"},
         ),
     ],
 )
 def test_log(
-    logging_service, user_settings, system_settings, route, func, kwargs, exec_info
+    logging_service,
+    user_settings,
+    system_settings,
+    route,
+    func,
+    kwargs,
+    exec_info,
+    custom_headers,
 ):
     with patch(
         "openbb_core.app.logs.logging_service.LoggingSettings",
@@ -148,6 +161,7 @@ def test_log(
                     func=func,
                     kwargs=kwargs,
                     exec_info=exec_info,
+                    custom_headers=custom_headers,
                 )
                 assert mock_log_startup.assert_called_once
 
@@ -165,6 +179,7 @@ def test_log(
                 func=mock_callable,
                 kwargs=kwargs,
                 exec_info=exec_info,
+                custom_headers=custom_headers,
             )
 
             message_label = "ERROR" if exec_info else "CMD"
@@ -173,6 +188,7 @@ def test_log(
                     "route": route,
                     "input": kwargs,
                     "error": str(exec_info[1]) if exec_info else None,
+                    "custom_headers": custom_headers,
                 }
             )
             log_message = f"{message_label}: {log_message}"