Merge pull request #406 from stacklok/issue-301

feat: add custom logger that includes the origin

Author: yrobla

src/codegate/cli.py

@@ -33,7 +33,7 @@ def __init__(self, config: UvicornConfig, server: Server):
         self._startup_complete = asyncio.Event()
         self._shutdown_event = asyncio.Event()
         self._should_exit = False
-        self.logger = structlog.get_logger("codegate")
+        self.logger = structlog.get_logger("codegate").bind(origin="generic_server")
 
     async def serve(self) -> None:
         """Start the uvicorn server and handle shutdown gracefully."""
@@ -84,8 +84,8 @@ async def cleanup(self) -> None:
 
 def validate_port(ctx: click.Context, param: click.Parameter, value: int) -> int:
     """Validate the port number is in valid range."""
-    logger = structlog.get_logger("codegate")
-    logger.debug(f"Validating port number: {value}")
+    cli_logger = structlog.get_logger("codegate").bind(origin="cli")
+    cli_logger.debug(f"Validating port number: {value}")
     if value is not None and not (1 <= value <= 65535):
         raise click.BadParameter("Port must be between 1 and 65535")
     return value
@@ -296,7 +296,7 @@ def serve(
 
         # Set up logging first
         setup_logging(cfg.log_level, cfg.log_format)
-        logger = structlog.get_logger("codegate")
+        logger = structlog.get_logger("codegate").bind(origin="cli")
 
         init_db_sync(cfg.db_path)
 
@@ -327,7 +327,7 @@ def serve(
         click.echo(f"Configuration error: {e}", err=True)
         sys.exit(1)
     except Exception as e:
-        logger = structlog.get_logger("codegate")
+        logger = structlog.get_logger("codegate").bind(origin="cli")
         logger.exception("Unexpected error occurred")
         click.echo(f"Error: {e}", err=True)
         sys.exit(1)
@@ -336,7 +336,7 @@ def serve(
 async def run_servers(cfg: Config, app) -> None:
     """Run the codegate server."""
     try:
-        logger = structlog.get_logger("codegate")
+        logger = structlog.get_logger("codegate").bind(origin="cli")
         logger.info(
             "Starting server",
             extra={

src/codegate/codegate_logging.py

@@ -48,6 +48,15 @@ def _missing_(cls, value: str) -> Optional["LogFormat"]:
             )
 
 
+def add_origin(logger, log_method, event_dict):
+    # Add 'origin' if it's bound to the logger but not explicitly in the event dict
+    if 'origin' not in event_dict and hasattr(logger, '_context'):
+        origin = logger._context.get('origin')
+        if origin:
+            event_dict['origin'] = origin
+    return event_dict
+
+
 def setup_logging(
     log_level: Optional[LogLevel] = None, log_format: Optional[LogFormat] = None
 ) -> logging.Logger:
@@ -74,6 +83,7 @@ def setup_logging(
     shared_processors = [
         structlog.processors.add_log_level,
         structlog.processors.TimeStamper(fmt="%Y-%m-%dT%H:%M:%S.%03dZ", utc=True),
+        add_origin,
         structlog.processors.CallsiteParameterAdder(
             [
                 structlog.processors.CallsiteParameter.MODULE,

src/codegate/providers/copilot/provider.py

@@ -1,11 +1,12 @@
 import asyncio
 import re
 import ssl
+from src.codegate.codegate_logging import setup_logging
+import structlog
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Tuple, Union
 from urllib.parse import unquote, urljoin, urlparse
 
-import structlog
 from litellm.types.utils import Delta, ModelResponse, StreamingChoices
 
 from codegate.ca.codegate_ca import CertificateAuthority
@@ -22,7 +23,8 @@
 )
 from codegate.providers.copilot.streaming import SSEProcessor
 
-logger = structlog.get_logger("codegate")
+setup_logging()
+logger = structlog.get_logger("codegate").bind(origin="copilot_proxy")
 
 # Constants
 MAX_BUFFER_SIZE = 10 * 1024 * 1024  # 10MB
@@ -637,7 +639,7 @@ async def get_target_url(path: str) -> Optional[str]:
         # Check for prefix match
         for route in VALIDATED_ROUTES:
             # For prefix matches, keep the rest of the path
-            remaining_path = path[len(route.path) :]
+            remaining_path = path[len(route.path):]
             logger.debug(f"Remaining path: {remaining_path}")
             # Make sure we don't end up with double slashes
             if remaining_path and remaining_path.startswith("/"):
@@ -791,7 +793,7 @@ def data_received(self, data: bytes) -> None:
                     self._proxy_transport_write(headers)
                     logger.debug(f"Headers sent: {headers}")
 
-                    data = data[header_end + 4 :]
+                    data = data[header_end + 4:]
 
             self._process_chunk(data)
 

tests/test_cli.py

@@ -190,28 +190,13 @@ def cli_runner():
     return CliRunner()
 
 
-@pytest.fixture
-def mock_logging(mocker):
-    return mocker.patch("your_cli_module.structlog.get_logger")
-
-
-@pytest.fixture
-def mock_setup_logging(mocker):
-    return mocker.patch("your_cli_module.setup_logging")
-
-
 def test_serve_default_options(cli_runner):
     """Test serve command with default options."""
     # Use patches for run_servers and logging setup
     with (
         patch("src.codegate.cli.run_servers") as mock_run,
-        patch("src.codegate.cli.structlog.get_logger") as mock_logging,
         patch("src.codegate.cli.setup_logging") as mock_setup_logging,
     ):
-
-        logger_instance = MagicMock()
-        mock_logging.return_value = logger_instance
-
         # Invoke the CLI command
         result = cli_runner.invoke(cli, ["serve"])
 
@@ -221,9 +206,6 @@ def test_serve_default_options(cli_runner):
         # Check if the logging setup was called with expected defaults
         mock_setup_logging.assert_called_once_with(LogLevel.INFO, LogFormat.JSON)
 
-        # Check if logging was done correctly
-        mock_logging.assert_called_with("codegate")
-
         # Validate run_servers was called once
         mock_run.assert_called_once()
 
@@ -232,13 +214,8 @@ def test_serve_custom_options(cli_runner):
     """Test serve command with custom options."""
     with (
         patch("src.codegate.cli.run_servers") as mock_run,
-        patch("src.codegate.cli.structlog.get_logger") as mock_logging,
         patch("src.codegate.cli.setup_logging") as mock_setup_logging,
     ):
-
-        logger_instance = MagicMock()
-        mock_logging.return_value = logger_instance
-
         # Invoke the CLI command with custom options
         result = cli_runner.invoke(
             cli,
@@ -271,9 +248,6 @@ def test_serve_custom_options(cli_runner):
         # Assert logging setup was called with the provided log level and format
         mock_setup_logging.assert_called_once_with(LogLevel.DEBUG, LogFormat.TEXT)
 
-        # Assert logger got called with the expected module name
-        mock_logging.assert_called_with("codegate")
-
         # Validate run_servers was called once
         mock_run.assert_called_once()
         # Retrieve the actual Config object passed to run_servers
@@ -332,20 +306,14 @@ def test_serve_with_config_file(cli_runner, temp_config_file):
     """Test serve command with config file."""
     with (
         patch("src.codegate.cli.run_servers") as mock_run,
-        patch("src.codegate.cli.structlog.get_logger") as mock_logging,
         patch("src.codegate.cli.setup_logging") as mock_setup_logging,
     ):
-
-        logger_instance = MagicMock()
-        mock_logging.return_value = logger_instance
-
         # Invoke the CLI command with the configuration file
         result = cli_runner.invoke(cli, ["serve", "--config", str(temp_config_file)])
 
         # Assertions to ensure the CLI ran successfully
         assert result.exit_code == 0
         mock_setup_logging.assert_called_once_with(LogLevel.DEBUG, LogFormat.JSON)
-        mock_logging.assert_called_with("codegate")
 
         # Validate that run_servers was called with the expected configuration
         mock_run.assert_called_once()
@@ -380,13 +348,8 @@ def test_serve_priority_resolution(cli_runner: CliRunner, temp_config_file: Path
     with (
         patch.dict(os.environ, {"LOG_LEVEL": "INFO", "PORT": "9999"}, clear=True),
         patch("src.codegate.cli.run_servers") as mock_run,
-        patch("src.codegate.cli.structlog.get_logger") as mock_logging,
         patch("src.codegate.cli.setup_logging") as mock_setup_logging,
     ):
-        # Set up mock logger
-        logger_instance = MagicMock()
-        mock_logging.return_value = logger_instance
-
         # Execute CLI command with specific options overriding environment and config file settings
         result = cli_runner.invoke(
             cli,
@@ -420,7 +383,6 @@ def test_serve_priority_resolution(cli_runner: CliRunner, temp_config_file: Path
 
         # Ensure logging setup was called with the highest priority settings (CLI arguments)
         mock_setup_logging.assert_called_once_with("ERROR", "TEXT")
-        mock_logging.assert_called_with("codegate")
 
         # Verify that the run_servers was called with the overridden settings
         config_arg = mock_run.call_args[0][0]  # Assuming Config is the first positional arg
@@ -448,13 +410,8 @@ def test_serve_certificate_options(cli_runner: CliRunner) -> None:
     """Test serve command with certificate options."""
     with (
         patch("src.codegate.cli.run_servers") as mock_run,
-        patch("src.codegate.cli.structlog.get_logger") as mock_logging,
         patch("src.codegate.cli.setup_logging") as mock_setup_logging,
     ):
-        # Set up mock logger
-        logger_instance = MagicMock()
-        mock_logging.return_value = logger_instance
-
         # Execute CLI command with certificate options
         result = cli_runner.invoke(
             cli,
@@ -478,7 +435,6 @@ def test_serve_certificate_options(cli_runner: CliRunner) -> None:
 
         # Ensure logging setup was called with expected arguments
         mock_setup_logging.assert_called_once_with("INFO", "JSON")
-        mock_logging.assert_called_with("codegate")
 
         # Verify that run_servers was called with the provided certificate options
         config_arg = mock_run.call_args[0][0]  # Assuming Config is the first positional arg