Skip to content

workflow: speak CBOR on the Vercel world's queue #265

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fantix wants to merge 1 commit into from
+152 −3
Open

workflow: speak CBOR on the Vercel world's queue #265

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions changes/vercel/queue-cbor-transport.bugfix.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Send and receive workflow queue messages as CBOR, mirroring `DualTransport` in
`@workflow/world-vercel`.
48 changes: 45 additions & 3 deletions src/vercel/_internal/workflow/worlds/vercel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import datetime
import json
import math
import os
import platform
import traceback
import urllib.parse
from collections.abc import Mapping
from collections.abc import AsyncIterator, Mapping
from typing import Any, TypeVar

import cbor2
Expand Down Expand Up @@ -34,6 +35,44 @@ def _cbor_filter_undefined(value: Mapping[Any, Any], shareable: bool = False) ->
return {k: None if v is cbor2.undefined else v for k, v in value.items()}


class _QueueTransport:
"""CBOR queue transport with a JSON fallback on receive.

Mirrors `DualTransport` in `@workflow/world-vercel`: workflow queue
messages are sent as CBOR (the transport every producer at
``specVersion >= 3`` uses) and decoded CBOR-first, so messages from a
JSON-only producer still arrive.
"""

content_type = "application/cbor"

def serialize(self, value: Any) -> bytes:
return cbor2.dumps(value)

async def deserialize(
self,
payload: AsyncIterator[bytes],
*,
content_type: str,
) -> Any:
chunks = bytearray()
async for chunk in payload:
chunks.extend(chunk)
body = bytes(chunks)
if "json" in content_type:
return json.loads(body)
try:
return cbor2.loads(body, tag_hook=_cbor_tag_hook, object_hook=_cbor_filter_undefined)
except cbor2.CBORDecodeError:
# A producer that predates the CBOR transport, or one that sent
# JSON without labelling it.
return json.loads(body)


# One instance: it is stateless, and both directions must agree.
_QUEUE_TRANSPORT = _QueueTransport()


# Events whose result the runtime reads back resolved (run/step entity fields);
# everything else is fetched lazily.
_EVENTS_NEEDING_RESOLVE = frozenset({"run_created", "run_started", "step_started"})
Expand Down Expand Up @@ -255,7 +294,9 @@ async def queue(
client = self._queue_client(deployment=deployment_id)
try:
message_id = await client.send(
w.get_physical_topic(queue_name),
# The topic carries the codec, so the send is CBOR whichever
# client happens to make it.
vqs.Topic(w.get_physical_topic(queue_name), transport=_QUEUE_TRANSPORT),
payload,
idempotency_key=idempotency_key,
delay=delay,
Expand Down Expand Up @@ -311,6 +352,7 @@ async def async_handler(message: vqs.Message[Any]) -> None:
vqs.subscribe(
topic=f"{topic_prefix}*",
consumer_group=w.QUEUE_CONSUMER_GROUP,
transport=_QUEUE_TRANSPORT,
)(async_handler)
self._queue_callbacks.append(async_handler)

Expand All @@ -320,7 +362,7 @@ async def http_handler(request: w.HTTPRequest) -> w.HTTPResponse:
# client the sends use, so they inherit the proxy base URL and
# token instead of falling back to an unconfigured default.
client = self._queue_client(deployment=vqs.ALL_DEPLOYMENTS)
await client.accept_and_handle(request)
await client.accept_and_handle(request, transport=_QUEUE_TRANSPORT)
except Exception:
traceback.print_exc()
raise
Expand Down
105 changes: 105 additions & 0 deletions src/vercel/tests/unit/test_workflow_queue_transport.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
"""Tests for the CBOR queue transport VercelWorld sends and receives with.

Mirrors ``DualTransport`` in ``@workflow/world-vercel``: every producer at
``specVersion >= 3`` puts workflow messages on the queue as CBOR, so the
default JSON transport cannot read a single delivery.
"""

from __future__ import annotations

import datetime
import json
from collections.abc import AsyncIterator
from typing import Any

import cbor2

from vercel._internal.core.polyfills import UTC
from vercel._internal.workflow import world as w
from vercel._internal.workflow.worlds import vercel as vercel_mod
from vercel.queue import MessageMetadata
from vercel.queue._internal.subscribers import infer_subscriber_transport

WRAPPER = {
"payload": {"runId": "wrun_1", "stepId": "step_1"},
"queueName": "__wkf_workflow_greet",
"deploymentId": "dpl_1",
}


async def _chunks(body: bytes) -> AsyncIterator[bytes]:
# Deliveries arrive as a stream; split so a transport that assumes a
# single chunk fails here rather than in production.
yield body[:3]
yield body[3:]


async def _decode(body: bytes, *, content_type: str = "application/cbor") -> Any:
return await vercel_mod._QueueTransport().deserialize(_chunks(body), content_type=content_type)


def test_serialize_writes_cbor() -> None:
transport = vercel_mod._QueueTransport()

assert transport.content_type == "application/cbor"
assert cbor2.loads(transport.serialize(WRAPPER)) == WRAPPER


async def test_deserialize_reads_cbor() -> None:
assert await _decode(cbor2.dumps(WRAPPER)) == WRAPPER


async def test_deserialize_resolves_undefined_and_typed_arrays() -> None:
"""The JS encoder emits both; neither has a Python equivalent."""
body = cbor2.dumps(
{
"payload": {"runId": "wrun_1", "stepId": cbor2.undefined},
"queueName": cbor2.CBORTag(64, b"\x00\x01"),
}
)

assert await _decode(body) == {
"payload": {"runId": "wrun_1", "stepId": None},
"queueName": b"\x00\x01",
}


async def test_deserialize_falls_back_to_json_for_an_unlabelled_body() -> None:
"""A producer that predates the CBOR transport sends JSON."""
assert await _decode(json.dumps(WRAPPER).encode()) == WRAPPER


async def test_deserialize_reads_json_when_the_content_type_says_so() -> None:
body = json.dumps(WRAPPER).encode()

assert await _decode(body, content_type="application/json") == WRAPPER


def test_dispatch_resolves_the_transport_without_help_from_the_client(
isolated_subscriptions: None,
) -> None:
"""The deployed entrypoint is a generated ``vercel.queue.asgi_app()``.

Its client is built by the platform and carries no transport, so what
matters is the transport dispatch resolves from the *subscription* — not
anything hanging off a client the world constructed. Asserting the latter
is what hid this: the world's own client never receives a delivery.
"""
world = vercel_mod.VercelWorld(token="tok")

async def handler(message, *, queue_name, attempt, message_id): # noqa: ANN001, ANN202
del message, queue_name, attempt, message_id
return None

world.create_queue_handler("__wkf_workflow_", handler)

metadata = MessageMetadata(
message_id="m1",
delivery_count=1,
created_at=datetime.datetime(2026, 1, 1, tzinfo=UTC),
topic="__wkf_workflow_greet",
consumer_group=w.QUEUE_CONSUMER_GROUP,
content_type="application/cbor",
)

assert infer_subscriber_transport(metadata) is vercel_mod._QUEUE_TRANSPORT