@salesforce/afv-skills 1.14.0 → 1.16.0

package/skills/investigating-agentforce-architecture/scripts/tests/test_rest_client.py

@@ -0,0 +1,1253 @@
+"""Tests for  (cross-host redirect strips Authorization) and
+ (redact_error scrubs bearer tokens).
+
+adds: `retry_on_401` decorator + `tooling_query` / `data_query`
+helpers. Tests cover both the HTTP-401 and the INVALID_SESSION_ID-in-body
+code paths, as well as the non-auth-error propagation guarantee.
+"""
+from __future__ import annotations
+
+import email.message
+import email.utils
+import io
+import json
+import time
+import unittest
+import unittest.mock as mock
+import urllib.error
+import urllib.request
+
+from . import _bootstrap  # noqa: F401
+
+import rest_client  # type: ignore
+
+
+class RedirectHeaderStripTests(unittest.TestCase):
+    """Authorization must be dropped on cross-host redirect."""
+
+    def setUp(self) -> None:
+        self.handler = rest_client.StripAuthOnCrossHostRedirect()
+        # Minimal stand-in for the `fp` argument — the handler ignores it
+        # except when constructing HTTPError on unsupported-scheme paths.
+        self.fp = io.BytesIO(b"")
+        # Headers are a Message-like; handler reads nothing from it for
+        # the 302 path. An empty dict-like suffices.
+        self.msg = {}
+
+    def _make_request(self, url: str, auth: str = "Bearer tok123"):
+        req = urllib.request.Request(url)
+        if auth:
+            req.add_header("Authorization", auth)
+        return req
+
+    def _has_authorization(self, req) -> bool:
+        for k in list(req.headers.keys()) + list(req.unredirected_hdrs.keys()):
+            if k.lower() == "authorization":
+                return True
+        return False
+
+    def test_same_host_redirect_preserves_authorization(self):
+        req = self._make_request("https://na1.salesforce.com/services/data/v60.0/query")
+        new_req = self.handler.redirect_request(
+            req, self.fp, 302, "Found", self.msg,
+            "https://na1.salesforce.com/services/data/v60.0/query?page=2",
+        )
+        self.assertIsNotNone(new_req)
+        self.assertTrue(
+            self._has_authorization(new_req),
+            "same-host redirect must preserve Authorization",
+        )
+
+    def test_cross_host_redirect_strips_authorization(self):
+        req = self._make_request("https://na1.salesforce.com/services/data/v60.0/query")
+        new_req = self.handler.redirect_request(
+            req, self.fp, 302, "Found", self.msg,
+            "https://attacker.example.com/steal",
+        )
+        self.assertIsNotNone(new_req)
+        self.assertFalse(
+            self._has_authorization(new_req),
+            "cross-host redirect must strip Authorization ",
+        )
+
+    def test_downgrade_to_http_cross_host_strips(self):
+        """A redirect from HTTPS→HTTP on a different host is the classic
+        MITM bait. Authorization must not follow."""
+        req = self._make_request("https://na1.salesforce.com/services/data/v60.0/query")
+        new_req = self.handler.redirect_request(
+            req, self.fp, 302, "Found", self.msg,
+            "http://attacker.example.com/",
+        )
+        self.assertIsNotNone(new_req)
+        self.assertFalse(self._has_authorization(new_req))
+
+    def test_case_insensitive_host_match(self):
+        """DNS is case-insensitive — `Foo.com` and `foo.com` are same host."""
+        req = self._make_request("https://Foo.Salesforce.com/path")
+        new_req = self.handler.redirect_request(
+            req, self.fp, 302, "Found", self.msg,
+            "https://foo.salesforce.com/other",
+        )
+        self.assertIsNotNone(new_req)
+        self.assertTrue(self._has_authorization(new_req))
+
+    def test_build_opener_wires_the_handler(self):
+        opener = rest_client.build_opener()
+        self.assertIsInstance(opener, urllib.request.OpenerDirector)
+        installed = [
+            h for h in opener.handlers
+            if isinstance(h, rest_client.StripAuthOnCrossHostRedirect)
+        ]
+        self.assertEqual(
+            len(installed), 1,
+            "build_opener must install exactly one StripAuthOnCrossHostRedirect",
+        )
+
+
+class RedactErrorTests(unittest.TestCase):
+    """redact_error must scrub bearer tokens from exception strings."""
+
+    def test_bearer_token_redacted(self):
+        # Scanner-inert placeholder — redactor only cares about the
+        # `Authorization: Bearer ` prefix; the value is opaque `\S+`.
+        raw = "Authorization: Bearer TESTONLY_HEADER.TESTONLY_PAYLOAD.TESTONLY_SIG"
+        exc = RuntimeError(raw)
+        out = rest_client.redact_error(exc)
+        self.assertIn("<redacted>", out)
+        self.assertNotIn("TESTONLY_HEADER", out)
+        self.assertNotIn("TESTONLY_PAYLOAD", out)
+
+    def test_bearer_token_case_insensitive(self):
+        exc = RuntimeError("authorization: BEARER TESTONLY_BEARER_VALUE")
+        out = rest_client.redact_error(exc)
+        self.assertNotIn("TESTONLY_BEARER_VALUE", out)
+
+    def test_access_token_querystring_redacted(self):
+        exc = RuntimeError("POST /oauth/token accessToken=TESTONLY_TOKEN&foo=1")
+        out = rest_client.redact_error(exc)
+        self.assertNotIn("TESTONLY_TOKEN", out)
+        self.assertIn("<redacted>", out)
+        # The non-sensitive tail is preserved — handy for debugging.
+        self.assertIn("foo=1", out)
+
+    def test_access_token_snake_case_querystring_redacted(self):
+        exc = RuntimeError("access_token=TESTONLY_SNAKE_TOKEN&other=ok")
+        out = rest_client.redact_error(exc)
+        self.assertNotIn("TESTONLY_SNAKE_TOKEN", out)
+        self.assertIn("other=ok", out)
+
+    def test_access_token_json_redacted(self):
+        exc = RuntimeError('{"accessToken":"TESTONLY_JSON_TOKEN","status":0}')
+        out = rest_client.redact_error(exc)
+        self.assertNotIn("TESTONLY_JSON_TOKEN", out)
+        self.assertIn("<redacted>", out)
+        # Non-sensitive JSON preserved.
+        self.assertIn("status", out)
+
+    def test_access_token_json_snake_case_redacted(self):
+        exc = RuntimeError('{"access_token":"tok.ABC"}')
+        out = rest_client.redact_error(exc)
+        self.assertNotIn("tok.ABC", out)
+
+    def test_output_includes_exception_type(self):
+        exc = ValueError("boring message")
+        out = rest_client.redact_error(exc)
+        self.assertTrue(out.startswith("ValueError:"))
+
+    def test_empty_exception_safe(self):
+        exc = RuntimeError("")
+        out = rest_client.redact_error(exc)
+        self.assertTrue(out.startswith("RuntimeError:"))
+
+    def test_plain_message_unchanged(self):
+        exc = RuntimeError("connection refused")
+        out = rest_client.redact_error(exc)
+        self.assertIn("connection refused", out)
+
+
+class RedactTextPublicSymbolTests(unittest.TestCase):
+    """`redact_text` is the public symbol; `_redact_text` stays as
+    a deprecated alias for backwards compatibility.
+    """
+
+    def test_redact_text_importable_at_module_top(self):
+        """Public name must be reachable via a plain `from` import — this
+        is what sf_cli now relies on at module-top import time.
+        """
+        from rest_client import redact_text  # noqa: F401
+        self.assertTrue(callable(redact_text))
+
+    def test_redact_text_same_behavior_as_alias(self):
+        """The deprecated alias must be bound to the same function object
+        so behavior drift is impossible.
+        """
+        self.assertIs(rest_client._redact_text, rest_client.redact_text)
+
+    def test_underscore_alias_still_works(self):
+        """Backwards-compat: any lingering `from rest_client import
+        _redact_text` must keep functioning until the alias is removed
+        in a follow-up batch.
+        """
+        from rest_client import _redact_text  # noqa: F401
+        out = _redact_text("Authorization: Bearer tokXYZ")
+        self.assertIn("<redacted>", out)
+        self.assertNotIn("tokXYZ", out)
+
+    def test_public_redact_text_scrubs_bearer(self):
+        out = rest_client.redact_text("Authorization: Bearer secrettok")
+        self.assertNotIn("secrettok", out)
+        self.assertIn("<redacted>", out)
+
+
+def _make_http_error(code: int, body: bytes = b"") -> urllib.error.HTTPError:
+    """Build a minimal HTTPError carrying `code` + `body`.
+
+    We don't need a real socket — the decorator only reads `.code` and
+    (optionally) the stringified form. Body is passed for parity with
+    the wire.
+    """
+    return urllib.error.HTTPError(
+        url="https://example.salesforce.com/x",
+        code=code,
+        msg="error",
+        hdrs={},  # type: ignore[arg-type]
+        fp=io.BytesIO(body),
+    )
+
+
+class RetryOn401Tests(unittest.TestCase):
+    """decorator refreshes the token on 401 and retries once."""
+
+    def test_401_then_success_calls_refresh_once(self):
+        """Wrapped fn raises 401 first, succeeds on retry.
+
+        refresh_fn MUST be called exactly once; wrapped fn returns normally.
+        """
+        call_count = {"n": 0}
+
+        def wrapped():
+            call_count["n"] += 1
+            if call_count["n"] == 1:
+                raise _make_http_error(401)
+            return {"ok": True}
+
+        refresh_calls = []
+
+        def refresh():
+            refresh_calls.append(1)
+            return ("https://example.salesforce.com", "new_tok")
+
+        decorated = rest_client.retry_on_401(refresh)(wrapped)
+        result = decorated()
+
+        self.assertEqual(result, {"ok": True})
+        self.assertEqual(len(refresh_calls), 1, "refresh_fn must be called exactly once")
+        self.assertEqual(call_count["n"], 2)
+
+    def test_401_twice_reraises_with_original_context(self):
+        """Second 401 → original 401 context surfaces (wrapped as RestClientError).
+
+        refresh_fn still called only once (no infinite retry).
+        """
+
+        def wrapped():
+            raise _make_http_error(401, body=b"session bad")
+
+        refresh_calls = []
+
+        def refresh():
+            refresh_calls.append(1)
+            return ("https://example.salesforce.com", "new_tok")
+
+        decorated = rest_client.retry_on_401(refresh)(wrapped)
+        with self.assertRaises(rest_client.RestClientError) as ctx:
+            decorated()
+
+        self.assertIn("auth error after refresh", str(ctx.exception))
+        self.assertEqual(len(refresh_calls), 1, "refresh_fn must be called only once")
+
+    def test_invalid_session_body_triggers_refresh(self):
+        """INVALID_SESSION_ID signalled via body (HTTP 200) → refresh + retry."""
+        call_count = {"n": 0}
+
+        def wrapped():
+            call_count["n"] += 1
+            if call_count["n"] == 1:
+                raise rest_client._InvalidSessionSignal("body had INVALID_SESSION_ID")
+            return {"ok": True}
+
+        refresh_calls = []
+
+        def refresh():
+            refresh_calls.append(1)
+            return ("https://example.salesforce.com", "new_tok")
+
+        decorated = rest_client.retry_on_401(refresh)(wrapped)
+        result = decorated()
+
+        self.assertEqual(result, {"ok": True})
+        self.assertEqual(len(refresh_calls), 1)
+
+    def test_non_auth_http_error_propagates_without_refresh(self):
+        """500 → NO refresh, error propagates as-is."""
+
+        def wrapped():
+            raise _make_http_error(500)
+
+        refresh_calls = []
+
+        def refresh():
+            refresh_calls.append(1)
+            return ("https://example.salesforce.com", "new_tok")
+
+        decorated = rest_client.retry_on_401(refresh)(wrapped)
+        with self.assertRaises(urllib.error.HTTPError) as ctx:
+            decorated()
+
+        self.assertEqual(ctx.exception.code, 500)
+        self.assertEqual(len(refresh_calls), 0, "500 must NOT trigger refresh")
+
+    def test_retry_error_messages_go_through_redact_error(self):
+        """A token baked into the original 401 body must not survive in the
+        RestClientError that gets re-raised on the second 401.
+        """
+
+        def wrapped():
+            # The HTTPError's str form echoes url + code. We bake a bearer
+            # token into the msg to exercise the redact path on the
+            # re-raise.
+            err = urllib.error.HTTPError(
+                url="https://example.salesforce.com",
+                code=401,
+                msg="Authorization: Bearer TESTONLY_HTTPERROR_TOKEN",
+                hdrs={},  # type: ignore[arg-type]
+                fp=io.BytesIO(b""),
+            )
+            raise err
+
+        def refresh():
+            return ("https://example.salesforce.com", "new_tok")
+
+        decorated = rest_client.retry_on_401(refresh)(wrapped)
+        with self.assertRaises(rest_client.RestClientError) as ctx:
+            decorated()
+        msg = str(ctx.exception)
+        self.assertNotIn("TESTONLY_HTTPERROR_TOKEN", msg)
+        # The redaction sentinel must appear somewhere — proof the
+        # substitution actually ran.
+        self.assertIn("<redacted>", msg)
+
+
+class QueryHelperTests(unittest.TestCase):
+    """tooling_query + data_query — happy paths with mocked opener."""
+
+    def _install_mock_opener(self, response_body: bytes):
+        """Monkeypatch build_opener to return an opener whose .open returns
+        a context manager yielding an object with .read() == response_body.
+        """
+        fake_resp = mock.MagicMock()
+        fake_resp.read.return_value = response_body
+        fake_resp.__enter__ = mock.MagicMock(return_value=fake_resp)
+        fake_resp.__exit__ = mock.MagicMock(return_value=False)
+
+        fake_opener = mock.MagicMock()
+        fake_opener.open = mock.MagicMock(return_value=fake_resp)
+        return mock.patch.object(
+            rest_client, "build_opener", return_value=fake_opener
+        )
+
+    def test_tooling_query_happy_path(self):
+        body = json.dumps({"records": [{"Id": "01p000000000001"}]}).encode("utf-8")
+        with self._install_mock_opener(body):
+            out = rest_client.tooling_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM ApexClass",
+                api_version="v60.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+        self.assertEqual(out["records"][0]["Id"], "01p000000000001")
+
+    def test_data_query_happy_path(self):
+        body = json.dumps({"totalSize": 0, "records": []}).encode("utf-8")
+        with self._install_mock_opener(body):
+            out = rest_client.data_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM Account LIMIT 1",
+                api_version="v60.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+        self.assertEqual(out["totalSize"], 0)
+
+    def test_tooling_query_apex_body_with_marker_does_not_trigger_refresh(self):
+        """BUGFIX 2026-05-03: 200 OK Tooling response where ApexClass.Body
+        references `INVALID_SESSION_ID` in source must pass through
+        unchanged — no refresh, no retry, no error wrapping."""
+        good_body = json.dumps({
+            "size": 1,
+            "totalSize": 1,
+            "done": True,
+            "queryLocator": None,
+            "entityTypeName": "ApexClass",
+            "records": [
+                {
+                    "attributes": {"type": "ApexClass"},
+                    "Id": "01pUv000003a1YaIAI",
+                    "Name": "XCSF_FlowFaultMessage",
+                    "Body": "// if (code == 'INVALID_SESSION_ID') rethrow;",
+                }
+            ],
+        }).encode("utf-8")
+
+        open_count = {"n": 0}
+
+        def fake_open(req):
+            open_count["n"] += 1
+            resp = mock.MagicMock()
+            resp.read.return_value = good_body
+            resp.__enter__ = mock.MagicMock(return_value=resp)
+            resp.__exit__ = mock.MagicMock(return_value=False)
+            return resp
+
+        fake_opener = mock.MagicMock()
+        fake_opener.open = fake_open
+
+        refresh_calls = []
+
+        def refresh():
+            refresh_calls.append(1)
+            return ("https://example.salesforce.com", "new_tok")
+
+        with mock.patch.object(rest_client, "build_opener", return_value=fake_opener):
+            out = rest_client.tooling_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id, Name, Body FROM ApexClass WHERE Id IN (...)",
+                api_version="v66.0",
+                on_401_refresh=refresh,
+            )
+
+        self.assertEqual(out["totalSize"], 1)
+        self.assertEqual(out["records"][0]["Name"], "XCSF_FlowFaultMessage")
+        self.assertEqual(open_count["n"], 1,
+                         "exactly one HTTP call — no spurious retry")
+        self.assertEqual(refresh_calls, [],
+                         "refresh must not fire on a 200 OK success payload")
+
+    def test_tooling_query_invalid_session_in_body_triggers_refresh(self):
+        """HTTP 200 with INVALID_SESSION_ID body → refresh closure invoked."""
+        bad_body = json.dumps([
+            {"errorCode": "INVALID_SESSION_ID", "message": "Session expired"}
+        ]).encode("utf-8")
+        good_body = json.dumps({"records": []}).encode("utf-8")
+
+        # Two-shot opener: first call returns bad_body, second returns good.
+        responses = [bad_body, good_body]
+
+        def fake_open(req):
+            resp = mock.MagicMock()
+            resp.read.return_value = responses.pop(0)
+            resp.__enter__ = mock.MagicMock(return_value=resp)
+            resp.__exit__ = mock.MagicMock(return_value=False)
+            return resp
+
+        fake_opener = mock.MagicMock()
+        fake_opener.open = fake_open
+
+        refresh_calls = []
+
+        def refresh():
+            refresh_calls.append(1)
+            return ("https://example.salesforce.com", "new_tok")
+
+        with mock.patch.object(rest_client, "build_opener", return_value=fake_opener):
+            out = rest_client.tooling_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM ApexClass",
+                api_version="v60.0",
+                on_401_refresh=refresh,
+            )
+
+        self.assertEqual(out, {"records": []})
+        self.assertEqual(len(refresh_calls), 1)
+
+
+class RetryOn401CredentialRefreshIntegrationTests(unittest.TestCase):
+    """Verify the retry carries the NEW token, not the stale one.
+
+    The previous design looked correct under unit tests (refresh_fn called
+    once, retry returns success) but was broken end-to-end: the closure
+    around `(instance_url, token)` was captured at decoration time, so
+    `refresh_fn`'s return value was thrown away and the retry saw the
+    SAME stale credentials. Any real 401 (bad token) would hit the retry
+    with the same bad token and 401 again.
+
+    This test exercises a full tooling_query through a mock opener that
+    401s for the stale token and succeeds for the refreshed token. A
+    fix that only changes decorator bookkeeping (without threading the
+    refresh into the retry's actual HTTP call) will fail this test.
+    """
+
+    def test_retry_on_401_uses_refreshed_credentials(self):
+        # creds cell — a single-slot list mutated by refresh_fn.
+        creds_cell: list = [("https://example.salesforce.com", "bad_token")]
+
+        def creds_provider():
+            return creds_cell[0]
+
+        def refresh_fn():
+            creds_cell[0] = ("https://example.salesforce.com", "good_token")
+            return creds_cell[0]
+
+        # Tracks the Authorization header seen on every request — this is
+        # what proves whether the retry carried the refreshed token or not.
+        auth_headers_seen: list[str] = []
+
+        def fake_open(req):
+            # urllib normalizes header names on add_header; read whichever
+            # case lands.
+            auth = None
+            for k, v in req.headers.items():
+                if k.lower() == "authorization":
+                    auth = v
+                    break
+            auth_headers_seen.append(auth or "")
+
+            if auth == "Bearer bad_token":
+                raise urllib.error.HTTPError(
+                    url=req.full_url,
+                    code=401,
+                    msg="Unauthorized",
+                    hdrs={},  # type: ignore[arg-type]
+                    fp=io.BytesIO(b""),
+                )
+            # Good token path — return a tiny success JSON.
+            resp = mock.MagicMock()
+            resp.read.return_value = json.dumps({"records": []}).encode("utf-8")
+            resp.__enter__ = mock.MagicMock(return_value=resp)
+            resp.__exit__ = mock.MagicMock(return_value=False)
+            return resp
+
+        fake_opener = mock.MagicMock()
+        fake_opener.open = fake_open
+
+        with mock.patch.object(rest_client, "build_opener", return_value=fake_opener):
+            out = rest_client.tooling_query(
+                creds_provider,
+                "SELECT Id FROM ApexClass",
+                api_version="v60.0",
+                on_401_refresh=refresh_fn,
+            )
+
+        self.assertEqual(out, {"records": []})
+        self.assertEqual(
+            len(auth_headers_seen), 2,
+            f"expected 2 attempts (401 + retry), got {len(auth_headers_seen)}",
+        )
+        self.assertEqual(
+            auth_headers_seen[0], "Bearer bad_token",
+            "first attempt should have used the stale token",
+        )
+        self.assertEqual(
+            auth_headers_seen[1], "Bearer good_token",
+            "retry MUST use the refreshed token — this is the bug",
+        )
+
+
+class TransientHttpWithRefreshIntegrationTests(unittest.TestCase):
+    """End-to-end stack test through `tooling_query`.
+
+    Drives the real decorator chain baked into `tooling_query`: first
+    response 429 (transient retry), second response 401 (refresh + retry),
+    third response 200. Proves that the stacking in the production
+    helper — NOT just the raw decorators — gives the intended behavior.
+    """
+
+    def test_tooling_query_429_then_401_then_success(self):
+        responses = [
+            ("429", None),
+            ("401", None),
+            ("ok", json.dumps({"records": [{"Id": "a"}]}).encode("utf-8")),
+        ]
+        creds_cell = [("https://example.salesforce.com", "bad_token")]
+        refresh_calls = []
+
+        def creds_provider():
+            return creds_cell[0]
+
+        def refresh():
+            refresh_calls.append(1)
+            creds_cell[0] = ("https://example.salesforce.com", "good_token")
+            return creds_cell[0]
+
+        def fake_open(req):
+            kind, body = responses.pop(0)
+            if kind == "429":
+                exc = urllib.error.HTTPError(
+                    url=req.full_url,
+                    code=429,
+                    msg="rate limited",
+                    hdrs=email.message.Message(),
+                    fp=io.BytesIO(b""),
+                )
+                exc.headers["Retry-After"] = "0"  # 0 → base_delay wins, but sleep is mocked
+                raise exc
+            if kind == "401":
+                raise urllib.error.HTTPError(
+                    url=req.full_url,
+                    code=401,
+                    msg="unauthorized",
+                    hdrs={},  # type: ignore[arg-type]
+                    fp=io.BytesIO(b""),
+                )
+            resp = mock.MagicMock()
+            resp.read.return_value = body
+            resp.__enter__ = mock.MagicMock(return_value=resp)
+            resp.__exit__ = mock.MagicMock(return_value=False)
+            return resp
+
+        fake_opener = mock.MagicMock()
+        fake_opener.open = fake_open
+
+        with mock.patch.object(rest_client, "build_opener", return_value=fake_opener), \
+             mock.patch.object(rest_client.time, "sleep"):
+            out = rest_client.tooling_query(
+                creds_provider,
+                "SELECT Id FROM ApexClass",
+                api_version="v60.0",
+                on_401_refresh=refresh,
+            )
+
+        self.assertEqual(out["records"][0]["Id"], "a")
+        self.assertEqual(len(refresh_calls), 1,
+                         "exactly one refresh — 429 must NOT trigger the refresh")
+        self.assertEqual(responses, [], "all three canned responses consumed")
+
+
+class AcceptEncodingAndForbiddenInvalidSessionTests(unittest.TestCase):
+    """Accept-Encoding: identity + 403+INVALID_SESSION_ID handling."""
+
+    def test_query_sends_accept_encoding_identity(self):
+        """The outbound request carries `Accept-Encoding: identity` so no
+        gzip body ever lands on the body-inspection path."""
+        body = json.dumps({"records": []}).encode("utf-8")
+
+        captured: list = []
+
+        def fake_open(req):
+            captured.append(req)
+            resp = mock.MagicMock()
+            resp.read.return_value = body
+            resp.__enter__ = mock.MagicMock(return_value=resp)
+            resp.__exit__ = mock.MagicMock(return_value=False)
+            return resp
+
+        fake_opener = mock.MagicMock()
+        fake_opener.open = fake_open
+
+        with mock.patch.object(rest_client, "build_opener", return_value=fake_opener):
+            rest_client.tooling_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM ApexClass",
+                api_version="v60.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+
+        self.assertEqual(len(captured), 1)
+        req = captured[0]
+        # urllib stores add_header'd names via .capitalize(); read any case.
+        accept_enc = None
+        for k, v in req.headers.items():
+            if k.lower() == "accept-encoding":
+                accept_enc = v
+                break
+        self.assertEqual(accept_enc, "identity",
+                         "Accept-Encoding must be 'identity'")
+
+    def test_403_invalid_session_triggers_refresh(self):
+        """HTTP 403 + INVALID_SESSION_ID body is treated as refreshable."""
+
+        call_count = {"n": 0}
+
+        def wrapped():
+            call_count["n"] += 1
+            if call_count["n"] == 1:
+                # 403 with INVALID_SESSION_ID body → should refresh.
+                raise urllib.error.HTTPError(
+                    url="https://example.salesforce.com/x",
+                    code=403,
+                    msg="Forbidden",
+                    hdrs={},  # type: ignore[arg-type]
+                    fp=io.BytesIO(
+                        b'[{"errorCode":"INVALID_SESSION_ID","message":"Session expired"}]'
+                    ),
+                )
+            return {"ok": True}
+
+        refresh_calls = []
+
+        def refresh():
+            refresh_calls.append(1)
+            return ("https://example.salesforce.com", "new_tok")
+
+        decorated = rest_client.retry_on_401(refresh)(wrapped)
+        result = decorated()
+
+        self.assertEqual(result, {"ok": True})
+        self.assertEqual(len(refresh_calls), 1,
+                         "403+INVALID_SESSION_ID must trigger refresh")
+
+    def test_403_non_invalid_session_does_not_trigger_refresh(self):
+        """HTTP 403 with some OTHER errorCode propagates unchanged."""
+
+        def wrapped():
+            raise urllib.error.HTTPError(
+                url="https://example.salesforce.com/x",
+                code=403,
+                msg="Forbidden",
+                hdrs={},  # type: ignore[arg-type]
+                fp=io.BytesIO(
+                    b'[{"errorCode":"INSUFFICIENT_ACCESS","message":"Denied"}]'
+                ),
+            )
+
+        refresh_calls = []
+
+        def refresh():
+            refresh_calls.append(1)
+            return ("https://example.salesforce.com", "new_tok")
+
+        decorated = rest_client.retry_on_401(refresh)(wrapped)
+        with self.assertRaises(urllib.error.HTTPError) as ctx:
+            decorated()
+
+        self.assertEqual(ctx.exception.code, 403)
+        self.assertEqual(len(refresh_calls), 0,
+                         "non-auth 403 must NOT burn a refresh")
+
+
+class RetryOnTransientHttpTests(unittest.TestCase):
+    """429/503 → exponential-backoff retry up to max_retries."""
+
+    def test_429_then_success_retries_once(self):
+        calls = {"n": 0}
+
+        def fn():
+            calls["n"] += 1
+            if calls["n"] == 1:
+                raise _make_http_error(429)
+            return {"ok": True}
+
+        with mock.patch.object(rest_client.time, "sleep") as slept:
+            decorated = rest_client.retry_on_transient_http()(fn)
+            out = decorated()
+
+        self.assertEqual(out, {"ok": True})
+        self.assertEqual(calls["n"], 2)
+        slept.assert_called_once()
+        # First retry: base_delay * 2**0 = 1.0
+        self.assertAlmostEqual(slept.call_args[0][0], 1.0)
+
+    def test_503_three_times_propagates_httperror(self):
+        """With max_retries=3, total attempts = 4. Four 503s must surface
+        the HTTPError after exactly 3 sleeps."""
+        def fn():
+            raise _make_http_error(503)
+
+        with mock.patch.object(rest_client.time, "sleep") as slept:
+            decorated = rest_client.retry_on_transient_http(max_retries=3)(fn)
+            with self.assertRaises(urllib.error.HTTPError) as ctx:
+                decorated()
+
+        self.assertEqual(ctx.exception.code, 503)
+        self.assertEqual(slept.call_count, 3,
+                         "3 sleeps for 3 retries before the final failure")
+
+    def test_retry_after_header_honored_when_larger(self):
+        """Retry-After: 10 overrides the 1s exponential schedule."""
+        def fn():
+            exc = urllib.error.HTTPError(
+                url="https://example.salesforce.com/x",
+                code=429,
+                msg="rate limited",
+                hdrs=email.message.Message(),
+                fp=io.BytesIO(b""),
+            )
+            exc.headers["Retry-After"] = "10"
+            raise exc
+
+        # fn always raises → exhaust retries; we just need to observe the
+        # first sleep value.
+        with mock.patch.object(rest_client.time, "sleep") as slept:
+            decorated = rest_client.retry_on_transient_http(max_retries=1)(fn)
+            with self.assertRaises(urllib.error.HTTPError):
+                decorated()
+
+        self.assertEqual(slept.call_count, 1)
+        self.assertAlmostEqual(slept.call_args_list[0][0][0], 10.0,
+                               msg="Retry-After (10) must win over base_delay (1)")
+
+    def test_retry_after_smaller_than_base_loses(self):
+        """Retry-After: 0.5 with base_delay=1.0 → sleep 1.0 (base_delay wins)."""
+        def fn():
+            exc = urllib.error.HTTPError(
+                url="https://example.salesforce.com/x",
+                code=429,
+                msg="rate limited",
+                hdrs=email.message.Message(),
+                fp=io.BytesIO(b""),
+            )
+            exc.headers["Retry-After"] = "0.5"
+            raise exc
+
+        with mock.patch.object(rest_client.time, "sleep") as slept:
+            decorated = rest_client.retry_on_transient_http(max_retries=1)(fn)
+            with self.assertRaises(urllib.error.HTTPError):
+                decorated()
+
+        self.assertAlmostEqual(slept.call_args_list[0][0][0], 1.0)
+
+    def test_500_propagates_without_sleep(self):
+        """Non-429/503 HTTP errors are out of scope — no retry, no sleep."""
+        def fn():
+            raise _make_http_error(500)
+
+        with mock.patch.object(rest_client.time, "sleep") as slept:
+            decorated = rest_client.retry_on_transient_http()(fn)
+            with self.assertRaises(urllib.error.HTTPError) as ctx:
+                decorated()
+
+        self.assertEqual(ctx.exception.code, 500)
+        slept.assert_not_called()
+
+    def test_401_propagates_without_sleep(self):
+        """401 is retry_on_401's job — this decorator must pass it through."""
+        def fn():
+            raise _make_http_error(401)
+
+        with mock.patch.object(rest_client.time, "sleep") as slept:
+            decorated = rest_client.retry_on_transient_http()(fn)
+            with self.assertRaises(urllib.error.HTTPError) as ctx:
+                decorated()
+
+        self.assertEqual(ctx.exception.code, 401)
+        slept.assert_not_called()
+
+    def test_stacks_correctly_with_retry_on_401(self):
+        """Integration: 429 → transient retry → 401 → refresh retry → success.
+
+        This verifies the documented ordering in `tooling_query` /
+        `data_query`: transient-HTTP is inner, 401-refresh is outer. If
+        the layering were reversed, the refresh would fire on the 429.
+        """
+        sequence = ["429", "401", "ok"]
+        refresh_calls = []
+
+        def fn():
+            code = sequence.pop(0)
+            if code == "429":
+                raise _make_http_error(429)
+            if code == "401":
+                raise _make_http_error(401)
+            return {"ok": True}
+
+        def refresh():
+            refresh_calls.append(1)
+            return ("https://example.salesforce.com", "new")
+
+        with mock.patch.object(rest_client.time, "sleep"):
+            # Inner: transient retry. Outer: 401 refresh. Matches the
+            # stack in tooling_query / data_query.
+            decorated = rest_client.retry_on_401(refresh)(
+                rest_client.retry_on_transient_http()(fn)
+            )
+            out = decorated()
+
+        self.assertEqual(out, {"ok": True})
+        self.assertEqual(len(refresh_calls), 1,
+                         "exactly one refresh — 429 must not burn a refresh")
+        self.assertEqual(sequence, [], "all three responses consumed")
+
+    def test_retry_after_http_date_parses(self):
+        """HTTP-date form of Retry-After is accepted.
+
+        We pass an RFC 7231 date roughly 5 seconds in the future. The
+        parsed delta should be >= 0 and, since base_delay=1.0, the delta
+        wins whenever it exceeds 1s. We accept any value >= 1.0 to avoid
+        flakiness on slow test machines.
+        """
+        future_ts = time.time() + 30  # 30s out, plenty of slack
+        http_date = email.utils.formatdate(future_ts, usegmt=True)
+
+        def fn():
+            exc = urllib.error.HTTPError(
+                url="https://example.salesforce.com/x",
+                code=503,
+                msg="unavailable",
+                hdrs=email.message.Message(),
+                fp=io.BytesIO(b""),
+            )
+            exc.headers["Retry-After"] = http_date
+            raise exc
+
+        with mock.patch.object(rest_client.time, "sleep") as slept:
+            decorated = rest_client.retry_on_transient_http(max_retries=1)(fn)
+            with self.assertRaises(urllib.error.HTTPError):
+                decorated()
+
+        self.assertEqual(slept.call_count, 1)
+        delay = slept.call_args_list[0][0][0]
+        # Sanity: must be larger than base_delay (1.0) and no bigger than
+        # the original offset (30) plus generous test-timing slack.
+        self.assertGreaterEqual(delay, 1.0)
+        self.assertLessEqual(delay, 35.0)
+
+    def test_malformed_retry_after_falls_back_to_base(self):
+        """`Retry-After: not-a-date` → base_delay wins."""
+        def fn():
+            exc = urllib.error.HTTPError(
+                url="https://example.salesforce.com/x",
+                code=429,
+                msg="rate limited",
+                hdrs=email.message.Message(),
+                fp=io.BytesIO(b""),
+            )
+            exc.headers["Retry-After"] = "not-a-date"
+            raise exc
+
+        with mock.patch.object(rest_client.time, "sleep") as slept:
+            decorated = rest_client.retry_on_transient_http(max_retries=1, base_delay=2.0)(fn)
+            with self.assertRaises(urllib.error.HTTPError):
+                decorated()
+
+        self.assertAlmostEqual(slept.call_args_list[0][0][0], 2.0,
+                               msg="malformed Retry-After falls back to base_delay")
+
+
+class InvalidSessionDetectionTests(unittest.TestCase):
+    """Unit coverage for the body-inspection helper."""
+
+    def test_list_with_invalid_session_detected(self):
+        body = [{"errorCode": "INVALID_SESSION_ID", "message": "foo"}]
+        self.assertTrue(rest_client._body_indicates_invalid_session(body))
+
+    def test_dict_without_invalid_session_not_detected(self):
+        body = {"records": [{"Id": "a"}]}
+        self.assertFalse(rest_client._body_indicates_invalid_session(body))
+
+    def test_none_safe(self):
+        self.assertFalse(rest_client._body_indicates_invalid_session(None))
+
+    def test_string_with_marker(self):
+        self.assertTrue(rest_client._body_indicates_invalid_session(
+            "...INVALID_SESSION_ID..."
+        ))
+
+    # --- BUGFIX 2026-05-03: false positive on ApexClass.Body substring -----
+    # Regression: _body_indicates_invalid_session used to walk ALL dict
+    # values and substring-match. A 200 OK Tooling Query response whose
+    # records carried Apex source referencing the string INVALID_SESSION_ID
+    # (e.g. `XCSF_FlowFaultMessage`, `SkillRulesMatchAction`) was
+    # misclassified as a session error, triggering spurious Wave B retries
+    # and a bogus `INVALID_SESSION_ID after refresh` entry in `_unresolved`.
+
+    def test_success_body_with_apex_source_substring_not_detected(self):
+        """200 OK Tooling response whose ApexClass.Body mentions
+        INVALID_SESSION_ID in source code must NOT be flagged."""
+        apex_source = (
+            "public class XCSF_FlowFaultMessage {\n"
+            "    // Handles INVALID_SESSION_ID by rethrowing a typed error.\n"
+            "    public static void raise(String code) {\n"
+            "        if (code == 'INVALID_SESSION_ID') {\n"
+            "            throw new HandledException('session expired');\n"
+            "        }\n"
+            "    }\n"
+            "}"
+        )
+        body = {
+            "size": 2,
+            "totalSize": 2,
+            "done": True,
+            "queryLocator": None,
+            "entityTypeName": "ApexClass",
+            "records": [
+                {
+                    "attributes": {
+                        "type": "ApexClass",
+                        "url": "/services/data/v66.0/tooling/sobjects/ApexClass/01pUv0",
+                    },
+                    "Id": "01pUv000003a1YaIAI",
+                    "Name": "XCSF_FlowFaultMessage",
+                    "Body": apex_source,
+                },
+                {
+                    "attributes": {"type": "ApexClass"},
+                    "Id": "01pUv000003a1YbIAI",
+                    "Name": "SkillRulesMatchAction",
+                    "Body": "// references INVALID_SESSION_ID constant",
+                },
+            ],
+        }
+        self.assertFalse(
+            rest_client._body_indicates_invalid_session(body),
+            "success payload with INVALID_SESSION_ID inside ApexClass.Body "
+            "must not be misclassified as an auth failure",
+        )
+
+    def test_error_envelope_under_errors_key_still_detected(self):
+        """Dict-wrapped error envelope under a known key still flags."""
+        body = {
+            "errors": [
+                {"errorCode": "INVALID_SESSION_ID", "message": "Session expired"}
+            ]
+        }
+        self.assertTrue(rest_client._body_indicates_invalid_session(body))
+
+    def test_success_body_with_marker_in_data_field_not_detected(self):
+        """A top-level data field whose string value mentions the marker
+        must not trigger — only error-envelope keys are walked."""
+        body = {"records": [{"Id": "a", "Description": "handles INVALID_SESSION_ID"}]}
+        self.assertFalse(rest_client._body_indicates_invalid_session(body))
+
+
+class QueryWireShapeTests(unittest.TestCase):
+    """`_query_once` must send GET with a urlencoded `q=`
+    querystring. The prior POST-with-JSON-body shape returned HTTP 405
+    on every real-org run — three reference fixtures failed
+    end-to-end. These tests pin the wire shape so a regression back to
+    POST would break CI before landing.
+    """
+
+    def _install_capturing_opener(self, captured: list, response_body: bytes):
+        def fake_open(req):
+            captured.append(req)
+            resp = mock.MagicMock()
+            resp.read.return_value = response_body
+            resp.__enter__ = mock.MagicMock(return_value=resp)
+            resp.__exit__ = mock.MagicMock(return_value=False)
+            return resp
+
+        fake_opener = mock.MagicMock()
+        fake_opener.open = fake_open
+        return mock.patch.object(
+            rest_client, "build_opener", return_value=fake_opener
+        )
+
+    def test_data_query_sends_get_with_urlencoded_querystring(self):
+        captured: list = []
+        body = json.dumps({"records": []}).encode("utf-8")
+        soql = "SELECT Id FROM BotDefinition WHERE DeveloperName = 'Foo'"
+
+        with self._install_capturing_opener(captured, body):
+            rest_client.data_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                soql,
+                api_version="v60.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+
+        self.assertEqual(len(captured), 1)
+        req = captured[0]
+        # method must be GET. Previously POST.
+        self.assertEqual(req.get_method(), "GET",
+                         "REST Query endpoint is GET-only; POST returns 405")
+        # GET has no body.
+        self.assertIsNone(req.data,
+                          "GET requests must carry no body")
+        # URL must carry the urlencoded querystring.
+        self.assertIn("/services/data/v60.0/query/?q=", req.full_url)
+        # urllib.parse.urlencode produces `+` for spaces — decoding
+        # should round-trip the SOQL exactly.
+        import urllib.parse
+        parsed = urllib.parse.urlparse(req.full_url)
+        qs = urllib.parse.parse_qs(parsed.query)
+        self.assertEqual(qs.get("q"), [soql],
+                         "urlencoded `q` must round-trip the SOQL exactly")
+
+    def test_tooling_query_sends_get_to_tooling_path(self):
+        captured: list = []
+        body = json.dumps({"records": []}).encode("utf-8")
+        soql = "SELECT Id, DeveloperName FROM BotVersion"
+
+        with self._install_capturing_opener(captured, body):
+            rest_client.tooling_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                soql,
+                api_version="v60.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+
+        self.assertEqual(len(captured), 1)
+        req = captured[0]
+        self.assertEqual(req.get_method(), "GET")
+        self.assertIsNone(req.data)
+        # Tooling path is distinct from Data API path.
+        self.assertIn("/services/data/v60.0/tooling/query/?q=", req.full_url)
+
+    def test_query_does_not_send_content_type_header(self):
+        """GET has no body — Content-Type is inappropriate and was
+        removed alongside the method change. Some edge proxies complain
+        about Content-Type on a bodyless GET; we just drop it."""
+        captured: list = []
+        body = json.dumps({"records": []}).encode("utf-8")
+
+        with self._install_capturing_opener(captured, body):
+            rest_client.data_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM Account LIMIT 1",
+                api_version="v60.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+
+        req = captured[0]
+        content_type = None
+        for k, v in req.headers.items():
+            if k.lower() == "content-type":
+                content_type = v
+                break
+        self.assertIsNone(content_type,
+                          "GET requests must not carry Content-Type")
+
+    def test_soql_with_special_chars_is_urlencoded(self):
+        """SOQL string literals contain single quotes, commas, spaces,
+        parentheses. urlencode must escape them all; a naive concat
+        would break either the URL parse or the SOQL semantics."""
+        captured: list = []
+        body = json.dumps({"records": []}).encode("utf-8")
+        # Exercise the full SOQL character classes that routinely appear.
+        soql = (
+            "SELECT Id, Name FROM ApexClass "
+            "WHERE Name IN ('Foo', 'Bar Baz') "
+            "AND NamespacePrefix IS NULL LIMIT 200"
+        )
+
+        with self._install_capturing_opener(captured, body):
+            rest_client.data_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                soql,
+                api_version="v60.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+
+        req = captured[0]
+        import urllib.parse
+        parsed = urllib.parse.urlparse(req.full_url)
+        qs = urllib.parse.parse_qs(parsed.query)
+        self.assertEqual(qs.get("q"), [soql])
+        # Spaces MUST be escaped — a raw space in a URL is malformed.
+        self.assertNotIn(" ", parsed.query,
+                         "querystring must contain no literal spaces")
+
+
+class ApiVersionRequiredKwargTests(unittest.TestCase):
+    """`api_version` is a REQUIRED keyword-only
+    argument on `tooling_query` + `data_query`.
+
+    The prior design hardcoded `v60.0`. Real orgs (my-org-alias,
+    my-perf-org-alias) run on v66 and expose fields v60 does not — confirmed
+    empirically: `BotDefinition.Description` resolves on v66, raises
+    `INVALID_FIELD` on v60 for the same org. Making `api_version`
+    required surfaces a missed call-site as a TypeError at call time,
+    not a silent regression back to a stale pinned version.
+    """
+
+    def test_tooling_query_requires_api_version(self):
+        with self.assertRaises(TypeError) as ctx:
+            rest_client.tooling_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM ApexClass",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+        self.assertIn("api_version", str(ctx.exception))
+
+    def test_data_query_requires_api_version(self):
+        with self.assertRaises(TypeError) as ctx:
+            rest_client.data_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM Account LIMIT 1",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+        self.assertIn("api_version", str(ctx.exception))
+
+
+class ApiVersionUrlSubstitutionTests(unittest.TestCase):
+    """the `api_version` kwarg lands in the URL path
+    unchanged. A caller passing `v66.0` must produce
+    `/services/data/v66.0/...` in the outbound Request, NOT the old
+    hardcoded `v60.0`.
+
+    Regression shield: if a future refactor re-pins the version string,
+    these tests fail at the wire-shape assertion even when unit-level
+    behavior of `tooling_query` still looks correct.
+    """
+
+    def _install_capturing_opener(self, captured: list, response_body: bytes):
+        def fake_open(req):
+            captured.append(req)
+            resp = mock.MagicMock()
+            resp.read.return_value = response_body
+            resp.__enter__ = mock.MagicMock(return_value=resp)
+            resp.__exit__ = mock.MagicMock(return_value=False)
+            return resp
+
+        fake_opener = mock.MagicMock()
+        fake_opener.open = fake_open
+        return mock.patch.object(
+            rest_client, "build_opener", return_value=fake_opener
+        )
+
+    def test_tooling_query_v66_in_url(self):
+        captured: list = []
+        body = json.dumps({"records": []}).encode("utf-8")
+        with self._install_capturing_opener(captured, body):
+            rest_client.tooling_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM ApexClass",
+                api_version="v66.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+        self.assertEqual(len(captured), 1)
+        url = captured[0].full_url
+        self.assertIn("/services/data/v66.0/tooling/query/?q=", url)
+        # And the old hardcoded version must NOT appear — otherwise the
+        # substitution silently failed.
+        self.assertNotIn("v60.0", url)
+
+    def test_data_query_v66_in_url(self):
+        captured: list = []
+        body = json.dumps({"records": []}).encode("utf-8")
+        with self._install_capturing_opener(captured, body):
+            rest_client.data_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM BotDefinition LIMIT 1",
+                api_version="v66.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+        url = captured[0].full_url
+        self.assertIn("/services/data/v66.0/query/?q=", url)
+        self.assertNotIn("v60.0", url)
+
+    def test_different_api_version_per_call(self):
+        """Two calls with different `api_version` land on distinct URLs —
+        proves the value is NOT cached in module-level state between
+        invocations."""
+        captured: list = []
+        body = json.dumps({"records": []}).encode("utf-8")
+        with self._install_capturing_opener(captured, body):
+            rest_client.tooling_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM ApexClass",
+                api_version="v60.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+            rest_client.tooling_query(
+                rest_client.static_creds("https://example.salesforce.com", "tok"),
+                "SELECT Id FROM ApexClass",
+                api_version="v66.0",
+                on_401_refresh=lambda: ("https://example.salesforce.com", "new"),
+            )
+        self.assertEqual(len(captured), 2)
+        self.assertIn("/services/data/v60.0/tooling/query/", captured[0].full_url)
+        self.assertIn("/services/data/v66.0/tooling/query/", captured[1].full_url)
+
+
+if __name__ == "__main__":
+    unittest.main()