Fix background-thread queue auth and error recovery

integrations/vercel-celery/tests/unit/test_celery_broker.py

@@ -8,6 +8,7 @@
 import json
 import logging
 import sys
+import threading
 from collections.abc import AsyncIterator, Iterator
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
@@ -27,13 +28,16 @@
 import vercel.integrations.celery._broker as vqs_celery
 from vercel.headers import get_headers, set_headers
 from vercel.queue import (
+    CommunicationError as QueueCommunicationError,
     Duration,
     Handoff,
     Message,
     MessageMetadata,
     RetryAfter,
     SanitizedName,
+    TokenResolutionError,
     Topic,
+    UnauthorizedError,
 )
 from vercel.queue._internal import subscribers as queue_subscribers
 
@@ -125,9 +129,11 @@ class FakeSyncQueueClient:
     def __init__(self, **kwargs: Any) -> None:
         self.kwargs = kwargs
         self.sent: list[dict[str, Any]] = []
+        self.send_headers: list[dict[str, str] | None] = []
         self.message_batches: list[list[FakeMessage]] = []
         self.acknowledged: list[MessageMetadata] = []
         self.ack_headers: list[dict[str, str] | None] = []
+        self.poll_headers: list[dict[str, str] | None] = []
         self.visibility_changes: list[tuple[MessageMetadata, Duration]] = []
         self.visibility_headers: list[dict[str, str] | None] = []
         self.lease_renewals: list[FakeLeaseRenewal] = []
@@ -138,6 +144,8 @@ def __init__(self, **kwargs: Any) -> None:
         FakeSyncQueueClient.instances.append(self)
 
     def send(self, topic: str, payload: dict[str, Any], **kwargs: Any) -> None:
