fix(single_brain): KI-244 — Z4 hardening — max_tokens, mark_recommendation gates, sticky-graceful retry

Three issues from Z2 5-persona test addressed:

1) Mid-session amnesia (Priya T3, Vikram T2/T4 returned "I'm sorry, I missed
a step" AFTER profile capture succeeded):
- single_brain.py: maxOutputTokens 1024 → 2048 (prose p95 ~600 + tool-call
JSON p95 ~800 + 20% margin = 1680, rounded up). Mirrors KI-150 fact_find
fix (420 → 700).
- Added finishReason=MAX_TOKENS warning log so future truncation surfaces.

2) Hallucinated closure (Priya T6 congratulated user despite zero rec;
Vikram T6 called mark_recommendation with empty policy_ids):
- brain_tools.mark_recommendation now gates BEFORE the Y2 record_policy_event
write:
(a) empty cleaned policy_ids → {recorded: False, error: 'no_policies_supplied'}
(b) non-empty ids but session.last_retrieved_chunks empty/missing →
{recorded: False, error: 'no_retrieval_history — call retrieve_policies first'}

3) Brain election bouncing (Priya bounced sales_brain → single_brain →
single_brain → sales_brain in one session):
- session_state.SessionState.single_brain_sticky: bool = False (new field).
- main.py: on first successful single_brain turn, stamp sticky=True.
On subsequent SingleBrainError after sticky, return a graceful
"Sorry, I'm having trouble — could you say that again?" reply with
brain_used="single_brain::sticky_graceful_retry" INSTEAD of falling
through to the legacy orchestrator. Non-sticky path (first turn) still
falls through.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

backend/brain_tools.py

