fix: harden composite routing and error handling

Author: Mateusz

docs/user_guide/features/auto-continue-removal.md

@@ -55,6 +55,29 @@ session:
 - Your workflow uses custom continue-like keywords that should reach the model.
 - You are testing context window behavior and need every message forwarded verbatim.
 
+## Usage Examples
+
+### Resume After A Network Interruption
+
+1. A coding session is interrupted mid-stream.
+2. The user sends `continue`.
+3. The proxy keeps the message in local history but does not forward it to the remote model.
+4. The next backend request reuses the meaningful prior context without wasting tokens on the mechanical resume prompt.
+
+### Preserve Literal Continue Prompts
+
+If you intentionally want the remote model to see `continue` or `proceed`, disable the feature first:
+
+```bash
+python -m src.core.cli --disable-auto-continue-removal
+```
+
+## Use Cases
+
+- Agentic coding sessions where transient disconnects are common and users resume by typing `continue`.
+- Long-running terminal workflows where preserving context window capacity matters.
+- Local debugging of session history, where the proxy should remember the resume command without sending it upstream.
+
 ## Logging
 
 When a message is tagged, the proxy logs at INFO level:

src/connectors/openai_codex/contracts.py

@@ -55,12 +55,15 @@ class ProcessedMessage(BaseModel):
     metadata: dict[str, object] | None = None
 
 
-class MessagePart(BaseModel):
-    """Content part in a multimodal message."""
-
-    type: str
-    text: str | None = None
-    data: object | None = None
+class MessagePart(BaseModel):
+    """Content part in a multimodal message."""
+
+    type: str
+    text: str | None = None
+    data: object | None = None
+
+
+ProcessedMessage.model_rebuild()
 
 
 class CodexInputItem(BaseModel):

src/core/common/exceptions.py

@@ -10,12 +10,35 @@
 import time
 from typing import TYPE_CHECKING, Any
 
-if TYPE_CHECKING:
-    from src.core.domain.client_termination import ClientTerminationReason
-    from src.core.domain.session_key import SessionKey
-
-
-class LLMProxyError(Exception):
+if TYPE_CHECKING:
+    from src.core.domain.client_termination import ClientTerminationReason
+    from src.core.domain.session_key import SessionKey
+
+
+def _is_json_scalar(value: Any) -> bool:
+    return isinstance(value, str | int | float | bool) or value is None
+
+
+def _to_json_safe_error_value(value: Any) -> Any:
+    if _is_json_scalar(value):
+        return value
+    if isinstance(value, dict):
+        return {
+            str(key): _to_json_safe_error_value(nested_value)
+            for key, nested_value in value.items()
+            if isinstance(key, str | int | float | bool)
+        }
+    if isinstance(value, list | tuple | set):
+        return [_to_json_safe_error_value(item) for item in value]
+
+    model_dump = getattr(value, "model_dump", None)
+    if callable(model_dump):
+        return _to_json_safe_error_value(model_dump(mode="json"))
+
+    return str(value)
+
+
+class LLMProxyError(Exception):
     """Base exception class for all LLM proxy errors."""
 
     __resilience_context__: dict[str, Any] | None
@@ -44,23 +67,25 @@ def __init__(
         for key, value in (kwargs or {}).items():
             setattr(self, key, value)
 
-    def to_dict(self) -> dict[str, Any]:
-        error_dict: dict[str, Any] = {
-            "message": self.message,
-            "type": self.__class__.__name__,
-            "details": self.details,
-        }
+    def to_dict(self) -> dict[str, Any]:
+        error_dict: dict[str, Any] = {
+            "message": self.message,
+            "type": self.__class__.__name__,
+            "details": _to_json_safe_error_value(self.details),
+        }
 
         # Include any additional attributes that were set via kwargs
         for attr_name in dir(self):
             if (
-                not attr_name.startswith("_")
-                and attr_name not in ["message", "details", "status_code", "args"]
-                and not callable(getattr(self, attr_name))
-            ):
-                error_dict[attr_name] = getattr(self, attr_name)
-
-        return {"error": error_dict}
+                not attr_name.startswith("_")
+                and attr_name not in ["message", "details", "status_code", "args"]
+                and not callable(getattr(self, attr_name))
+            ):
+                error_dict[attr_name] = _to_json_safe_error_value(
+                    getattr(self, attr_name)
+                )
+
+        return {"error": error_dict}
 
 
 class AuthenticationError(LLMProxyError):

src/core/services/backend_model_resolver.py

@@ -34,12 +34,12 @@
 from src.core.interfaces.planning_phase_manager_interface import IPlanningPhaseManager
 from src.core.interfaces.session_service_interface import ISessionService
 from src.core.services.backend_routing_service import BackendRoutingService
