From cc93053b42ea98713cf46192ec340682a3728862 Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Sat, 23 May 2026 02:45:25 -0700
Subject: [PATCH] fix(xai-oauth): apply WKE disambiguator to recovery-path
 catch-all (#29344)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_recover_with_credential_pool had a second classification site that blanket-
treated any 403 against xai-oauth as entitlement (defense-in-depth for
#26847).  That override defeated the new _is_entitlement_failure
disambiguator from the parent commit — bad-credentials 403s still
short-circuited the refresh path.

Apply the same WKE-unauthenticated / OAuth2-validation-phrase guard at
the override site so xAI's authoritative 'this is auth, not entitlement'
signal wins there too.  The #26847 catch-all still triggers for genuine
entitlement bodies that don't carry the disambiguator.

Closes the end-to-end gap exposed by
test_recover_with_credential_pool_refreshes_on_xai_bad_credentials_403.
---
 agent/agent_runtime_helpers.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py
index 6fee1baeca9..4175f3e1898 100644
--- a/agent/agent_runtime_helpers.py
+++ b/agent/agent_runtime_helpers.py
@@ -617,9 +617,28 @@ def recover_with_credential_pool(
         # existing entitlement keyword set in ``_is_entitlement_failure``.
         # Any 403 against ``xai-oauth`` is treated as entitlement here so
         # the refresh loop can't spin in those cases either.
+        #
+        # Exception (#29344): xAI's ``[WKE=unauthenticated:...]`` suffix and
+        # the ``OAuth2 access token could not be validated`` phrasing are
+        # xAI's authoritative "this is a stale token, not entitlement"
+        # signal.  When either fires we must NOT apply the catch-all
+        # override — refresh is the recoverable path for these bodies, and
+        # blanket-classifying them as entitlement was the bug that left
+        # long-running TUI sessions stuck on stale tokens until the user
+        # exited and reopened.
         is_entitlement = agent._is_entitlement_failure(error_context, status_code)
         if not is_entitlement and status_code == 403 and (agent.provider or "") == "xai-oauth":
-            is_entitlement = True
+            _disambiguator_haystack = " ".join(
+                str(error_context.get(k) or "").lower()
+                for k in ("message", "reason", "code", "error")
+                if isinstance(error_context, dict)
+            )
+            _is_xai_auth_failure = (
+                "[wke=unauthenticated:" in _disambiguator_haystack
+                or "oauth2 access token could not be validated" in _disambiguator_haystack
+            )
+            if not _is_xai_auth_failure:
+                is_entitlement = True
         if is_entitlement:
             _ra().logger.info(
                 "Credential %s — entitlement-shaped 403 from %s; "