Fix speaker sample extraction after WS reconnect

[beastoin]

Summary

Fixes #5949 (scoped to the real backend bug — see analysis below).

• Bug: current_session_segments dict in transcribe.py is declared inside _stream_handler() and resets on every new WebSocket connection. After WS reconnect, can_assign returns false because segment IDs from the previous session are gone, blocking speaker sample extraction for all users who experience WS reconnect during recording (not just TestFlight).
• Fix: Extract rehydrate_session_segments() to utils/speaker_assignment.py and call it from transcribe.py when resuming a conversation after reconnect. Uses defensive reads (seg.get('id'), seg.get('speech_profile_processed', True)) and logs the rehydration count.
• Scope note: The issue's core narrative (prod/staging backend split on reconnect) was verified as not reproducible from code — Env.overrideApiBaseUrl is set once at init, no mid-session mutation exists. The staging banner claim was already fixed by commit 2f2b2e01d. This PR fixes the one real bug: session-local segment state loss on reconnect.

Changes

File	Change
backend/utils/speaker_assignment.py	New rehydrate_session_segments() shared utility function
backend/routers/transcribe.py	Import + call rehydrate_session_segments() in _prepare_in_progess_conversations() resume path
backend/tests/unit/test_speaker_segment_reconnect.py	8 unit tests: dict segments, empty, missing fields, can_assign, cross-convo leakage, large list, TranscriptSegment objects
backend/test.sh	Add new test file to test runner

Testing

$ cd backend && pytest tests/unit/test_speaker_segment_reconnect.py tests/unit/test_speaker_assignment.py -v
22 passed in 0.11s

Review Cycle

• Round 1: Reviewer flagged tests re-implementing logic locally → extracted to shared utility
• Round 2: Approved (no blocking issues, 22 tests pass)
• Tester: Added TranscriptSegment object test for getattr branch → 8 tests, TESTS_APPROVED

Risks / Edge Cases

1. Cross-conversation leakage: Not a risk — current_session_segments is initialized empty at _stream_handler() entry; rehydration only on resume
2. Missing speech_profile_processed: Defaults to True (matching model default)
3. Large conversations: O(n) pass acceptable — conversation already in memory
4. Speaker map continuity: NOT rehydrating speaker_to_person_map / segment_person_assignment_map — speaker_id continuity not guaranteed across reconnect

Issue Analysis (Codex-verified)

Issue Claim	Verified
testFlightApiEnvironment defaults to staging	TRUE
main.dart overrides apiBaseUrl for TestFlight	TRUE
WS reconnect uses current apiBaseUrl	TRUE
Prod→staging split on reconnect	FALSE — no mid-session URL mutation
current_session_segments resets on reconnect	TRUE — THIS PR
Staging banner broken (isUsingStagingApi)	FALSE — already fixed

Closes #5949

---

by AI for @beastoin

[greptile-apps]

Greptile Summary

This PR fixes a real bug where current_session_segments — a local dict in _stream_handler() that tracks which transcript segments have been processed against a speech profile — was silently reset to empty on every WebSocket reconnect. After reconnect, the can_assign guard in the speaker_assigned event handler could never find any segment IDs from the prior session, blocking speaker sample extraction for any user who experienced a WS drop during recording.

The fix is a small, well-scoped addition inside _prepare_in_progess_conversations(): when an existing in-progress conversation is resumed (i.e. within the timeout window), its persisted transcript_segments are iterated and their id / speech_profile_processed values are loaded back into current_session_segments. Defensive reads (seg.get(...) vs getattr(...)) handle both raw-dict and model-object representations; missing speech_profile_processed defaults to True (matching the model default). No nonlocal is required because the code mutates the dict rather than rebinding the variable.

Key points:

• Rehydration is correctly limited to the resume path; new conversations and timed-out conversations start with an empty dict as before, so there is no cross-conversation leakage.
• speaker_to_person_map and segment_person_assignment_map are intentionally NOT rehydrated because speaker_id continuity across reconnect is not guaranteed.
• 7 new unit tests cover the core scenarios; the test helper _rehydrate() mirrors the production logic exactly. The object-attribute branch of the isinstance check is not yet tested (see inline comment).
• test.sh is updated to include the new test file in CI.

Confidence Score: 5/5

• This PR is safe to merge — the fix is minimal, correctly scoped, and backed by passing tests with no regressions.
• The change is a small dict-population loop in a single code path (conversation resume on reconnect). It uses correct Python closure semantics (mutation vs rebinding), adds defensive field access, and includes a logging statement. The root cause analysis in the PR description is accurate. All existing tests pass, and 7 new unit tests directly validate the fix. The only open items are P2 style suggestions (stale line number comment, missing object-path test coverage) that do not block correctness.
• No files require special attention.

Important Files Changed