+        headers = get_headers()
+        self.send_headers.append(dict(headers) if headers is not None else None)
         self.sent.append({"topic": topic, "payload": payload, "kwargs": kwargs})
 
     def poll(
@@ -146,6 +154,8 @@ def poll(
         consumer_group: str,
         **kwargs: Any,
     ) -> Iterator[FakeDelivery]:
+        headers = get_headers()
+        self.poll_headers.append(dict(headers) if headers is not None else None)
         batch = self.message_batches.pop(0) if self.message_batches else []
         self.last_topic = topic
         self.last_consumer_group = consumer_group
@@ -172,13 +182,18 @@ def extend_lease(
         self.visibility_headers.append(dict(headers) if headers is not None else None)
         self.visibility_changes.append((metadata, duration))
 
+    def retry_after(
+        self,
+        message: Message[dict[str, Any]] | MessageMetadata,
+        delay: Duration,
+    ) -> None:
+        self.extend_lease(message, delay)
+
     def run_lease_renewal(
         self,
         message: Message[dict[str, Any]],
         lease_duration: Duration | None = None,
-        headers_context: Any | None = None,
     ) -> FakeLeaseRenewal:
-        del headers_context
         renewal = FakeLeaseRenewal(message=message, lease_duration=lease_duration)
         self.lease_renewals.append(renewal)
         return renewal
@@ -257,6 +272,7 @@ def clean_celery_integration_state() -> Iterator[None]:
         vqs_celery._push_channels.clear()
         vqs_celery._finalize_hook_state.installed = False
         vqs_celery._finalize_hook_state.register_queues = False
+        set_headers(None)
 
 
 @pytest.fixture(autouse=True)
@@ -311,7 +327,6 @@ def track(channel: vqs_celery._BaseChannel, tag: str, queued_message: FakeMessag
         message=queued_message,
         lease_renewal=renewal,
         queue_client=channel._queue_client,
-        headers_context=vqs_celery.get_headers_context(),
     )
 
 
@@ -1064,6 +1079,48 @@ def start(self) -> None:
     assert embedded.thread.daemon is True
 
 
+def test_start_embedded_worker_thread_does_not_inherit_header_context(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    seen_headers: list[dict[str, str] | None] = []
+    worker_started = threading.Event()
+
+    class FakeWorkController:
+        consumer = object()
+
+        def __init__(self, **kwargs: Any) -> None:
+            pass
+
+        def start(self) -> None:
+            headers = get_headers()
+            seen_headers.append(dict(headers) if headers is not None else None)
+            worker_started.set()
+
+    app = CeleryApp("worker-context")
+    monkeypatch.setattr(app, "WorkController", FakeWorkController)
+    monkeypatch.setattr(vqs_celery, "_wait_for_embedded_worker_channel", lambda worker: None)
+
+    set_headers({"x-vercel-oidc-token": "boot"})
+    try:
+        _REAL_START_EMBEDDED_WORKER(app)
+    finally:
+        set_headers(None)
+
+    assert worker_started.wait(5)
+    assert seen_headers == [None]
+
+
+def test_transports_treat_auth_and_network_errors_as_recoverable() -> None:
+    for transport in (
+        public_vqs_celery.VercelQueueTransport,
+        public_vqs_celery.VercelQueuePollTransport,
+        public_vqs_celery.VercelQueuePushTransport,
+    ):
+        assert TokenResolutionError in transport.connection_errors
+        assert UnauthorizedError in transport.connection_errors
+        assert QueueCommunicationError in transport.connection_errors
+
+
 def test_embedded_worker_ready_requires_channel_for_same_app() -> None:
     @dataclass
     class FakeWorker:
@@ -1485,8 +1542,8 @@ def test_registered_callback_delivers_to_active_push_channel() -> None:
     assert renewal.entered == 1
 
 
-def test_push_delivery_follow_ups_use_delivery_header_contexts() -> None:
-    app = CeleryApp("callback-context")
+def test_push_delivery_follow_ups_do_not_replay_delivery_headers() -> None:
+    app = CeleryApp("callback-no-context-replay")
     configure_push_broker(app)
     app.conf.task_queues = (Queue("emails"),)
     channel = make_push_channel(requeue_delay_seconds=7)
@@ -1509,12 +1566,85 @@ def test_push_delivery_follow_ups_use_delivery_header_contexts() -> None:
     channel.basic_reject(second_tag, requeue=True)
 
     fake_client = cast("FakeSyncQueueClient", channel._queue_client)
-    assert fake_client.ack_headers == [{"x-vercel-oidc-token": "token-1"}]
-    assert fake_client.visibility_headers == [{"x-vercel-oidc-token": "token-2"}]
+    assert fake_client.ack_headers == [{"x-vercel-oidc-token": "current"}]
+    assert fake_client.visibility_headers == [{"x-vercel-oidc-token": "current"}]
     assert get_headers() == {"x-vercel-oidc-token": "current"}
     assert len(FakeSyncQueueClient.instances) == 1
 
 
+def deliver_push_message(headers: dict[str, str]) -> None:
+    set_headers(headers)
+    try:
+        with pytest.raises(Handoff):
+            registered_callback()(FakeMessage(message(), topic="emails"))
+    finally:
+        set_headers(None)
+
+
+def test_publish_without_request_context_does_not_install_delivery_headers() -> None:
+    app = CeleryApp("no-context-fallback-put")
+    configure_push_broker(app)
+    app.conf.task_queues = (Queue("emails"),)
+    channel = make_push_channel()
+    received: list[Any] = []
+    channel.basic_consume("emails", no_ack=False, callback=received.append, consumer_tag="ctag")
+    vqs_celery.register_celery_app_queues(app)
+
+    deliver_push_message({"x-vercel-oidc-token": "token-1"})
+    channel._put("emails", message())
+
+    deliver_push_message({"x-vercel-oidc-token": "token-2"})
+    channel._put("emails", message())
+
+    fake_client = cast("FakeSyncQueueClient", channel._queue_client)
+    assert fake_client.send_headers == [None, None]
+
+
+def test_publish_preserves_ambient_request_context() -> None:
+    app = CeleryApp("context-ambient")
+    configure_push_broker(app)
+    app.conf.task_queues = (Queue("emails"),)
+    channel = make_push_channel()
+    received: list[Any] = []
+    channel.basic_consume("emails", no_ack=False, callback=received.append, consumer_tag="ctag")
+    vqs_celery.register_celery_app_queues(app)
+
+    deliver_push_message({"x-vercel-oidc-token": "stale"})
+    set_headers({"x-vercel-oidc-token": "current"})
+    try:
+        channel._put("emails", message())
+    finally:
+        set_headers(None)
+
+    fake_client = cast("FakeSyncQueueClient", channel._queue_client)
+    assert fake_client.send_headers == [{"x-vercel-oidc-token": "current"}]
+
+
+def test_poll_without_request_context_does_not_install_delivery_headers() -> None:
+    app = CeleryApp("no-context-fallback-poll")
+    configure_push_broker(app)
+    app.conf.task_queues = (Queue("emails"),)
+    push_channel = make_push_channel()
+    received: list[Any] = []
+    push_channel.basic_consume(
+        "emails",
+        no_ack=False,
+        callback=received.append,
+        consumer_tag="ctag",
+    )
+    vqs_celery.register_celery_app_queues(app)
+    deliver_push_message({"x-vercel-oidc-token": "token-1"})
+
+    channel = make_poll_channel()
+    fake_client = cast("FakeSyncQueueClient", channel._queue_client)
+    fake_client.message_batches.append([FakeMessage(message())])
+    payload = channel._get("emails")
+    channel.basic_ack(payload["properties"]["delivery_tag"])
+
+    assert fake_client.poll_headers == [None]
+    assert fake_client.ack_headers == [None]
+
+
 def test_registered_callback_matches_push_channel_consumer_group() -> None:
     app = CeleryApp("callback-groups")
     configure_push_broker(app)

integrations/vercel-celery/vercel/integrations/celery/__init__.py

@@ -4,6 +4,7 @@
 
 from celery.app import backends as celery_backends
 from celery.app.defaults import DEFAULTS as CELERY_DEFAULTS
+from vercel.queue import CommunicationError, TokenResolutionError, UnauthorizedError
 
 from ._broker import (
     AutoChannel,
@@ -20,6 +21,21 @@
 from .version import __version__
 
 _DEFAULT_CONNECTION_PARAMS = {"hostname": None, "port": None}
+# Token resolution depends on request-scoped OIDC context that may be
+# momentarily unavailable on Celery-owned threads (and refreshes with the next
+# push delivery). A cached token can also expire between local resolution and the
+# remote request, producing UnauthorizedError once before the next resolution
+# evicts it. CommunicationError is ordinary network trouble. Treat all three as
+# recoverable connection failures so the Celery consumer restarts instead of
+# shutting the worker down permanently. Kombu's own connection_errors use
+# amqp.exceptions.ConnectionError, not the builtin, so CommunicationError must be
+# listed explicitly.
+_CONNECTION_ERRORS = (
+    *virtual.Transport.connection_errors,
+    TokenResolutionError,
+    UnauthorizedError,
+    CommunicationError,
+)
 
 
 class VercelQueueTransport(virtual.Transport):
@@ -33,6 +49,7 @@ class VercelQueueTransport(virtual.Transport):
     driver_type = "vercel"
     driver_name = "Vercel Queue Service"
     default_connection_params = _DEFAULT_CONNECTION_PARAMS
+    connection_errors = _CONNECTION_ERRORS
 
 
 class VercelQueuePollTransport(virtual.Transport):
@@ -46,6 +63,7 @@ class VercelQueuePollTransport(virtual.Transport):
     driver_type = "vercel"
     driver_name = "Vercel Queue Service (poll)"
     default_connection_params = _DEFAULT_CONNECTION_PARAMS
+    connection_errors = _CONNECTION_ERRORS
 
 
 class VercelQueuePushTransport(virtual.Transport):
@@ -59,6 +77,7 @@ class VercelQueuePushTransport(virtual.Transport):
     driver_type = "vercel"
     driver_name = "Vercel Queue Service (push)"
     default_connection_params = _DEFAULT_CONNECTION_PARAMS
+    connection_errors = _CONNECTION_ERRORS
 
 
 def install_vercel_celery_integration(