feat: WebSocket auth hardening and e2e tests#47
Conversation
Harden WebSocket authentication and add comprehensive testing: - Rewrite WebSocket auth to use denial responses before accepting - Add separate auth checks for missing and invalid tokens - Add WebSocket bypass/alias handling with mock response streaming - Add request context and model alias restoration for WebSocket events - Add WebSocket payload sanitization through adapter pipeline - Simplify detection prompt merging to always merge with fallback - Add use_mock_adapter_in_bypass_mode flag for per-plugin control - Pass mock_handler to adapter kwargs for integrated bypass handling - Add WebSocket e2e tests, route tests, and detection alias tests - Remove agent-framework dependency
Extract _sanitize_provider_body from prepare_provider_request so both the real and mock WebSocket paths apply the same Codex-specific cleanup: removing unsupported keys, filtering item_reference inputs, and stripping _-prefixed metadata fields.
… errors _make_websocket_terminal_event was always emitting type "response.completed" even for errors, and omitting sequence_number. Per the OpenAI Responses API spec (ResponseFailedEvent model), error terminal events must use type "response.failed" and all stream events require a sequence_number field.
There was a problem hiding this comment.
Pull request overview
This PR hardens the Codex WebSocket authentication flow (denying unauthorized handshakes before accepting) and expands WebSocket support in bypass/mock mode, while adding unit/integration/e2e test coverage for streaming, warmup, error, and multi-message WebSocket behaviors.
Changes:
- Update Codex WebSocket auth to send denial responses (401) before
accept(), plus helper plumbing for settings resolution. - Add mock/bypass-mode WebSocket streaming support and model-alias restoration for WS events, aligning payload sanitization with HTTP.
- Add comprehensive unit/integration/e2e tests for Codex WebSocket routes and shared WS validation/test-data helpers.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
ccproxy/plugins/codex/routes.py |
WebSocket auth denial-before-accept, WS payload sanitization/model alias restore, and bypass/mock WS streaming path. |
ccproxy/plugins/codex/adapter.py |
Route requests through MockAdapter when mock_handler is injected; unify sanitization via _sanitize_provider_body. |
ccproxy/plugins/codex/detection_service.py |
Merge partial detected prompt caches with fallback prompt data. |
ccproxy/plugins/codex/plugin.py |
Disable MockAdapter replacement for Codex in bypass mode (inject mock handler instead). |
ccproxy/core/plugins/factories.py |
Add use_mock_adapter_in_bypass_mode flag and inject mock_handler into adapters in bypass mode when configured. |
ccproxy/services/adapters/base.py |
Persist mock_handler on adapters via BaseAdapter kwargs. |
ccproxy/services/adapters/mock_adapter.py |
Tighten typing around streaming mock handler returns. |
tests/unit/plugins/test_codex_detection.py |
Add test for merging partial prompt cache with fallback prompts. |
tests/unit/core/test_provider_factory_bypass.py |
Update bypass-mode factory test to expect real adapter w/ injected mock handler for Codex. |
tests/plugins/codex/unit/test_routes.py |
New unit tests for WS route helpers (sanitization, alias restore, settings resolution). |
tests/plugins/codex/integration/test_codex_websocket.py |
New integration tests covering Codex WS streaming, warmup, upstream errors, bypass mode, and auth denial. |
tests/integration/test_websocket_e2e.py |
New parameterized e2e suite for Codex WS endpoints (including optional live-server tests). |
tests/helpers/test_data.py |
Add WS endpoint configurations + WS request builders and terminal type constants. |
tests/helpers/e2e_validation.py |
Add WS-specific validation helpers for event sequences, streaming content, warmups, and error events. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if not saw_terminal_event: | ||
| await websocket.send_text( | ||
| json.dumps( | ||
| _make_websocket_terminal_event( | ||
| provider_payload, | ||
| error={ | ||
| "type": "server_error", | ||
| "message": "WebSocket stream ended before response.completed", | ||
| }, | ||
| _restore_websocket_event_models( | ||
| _make_websocket_terminal_event( | ||
| provider_payload, | ||
| error={ | ||
| "type": "server_error", | ||
| "message": "WebSocket stream ended before response.completed", | ||
| }, | ||
| ), | ||
| request_context, | ||
| ), | ||
| separators=(",", ":"), | ||
| ) | ||
| ) |
There was a problem hiding this comment.
The fallback terminal event emitted when the upstream stream ends without a terminal event always uses sequence_number=0 (default in _make_websocket_terminal_event). If some events were already forwarded, this produces a non-monotonic sequence_number and can break clients that rely on ordering. Track the last seen sequence_number (or increment a local counter when forwarding events) and pass an appropriate value into _make_websocket_terminal_event for this error case.
| if not saw_terminal_event: | ||
| await _send_websocket_event( | ||
| websocket, | ||
| _make_websocket_terminal_event( | ||
| provider_payload, | ||
| error={ | ||
| "type": "server_error", | ||
| "message": "Mock WebSocket stream ended before response.completed", | ||
| }, | ||
| ), | ||
| request_context, | ||
| ) |
There was a problem hiding this comment.
Same as the upstream path: when the mock WS SSE stream ends without a terminal event, the locally generated response.failed terminal event uses the default sequence_number=0. If the mock handler already emitted events with increasing sequence numbers, this final error event will be out of order. Track the last forwarded sequence number and pass it through to _make_websocket_terminal_event here as well.
Summary
Replaces #45 (auto-closed when #44 merged). PR 3 of 3 from original #41.
Test plan