FastStream Kafka RPC

serve.py

import asyncio
from typing import Annotated, Any
from uuid import uuid4

from faststream import FastStream, Context, context
from faststream.kafka import KafkaBroker


context.set_global("responses", {})

Responses = Annotated[
    dict[str, asyncio.Future[Any]],
    Context("responses"),
]

CorrelationId = Annotated[str, Context("message.correlation_id")]


broker = KafkaBroker()
app = FastStream(broker)


@broker.subscriber("responses")
async def handle_responses(
    msg: Any,
    cor_id: CorrelationId,
    responses: Responses,
):
    if (future := responses.get(cor_id)):
        future.set_result(msg)


@broker.subscriber("in")
async def just_handler():
    return "Response"


@app.after_startup
async def rpc_pub(responses: Responses):
    cor_id = str(uuid4())

    future = responses[cor_id] = asyncio.Future()

    await broker.publish(
        message="test",
        topic="in",
        reply_to="responses",
        correlation_id=cor_id,
    )

    result = await future
    del responses[cor_id]

    assert result == "Response"