diff --git a/src/codegate/cli.py b/src/codegate/cli.py index 01f10945..e1e30634 100644 --- a/src/codegate/cli.py +++ b/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={ diff --git a/src/codegate/codegate_logging.py b/src/codegate/codegate_logging.py index d99f55ae..5a15dec1 100644 --- a/src/codegate/codegate_logging.py +++ b/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, diff --git a/src/codegate/providers/copilot/provider.py b/src/codegate/providers/copilot/provider.py index 3da8a4b5..311c4a93 100644 --- a/src/codegate/providers/copilot/provider.py +++ b/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) diff --git a/tests/test_cli.py b/tests/test_cli.py index caa3a215..fa85b3ea 100644 --- a/tests/test_cli.py +++ b/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