Improve logging options and exception handling

.cursor/rules/run-tests.mdc

@@ -0,0 +1,16 @@
+---
+description:
+globs:
+alwaysApply: true
+---
+# Running tests
+
+After finishing any big code changes, and before yielding to the user, run with your terminal tool the following comand:
+
+```
+source .venv/bin/activate && pytest -k ""
+```
+
+You don't need to preface it with /bin/bash -c, nor add | cat at the end. Just run it as is.
+
+If you want to run linting and tests, perform 2 parallel function calls

.env.example

@@ -4,6 +4,8 @@ FLASK_DEBUG=1
 FLASK_ENV=development
 GUNICORN_WORKERS=1
 LOG_LEVEL=debug
+LOG_CONTEXT=on
+LOG_COMPLETION=on
 RECORD_TRAFFIC=off
 
 # Arbitrary API key to protect your service.

README.md

@@ -5,7 +5,6 @@
 [![Pytest](https://img.shields.io/badge/Pytest-fff?logo=pytest&logoColor=000)](#)
 [![Docker](https://img.shields.io/badge/Docker-2496ED?logo=docker&logoColor=fff)](#)
 ![GitHub License](https://img.shields.io/github/license/gabrii/Cursor-Azure-GPT-5)
-![GitHub Actions Workflow Status](https://img.shields.io/github/actions/workflow/status/gabrii/Cursor-Azure-GPT-5/lint.yml?label=lint)
 ![GitHub Actions Workflow Status](https://img.shields.io/github/actions/workflow/status/gabrii/Cursor-Azure-GPT-5/lint.yml?label=test)
 ![Codecov](https://img.shields.io/codecov/c/github/gabrii/Cursor-Azure-GPT-5)
 
@@ -26,8 +25,6 @@ This project originates from Cursor's lack of support for Azure models that are
 
 Upcoming features:
 - Multimodal: Will be implemented as soon as better testing is in place and there is demand (PRs welcome).
-- Multiple models simultaneously: Even Cursor's Azure configuration only supports a single deployment at a time. It would be fairly easy to implement support for multiple models in this service, covering even more needs.
-- Full test coverage: See the [Testing](#testing) section for an explanation of the current low coverage.
 
 Feel free to create or vote on any [project issues](https://github.com/gabrii/Cursor-Azure-GPT-5/issues), and star the project to show your support.
 
@@ -59,6 +56,8 @@ Alternatively, you can pass them through the environment where you run the appli
 | `AZURE_TRUNCATION`  | Truncation strategy for long inputs.                                   | `auto`               |
 | `FLASK_ENV`         | Flask environment. Use `development` for dev or `production` for prod. | `production`         |
 | `RECORD_TRAFFIC`    | Toggle writing request/response traffic to `recordings/`               | `off`                |
+| `LOG_CONTEXT`       | Enable rich pretty-printing of request context to console.             | `on`                 |
+| `LOG_COMPLETION`    | Enable logging of completion responses (not yet implemented).          | `on`                 |
 
 </details>
 

app/azure/response_adapter.py

@@ -11,9 +11,12 @@ import time
 from string import ascii_letters, digits
 from typing import Any, Dict, Iterable, Optional
 
-from flask import Response, stream_with_context
+from flask import Response, current_app, stream_with_context
+from rich.live import Live
 
+from ..common.logging import console, create_message_panel
 from ..common.sse import chunks_to_sse, sse_to_events
+from ..exceptions import ClientClosedConnection
 
 # Centralized events that should end a <think> block before handling
 THINKING_STOP_EVENTS = {"response.output_text.delta", "response.output_item.added"}
@@ -158,10 +161,27 @@ class ResponseAdapter:
             self._tool_calls = 0
 
             def gen_dicts() -> Iterable[Dict[str, Any]]:
-                try:
+                # Initialize message object for completion logging
+                completion_msg: Dict[str, Any] = {
+                    "role": "assistant",
+                    "content": "",
+                    "tool_calls": [],
+                }
+
+                events = 0
+                with Live(
+                    None,
+                    console=console,
+                    refresh_per_second=2,
+                ) as live:  # update 4 times a second to feel fluid
                     for ev in sse_to_events(
                         upstream_resp.iter_content(chunk_size=8192)
                     ):
+                        if current_app.config["LOG_COMPLETION"]:
+                            if events > 1:
+                                live.update(create_message_panel(completion_msg, 1, 1))
+                            events += 1
+
                         handler_name = "_" + (ev.event or "").replace(
                             "response.", ""
                         ).replace(".", "__")
@@ -176,19 +196,61 @@ class ResponseAdapter:
                             )
                             self._thinking = False
 
+                            if current_app.config["LOG_COMPLETION"]:
+                                completion_msg["content"] += "</think>\n\n"
+
                         res = handler(ev.json)
                         if res is not None:
                             yield res
-                finally:
+
-                    # Emit finish reason at the end of stream
+                            if current_app.config["LOG_COMPLETION"]:
+                                delta = res.get("choices", [{}])[0].get("delta", {})
+                                content = delta.get("content")
+
+                                if content is not None:
+                                    # Append content to the message
+                                    completion_msg["content"] += content
+                                else:
+                                    # Handle tool calls
+                                    tool_calls_delta = delta.get("tool_calls", [])
+                                    for tool_call_delta in tool_calls_delta:
+                                        function = tool_call_delta.get("function", {})
+                                        name = function.get("name")
+                                        arguments = function.get("arguments", "")
+
+                                        if name:
+                                            # New tool call - add to the list
+                                            completion_msg["tool_calls"].append(
+                                                {
+                                                    "id": tool_call_delta.get("id", ""),
+                                                    "type": "function",
+                                                    "function": {
+                                                        "name": name,
+                                                        "arguments": arguments,
+                                                    },
+                                                }
+                                            )
+                                        else:
+                                            # Append arguments to the last tool call
+                                            completion_msg["tool_calls"][-1][
+                                                "function"
+                                            ]["arguments"] += arguments
+
                     if self._tool_calls > 0:
                         yield self._build_completion_chunk(finish_reason="tool_calls")
                     else:
                         yield self._build_completion_chunk(finish_reason="stop")
+                    if current_app.config["LOG_COMPLETION"]:
+                        live.update(create_message_panel(completion_msg, 1, 1))
 
-            # Wrap as SSE with [DONE]
             try:
                 yield from chunks_to_sse(gen_dicts())
+            except GeneratorExit:
+                # Downstream client closed the connection mid-stream
+                # Translate to a clearer exception; upstream will be closed in finally
+                raise ClientClosedConnection(
+                    "Client closed connection during streaming response"
+                ) from None
             finally:
                 upstream_resp.close()
 

app/blueprint.py

@@ -4,7 +4,7 @@ This module defines the application blueprint, configures logging, and
 forwards incoming HTTP requests to the configured backend implementation.
 """
 
-from flask import Blueprint, jsonify, request
+from flask import Blueprint, current_app, jsonify, request
 
 from .auth import require_auth
 from .azure.adapter import AzureAdapter
@@ -46,7 +46,8 @@ def catch_all(path: str):
     implementation, returning the backend's response. If forwarding fails,
     returns a 502 JSON error payload.
     """
-    log_request(request)
+    if current_app.config.get("LOG_CONTEXT"):
+        log_request(request)
     init_last_recording()
     increment_last_recording()
     record_payload(request.json, "downstream_request")

app/common/logging.py

@@ -105,6 +105,82 @@ def escape_tags(text: str) -> str:
     ).replace(">`\n\n\n`<", ">`\n\n`<")
 
 
+def create_message_panel(msg: Dict[str, Any], idx: int, total: int) -> Panel:
+    """Create a Rich Panel for displaying a message.
+
+    Args:
+        msg: Message object with 'role', 'content', optional 'name', 'tool_call_id', 'tool_calls'
+        idx: Current message index (1-based)
+        total: Total number of messages
+
+    Returns:
+        A Rich Panel object ready to be printed
+    """
+    role = str(msg.get("role", ""))
+    content_val = msg.get("content", "")
+    name = msg.get("name")
+    tool_call_id = msg.get("tool_call_id")
+    message_title = (
+        f"[italic]{idx}/{total}[/italic] [bold]<{role}>[/bold]"
+        if not name
+        else f"[italic]{idx}/{total}[/italic] [bold]<{role} name={name} id={tool_call_id}>[/bold]"
+    )
+    message_elements = []
+    message_elements.append(
+        Padding(
+            Markdown(
+                escape_tags(content_val),
+            ),
+            (1, 0),
+        )
+    )
+    tool_calls = msg.get("tool_calls", [])
+    for tool_call in tool_calls:
+        function = tool_call.get("function", {})
+        arguments = function.get("arguments")
+        tool_elements = []
+        tool_title = f"[bold]<tool_call id={tool_call.get('id')}>[/bold]"
+        tool_elements.append(
+            f"[bold][magenta]{function.get('name')}[/magenta] [blue]([/blue][/bold]",
+        )
+        try:
+            tool_elements.append(Padding(JSON(arguments), (0, 4)))
+        except json.JSONDecodeError:
+            tool_elements.append(
+                Padding("[red]Invalid JSON generated by the model:[/red]", (0, 4))
+            )
+            tool_elements.append(arguments)
+        tool_elements.append("[bold][blue])[/blue][/bold]")
+        message_elements.append(
+            Panel(
+                Group(*tool_elements),
+                title=tool_title,
+                border_style=ROLE_COLORS["tool"],
+            )
+        )
+    message_style = ROLE_COLORS.get(role, "red")
+    message_subtitle = message_title[message_title.find("<") :].replace("<", "</")
+    message_subtitle = (
+        "[bold]"
+        + message_subtitle[
+            : (
+                message_subtitle.find(" ")
+                if " " in message_subtitle
+                else message_subtitle.find(">")
+            )
+        ]
+        + ">[/bold]"
+    )
+    return Panel(
+        Group(*message_elements),
+        title=message_title,
+        title_align="left",
+        subtitle=message_subtitle,
+        subtitle_align="right",
+        border_style=message_style,
+    )
+
+
 def log_request(req: Request) -> str:
     """Pretty-print a Flask request using Rich and return the request id."""
     request_id = uuid.uuid4().hex[:8]
@@ -131,130 +207,65 @@ def log_request(req: Request) -> str:
     messages = json_payload.get("messages", [])
     tools = json_payload.get("tools", [])
 
-    for idx, tool in enumerate(tools, start=1):
+    # Render tools section once (no duplicate panels)
-        table = Table(
+    table = Table(
-            caption="[italic]Required fields are marked with *[/italic]",
+        caption="[italic]Required fields are marked with *[/italic]",
-            pad_edge=False,
+        pad_edge=False,
-            box=box.SIMPLE,
+        box=box.SIMPLE,
-            leading=2,
+        leading=2,
+    )
+    table.add_column("Name", justify="right", style="cyan", no_wrap=True)
+    table.add_column("Description", style="white")
+
+    for tool in tools:
+        function = tool.get("function")
+        parameters = function.get("parameters", {}) or {}
+        required = parameters.get("required", []) or []
+        props = parameters.get("properties", {}) or {}
+
+        params_table = Table(
+            show_header=False,
+            show_edge=True,
+            show_lines=True,
+            title="Parameters",
+            title_justify="left",
+            expand=True,
+            leading=3,
+        )
+        params_table.add_column("Name", justify="right", no_wrap=True, style="cyan")
+        params_table.add_column("Description", style="white")
+        for param_name, param_value in props.items():
+            param_type = param_value.get("type")
+            if param_type == "array":
+                param_type += f"({param_value.get('items').get('type')})"
+            if param_name in required:
+                param_type = f"[bold]*{param_type}[/bold]"
+            param_type = f"[magenta]{param_type}[/magenta]"
+            params_table.add_row(
+                Group(param_name, param_type),
+                f"{param_value.get('description')}",
+            )
+        table.add_row(
+            f"{function.get('name')}",
+            Group(
+                Markdown(escape_tags(function.get("description"))),
+                params_table,
+            ),
         )
 
-        table.add_column("Name", justify="right", style="cyan", no_wrap=True)
+    console.print(
-        table.add_column("Description", style="white")
+        Panel(
-
+            table,
-        for tool in tools:
+            title=f"[italic]{0}/{len(messages)}[/italic] [bold]<tools>[/bold]",
-            function = tool.get("function")
+            title_align="left",
-            required = function.get("parameters").get("required")
+            subtitle="[bold]</tools>[/bold]",
-
+            subtitle_align="right",
-            params_table = Table(
+            border_style=ROLE_COLORS["tool"],
-                show_header=False,
-                show_edge=True,
-                show_lines=True,
-                title="Parameters",
-                title_justify="left",
-                expand=True,
-                leading=3,
-            )
-            params_table.add_column("Name", justify="right", no_wrap=True, style="cyan")
-            params_table.add_column("Description", style="white")
-            for param_name, param_value in (
-                function.get("parameters").get("properties").items()
-            ):
-                param_type = param_value.get("type")
-                if param_type == "array":
-                    param_type += f"({param_value.get('items').get('type')})"
-                if param_name in required:
-                    param_type = f"[bold]*{param_type}[/bold]"
-                param_type = f"[magenta]{param_type}[/magenta]"
-                param_name = f"{param_name}"
-                params_table.add_row(
-                    Group(param_name, param_type),
-                    f"{param_value.get("description")}",
-                )
-            table.add_row(
-                f"{function.get('name')}",
-                Group(
-                    Markdown(escape_tags(function.get("description"))),
-                    params_table,
-                ),
-            )
-
-        console.print(
-            Panel(
-                table,
-                title=f"[italic]{0}/{len(messages)}[/italic] [bold]<tools>[/bold]",
-                title_align="left",
-                subtitle="[bold]</tools>[/bold]",
-                subtitle_align="right",
-                border_style=ROLE_COLORS["tool"],
-            )
         )
+    )
 
     for idx, msg in enumerate(messages, start=1):
-        role = str(msg.get("role", ""))
+        panel = create_message_panel(msg, idx, len(messages))
-        content_val = msg.get("content", "")
+        console.print(panel)
-        name = msg.get("name")
-        tool_call_id = msg.get("tool_call_id")
-        message_title = (
-            f"[italic]{idx}/{len(messages)}[/italic] [bold]<{role}>[/bold]"
-            if not name
-            else f"[italic]{idx}/{len(messages)}[/italic] [bold]<{role} name={name} id={tool_call_id}>[/bold]"
-        )
-        message_elements = []
-        message_elements.append(
-            Padding(
-                Markdown(
-                    escape_tags(content_val),
-                ),
-                (1, 0),
-            )
-        )
-        tool_calls = msg.get("tool_calls", [])
-        for tool_call in tool_calls:
-            function = tool_call.get("function", {})
-            arguments = function.get("arguments")
-            tool_elements = []
-            tool_title = f"[bold]<tool_call id={tool_call.get('id')}>[/bold]"
-            tool_elements.append(
-                f"[bold][magenta]{function.get('name')}[/magenta] [blue]([/blue][/bold]",
-            )
-            try:
-                tool_elements.append(Padding(JSON(arguments), (0, 4)))
-            except json.JSONDecodeError:
-                tool_elements.append(
-                    Padding("[red]Invalid JSON generated by the model:[/red]", (0, 4))
-                )
-                tool_elements.append(arguments)
-            tool_elements.append("[bold][blue])[/blue][/bold]")
-            message_elements.append(
-                Panel(
-                    Group(*tool_elements),
-                    title=tool_title,
-                    border_style=ROLE_COLORS["tool"],
-                )
-            )
-        message_style = ROLE_COLORS.get(role, "red")
-        message_subtitle = message_title[message_title.find("<") :].replace("<", "</")
-        message_subtitle = (
-            "[bold]"
-            + message_subtitle[
-                : (
-                    message_subtitle.find(" ")
-                    if " " in message_subtitle
-                    else message_subtitle.find(">")
-                )
-            ]
-            + ">[/bold]"
-        )
-        console.print(
-            Panel(
-                Group(*message_elements),
-                title=message_title,
-                title_align="left",
-                subtitle=message_subtitle,
-                subtitle_align="right",
-                border_style=message_style,
-            )
-        )
 
     return request_id

app/common/sse.py

@@ -155,10 +155,10 @@ def chunks_to_sse(chunks: Iterable[Dict[str, Any]]) -> Iterator[bytes]:
             sse = encode_sse_json(obj)
             buffer += sse
             yield sse
-    finally:
         sse = done_event_bytes()
         buffer += sse
         yield sse
+    finally:
         record_sse(buffer, "downstream_response")
 
 

app/exceptions.py

@@ -23,3 +23,10 @@ class CursorConfigurationError(ConfigurationError):
     """Exception raised for configuration errors in Cursor configuration."""
 
     preamble = "Cursor configuration error, check your Cursor settings."
+
+
+class ClientClosedConnection(Exception):  # noqa: N818
+    """Raised when the downstream client closes the HTTP connection mid-stream.
+
+    This helps distinguish client disconnects from other server-side errors.
+    """

app/settings.py

@@ -14,6 +14,8 @@ env.read_env()
 ENV = env.str("FLASK_ENV", default="production")
 DEBUG = ENV == "development"
 RECORD_TRAFFIC = env.bool("RECORD_TRAFFIC", False)
+LOG_CONTEXT = env.bool("LOG_CONTEXT", True)
+LOG_COMPLETION = env.bool("LOG_COMPLETION", True)
 
 SERVICE_API_KEY = env.str("SERVICE_API_KEY", "change-me")
 

tests/settings.py

@@ -14,6 +14,8 @@ AZURE_VERBOSITY_LEVEL = "medium"
 AZURE_TRUNCATION = "auto"
 
 RECORD_TRAFFIC = False
+LOG_CONTEXT = True
+LOG_COMPLETION = True
 
 
 AZURE_RESPONSES_API_URL = (

tests/test_client_closed_connection.py

@@ -0,0 +1,51 @@
+"""Tests for client-closed connection during streaming.
+
+This test simulates a downstream client closing the connection mid-stream by
+forcing a GeneratorExit inside the streaming pipeline. The outer boundary in
+ResponseAdapter should translate that into ClientClosedConnection.
+"""
+
+import pytest
+
+from app.exceptions import ClientClosedConnection
+from tests.replay_base import ReplyBase
+
+
+def _raise_generator_exit(
+    _chunks,
+):  # pragma: no cover - behavior is validated via exception
+    """Stub encoder that simulates generator shutdown.
+
+    Raising GeneratorExit here mimics the server-side observation that the
+    response iterable was closed. We don't yield any bytes.
+    """
+    raise GeneratorExit
+
+
+class TestClientClosedConnection(ReplyBase):
+    """Ensure a client disconnect is surfaced as ClientClosedConnection."""
+
+    # Use any existing recording that yields a normal streaming response
+    recording: str = "one_ping_pong"
+
+    def test(self, testapp, requests_mock, monkeypatch):
+        """Test chunks_to_see raising GeneratorExit."""
+        # Patch the exact symbol used inside ResponseAdapter to encode SSE
+        monkeypatch.setattr(
+            "app.azure.response_adapter.chunks_to_sse", _raise_generator_exit
+        )
+
+        # Mock upstream with recorded SSE so the adapter enters streaming code
+        mock = self.mock_upstream(requests_mock)
+
+        # When the streaming pipeline is interrupted, the adapter should raise
+        # our domain-specific exception. No response bytes should be produced.
+        with pytest.raises(ClientClosedConnection):
+            testapp.post(
+                "/chat/completions",
+                params=self.downstream_request_body,
+                headers=self.downstream_request_headers,
+            )
+
+        # Sanity check: upstream was called with the expected payload
+        self.assert_upstream_request(mock)

tests/test_logging.py

@@ -50,3 +50,45 @@ def test_should_not_refact(mocker):
     }
     redacted_headers = redact_headers(headers)
     assert redacted_headers == headers
+
+
+class TestLogContextEnabled(ReplyBase):
+    """Test that log_request is called when LOG_CONTEXT=True."""
+
+    def modify_settings(self, app) -> None:
+        """Ensure LOG_CONTEXT is enabled."""
+        app.config["LOG_CONTEXT"] = True
+
+    def test(self, testapp, requests_mock, mocker):
+        """Test that log_request is called when LOG_CONTEXT is True."""
+        log_request_mock = mocker.patch("app.blueprint.log_request")
+        super().test(testapp, requests_mock)
+        log_request_mock.assert_called_once()
+
+
+class TestLogContextDisabled(ReplyBase):
+    """Test that log_request is NOT called when LOG_CONTEXT=False."""
+
+    def modify_settings(self, app) -> None:
+        """Disable LOG_CONTEXT."""
+        app.config["LOG_CONTEXT"] = False
+
+    def test(self, testapp, requests_mock, mocker):
+        """Test that log_request is NOT called when LOG_CONTEXT is False."""
+        log_request_mock = mocker.patch("app.blueprint.log_request")
+        super().test(testapp, requests_mock)
+        log_request_mock.assert_not_called()
+
+
+class TestLogCompletionDisabled(ReplyBase):
+    """Test that logging of create_message_panel is NOT called when LOG_COMPLETION=False."""
+
+    def modify_settings(self, app) -> None:
+        """Disable LOG_COMPLETION."""
+        app.config["LOG_COMPLETION"] = False
+
+    def test(self, testapp, requests_mock, mocker):
+        """Test that Rich.Live.update is NOT called when LOG_COMPLETION is False."""
+        live_update_mock = mocker.patch("rich.live.Live.update")
+        super().test(testapp, requests_mock)
+        live_update_mock.assert_not_called()