fix(routing): tighten router contracts after review

Remove direct router access to broker private task storage and add an explicit
broker-owned storage boundary for router-managed task registration.

Clarify prepared route snapshots, subscription task-name semantics and setup-time
router mutation contracts. Extend routing docs with flow resolution precedence
and custom task base class middleware behavior.

Add regression coverage for shared task local registration, custom base_cls
middleware hooks and scheduler fallback when a route is removed before dispatch.

Author: GefMar

docs/guide/routing-and-flows.md

@@ -150,6 +150,16 @@ keep a snapshot of the route that was resolved at prepare time, so later changes
 to the mutable kicker object or router route table do not change that prepared
 send.
 
+Flow selection is resolved in this order:
+
+| Source | Effect |
+| --- | --- |
+| Explicit route | `.with_route(route)` uses that route's broker and flow. A later `.with_flow(flow)` updates the route flow, while a later `.with_broker(broker)` clears the explicit route. |
+| Explicit broker | `.with_broker(broker)` sends through that broker. If `.with_flow(flow)` is also set, that flow wins; otherwise Taskiq uses the same-broker task route flow, then the broker `default_flow`. |
+| Explicit flow | `.with_flow(flow)` without a broker override replaces the flow of the resolved route or broker default route. |
+| Router task route | `router.route_task(task, broker=..., flow=...)` is the default outbound route for that task. |
+| Broker default flow | `InMemoryBroker(..., default_flow=...)` and other brokers with `default_flow` provide the fallback flow for that broker. |
+
 ## Subscriptions
 
 Routers also keep inbound flow subscriptions. Routing and subscribing are
@@ -179,6 +189,11 @@ The deprecated `route_task(..., subscribe=True)` shim still performs this
 subscription when a flow is resolved, but new code should call `subscribe()`
 directly.
 
+`TaskiqSubscription.task_names` is diagnostic listen-plan metadata. It records
+which task names caused the broker to listen to a flow, but it does not make the
+flow an inbound task router. Workers still execute messages by
+`TaskiqMessage.task_name`.
+
 Existing brokers can keep implementing `listen()` as before. New flow-aware
 brokers may use `get_subscribed_flows()` to subscribe to queues, topics,
 subjects or streams while the routing rules stay in the router.
@@ -187,6 +202,10 @@ subjects or streams while the routing rules stay in the router.
 subscriptions. It deduplicates by flow identity and checks broker options for
 conflicts.
 
+Router routes and subscriptions are mutable setup-time configuration. They are
+not thread-safe runtime coordination primitives; configure them before starting
+concurrent worker or client activity.
+
 ## Scheduler and requeue
 
 `ScheduledTask` remains a transport-neutral invocation payload. It stores task
@@ -277,6 +296,9 @@ When the final application registers this definition, Taskiq creates the bound
 task using `TracingTask`. If `base_cls` is not provided, Taskiq uses the native
 decorated task class.
 
+Custom task classes do not bypass broker middleware. Send and execute lifecycle
+hooks are still owned by the selected broker and worker path.
+
 For low-level integrations, `TaskDefinition.message(...)` builds a
 `TaskiqMessage` without binding the task. It uses the same argument and label
 preparation contract as a normal `.kicker().prepare(...)` invocation.

taskiq/abc/broker.py

@@ -439,6 +439,19 @@ def bind_task_definition(
             register=register,
         )
 