Filename	Overview
backend/routers/transcribe.py	Rehydrates current_session_segments from persisted transcript_segments on conversation resume; uses safe defensive reads and correct dict-mutation semantics (no nonlocal needed). The fix is minimal and correctly scoped to the reconnect path only.
backend/tests/unit/test_speaker_segment_reconnect.py	7 unit tests covering dict segments, empty list, missing fields, can_assign simulation, cross-conversation isolation, and large lists. Missing coverage for the object-attribute branch of the isinstance check; stale line number in one comment.
backend/test.sh	New test file correctly registered in the unit test runner; no issues.

Sequence Diagram

sequenceDiagram
    participant Client
    participant StreamHandler as _stream_handler()
    participant PrepareConv as _prepare_in_progress_conversations()
    participant DB as conversations_db

    Note over StreamHandler: current_session_segments = {} (fresh on every WS connect)

    Client->>StreamHandler: WebSocket reconnect
    StreamHandler->>PrepareConv: await _prepare_in_progess_conversations()
    PrepareConv->>DB: retrieve_in_progress_conversation(uid)
    DB-->>PrepareConv: existing_conversation (with transcript_segments)
    
    alt seconds_since_last_segment < timeout (resume path)
        PrepareConv->>StreamHandler: current_conversation_id = existing_conversation['id']
        loop for each seg in transcript_segments (NEW in this PR)
            PrepareConv->>StreamHandler: current_session_segments[seg.id] = seg.speech_profile_processed
        end
        PrepareConv->>PrepareConv: log rehydration count
    else timed out
        PrepareConv->>StreamHandler: _create_new_in_progress_conversation()
        Note over StreamHandler: current_session_segments stays empty (correct)
    end

    Client->>StreamHandler: speaker_assigned event (segment_ids)
    StreamHandler->>StreamHandler: can_assign check against current_session_segments
    Note over StreamHandler: Now finds pre-reconnect segment IDs → can_assign = True

Loading

_{Reviews (1): Last reviewed commit: "Add speaker segment reconnect tests to t..." | Re-trigger Greptile}

[greptile-apps]

Stale line number in comment

The comment references transcribe.py:2509-2514, but the can_assign logic is currently around lines 2536–2541. This is a minor stale reference that may confuse future readers.

Suggested change

	# Simulate the can_assign check from transcribe.py:2509-2514
	# Simulate the can_assign check from transcribe.py (speaker_assigned handler)

[greptile-apps]

No test coverage for object-attribute segment path

The production code in transcribe.py handles both dict segments and model-object segments via the isinstance(seg, dict) check:

sid = seg.get('id') if isinstance(seg, dict) else getattr(seg, 'id', None)

All seven tests only exercise the dict path. The getattr branch (reached when transcript_segments contains Pydantic/dataclass model instances) is entirely untested. Adding a simple object-based fixture (e.g. a SimpleNamespace or a minimal stub) would complete coverage of that branch and guard against regressions if the storage layer ever returns model objects instead of raw dicts.

[beastoin]

@beastoin I re-reviewed the latest diff for #5989 and the previous concern is addressed: rehydrate_session_segments is now shared production code in utils/speaker_assignment.py, and the reconnect tests import that function directly instead of duplicating logic; I also ran pytest tests/unit/test_speaker_assignment.py tests/unit/test_speaker_segment_reconnect.py -v locally and all 21 tests passed. I have no blocking issues on this update; an extra object-based segment test for the getattr(...) branch would be optional but non-blocking. Please merge when ready.

---

by AI for @beastoin

[beastoin]

CP9 Changed-Path Coverage Checklist

Path ID	Sequence ID(s)	Changed path	Happy-path test	Non-happy-path test	L1 result	L2 result	If untested: justification
P1	N/A	transcribe.py:843 rehydration call in resume branch	test_can_assign_after_rehydration	test_rehydrate_empty_segments	PASS (22/22 tests)	UNTESTED	WS reconnect + resume requires live WS session + network disruption; rehydration logic fully covered by unit tests
P2	N/A	transcribe.py:849 log on rehydration	test_rehydrate_from_dict_segments (non-empty)	test_rehydrate_empty_segments (no log)	PASS	UNTESTED	Logging branch covered by unit test of function return value
P3	N/A	speaker_assignment.py:rehydrate_session_segments dict branch	test_rehydrate_from_dict_segments	test_rehydrate_missing_speech_profile_processed, test_rehydrate_missing_id	PASS	UNTESTED	Utility function; no integration boundary
P4	N/A	speaker_assignment.py:rehydrate_session_segments object branch	test_rehydrate_from_transcript_segment_objects	N/A (objects always have id)	PASS	UNTESTED	Utility function; no integration boundary
P5	N/A	speaker_assignment.py:rehydrate_session_segments skip-no-id branch	test_rehydrate_missing_id	N/A	PASS	UNTESTED	Utility function; no integration boundary

L1 Synthesis

All 5 changed paths (P1-P5) proven by 8 unit tests (N/A sequences, path-only mode). Non-happy paths proven: empty segments, missing id, missing speech_profile_processed, can_assign false for unknown/unprocessed segments. 22/22 tests pass including pre-existing speaker_assignment tests. Backend utility imports and tests cleanly.

L2 Synthesis