-from src.core.services.composite_routing_state import (
-    COMPOSITE_LEAF_RESOLUTION_EXTRA_BODY_KEY,
-    COMPOSITE_LEAF_RESOLUTION_FLAG,
-    is_composite_selector,
-    resolve_composite_routing_surface,
-)
+from src.core.services.composite_routing_state import (
+    COMPOSITE_LEAF_RESOLUTION_EXTRA_BODY_KEY,
+    COMPOSITE_LEAF_RESOLUTION_FLAG,
+    COMPOSITE_LEAF_SELECTOR_EXTRA_BODY_KEY,
+    resolve_composite_routing_surface,
+)
 from src.core.services.replacement_compatibility_bridge import (
     ReplacementCompatibilityBridge,
 )
@@ -476,21 +476,18 @@ def _is_composite_leaf_resolution(
         context: RequestContext | None,
         request: ChatRequest,
     ) -> bool:
-        model = request.model
-        if is_composite_selector(model):
-            return False
+        if context is not None and bool(
+            context.extensions.get(COMPOSITE_LEAF_RESOLUTION_FLAG)
+        ):
+            return True
 
-        if context is None:
-            extra_body = request.extra_body
-            if isinstance(extra_body, dict):
-                return bool(extra_body.get(COMPOSITE_LEAF_RESOLUTION_EXTRA_BODY_KEY))
+        extra_body = request.extra_body
+        if not isinstance(extra_body, dict):
             return False
-        if bool(context.extensions.get(COMPOSITE_LEAF_RESOLUTION_FLAG)):
-            return True
-        extra_body = request.extra_body
-        if isinstance(extra_body, dict):
-            return bool(extra_body.get(COMPOSITE_LEAF_RESOLUTION_EXTRA_BODY_KEY))
-        return False
+        if not bool(extra_body.get(COMPOSITE_LEAF_RESOLUTION_EXTRA_BODY_KEY)):
+            return False
+        leaf_selector = extra_body.get(COMPOSITE_LEAF_SELECTOR_EXTRA_BODY_KEY)
+        return isinstance(leaf_selector, str) and leaf_selector == request.model
 
     @staticmethod
     def _normalize_uri_params(raw_value: Any) -> dict[str, JsonValue]:

src/core/services/composite_leaf_target_resolver_adapter.py

@@ -13,6 +13,7 @@
 from src.core.services.composite_routing_state import (
     COMPOSITE_LEAF_RESOLUTION_EXTRA_BODY_KEY,
     COMPOSITE_LEAF_RESOLUTION_FLAG,
+    COMPOSITE_LEAF_SELECTOR_EXTRA_BODY_KEY,
 )
 
 __all__ = ["CompositeLeafTargetResolverAdapter"]
@@ -35,6 +36,7 @@ async def resolve_leaf(
     ) -> BackendTarget:
         leaf_extra_body = dict(request.extra_body or {})
         leaf_extra_body[COMPOSITE_LEAF_RESOLUTION_EXTRA_BODY_KEY] = True
+        leaf_extra_body[COMPOSITE_LEAF_SELECTOR_EXTRA_BODY_KEY] = leaf_selector
 
         parsed_leaf = parse_model_with_params(leaf_selector, default_backend="")
         if has_explicit_backend_selector(leaf_selector):

src/core/services/composite_routing_state.py

@@ -11,6 +11,7 @@
 __all__ = [
     "COMPOSITE_LEAF_RESOLUTION_EXTRA_BODY_KEY",
     "COMPOSITE_LEAF_RESOLUTION_FLAG",
+    "COMPOSITE_LEAF_SELECTOR_EXTRA_BODY_KEY",
     "COMPOSITE_ROUTING_STATE_KEY",
     "COMPOSITE_ROUTING_SURFACE_KEY",
     "FAILOVER_MODE",
@@ -29,6 +30,7 @@
 COMPOSITE_ROUTING_SURFACE_KEY = "composite_routing_surface"
 COMPOSITE_LEAF_RESOLUTION_FLAG = "composite_leaf_resolution"
 COMPOSITE_LEAF_RESOLUTION_EXTRA_BODY_KEY = "_composite_leaf_resolution"
+COMPOSITE_LEAF_SELECTOR_EXTRA_BODY_KEY = "_composite_leaf_selector"
 FAILOVER_MODE = "failover"
 WEIGHTED_RETRY_MODE = "weighted_retry"
 

src/core/services/composite_selector_parser.py

@@ -62,7 +62,7 @@ def parse(self, routing_input: CompositeRoutingInput) -> CompositeRoutePlan:
         parts = self._split_top_level(selector, operator)
         other_operator = "^" if operator == "|" else "|"
         has_mixed_operator = any(
-            self._contains_operator_in_route_segment(segment, other_operator)
+            self._contains_operator_outside_brackets(segment, other_operator)
             for segment in parts
         )
         if has_mixed_operator:
@@ -132,11 +132,6 @@ def _detect_primary_operator(selector: str) -> str | None:
                 return char
         return None
 
-    @classmethod
-    def _contains_operator_in_route_segment(cls, segment: str, operator: str) -> bool:
-        route_segment, _, _ = segment.partition("?")
-        return cls._contains_operator_outside_brackets(route_segment, operator)
-
     @staticmethod
     def _contains_operator_outside_brackets(selector: str, operator: str) -> bool:
         bracket_depth = 0

src/core/services/streaming/error_mapping.py

@@ -21,12 +21,29 @@
     LLMProxyError,
     ParsingError,
     RateLimitExceededError,
+    RoutingError,
 )
 from src.core.domain.streaming.contracts import StreamingErrorInfo
 
 logger = logging.getLogger(__name__)
 
 
