fix scope of codegate logger

Author: yrobla

src/codegate/cli.py

@@ -7,6 +7,7 @@
 from typing import Dict, Optional
 
 import click
+import structlog
 from uvicorn.config import Config as UvicornConfig
 from uvicorn.server import Server
 
@@ -19,7 +20,6 @@
 from codegate.providers.copilot.provider import CopilotProvider
 from codegate.server import init_app
 from codegate.storage.utils import restore_storage_backup
-from codegate.logger.logger import OriginLogger
 
 
 class UvicornServer:
@@ -33,9 +33,7 @@ def __init__(self, config: UvicornConfig, server: Server):
         self._startup_complete = asyncio.Event()
         self._shutdown_event = asyncio.Event()
         self._should_exit = False
-
-        logger_obj = OriginLogger("generic_server")
-        self.logger = logger_obj.logger
+        self.logger = structlog.get_logger("codegate").bind(origin="generic_server")
 
     async def serve(self) -> None:
         """Start the uvicorn server and handle shutdown gracefully."""
@@ -86,9 +84,7 @@ 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."""
-    cli_logger_obj = OriginLogger("cli")
-    cli_logger = cli_logger_obj.logger
-
+    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")
@@ -300,8 +296,7 @@ def serve(
 
         # Set up logging first
         setup_logging(cfg.log_level, cfg.log_format)
-        cli_logger_obj = OriginLogger("cli")
-        logger = cli_logger_obj.logger
+        logger = structlog.get_logger("codegate").bind(origin="cli")
 
         init_db_sync(cfg.db_path)
 
@@ -332,9 +327,7 @@ def serve(
         click.echo(f"Configuration error: {e}", err=True)
         sys.exit(1)
     except Exception as e:
-        cli_logger_obj = OriginLogger("cli")
-        logger = cli_logger_obj.logger
-
+        logger = structlog.get_logger("codegate").bind(origin="cli")
         logger.exception("Unexpected error occurred")
         click.echo(f"Error: {e}", err=True)
         sys.exit(1)
@@ -343,9 +336,7 @@ def serve(
 async def run_servers(cfg: Config, app) -> None:
     """Run the codegate server."""
     try:
-        cli_logger_obj = OriginLogger("cli")
-        logger = cli_logger_obj.logger
-
+        logger = structlog.get_logger("codegate").bind(origin="cli")
         logger.info(
             "Starting server",
             extra={

src/codegate/logger/logger.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
 
-from src.codegate.logger.logger import OriginLogger
 from litellm.types.utils import Delta, ModelResponse, StreamingChoices
 
 from codegate.ca.codegate_ca import CertificateAuthority
@@ -22,8 +23,8 @@
 )
 from codegate.providers.copilot.streaming import SSEProcessor
 
-logger_obj = OriginLogger("copilot_proxy")
-logger = logger_obj.logger
+setup_logging()
+logger = structlog.get_logger("codegate").bind(origin="copilot_proxy")
 
 # Constants
 MAX_BUFFER_SIZE = 10 * 1024 * 1024  # 10MB

src/codegate/providers/copilot/provider.py

@@ -195,13 +195,8 @@ def test_serve_default_options(cli_runner):
     # Use patches for run_servers and logging setup
     with (
         patch("src.codegate.cli.run_servers") as mock_run,
-        patch("src.codegate.cli.OriginLogger") as mock_origin_logger,
         patch("src.codegate.cli.setup_logging") as mock_setup_logging,
     ):
-
-        logger_instance = MagicMock()
-        mock_origin_logger.return_value = logger_instance
-
         # Invoke the CLI command
         result = cli_runner.invoke(cli, ["serve"])
 
@@ -211,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_origin_logger.assert_called_with("cli")
-
         # Validate run_servers was called once
         mock_run.assert_called_once()
 
@@ -222,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.OriginLogger") as mock_origin_logger,
         patch("src.codegate.cli.setup_logging") as mock_setup_logging,
     ):
-
-        logger_instance = MagicMock()
-        mock_origin_logger.return_value = logger_instance
-
         # Invoke the CLI command with custom options
         result = cli_runner.invoke(
             cli,
@@ -261,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_origin_logger.assert_called_with("cli")
-
         # Validate run_servers was called once
         mock_run.assert_called_once()
         # Retrieve the actual Config object passed to run_servers
@@ -322,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.OriginLogger") as mock_origin_logger,
         patch("src.codegate.cli.setup_logging") as mock_setup_logging,
     ):
-
-        logger_instance = MagicMock()
-        mock_origin_logger.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_origin_logger.assert_called_with("cli")
 
         # Validate that run_servers was called with the expected configuration
         mock_run.assert_called_once()
@@ -370,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.OriginLogger") as mock_origin_logger,
         patch("src.codegate.cli.setup_logging") as mock_setup_logging,
     ):
-        # Set up mock logger
-        logger_instance = MagicMock()
-        mock_origin_logger.return_value = logger_instance
-
         # Execute CLI command with specific options overriding environment and config file settings
         result = cli_runner.invoke(
             cli,
@@ -410,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_origin_logger.assert_called_with("cli")
 
         # 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
@@ -438,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.OriginLogger") as mock_origin_logger,
         patch("src.codegate.cli.setup_logging") as mock_setup_logging,
     ):
-        # Set up mock logger
-        logger_instance = MagicMock()
-        mock_origin_logger.return_value = logger_instance
-
         # Execute CLI command with certificate options
         result = cli_runner.invoke(
             cli,
@@ -468,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_origin_logger.assert_called_with("cli")
 
         # 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