@@ -328,6 +328,25 @@ def mark_recommendation(
             seen.add(s)
             cleaned.append(s)
 
+    # Z2 fix — Issue 2 (hallucinated closure). Vikram T6 saw the LLM emit
+    # mark_recommendation with an empty policy_ids list, the tool silently
+    # returned {"recorded": True, "policy_ids": []}, and the bot then said
+    # "I'm glad we found a good fit" despite ZERO cards shown. Two
+    # preconditions, gated BEFORE any session write so we don't poison
+    # last_recommendation_ids / shown_policies events with junk:
+    #   (a) empty (after dedup)  → no_policies_supplied
+    #   (b) non-empty BUT no retrieval history this session → caller
+    #       must run retrieve_policies first (Y2 cache check via
+    #       session.last_retrieved_chunks)
+    if not cleaned:
+        return {"recorded": False, "error": "no_policies_supplied"}
+    _retrieval_cache = getattr(session, "last_retrieved_chunks", None)
+    if not _retrieval_cache:
+        return {
+            "recorded": False,
+            "error": "no_retrieval_history — call retrieve_policies first",
+        }
+
     try:
         session.last_recommendation_ids = cleaned
     except Exception as e:  # noqa: BLE001

backend/main.py

@@ -445,6 +445,14 @@ async def chat(req: ChatRequest, request: Request):
             from backend.session_state import get_session
 
             _sb_session = get_session(session_id)
+            # Z2 fix — Issue 3 (brain election bouncing). Once a session
+            # has had ANY successful single_brain turn, it must stay on
+            # single_brain for the rest of its lifetime. Falling back to
+            # the legacy orchestrator mid-stream (which is what Priya saw
+            # 6 turns in a row) discards everything single_brain captured
+            # in last_recommendation_ids / last_retrieved_chunks /
+            # slug_to_insurer and confuses the user. Sticky check below.
+            _sb_was_sticky = getattr(_sb_session, "single_brain_sticky", False)
             try:
                 turn = await asyncio.wait_for(
                     single_brain.handle_turn(
@@ -454,23 +462,58 @@ async def chat(req: ChatRequest, request: Request):
                     ),
                     timeout=45.0,
                 )
+                # First successful single_brain turn stamps the flag so
+                # every subsequent turn on this session is locked in.
+                try:
+                    _sb_session.single_brain_sticky = True
+                except Exception:  # noqa: BLE001
+                    pass
             except single_brain.SingleBrainError as _sb_err:
-                logging.warning(
-                    "single_brain failed, falling back to orchestrator "
-                    "(session=%s): %s",
-                    session_id, _sb_err,
-                )
-                turn = await asyncio.wait_for(
-                    handle_turn(
-                        user_text=req.user_text,
-                        chat_history=req.chat_history,
-                        user_profile=req.profile,
-                        policy_filter_ids=req.policy_filter_ids,
-                        session_id=session_id,
-                        view_context=req.view_context,
-                    ),
-                    timeout=45.0,
-                )
+                if _sb_was_sticky:
+                    # Z2 belt+suspenders — session already had a clean
+                    # single_brain turn. Do NOT cross-fade to the legacy
+                    # orchestrator (loses turn state + frontend sees the
+                    # brain hop). Emit a graceful retry prompt instead.
+                    logging.warning(
+                        "single_brain failed on STICKY session (session=%s); "
+                        "emitting graceful retry, NOT falling back to "
+                        "orchestrator: %s",
+                        session_id, _sb_err,
+                    )
+                    turn = single_brain.TurnResult(
+                        reply_text=(
+                            "Sorry, I'm having trouble — could you say "
+                            "that again?"
+                        ),
+                        citations=[],
+                        retrieved_chunk_ids=[],
+                        brain_used="single_brain::sticky_graceful_retry",
+                        intent="qa",
+                        language="en",
+                        latency_ms=int((time.time() - t_chat0) * 1000),
+                        raw_reply=f"SingleBrainError: {_sb_err}",
+                        faithfulness_passed=True,
+                        faithfulness_reasons=[],
+                        blocked=False,
+                        profile_updates={},
+                    )
+                else:
+                    logging.warning(
+                        "single_brain failed, falling back to orchestrator "
+                        "(session=%s): %s",
+                        session_id, _sb_err,
+                    )
+                    turn = await asyncio.wait_for(
+                        handle_turn(
+                            user_text=req.user_text,
+                            chat_history=req.chat_history,
+                            user_profile=req.profile,
+                            policy_filter_ids=req.policy_filter_ids,
+                            session_id=session_id,
+                            view_context=req.view_context,
+                        ),
+                        timeout=45.0,
+                    )
         else:
             turn = await asyncio.wait_for(
                 handle_turn(

backend/session_state.py

@@ -74,6 +74,16 @@ class SessionState:
     # "Conversation turn" column in the admin Recommendation History panel
     # (previously showed "—" because no caller populated the field).
     turn_idx: int = 0
+    # Z2 fix — Issue 3 (brain election bouncing). Priya's session hopped
+    # sales_brain → single_brain → single_brain → sales_brain across 6
+    # turns even though USE_SINGLE_BRAIN=true; the Gemini 503 fallback
+    # (Z1 retry already softens) was dropping the session back onto the
+    # legacy orchestrator mid-stream. Belt-and-suspenders: main.py stamps
+    # this True after the FIRST successful single_brain turn, and
+    # subsequent SingleBrainError responses on the same session must NOT
+    # fall through to the orchestrator — they emit a graceful retry
+    # prompt so the session stays sticky on single_brain.
+    single_brain_sticky: bool = False
 
     def _flush(self) -> None:
         """No-op since KI-118 (2026-05-15). Disk persistence was removed; the

backend/single_brain.py

@@ -402,8 +402,17 @@ async def _gemini_call(
         "tools": [{"functionDeclarations": tools}],
         "toolConfig": {"functionCallingConfig": {"mode": "AUTO"}},
         "generationConfig": {
+            # Z2 fix — Issue 1 (mid-session amnesia). Priya T3 + Vikram T2/T4
+            # came back with the "I lost my train of thought" template even
+            # though slot capture succeeded. Root cause matches KI-150
+            # (fact_find LLM, 420 → 700): when Gemini must emit prose AND a
+            # tool-call trailer in the same response, the model hits
+            # maxOutputTokens mid-emission, the trailer truncates, and the
+            # caller falls through to the defensive reply. Budget breakdown
+            # at p95: prose ~600 tok + tool-call JSON ~800 tok + 20% margin
+            # ⇒ 1680, rounded up to a safe power-of-two-ish 2048.
             "temperature": 0.4,
-            "maxOutputTokens": 1024,
+            "maxOutputTokens": 2048,
         },
     }
     headers = {"Content-Type": "application/json"}
@@ -477,10 +486,35 @@ async def _gemini_call(
             raise SingleBrainError(last_err)
 
         try:
-            return resp.json()
+            _payload = resp.json()
         except Exception as e:  # noqa: BLE001
             raise SingleBrainError(f"Gemini malformed JSON: {e}") from e
 
+        # Z2 fix — Issue 1 truncation detector. If Gemini hit our
+        # maxOutputTokens budget the candidate's finishReason will be
+        # "MAX_TOKENS" and the tool-call trailer (if any) is likely
+        # truncated → caller will degrade to the defensive "I lost my
+        # train of thought" reply. Log a WARNING (not raise) so the turn
+        # still flows, but ops can detect a future budget regression by
+        # alerting on this log line. Swallow any shape errors — this is
+        # purely observational.
+        try:
+            _cands = _payload.get("candidates") or []
+            if _cands:
+                _fr = (_cands[0].get("finishReason") or "").upper()
+                if _fr == "MAX_TOKENS":
+                    _log.warning(
+                        "single_brain Gemini finishReason=MAX_TOKENS "
+                        "(model=%s, budget=%d) — prose+tool-call trailer "
+                        "may be truncated; raise maxOutputTokens if this "
+                        "recurs",
+                        model, body["generationConfig"]["maxOutputTokens"],
+                    )
+        except Exception:  # noqa: BLE001
+            pass
+
+        return _payload
+
     # Defensive — loop fell through without returning or raising. Should
     # be unreachable, but raise so we never silently return None.
     raise SingleBrainError(