+def _resolve_routing_error_status(error: RoutingError) -> int:
+    details = getattr(error, "details", None)
+    if isinstance(details, dict):
+        code = details.get("code")
+        if code == "unknown_model":
+            return 404
+        if code == "unsupported_on_instance":
+            return 400
+        if code == "temporarily_unavailable":
+            return 503
+        if code == "policy_rejected":
+            return 403
+    status_code = getattr(error, "status_code", None)
+    return status_code if isinstance(status_code, int) else 500
+
+
 def _merge_provider_retry_metadata(
     details: dict[str, str], detail_payload: dict[str, Any]
 ) -> None:
@@ -240,7 +257,9 @@ async def handle_streaming_error(
 
     # Extract status_code if available
     status_code: int | None = None
-    if hasattr(mapped_error, "status_code"):
+    if isinstance(mapped_error, RoutingError):
+        status_code = _resolve_routing_error_status(mapped_error)
+    elif hasattr(mapped_error, "status_code"):
         status_code = mapped_error.status_code
 
     # For quota_exceeded errors, use 503 instead of 429

tests/unit/app/middleware/test_exception_middleware.py

@@ -5,6 +5,11 @@
 from fastapi.testclient import TestClient
 from src.core.app.middleware.exception_middleware import DomainExceptionMiddleware
 from src.core.common.exceptions import DuplicateRequestError, RateLimitExceededError
+from src.core.domain.composite_routing import (
+    CompositeSelectorValidationError,
+    CompositeValidationErrorCode,
+    CompositeValidationErrorEnvelope,
+)
 
 
 def test_domain_exception_middleware_sets_retry_after_header(monkeypatch):
@@ -62,6 +67,36 @@ async def duplicate_endpoint() -> None:
     assert body["error"]["type"] == "DuplicateRequestError"
 
 
+def test_domain_exception_middleware_serializes_composite_validation_envelope():
+    app = FastAPI()
+    app.add_middleware(DomainExceptionMiddleware)
+
+    @app.get("/composite-validation")
+    async def composite_validation_endpoint() -> None:
+        raise CompositeSelectorValidationError(
+            CompositeValidationErrorEnvelope(
+                code=CompositeValidationErrorCode.UNSUPPORTED_CONSTRUCT,
+                message="Mixed operators are not supported.",
+                selector_echo="openai:gpt-4|anthropic:claude^gemini:flash",
+            )
+        )
+
+    _ = composite_validation_endpoint
+
+    with TestClient(app) as client:
+        response = client.get("/composite-validation")
+
+    assert response.status_code == 400
+    body = response.json()
+    assert body["error"]["type"] == "CompositeSelectorValidationError"
+    assert body["error"]["details"]["composite_validation"]["code"] == (
+        "unsupported_construct"
+    )
+    assert body["error"]["envelope"]["selector_echo"] == (
+        "openai:gpt-4|anthropic:claude^gemini:flash"
+    )
+
+
 async def test_domain_exception_middleware_reraises_transport_error_on_final_body_send():
     """Transport failures on the final body write must not be swallowed."""
 

tests/unit/core/services/test_backend_completion_flow_streaming_error_envelope.py

@@ -230,3 +230,33 @@ async def test_streaming_error_envelope_enters_failover_recovery_when_enabled()
         "category": "validation",
         "retryable": False,
     }
+
+
+@pytest.mark.asyncio
+async def test_streaming_terminal_routing_error_uses_unknown_model_http_status() -> (
+    None
+):
+    flow = _build_flow_with_erroring_backend()
+
+    routing_error = RoutingError(
+        message="Missing extracted backend",
+        details={
+            "code": "unknown_model",
+            "category": "validation",
+            "retryable": False,
+        },
+    )
+
+    envelope = await flow._build_terminal_error_stream_envelope(
+        error=routing_error,
+        provider="gemini-oauth-plan",
+    )
+
+    assert envelope.status_code == 404
+
+    assert envelope.content is not None
+    chunks = [chunk async for chunk in envelope.content]
+    assert len(chunks) == 1
+    error_payload = chunks[0].metadata.get("error")
+    assert isinstance(error_payload, dict)
+    assert error_payload.get("status_code") == 404