Tighten WebSocket auth close codes

[tianmind-studio]

Summary

• distinguish WebSocket auth failures so clients can refresh expired/certificate-stale tokens or force re-login for revoked tokens
• keep generic invalid auth on policy violation code 1008
• add client-side socket close-code labels for 1008, 4001, and 4004
• add backend coverage for expired, revoked, certificate-key, certificate-fetch, and generic key auth failures
• keep the backend auth test isolated on Windows by stubbing Firebase/database imports before loading utils.other.endpoints

Current status

• Rebased on main at 1a5824403b68ce47c3b0909577cadc1242ba0d3f
• Head refreshed to 41859b19a479831be67953776b910d770e4a3dd2
• Existing review threads are resolved/outdated after the certificate-fetch and generic-key follow-ups

Verification

• D:\codex-omi-work\.venvs\omi-backend-vad-refresh\Scripts\python.exe -m pytest backend\tests\unit\test_ws_auth_handshake.py -q --tb=short
  • 23 passed, 1 warning
• D:\codex-omi-work\.venvs\omi-backend-vad-refresh\Scripts\python.exe -m py_compile backend\utils\other\endpoints.py backend\tests\unit\test_ws_auth_handshake.py
• D:\codex-omi-work\.venvs\omi-backend-vad-refresh\Scripts\python.exe -m black --check backend\utils\other\endpoints.py backend\tests\unit\test_ws_auth_handshake.py
  • 2 files would be left unchanged
• D:\codex-omi-work\.tools\dart-sdk\bin\dart.exe format --set-exit-if-changed app\lib\services\sockets\pure_socket.dart
  • Formatted 1 file (0 changed). Dart still reports a flutter_lints package-resolution warning because app package dependencies are not installed in this Windows worktree.
• git diff --check origin/main...HEAD
• scripts/pre-commit with the backend Windows venv and local Dart SDK on PATH

[greptile-apps]

Greptile Summary

This PR differentiates WebSocket authentication close codes by mapping Firebase auth errors to 4001 (token refresh needed), 4004 (re-login required), and 1008 (generic failure), and adds matching telemetry labels on the Dart client side.

• Backend (endpoints.py): A new _get_ws_auth_close helper uses isinstance checks for ExpiredIdTokenError/RevokedIdTokenError/CertificateFetchError, then falls back to substring matching. The CertificateFetchError isinstance branch is unreachable because that class inherits from UnknownError, not InvalidIdTokenError.
• Tests (test_ws_auth_handshake.py): New tests cover expired/revoked token paths, but the certificate-error test uses InvalidIdTokenError('Certificate key not found') rather than an actual CertificateFetchError, leaving the dead isinstance branch undetected.
• Dart (pure_socket.dart): Adds human-readable labels for 1008, 4001, and 4004 in _getCloseCodeReason for telemetry only; no reconnect behavior changes.

Confidence Score: 3/5

Safe for the happy-path and string-fallback cases, but a real CertificateFetchError silently returns 1008 instead of the intended 4001, so clients will not know to refresh their token in that scenario.

The core logic routing ExpiredIdTokenError to 4001 and RevokedIdTokenError to 4004 is correct and tested. However, CertificateFetchError inherits from UnknownError not InvalidIdTokenError, making its isinstance branch dead code. The test that appears to cover this path exercises a completely different branch via string matching, masking the gap.

backend/utils/other/endpoints.py — the _verify_ws_auth and _get_ws_auth_close functions need attention for the unreachable CertificateFetchError branch.

Important Files Changed

Filename	Overview
backend/utils/other/endpoints.py	Introduces _get_ws_auth_close to map Firebase auth errors to distinct WS close codes; contains a dead isinstance(CertificateFetchError) check because that class does not inherit from InvalidIdTokenError, and an overly broad 'key' substring match in the fallback.
backend/tests/unit/test_ws_auth_handshake.py	Adds tests for 4001/4004 close codes; the CertificateFetchError test case exercises the string-matching fallback rather than an actual CertificateFetchError instance, leaving the dead isinstance branch untested.
app/lib/services/sockets/pure_socket.dart	Adds telemetry labels for close codes 1008, 4001, and 4004 in _getCloseCodeReason; straightforward and correct for its stated purpose.

Sequence Diagram

sequenceDiagram
    participant Client as Flutter Client
    participant Auth as _verify_ws_auth
    participant Close as _get_ws_auth_close
    participant FB as firebase_admin

    Client->>Auth: WS connect (Authorization header)
    alt No or malformed header
        Auth-->>Client: close 1008
    else Valid format
        Auth->>FB: verify_token(token)
        alt RevokedIdTokenError (subclass of InvalidIdTokenError)
            FB-->>Auth: raises RevokedIdTokenError
            Auth->>Close: _get_ws_auth_close(e)
            Close-->>Auth: 4004
            Auth-->>Client: close 4004 (re-login required)
        else ExpiredIdTokenError (subclass of InvalidIdTokenError)
            FB-->>Auth: raises ExpiredIdTokenError
            Auth->>Close: _get_ws_auth_close(e)
            Close-->>Auth: 4001
            Auth-->>Client: close 4001 (token refresh)
        else Generic InvalidIdTokenError
            FB-->>Auth: raises InvalidIdTokenError
            Auth->>Close: _get_ws_auth_close(e)
            Close-->>Auth: 1008 or 4001 via string match
            Auth-->>Client: close 1008 or 4001
        else CertificateFetchError (NOT subclass of InvalidIdTokenError)
            FB-->>Auth: raises CertificateFetchError
            Note over Auth: Caught by except Exception, not InvalidIdTokenError
            Auth-->>Client: close 1008 instead of intended 4001
        else Success
            FB-->>Auth: uid string
            Auth-->>Client: connection accepted
        end
    end

Loading

Comments Outside Diff (1)

1. backend/utils/other/endpoints.py, line 128-132 (link)

   The docstring still says the function only raises WebSocketException(code=1008), but it can now raise with codes 4001 or 4004 as well.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "Differentiate websocket auth close codes" | Re-trigger Greptile}

[kodjima33]

Differentiated WebSocket auth close codes — mixed backend+app/dart, approve only per policy.

[tianmind-studio]

Rebased and force-pushed onto current main.

Greptile follow-up addressed in the current branch:

• real CertificateFetchError is caught before the generic Exception handler and is covered by test_certificate_fetch_error_sends_close_4001
• the fallback no longer treats every key message as refreshable; only expired/certificate map to 4001, with API key invalid covered as 1008
• the _verify_ws_auth docstring now documents 1008, 4001, and 4004

Backend verification passes: python -m pytest tests\unit\test_ws_auth_handshake.py -q -> 23 passed, 1 warning. I could not run Dart/Flutter formatting locally because this Windows host does not have dart or flutter installed.

[kodjima33]

re-approved on new sha: backend+dart-client (distinct WS auth close codes + close-code labels) — approve only per policy

[kodjima33]

Re-approved on new commits — backend (approve-only area).

[kodjima33]

Tighten WebSocket auth close codes — approve only, auth area (re-approve)