P1-P5 are UNTESTED at L2 level. The rehydration code path triggers only on WS reconnection that resumes an existing conversation — requires simulating network disruption during active recording, which is not feasible in local integrated testing without infrastructure support. The fix is a pure data transformation (read segments → populate dict) with no cross-service boundary change — the WS protocol, Pusher communication, and Firestore schema are all unchanged. L1 unit tests provide complete coverage of the rehydration logic.

---

by AI for @beastoin

[beastoin]

CP8 Test Detail Table

Sequence ID	Path ID	Scenario ID	Changed path	Exact test command	Test name(s)	Assertion intent	Result	Evidence
N/A	P3	S1	speaker_assignment.py:rehydrate_session_segments dict happy	cd backend && pytest tests/unit/test_speaker_segment_reconnect.py::TestRehydrateSessionSegments::test_rehydrate_from_dict_segments -v	test_rehydrate_from_dict_segments	Dict segments produce correct {id: spp} mapping	PASS	PR comment test output
N/A	P3	S2	speaker_assignment.py:rehydrate_session_segments empty	cd backend && pytest tests/unit/test_speaker_segment_reconnect.py::TestRehydrateSessionSegments::test_rehydrate_empty_segments -v	test_rehydrate_empty_segments	Empty input → empty dict	PASS	PR comment test output
N/A	P3	S3	speaker_assignment.py:rehydrate_session_segments missing spp	cd backend && pytest tests/unit/test_speaker_segment_reconnect.py::TestRehydrateSessionSegments::test_rehydrate_missing_speech_profile_processed -v	test_rehydrate_missing_speech_profile_processed	Missing spp defaults to True	PASS	PR comment test output
N/A	P5	S4	speaker_assignment.py:rehydrate_session_segments missing id	cd backend && pytest tests/unit/test_speaker_segment_reconnect.py::TestRehydrateSessionSegments::test_rehydrate_missing_id -v	test_rehydrate_missing_id	Segments without id skipped	PASS	PR comment test output
N/A	P1	S5	transcribe.py:843 can_assign after rehydrate	cd backend && pytest tests/unit/test_speaker_segment_reconnect.py::TestRehydrateSessionSegments::test_can_assign_after_rehydration -v	test_can_assign_after_rehydration	can_assign True for processed, False for unprocessed/unknown	PASS	PR comment test output
N/A	P1	S6	cross-convo leakage prevention	cd backend && pytest tests/unit/test_speaker_segment_reconnect.py::TestRehydrateSessionSegments::test_no_cross_conversation_leakage -v	test_no_cross_conversation_leakage	New conversation gets empty dict	PASS	PR comment test output
N/A	P3	S7	large segment list	cd backend && pytest tests/unit/test_speaker_segment_reconnect.py::TestRehydrateSessionSegments::test_rehydrate_large_segment_list -v	test_rehydrate_large_segment_list	1000 segments rehydrated correctly	PASS	PR comment test output
N/A	P4	S8	TranscriptSegment objects	cd backend && pytest tests/unit/test_speaker_segment_reconnect.py::TestRehydrateSessionSegments::test_rehydrate_from_transcript_segment_objects -v	test_rehydrate_from_transcript_segment_objects	getattr branch works with real objects	PASS	PR comment test output

---

by AI for @beastoin

[beastoin]

✅ All Checkpoints Passed — Ready for Merge

Checkpoint	Status
CP0 Preflight	✅
CP1 Issue understood	✅
CP2 Workspace setup	✅
CP3 Exploration	✅
CP4 Codex consult	✅
CP5 Implementation	✅
CP6 PR created	✅
CP7 Reviewer approved	✅ (PR_APPROVED_LGTM, 2 rounds)
CP8 Tester approved	✅ (TESTS_APPROVED, 8 unit tests)
CP9A L1 live test	✅ (22/22 tests pass)
CP9B L2 integrated	✅ (UNTESTED with justification — pure data transform, no cross-service boundary)
CP9C L3 remote	N/A (level3_required=false)

Summary: Fixes #5949 — rehydrates current_session_segments from persisted transcript_segments on WS reconnect so can_assign works correctly for speaker sample extraction.

Changes: 4 files, 163 additions — utils/speaker_assignment.py (shared utility), routers/transcribe.py (rehydration call), test file + test.sh registration.

Awaiting human merge approval.

by AI for @beastoin

[github-actions]

Hey @beastoin 👋

Thank you so much for taking the time to contribute to Omi! We truly appreciate you putting in the effort to submit this pull request.

After careful review, we've decided not to merge this particular PR. Please don't take this personally — we genuinely try to merge as many contributions as possible, but sometimes we have to make tough calls based on:

• Project standards — Ensuring consistency across the codebase
• User needs — Making sure changes align with what our users need
• Code best practices — Maintaining code quality and maintainability
• Project direction — Keeping aligned with our roadmap and vision

Your contribution is still valuable to us, and we'd love to see you contribute again in the future! If you'd like feedback on how to improve this PR or want to discuss alternative approaches, please don't hesitate to reach out.

Thank you for being part of the Omi community! 💜