+    def store_registered_task(
+        self,
+        task: AsyncTaskiqDecoratedTask[Any, Any],
+    ) -> None:
+        """
+        Store a task after router-owned registration accepted it.
+
+        Application code should use `task` or `register_task`. This method is
+        the broker/router integration boundary for cases where the router owns
+        registration ordering and the broker owns local task storage.
+        """
+        self._store_task(task.task_name, task)
+
     def _decorate_task(
         self,
         func: Callable[_FuncParams, _ReturnType],

taskiq/kicker.py

@@ -321,7 +321,14 @@ async def _legacy_kiq(
     def _prepare_route_snapshot(
         self,
     ) -> tuple[AsyncBroker, TaskiqRoute | None, FlowProtocol | None]:
-        """Resolve the route that a prepared invocation must keep."""
+        """
+        Resolve route state that a prepared invocation must keep.
+
+        The returned tuple is `(broker, route, explicit_flow_override)`.
+        When a route is present, the route already carries its flow, so the
+        separate flow override is `None`. A separate flow is returned only for
+        legacy broker paths that cannot snapshot a `TaskiqRoute`.
+        """
         router = getattr(self.broker, "router", None)
         if not isinstance(router, TaskiqRouter):
             return self.broker, None, self.route_flow

taskiq/router.py

@@ -132,11 +132,7 @@ def register_task(
                 broker=target_broker,
                 flow=flow,
             )
-            # Router and broker share this internal registration boundary.
-            target_broker._store_task(  # noqa: SLF001
-                registered_task.task_name,
-                registered_task,
-            )
+            target_broker.store_registered_task(registered_task)
             return registered_task
 
         return self._register_bound_task(task, broker=broker, flow=flow)
@@ -178,7 +174,7 @@ def register(
             func: Callable[_FuncParams, _ReturnType],
         ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]:
             target_broker = self._brokers.resolve(broker)
-            real_task_name = task_name if not callable(task_name) else None
+            real_task_name: str | None = None if callable(task_name) else task_name
             task = target_broker.task(task_name=real_task_name, **labels)(func)
             if flow is not None:
                 self.route_task(task, broker=target_broker, flow=flow)

taskiq/routing/models.py

@@ -26,7 +26,13 @@ def broker_name(self) -> str:
 
 @dataclass(frozen=True, slots=True)
 class TaskiqSubscription:
-    """Inbound flow subscription owned by a router."""
+    """
+    Inbound flow subscription owned by a router.
+
+    `task_names` is diagnostic listen-plan metadata. It records which task
+    declarations caused a broker to listen to this flow; worker execution still
+    resolves inbound messages by `TaskiqMessage.task_name`.
+    """
 
     broker: AsyncBroker
     flow: FlowProtocol

taskiq/routing/subscriptions.py

@@ -15,7 +15,12 @@
 
 
 class SubscriptionPlan:
-    """Inbound listen-plan subscriptions for flow-aware brokers."""
+    """
+    Inbound listen-plan subscriptions for flow-aware brokers.
+
+    This is mutable setup-time configuration, not a thread-safe runtime
+    coordination primitive.
+    """
 
     def __init__(self, brokers: BrokerRegistry) -> None:
         self.brokers = brokers

tests/routing/test_shared_tasks.py

@@ -1,7 +1,7 @@
 import pytest
 
 from taskiq import AsyncTaskiqDecoratedTask, Flow, task_builder
-from tests.routing.models import TracingTask
+from tests.routing.models import RecordingMiddleware, TracingTask
 from tests.utils import RecordingBroker
 
 
@@ -41,6 +41,7 @@ async def shared_task() -> None:
 
     assert source.sent == []
     assert target.find_task("shared.routed") is registered
+    assert target.local_task_registry["shared.routed"] is registered
     assert target.sent[0][0].task_name == "shared.routed"
     assert target.sent[0][1] == flow
 
@@ -64,6 +65,27 @@ async def traced(value: int) -> int:
     assert broker.sent[0][0].task_name == "shared.traced"
 
 
+async def test_task_builder_custom_base_cls_uses_broker_middleware() -> None:
+    broker = RecordingBroker()
+    events: list[tuple[str, str, str]] = []
+    broker.add_middlewares(RecordingMiddleware("broker", events))
+
+    @task_builder("shared.traced.middleware", base_cls=TracingTask)
+    async def traced(value: int) -> int:
+        return value + 1
+
+    registered = broker.register_task(traced)
+
+    assert isinstance(registered, TracingTask)
+
+    await registered.kiq(1)
+
+    assert events == [
+        ("broker", "pre_send", "shared.traced.middleware"),
+        ("broker", "post_send", "shared.traced.middleware"),
+    ]
+
+
 def test_task_definition_default_flow_does_not_create_subscription() -> None:
     flow = Flow("shared.default")
     broker = RecordingBroker(default_flow=flow)

tests/scheduler/test_scheduler.py

@@ -114,6 +114,40 @@ async def test_scheduler_without_route_uses_scheduler_broker() -> None:
     assert scheduled_task.labels == {"source": "database"}
 
 
+async def test_scheduler_removed_route_uses_scheduler_broker() -> None:
+    router = TaskiqRouter()
+    scheduler_flow = Flow("scheduler.default")
+    scheduler_broker = RecordingBroker(
+        router=router,
+        broker_name="scheduler",
+        default_flow=scheduler_flow,
+    )
+    target_broker = RecordingBroker(router=router, broker_name="target")
+    source = RecordingScheduleSource()
+    scheduler = TaskiqScheduler(broker=scheduler_broker, sources=[source])
+
+    @scheduler_broker.task(task_name="demo.task")
+    async def demo_task() -> None:
+        return None
+
+    router.route_task(demo_task, broker=target_broker, flow=Flow("target"))
+    scheduled_task = ScheduledTask(
+        task_name="demo.task",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+    del router.routes["demo.task"]
+
+    await scheduler.on_ready(source, scheduled_task)
+
+    sent_message, sent_flow = scheduler_broker.sent[0]
+    assert target_broker.sent == []
+    assert sent_message.task_name == "demo.task"
+    assert sent_flow == scheduler_flow
+
+
 async def test_created_schedule_kiq_does_not_mutate_schedule_payload() -> None:
     router = TaskiqRouter()
     source_broker = RecordingBroker(router=router, broker_name="source")