diff --git a/README.md b/README.md index 5601eec35..27f7eb581 100644 --- a/README.md +++ b/README.md @@ -469,6 +469,13 @@ channels: # Gateway API URL (default: http://localhost:8001) gateway_url: http://localhost:8001 + # Maximum queued or provider-reserved inbound messages (default: 1000) + inbound_queue_maxsize: 1000 + # Fixed number of long-lived inbound handler workers (default: 5) + max_concurrency: 5 + # Seconds to drain accepted work before cancelling active handlers (default: 3) + shutdown_grace_period_seconds: 3 + # Optional: global session defaults for all mobile channels session: assistant_id: lead_agent # or a custom agent name; custom agents are routed via lead_agent + agent_name @@ -546,6 +553,7 @@ Notes: - `assistant_id: lead_agent` calls the default LangGraph assistant directly. - If `assistant_id` is set to a custom agent name, DeerFlow still routes through `lead_agent` and injects that value as `agent_name`, so the custom agent's SOUL/config takes effect for IM channels. - IM channel workers call Gateway's LangGraph-compatible API internally and automatically attach process-local internal auth plus the CSRF cookie/header pair required for thread and run creation. +- Inbound work is bounded to `inbound_queue_maxsize` pending messages plus `max_concurrency` active workers. When capacity is exhausted, socket/polling providers drop new messages before sending DeerFlow's working acknowledgment and emit a rate-limited warning. Buzz leaves its replay cursor unchanged and reconnects for relay replay; GitHub webhooks return `503`, marking the delivery failed for manual/API redelivery. Shutdown closes admission immediately, keeps channel transports available while accepted messages drain for up to `shutdown_grace_period_seconds`, then cancels and awaits active handlers before closing provider resources; the Gateway's outer timeout can cancel an incomplete shutdown without detaching those resources. - Feishu/Lark now queues rapid follow-up messages per mapped DeerFlow `thread_id` instead of immediately surfacing the generic busy reply, and topic replies keep a per-message card with a compact source-message preview across queued/running/final patches. Set the corresponding API keys in your `.env` file: diff --git a/backend/app/channels/AGENTS.md b/backend/app/channels/AGENTS.md index f1210ad56..243e64b5b 100644 --- a/backend/app/channels/AGENTS.md +++ b/backend/app/channels/AGENTS.md @@ -10,8 +10,8 @@ Bridges external messaging platforms (Feishu, Slack, Telegram, Discord, DingTalk - `manager.py` - Core dispatcher: creates threads via `client.threads.create()`, routes commands including `/goal` (setting a goal persists it through Gateway and then routes the objective as a chat turn), keeps Slack/Discord on `client.runs.wait()`, uses `client.runs.stream(["messages-tuple", "values"])` for Feishu/Telegram incremental outbound updates, serializes same-thread Feishu turns in-manager when the channel's `ChannelRunPolicy.serialize_thread_runs=True` so rapid follow-ups queue instead of tripping the runtime busy reply, and switches to `client.runs.create()` (fire-and-forget, returns once the run is `pending`) for channels whose `ChannelRunPolicy.fire_and_forget=True` so long autonomous runs do not hit the SDK default 300s `httpx.ReadTimeout` A swallowed streaming failure publishes its final outbound before releasing the inbound dedupe key, so a provider redelivery can retry without overtaking the terminal reply. **What may be published from the stream is an allowlist, not a denylist** (`_accumulate_stream_text`): only assistant message types — LangChain serializes `AIMessage.type` as `"ai"` and `AIMessageChunk.type` as `"AIMessageChunk"`, plus the OpenAI-style `"assistant"` spelling for foreign runtimes — become displayable text. The previous rule rejected only payloads whose `type` contained `"tool"` and therefore published everything else, which leaked DeerFlow's hidden model context to every streaming IM channel: `DynamicContextMiddleware` injects recalled memory as a hidden `HumanMessage` (`type == "human"`) and rewrites the user's own turn into a new `HumanMessage`, `DurableContextMiddleware` injects a hidden `` `HumanMessage`, and LangGraph fans those state writes out on the `messages-tuple` stream. Proved live on a Buzz relay, which published a `` fact block and, in another run, a verbatim echo of the user's own message as the assistant's reply. Matching is by prefix (`ai` / `assistant`), never substring, because ordinary words contain `"ai"` (`chain`, `domain`). The message type is resolved by `_stream_payload_type`, which handles both the `model_dump()` shape DeerFlow's own gateway emits and LangChain's `to_json()` constructor shape (whose top-level `type` is the literal `"constructor"`, with the class name at the tail of the `id` path). A bare `str` payload is no longer accepted at all: it carries no type information, so it cannot be attributed to the assistant, and nothing in DeerFlow produces one (`runtime/serialization.py::serialize_messages_tuple` always emits `[message_dict, metadata]`). -- `base.py` - Abstract `Channel` base class (start/stop/send lifecycle) -- `service.py` - Manages lifecycle of all configured channels from `config.yaml` +- `base.py` - Abstract `Channel` base class (start/stop/send lifecycle). Provider callbacks that submit coroutines from SDK threads through `run_coroutine_threadsafe()` must use `_submit_threadsafe_coroutine()`: submission is closed atomically with shutdown, and `stop()` must call `_close_and_drain_threadsafe_futures()` before tearing down SDK resources. +- `service.py` - Manages lifecycle of all configured channels from `config.yaml`. Shutdown closes manager admission first and keeps transports alive until every manager worker/follow-up watcher has exited. A successful manager stop therefore owns no live handler; if the Gateway's outer timeout cancels shutdown, the service retains its channel objects and global singleton so unfinished resources are not detached and cleanup can be retried. - `slack.py` / `feishu.py` / `telegram.py` / `discord.py` / `dingtalk.py` - Platform-specific implementations (`feishu.py` tracks the running card `message_id` in memory and patches the same card in place; `telegram.py` accepts inbound text/photos/documents, preserves media captions, hands token-free attachment bytes to the shared upload pipeline, edits the "Working on it..." stream target in place via `editMessageText`, and can optionally send final Markdown replies as Rich Messages through `channels.telegram.rich_messages`; `discord.py` registers typing loops before inbound handling yields and `_start_typing()` refuses work once `_running` is false; because `stop()` runs on the main loop while typing tasks belong to `_discord_loop`, normal cross-thread cancellation, awaiting, and map cleanup are scheduled there with a bounded wait, while `_run_client()` drains tasks in its `finally` block before an exception/disconnect can make that loop unusable; an already-stopped foreign loop must never have its tasks awaited from the main loop; `dingtalk.py` optionally uses AI Card streaming for in-place updates when `card_template_id` is configured, and overrides `receive_file` to download inbound images (`picture`/`richText`) and documents (`file`) by `downloadCode` into the thread uploads bucket, mirroring `feishu.py`) - `buzz.py` - Buzz (Nostr relay) implementation: one NIP-42-authenticated websocket, pubkey-allowlist + mention/DM/thread-follow gating, streaming replies via in-place kind-40003 edits; requires the `buzz` dependency extra. **Subscription model** (operator-facing version in [IM_CHANNEL_CONNECTIONS.md](../../docs/IM_CHANNEL_CONNECTIONS.md#buzz-subscription-model)): the relay fans kind-9 chat events out **only** to channel-scoped subscriptions, proved against a live relay — `REQ {"kinds":[9]}` is accepted and answered with `EOSE` but never receives an event (the connector authenticated and then sat silent forever), `REQ {"kinds":[9],"#h":[uuid]}` works, and a multi-value `#h` matches nothing, so it is strictly one REQ per channel (the same shape as Buzz's own `buzz-acp` harness). Every connection therefore rebuilds three kinds of subscription after NIP-42 auth: `buzz-discovery` (`{"kinds":[39000]}`, a historical query returning exactly the channels this identity belongs to, one stored event each then `EOSE` — adding `#p` returns zero, do not "narrow" it), `buzz-membership` (`{"kinds":[44100,44101],"#p":[us],"since":}`, the relay-signed member-added/member-removed notifications whose `p` tag names the affected member and `h` tag the channel), and one `buzz-chat-` per discovered channel. Chat subscriptions open as each kind-39000 arrives; the discovery `EOSE` is the completeness barrier that retries any that failed and warns when discovery found nothing. A kind-44100 for our pubkey subscribes to the new channel **live** (then re-issues discovery so its name/type reach the DM-detection cache, but only when that channel's metadata is actually missing — an unconditional refresh is how a burst of 44100s multiplied into one discovery pass each); a kind-44101 issues `buzz_nostr.close_frame` for exactly that channel's subscription and drops its metadata. **The membership subscription is scoped to LIVE events** and this is load-bearing: buzz-relay *stores* 44100/44101 and serves history newest-first (default limit 2000), so an unscoped filter replayed the whole membership history on every connect — every stored add read as live, re-running discovery once each (M+1 discovery passes × N stored kind-39000 events per connect, observed live as two `channel discovery complete` lines and one channel logged `` because the 44100 path subscribed before its metadata arrived), re-subscribing channels we have since been removed from, and letting a stored 44101 transiently unsubscribe a channel we are still in. `since` is anchored at the moment the socket opened (`_session_started_at`) minus `MEMBERSHIP_LOOKBACK_SECONDS` (60) of slack, which covers both relay clock skew and a membership change published *during* the connect/auth handshake; the slack can only cost an idempotent replay of the last minute. **A relay `CLOSED` frame is recovered, not merely forgotten** — every subscription on the socket fails silently when dropped, so `_handle_closed` re-issues it, bounded by `MAX_RESUBSCRIBE_ATTEMPTS` (3) per subscription id per connection *and per auth epoch* (`_resubscribe_attempts` is reset on session start, on re-auth, and by `stop()`, so pre-auth rejections — which the auth branch already recovers wholesale — never spend the authenticated session's budget). The first retry is immediate (the common case is a one-off hiccup); later ones back off 1s then 2s, awaited inline in the read loop rather than in a background task that could outlive its own socket. **`auth-required:` before this socket has completed its NIP-42 handshake is the expected bootstrap sequence, not a refusal**, and `_handle_closed` short-circuits it ahead of every recovery path: the connector opens its control REQs immediately in case the relay serves unauthenticated reads, a closed relay answers `auth-required:` plus an `AUTH` challenge, and the auth branch re-opens everything. That case is logged at DEBUG and consumes neither the permanent-refusal branch nor the retry budget (a chat subscription is still dropped from `_chat_subscriptions`, since it genuinely is not subscribed and discovery is what re-opens it). Treating it as permanent produced an operator-facing warning claiming discovery/membership tracking was DOWN in the same run where discovery then completed, every channel was subscribed, and a brand-new channel's kind-44100 was picked up one second later. The boundary is the per-socket `_auth_completed` flag, set once the signed AUTH event has been sent and cleared on session entry, session exit, and `stop()`; the same reason *after* that stays loud, and says the subscription is down until the relay's next AUTH challenge or the next reconnect rather than borrowing the non-auth wording. Then `_is_transient_close` decides whether to retry at all: NIP-01/NIP-42 `auth-required:`/`restricted:`/`blocked:`/`mute:`/`invalid:`/`pow:` prefixes and buzz-relay's own removal/revocation prose are permanent (do not fight the relay over a channel that is no longer ours; a post-auth `auth-required:` is recovered by the AUTH branch, not by re-issuing), `rate-limited:`/`error:`/no-reason-at-all and anything unrecognized are transient — the default resolves toward "keep listening" because going silently deaf is the failure this exists to remove, and the attempt budget bounds a wrong guess. A chat `CLOSED` is only ever recovered for a channel already in `_chat_subscriptions`: a `CLOSED` is relay-supplied, so acting on an unknown one would let a relay induce a subscription just by naming a channel. Every subscription that goes unlistened is logged at WARNING, never INFO. Subscription ids are deterministic per channel precisely so one can be replaced or closed without disturbing the others on the socket, and `_chat_subscriptions` is per-socket state cleared on session end, on re-auth (a pre-auth REQ may have been rejected), and by `stop()`. Three bounds on remote-fed state: `MAX_CACHED_CHANNELS` (512) caps the kind-39000 metadata cache, the watermark map, and the resubscribe-attempt map; `MAX_CHANNEL_SUBSCRIPTIONS` (256, well under buzz-relay's own 1024-per-connection ceiling) caps live chat subscriptions — at the cap new channels are refused and named in a warning rather than evicting a working subscription. **Known bound (documented, not fixed):** the relay caps historical delivery at 2000 events per subscription, newest-first, even with a `since`, so >2000 unread messages in a *single* channel across a disconnect loses the oldest — the relay never sends them and the watermark advances past them. That is the one remaining path that can skip; everything else fails toward replay. **Trust model** (operator-facing version in [IM_CHANNEL_CONNECTIONS.md](../../docs/IM_CHANNEL_CONNECTIONS.md#buzz-trust-model)): every inbound `EVENT` is authenticated at the single `handle_relay_frame` choke point — the NIP-01 id is recomputed from the delivered payload and the BIP-340 Schnorr signature verified against the claimed `pubkey` (`buzz_nostr.verify_event`, pure and total: malformed input returns `False`, never raises) — so `ev["pubkey"]`, the authorization principal for both the allowlist and the `/connect` bind, cannot be forged by a relay the DeerFlow operator does not run. What remains trusted is the *authorship* of kind-39000 channel metadata: any member can sign one, and because per-channel subscriptions are now driven by discovery, a forged kind-39000 has two effects rather than one — it can mark a channel `type: "dm"` (relaxing `require_mention` for that channel) **and** it can induce a chat subscription for a channel of the forger's choosing, since the channels we listen to are exactly the channels we hold metadata for. Neither makes anything be *acted on*: `allowed_users` and per-event signature verification are independent gates, so an induced subscription only means the relay reads its own traffic back to a subscriber that drops it, bounded by `MAX_CHANNEL_SUBSCRIPTIONS` (which refuses rather than evicts, so it cannot displace a real channel). Same for a forged kind-44100, except its `p` tag is re-checked locally so it must at least name us. Closing this needs a configured trusted relay pubkey, which `relay_url` is not. `allowed_users` is deny-by-default (empty = nobody, unlike siblings' empty = everyone), so `start()` logs a WARNING when it is empty and each drop logs at DEBUG. The resubscribe cursor (`since`) is **per channel**, advances only for events that were actually processed, and never past `now + MAX_FUTURE_SKEW_SECONDS`, because it is peer-supplied (`created_at`) and a single future-dated event otherwise made the connector permanently deaf. Per channel rather than global is the safety-critical half: subscriptions are per channel, so one shared cursor is the newest event seen in *any* channel and a busy channel would drag it past a quiet channel's unread messages, skipping them on the next reconnect — measured on a live relay, three channels of one identity sat ~28h apart. Per-channel cursors can only ever cost duplicate delivery (absorbed by the manager's `event_id` dedupe), and an evicted cursor degrades to "no `since`", i.e. the relay's default backlog — both fail toward replay, never toward a miss. Streaming tracks every oversize chunk index (`_stream_targets` for chunk 0, `_stream_tails` for the rest), since the manager republishes cumulative text and reposting `chunks[1:]` per update flooded the channel; all of it is per-connection state cleared by `stop()`, and the remote-fed kind-39000 cache is capped at `MAX_CACHED_CHANNELS`. **`send()` refuses outright to publish text carrying a hidden model-context wrapper** (``, ``, `` — `_HIDDEN_CONTEXT_MARKERS`), logging at ERROR and clearing the stream bookkeeping on a blocked `is_final`. This is defense in depth behind the manager's allowlist, and it lives here rather than in a sibling connector because on Buzz a leak is permanent: every streaming update is an immutable public Nostr event, so a corrective edit only changes what clients render while the original leaked event stays on the relay. Matching is on the literal opening tag, so a reply that merely talks about memory is still published - `github.py` - Webhook-driven GitHub channel. Inbound messages come from `POST /api/webhooks/github`; outbound is log-only because GitHub agents post explicitly with `gh` from their sandbox when they choose to comment or create a PR diff --git a/backend/app/channels/base.py b/backend/app/channels/base.py index ecdeda1b2..62f0a3bf6 100644 --- a/backend/app/channels/base.py +++ b/backend/app/channels/base.py @@ -4,13 +4,25 @@ from __future__ import annotations import asyncio import logging +import threading from abc import ABC, abstractmethod -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Coroutine from concurrent.futures import CancelledError as FutureCancelledError +from concurrent.futures import Future from typing import Any, TypeVar from app.channels.commands import extract_connect_code -from app.channels.message_bus import InboundMessage, InboundMessageType, MessageBus, OutboundMessage, ResolvedAttachment +from app.channels.message_bus import ( + InboundMessage, + InboundMessageType, + InboundQueueClosedError, + InboundQueueFullError, + InboundReservation, + InboundReservationExpiredError, + MessageBus, + OutboundMessage, + ResolvedAttachment, +) logger = logging.getLogger(__name__) @@ -33,6 +45,12 @@ class Channel(ABC): self.config = config self._running = False self._connection_repo: Any = config.get("connection_repo") + # Provider SDK callbacks often run on a dedicated thread and submit + # preparation work to the Gateway loop. Submission and shutdown share + # this lock so stop() cannot miss a future created concurrently. + self._threadsafe_futures: set[Future[Any]] = set() + self._threadsafe_futures_lock = threading.Lock() + self._threadsafe_future_intake_open = True @property def is_running(self) -> bool: @@ -119,6 +137,78 @@ class Channel(ABC): if exc: logger.error("[%s] %s failed for msg_id=%s: %s", self.name, name, msg_id, exc) + def _open_threadsafe_future_intake(self) -> None: + """Allow a newly started provider to submit work to its main loop.""" + with self._threadsafe_futures_lock: + pending = [future for future in self._threadsafe_futures if not future.done()] + if pending: + raise RuntimeError(f"cannot restart {self.name} while cross-thread work is still running") + self._threadsafe_futures.clear() + self._threadsafe_future_intake_open = True + + def _submit_threadsafe_coroutine( + self, + coroutine: Coroutine[Any, Any, T], + loop: asyncio.AbstractEventLoop | None, + *, + name: str, + msg_id: Any, + reservation: InboundReservation | None = None, + ) -> Future[T] | None: + """Submit and retain provider-thread work until completion or shutdown.""" + + with self._threadsafe_futures_lock: + if not self._threadsafe_future_intake_open or loop is None or not loop.is_running(): + coroutine.close() + if reservation is not None: + reservation.release() + return None + try: + future = asyncio.run_coroutine_threadsafe(coroutine, loop) + except RuntimeError: + coroutine.close() + if reservation is not None: + reservation.release() + return None + self._threadsafe_futures.add(future) + + future.add_done_callback( + lambda completed: self._finalize_threadsafe_future( + completed, + name=name, + msg_id=msg_id, + reservation=reservation, + ) + ) + return future + + def _finalize_threadsafe_future( + self, + future: Future[Any], + *, + name: str, + msg_id: Any, + reservation: InboundReservation | None, + ) -> None: + with self._threadsafe_futures_lock: + self._threadsafe_futures.discard(future) + if reservation is not None: + reservation.release() + self._log_future_error(future, name, msg_id) + + async def _close_and_drain_threadsafe_futures(self) -> None: + """Close cross-thread submission, then cancel and await owned futures.""" + with self._threadsafe_futures_lock: + self._threadsafe_future_intake_open = False + futures = tuple(self._threadsafe_futures) + + for future in futures: + future.cancel() + if futures: + await asyncio.gather(*(asyncio.wrap_future(future) for future in futures), return_exceptions=True) + with self._threadsafe_futures_lock: + self._threadsafe_futures.difference_update(future for future in futures if future.done()) + def _pending_connect_code(self, text: str) -> str | None: """Return the one-time bind code if *text* is a ``/connect `` command and channel connections are configured, else ``None``. @@ -155,6 +245,48 @@ class Channel(ABC): metadata=metadata or {}, ) + def _reserve_inbound(self, msg: InboundMessage) -> InboundReservation | None: + """Reserve bounded intake capacity or explicitly drop under overload. + + Real-time socket/polling providers do not expose a reliable delivery + retry contract to this adapter layer. They therefore drop a message + that cannot be admitted immediately. ``MessageBus`` emits a + rate-limited warning with the cumulative rejection count. + """ + try: + return self.bus.reserve_inbound(msg) + except InboundQueueFullError: + return None + except (InboundQueueClosedError, InboundReservationExpiredError): + logger.debug("[%s] inbound ignored because channel intake is closed", self.name) + return None + + def _commit_reserved_inbound( + self, + reservation: InboundReservation, + msg: InboundMessage, + ) -> bool: + """Commit a reservation on the MessageBus loop, releasing on failure.""" + try: + reservation.commit(msg) + return True + except (InboundQueueClosedError, InboundReservationExpiredError): + logger.debug("[%s] inbound reservation expired during shutdown", self.name) + return False + finally: + reservation.release() + + async def _publish_inbound_or_drop(self, msg: InboundMessage) -> bool: + """Publish from an already-serialized provider loop without waiting.""" + try: + await self.bus.publish_inbound(msg) + return True + except InboundQueueFullError: + return False + except InboundQueueClosedError: + logger.debug("[%s] inbound ignored because channel intake is closed", self.name) + return False + async def _on_outbound(self, msg: OutboundMessage) -> None: """Outbound callback registered with the bus. diff --git a/backend/app/channels/buzz.py b/backend/app/channels/buzz.py index 212e12fbb..fa65def29 100644 --- a/backend/app/channels/buzz.py +++ b/backend/app/channels/buzz.py @@ -52,7 +52,7 @@ from app.channels import buzz_nostr from app.channels.base import Channel from app.channels.commands import is_known_channel_command from app.channels.connection_identity import attach_connection_identity -from app.channels.message_bus import InboundMessage, InboundMessageType, MessageBus, OutboundMessage +from app.channels.message_bus import InboundMessage, InboundMessageType, InboundQueueFullError, MessageBus, OutboundMessage logger = logging.getLogger(__name__) @@ -941,6 +941,11 @@ class BuzzChannel(Channel): await self._handle_chat_event(ev) elif ev_kind in (buzz_nostr.KIND_MEMBER_ADDED, buzz_nostr.KIND_MEMBER_REMOVED): await self._handle_membership_event(ev) + except InboundQueueFullError: + # This is not malformed relay input. Escape to _run_loop so the + # connection is reopened with the last processed watermark and + # relay history can retry the message once intake has capacity. + raise except Exception: # Defense in depth against malformed tags/timestamps inside an # otherwise well-shaped EVENT frame (e.g. a non-integer created_at, diff --git a/backend/app/channels/dingtalk.py b/backend/app/channels/dingtalk.py index 57cad0524..2dff8c339 100644 --- a/backend/app/channels/dingtalk.py +++ b/backend/app/channels/dingtalk.py @@ -16,7 +16,7 @@ import httpx from app.channels.base import Channel from app.channels.commands import is_known_channel_command, strip_leading_mentions from app.channels.connection_identity import attach_connection_identity -from app.channels.message_bus import InboundMessage, InboundMessageType, MessageBus, OutboundMessage, ResolvedAttachment +from app.channels.message_bus import InboundMessage, InboundMessageType, InboundReservation, MessageBus, OutboundMessage, ResolvedAttachment from deerflow.config.paths import VIRTUAL_PATH_PREFIX, get_paths from deerflow.runtime.user_context import get_effective_user_id from deerflow.sandbox.sandbox_provider import get_sandbox_provider @@ -188,6 +188,7 @@ class DingTalkChannel(Channel): if self._card_template_id: logger.info("[DingTalk] AI Card mode enabled (template=%s)", self._card_template_id) + self._open_threadsafe_future_intake() self._running = True self.bus.subscribe_outbound(self._on_outbound) @@ -202,6 +203,7 @@ class DingTalkChannel(Channel): async def stop(self) -> None: self._running = False self.bus.unsubscribe_outbound(self._on_outbound) + await self._close_and_drain_threadsafe_futures() stream_client = self._stream_client if stream_client is not None: @@ -396,7 +398,7 @@ class DingTalkChannel(Channel): connect_code = self._pending_connect_code(text) if connect_code: if self._main_loop and self._main_loop.is_running(): - fut = asyncio.run_coroutine_threadsafe( + future = self._submit_threadsafe_coroutine( self._bind_connection_from_connect_code( conversation_type=conversation_type, sender_staff_id=sender_staff_id, @@ -405,8 +407,11 @@ class DingTalkChannel(Channel): code=connect_code, ), self._main_loop, + name="bind_connection", + msg_id=msg_id, ) - fut.add_done_callback(lambda f, mid=msg_id: self._log_future_error(f, "bind_connection", mid)) + if future is None: + logger.info("[DingTalk] main loop stopped before channel connection bind could be scheduled") else: logger.warning("[DingTalk] main loop not running, cannot bind channel connection") return @@ -478,18 +483,24 @@ class DingTalkChannel(Channel): ) inbound.topic_id = topic_id - if self._card_template_id: - source_key = self._make_card_source_key(inbound) - with self._incoming_messages_lock: - self._incoming_messages[source_key] = message - if self._main_loop and self._main_loop.is_running(): + reservation = self._reserve_inbound(inbound) + if reservation is None: + return + if self._card_template_id: + source_key = self._make_card_source_key(inbound) + with self._incoming_messages_lock: + self._incoming_messages[source_key] = message logger.info("[DingTalk] publishing inbound message to bus (type=%s, msg_id=%s)", msg_type.value, msg_id) - fut = asyncio.run_coroutine_threadsafe( - self._prepare_inbound(chat_id, inbound), + future = self._submit_threadsafe_coroutine( + self._prepare_inbound(chat_id, inbound, reservation=reservation), self._main_loop, + name="prepare_inbound", + msg_id=msg_id, + reservation=reservation, ) - fut.add_done_callback(lambda f, mid=msg_id: self._log_future_error(f, "prepare_inbound", mid)) + if future is None: + logger.info("[DingTalk] main loop stopped before reserved inbound could be scheduled") else: logger.warning("[DingTalk] main loop not running, cannot publish inbound message") except Exception: @@ -735,12 +746,25 @@ class DingTalkChannel(Channel): logger.exception("[DingTalk] failed to download file by code") return None - async def _prepare_inbound(self, chat_id: str, inbound: InboundMessage) -> None: - inbound = await self._attach_connection_identity(inbound) - # Running reply must finish before publish_inbound so AI card tracks are - # registered before the manager emits streaming outbounds. - await self._send_running_reply(chat_id, inbound) - await self.bus.publish_inbound(inbound) + async def _prepare_inbound( + self, + chat_id: str, + inbound: InboundMessage, + *, + reservation: InboundReservation | None = None, + ) -> None: + try: + inbound = await self._attach_connection_identity(inbound) + # Running reply must finish before commit so AI card tracks are + # registered before the manager emits streaming outbounds. + await self._send_running_reply(chat_id, inbound) + if reservation is None: + await self.bus.publish_inbound(inbound) + else: + self._commit_reserved_inbound(reservation, inbound) + finally: + if reservation is not None: + reservation.release() @staticmethod def _connection_workspace_id(conversation_type: str, conversation_id: str) -> str | None: diff --git a/backend/app/channels/discord.py b/backend/app/channels/discord.py index e00e5d6ea..977cebf1c 100644 --- a/backend/app/channels/discord.py +++ b/backend/app/channels/discord.py @@ -13,7 +13,7 @@ from typing import Any from app.channels.base import Channel from app.channels.commands import is_known_channel_command from app.channels.connection_identity import attach_connection_identity -from app.channels.message_bus import InboundMessage, InboundMessageType, MessageBus, OutboundMessage, ResolvedAttachment +from app.channels.message_bus import InboundMessage, InboundMessageType, InboundReservation, MessageBus, OutboundMessage, ResolvedAttachment logger = logging.getLogger(__name__) @@ -373,64 +373,135 @@ class DiscordChannel(Channel): if connect_code and await self._bind_connection_from_connect_code(message, connect_code): return - # --- Determine thread/channel routing and typing target --- - thread_id = None - chat_id = None - typing_target = None # The Discord object to type into + # /connect is a local control-plane command and is consumed above. For + # ordinary messages, reserve the shared intake slot before any Discord + # API call, thread mapping write, or connection-identity query. The + # provisional routing fields are used only for overload logging; the + # final message is built after thread routing completes. + msg_type = InboundMessageType.COMMAND if is_known_channel_command(text) else InboundMessageType.CHAT + provisional = self._make_inbound( + chat_id=str(message.channel.id), + user_id=str(message.author.id), + text=text, + msg_type=msg_type, + metadata={ + "guild_id": str(guild.id) if guild else None, + "channel_id": str(message.channel.id), + "message_id": str(message.id), + }, + ) + reservation = self._reserve_inbound(provisional) + if reservation is None: + return - if isinstance(message.channel, self._discord_module.Thread): - # --- Message already inside a thread --- - thread_obj = message.channel - thread_id = str(thread_obj.id) - chat_id = str(thread_obj.parent_id or thread_obj.id) - typing_target = thread_obj + await self._handle_admitted_message( + message, + text=text, + msg_type=msg_type, + guild=guild, + has_mention=bool(has_mention), + reservation=reservation, + ) - # If this is a known active thread, process normally - if thread_id in self._active_thread_ids: - msg_type = InboundMessageType.COMMAND if is_known_channel_command(text) else InboundMessageType.CHAT - inbound = self._make_inbound( - chat_id=chat_id, - user_id=str(message.author.id), - text=text, - msg_type=msg_type, - thread_ts=thread_id, - metadata={ - "guild_id": str(guild.id) if guild else None, - "channel_id": str(message.channel.id), - "message_id": str(message.id), - }, - ) - inbound.topic_id = thread_id - inbound = await self._attach_connection_identity(inbound, guild_id=str(guild.id) if guild else None) - self._publish(inbound) - # Start typing indicator in the thread - if typing_target: - await self._start_typing(typing_target, chat_id, thread_id) - asyncio.create_task(self._add_reaction(message)) - return - - # Thread not tracked (orphaned) — create new thread and handle below - logger.debug("[Discord] message in orphaned thread %s, will create new thread", thread_id) + async def _handle_admitted_message( + self, + message, + *, + text: str, + msg_type: InboundMessageType, + guild, + has_mention: bool, + reservation: InboundReservation, + ) -> None: + """Finish Discord routing while owning one bounded intake slot.""" + reservation_transferred = False + try: + # --- Determine thread/channel routing and typing target --- thread_id = None - typing_target = None + chat_id = None + typing_target = None # The Discord object to type into - # At this point we're guaranteed to be in a channel, not a thread - # (the Thread case is handled above). Apply mention_only for all - # non-thread messages — no special case needed. - channel_id = str(message.channel.id) + if isinstance(message.channel, self._discord_module.Thread): + # --- Message already inside a thread --- + thread_obj = message.channel + thread_id = str(thread_obj.id) + chat_id = str(thread_obj.parent_id or thread_obj.id) + typing_target = thread_obj - # Check if there's an active thread for this channel - if channel_id in self._active_threads: - # respect mention_only: if enabled, only process messages that mention the bot - # (unless the channel is in allowed_channels) - # Messages within a thread are always allowed through (continuation). - # At this code point we know the message is in a channel, not a thread - # (Thread case handled above), so always apply the check. - if self._mention_only and not has_mention and channel_id not in self._allowed_channels: - logger.debug("[Discord] skipping no-@ message in channel %s (not in thread)", channel_id) + # If this is a known active thread, process normally + if thread_id in self._active_thread_ids: + inbound = self._make_inbound( + chat_id=chat_id, + user_id=str(message.author.id), + text=text, + msg_type=msg_type, + thread_ts=thread_id, + metadata={ + "guild_id": str(guild.id) if guild else None, + "channel_id": str(message.channel.id), + "message_id": str(message.id), + }, + ) + inbound.topic_id = thread_id + inbound = await self._attach_connection_identity(inbound, guild_id=str(guild.id) if guild else None) + if not self._publish_reserved(inbound, reservation): + return + reservation_transferred = True + # Start typing indicator in the thread + if typing_target: + await self._start_typing(typing_target, chat_id, thread_id) + asyncio.create_task(self._add_reaction(message)) + return + + # Thread not tracked (orphaned) — create new thread and handle below + logger.debug("[Discord] message in orphaned thread %s, will create new thread", thread_id) + thread_id = None + typing_target = None + + # At this point we're guaranteed to be in a channel, not a thread + # (the Thread case is handled above). Apply mention_only for all + # non-thread messages — no special case needed. + channel_id = str(message.channel.id) + + # Check if there's an active thread for this channel + if channel_id in self._active_threads: + # respect mention_only: if enabled, only process messages that mention the bot + # (unless the channel is in allowed_channels) + # Messages within a thread are always allowed through (continuation). + # At this code point we know the message is in a channel, not a thread + # (Thread case handled above), so always apply the check. + if self._mention_only and not has_mention and channel_id not in self._allowed_channels: + logger.debug("[Discord] skipping no-@ message in channel %s (not in thread)", channel_id) + return + # mention_only + fresh @ → create new thread instead of routing to existing one + if self._mention_only and has_mention: + thread_obj = await self._create_thread(message) + if thread_obj is not None: + target_thread_id = str(thread_obj.id) + self._record_thread_mapping(channel_id, target_thread_id) + await asyncio.to_thread(self._persist_thread_mappings) + thread_id = target_thread_id + chat_id = channel_id + typing_target = thread_obj + logger.info("[Discord] created new thread %s in channel %s on mention (replacing existing thread)", target_thread_id, channel_id) + else: + logger.info("[Discord] thread creation failed in channel %s, falling back to channel replies", channel_id) + thread_id = channel_id + chat_id = channel_id + typing_target = message.channel + else: + # Existing session → route to the existing thread + target_thread_id = self._active_threads[channel_id] + logger.debug("[Discord] routing message in channel %s to existing thread %s", channel_id, target_thread_id) + thread_id = target_thread_id + chat_id = channel_id + typing_target = await self._get_channel_or_thread(target_thread_id) + elif self._mention_only and not has_mention and channel_id not in self._allowed_channels: + # Not mentioned and not in an allowed channel → skip + logger.debug("[Discord] skipping message without mention in channel %s", channel_id) return - # mention_only + fresh @ → create new thread instead of routing to existing one - if self._mention_only and has_mention: + elif self._mention_only and has_mention: + # First mention in this channel → create thread thread_obj = await self._create_thread(message) if thread_obj is not None: target_thread_id = str(thread_obj.id) @@ -438,91 +509,74 @@ class DiscordChannel(Channel): await asyncio.to_thread(self._persist_thread_mappings) thread_id = target_thread_id chat_id = channel_id - typing_target = thread_obj - logger.info("[Discord] created new thread %s in channel %s on mention (replacing existing thread)", target_thread_id, channel_id) + typing_target = thread_obj # Type into the new thread + logger.info("[Discord] created thread %s in channel %s for user %s", target_thread_id, channel_id, message.author.display_name) else: + # Fallback: thread creation failed (disabled/permissions), reply in channel logger.info("[Discord] thread creation failed in channel %s, falling back to channel replies", channel_id) thread_id = channel_id chat_id = channel_id - typing_target = message.channel + typing_target = message.channel # Type into the channel + elif self._thread_mode: + # thread_mode but mention_only is False → create thread anyway for conversation grouping + thread_obj = await self._create_thread(message) + if thread_obj is None: + # Thread creation failed (disabled/permissions), fall back to channel replies + logger.info("[Discord] thread creation failed in channel %s, falling back to channel replies", channel_id) + thread_id = channel_id + chat_id = channel_id + typing_target = message.channel # Type into the channel + else: + target_thread_id = str(thread_obj.id) + self._record_thread_mapping(channel_id, target_thread_id) + await asyncio.to_thread(self._persist_thread_mappings) + thread_id = target_thread_id + chat_id = channel_id + typing_target = thread_obj # Type into the new thread else: - # Existing session → route to the existing thread - target_thread_id = self._active_threads[channel_id] - logger.debug("[Discord] routing message in channel %s to existing thread %s", channel_id, target_thread_id) - thread_id = target_thread_id - chat_id = channel_id - typing_target = await self._get_channel_or_thread(target_thread_id) - elif self._mention_only and not has_mention and channel_id not in self._allowed_channels: - # Not mentioned and not in an allowed channel → skip - logger.debug("[Discord] skipping message without mention in channel %s", channel_id) - return - elif self._mention_only and has_mention: - # First mention in this channel → create thread - thread_obj = await self._create_thread(message) - if thread_obj is not None: - target_thread_id = str(thread_obj.id) - self._record_thread_mapping(channel_id, target_thread_id) - await asyncio.to_thread(self._persist_thread_mappings) - thread_id = target_thread_id - chat_id = channel_id - typing_target = thread_obj # Type into the new thread - logger.info("[Discord] created thread %s in channel %s for user %s", target_thread_id, channel_id, message.author.display_name) - else: - # Fallback: thread creation failed (disabled/permissions), reply in channel - logger.info("[Discord] thread creation failed in channel %s, falling back to channel replies", channel_id) + # No threading — reply directly in channel thread_id = channel_id chat_id = channel_id typing_target = message.channel # Type into the channel - elif self._thread_mode: - # thread_mode but mention_only is False → create thread anyway for conversation grouping - thread_obj = await self._create_thread(message) - if thread_obj is None: - # Thread creation failed (disabled/permissions), fall back to channel replies - logger.info("[Discord] thread creation failed in channel %s, falling back to channel replies", channel_id) - thread_id = channel_id - chat_id = channel_id - typing_target = message.channel # Type into the channel - else: - target_thread_id = str(thread_obj.id) - self._record_thread_mapping(channel_id, target_thread_id) - await asyncio.to_thread(self._persist_thread_mappings) - thread_id = target_thread_id - chat_id = channel_id - typing_target = thread_obj # Type into the new thread - else: - # No threading — reply directly in channel - thread_id = channel_id - chat_id = channel_id - typing_target = message.channel # Type into the channel - msg_type = InboundMessageType.COMMAND if is_known_channel_command(text) else InboundMessageType.CHAT - inbound = self._make_inbound( - chat_id=chat_id, - user_id=str(message.author.id), - text=text, - msg_type=msg_type, - thread_ts=thread_id, - metadata={ - "guild_id": str(guild.id) if guild else None, - "channel_id": str(message.channel.id), - "message_id": str(message.id), - }, - ) - inbound.topic_id = thread_id - inbound = await self._attach_connection_identity(inbound, guild_id=str(guild.id) if guild else None) + inbound = self._make_inbound( + chat_id=chat_id, + user_id=str(message.author.id), + text=text, + msg_type=msg_type, + thread_ts=thread_id, + metadata={ + "guild_id": str(guild.id) if guild else None, + "channel_id": str(message.channel.id), + "message_id": str(message.id), + }, + ) + inbound.topic_id = thread_id + inbound = await self._attach_connection_identity(inbound, guild_id=str(guild.id) if guild else None) - # Start typing indicator in the correct target (thread or channel) - if typing_target: - await self._start_typing(typing_target, chat_id, thread_id) + if not self._publish_reserved(inbound, reservation): + return + reservation_transferred = True - self._publish(inbound) - asyncio.create_task(self._add_reaction(message)) + # Start typing/reaction only after bounded admission succeeds. + if typing_target: + await self._start_typing(typing_target, chat_id, thread_id) + asyncio.create_task(self._add_reaction(message)) + finally: + if not reservation_transferred: + reservation.release() - def _publish(self, inbound) -> None: - """Publish an inbound message to the main event loop.""" + def _publish_reserved(self, inbound: InboundMessage, reservation: InboundReservation) -> bool: + """Transfer an already-reserved message to the Gateway loop.""" if self._main_loop and self._main_loop.is_running(): - future = asyncio.run_coroutine_threadsafe(self.bus.publish_inbound(inbound), self._main_loop) - future.add_done_callback(lambda f: logger.exception("[Discord] publish_inbound failed", exc_info=f.exception()) if f.exception() else None) + try: + # This is called from discord.py's private event-loop thread. + self._main_loop.call_soon_threadsafe(self._commit_reserved_inbound, reservation, inbound) + return True + except RuntimeError: + logger.info("[Discord] main loop stopped before reserved inbound could be scheduled") + reservation.release() + return False async def _attach_connection_identity(self, inbound: InboundMessage, guild_id: str | None = None) -> InboundMessage: return await attach_connection_identity( diff --git a/backend/app/channels/feishu.py b/backend/app/channels/feishu.py index 72939deb4..417bbe9c3 100644 --- a/backend/app/channels/feishu.py +++ b/backend/app/channels/feishu.py @@ -18,6 +18,7 @@ from app.channels.message_bus import ( RESOLVED_FROM_PENDING_CLARIFICATION_METADATA_KEY, InboundMessage, InboundMessageType, + InboundReservation, MessageBus, OutboundMessage, ResolvedAttachment, @@ -200,6 +201,7 @@ class FeishuChannel(Channel): logger.info("[Feishu] using domain: %s", domain) self._main_loop = asyncio.get_event_loop() + self._open_threadsafe_future_intake() self._running = True self.bus.subscribe_outbound(self._on_outbound) @@ -259,11 +261,17 @@ class FeishuChannel(Channel): async def stop(self) -> None: self._running = False self.bus.unsubscribe_outbound(self._on_outbound) - for task in list(self._background_tasks): + await self._close_and_drain_threadsafe_futures() + + background_tasks = tuple(self._background_tasks) + running_card_tasks = tuple(self._running_card_tasks.values()) + for task in background_tasks: task.cancel() + for task in running_card_tasks: + task.cancel() + if background_tasks or running_card_tasks: + await asyncio.gather(*background_tasks, *running_card_tasks, return_exceptions=True) self._background_tasks.clear() - for task in list(self._running_card_tasks.values()): - task.cancel() self._running_card_tasks.clear() if self._thread: self._thread.join(timeout=5) @@ -832,19 +840,37 @@ class FeishuChannel(Channel): source_message_ids: list[str] | None = None, ) -> None: if self._main_loop and self._main_loop.is_running(): + reservation = self._reserve_inbound(inbound) + if reservation is None: + return logger.info("[Feishu] publishing inbound message to bus (type=%s, msg_id=%s)", inbound.msg_type.value, msg_id) - fut = asyncio.run_coroutine_threadsafe( - self._prepare_inbound(msg_id, inbound, source_message_ids=source_message_ids), + future = self._submit_threadsafe_coroutine( + self._prepare_inbound( + msg_id, + inbound, + source_message_ids=source_message_ids, + reservation=reservation, + ), self._main_loop, + name="prepare_inbound", + msg_id=msg_id, + reservation=reservation, ) - fut.add_done_callback(lambda f, mid=msg_id: self._log_future_error(f, "prepare_inbound", mid)) + if future is None: + logger.info("[Feishu] main loop stopped before reserved inbound could be scheduled") else: logger.warning("[Feishu] main loop not running, cannot publish inbound message") def _schedule_batch_flush(self, key: tuple[str, str], source_message_id: str) -> None: if self._main_loop and self._main_loop.is_running(): - fut = asyncio.run_coroutine_threadsafe(self._flush_pending_inbound_batch_after(key, source_message_id), self._main_loop) - fut.add_done_callback(lambda f, mid=source_message_id: self._log_future_error(f, "flush_inbound_batch", mid)) + future = self._submit_threadsafe_coroutine( + self._flush_pending_inbound_batch_after(key, source_message_id), + self._main_loop, + name="flush_inbound_batch", + msg_id=source_message_id, + ) + if future is None: + logger.info("[Feishu] main loop stopped before inbound batch flush could be scheduled") else: logger.warning("[Feishu] main loop not running, cannot flush inbound batch") @@ -909,6 +935,9 @@ class FeishuChannel(Channel): if not batch: return anchor_message_id, inbound, source_message_ids = batch + reservation = self._reserve_inbound(inbound) + if reservation is None: + return logger.info( "[Feishu] flushing inbound file batch: chat_id=%s user_id=%s anchor=%s messages=%d files=%d", inbound.chat_id, @@ -917,7 +946,12 @@ class FeishuChannel(Channel): len(source_message_ids), len(inbound.files), ) - await self._prepare_inbound(anchor_message_id, inbound, source_message_ids=source_message_ids) + await self._prepare_inbound( + anchor_message_id, + inbound, + source_message_ids=source_message_ids, + reservation=reservation, + ) @staticmethod def _log_task_error(task: asyncio.Task, name: str, msg_id: str) -> None: @@ -931,15 +965,29 @@ class FeishuChannel(Channel): except Exception: pass - async def _prepare_inbound(self, msg_id: str, inbound, *, source_message_ids: list[str] | None = None) -> None: + async def _prepare_inbound( + self, + msg_id: str, + inbound, + *, + source_message_ids: list[str] | None = None, + reservation: InboundReservation | None = None, + ) -> None: """Kick off Feishu side effects without delaying inbound dispatch.""" - inbound = await self._attach_connection_identity(inbound) - reaction_message_ids = source_message_ids or [msg_id] - for reaction_message_id in reaction_message_ids: - reaction_task = asyncio.create_task(self._add_reaction(reaction_message_id, "OK")) - self._track_background_task(reaction_task, name="add_reaction", msg_id=reaction_message_id) - self._ensure_running_card_started(msg_id, metadata=inbound.metadata) - await self.bus.publish_inbound(inbound) + try: + inbound = await self._attach_connection_identity(inbound) + reaction_message_ids = source_message_ids or [msg_id] + for reaction_message_id in reaction_message_ids: + reaction_task = asyncio.create_task(self._add_reaction(reaction_message_id, "OK")) + self._track_background_task(reaction_task, name="add_reaction", msg_id=reaction_message_id) + self._ensure_running_card_started(msg_id, metadata=inbound.metadata) + if reservation is None: + await self.bus.publish_inbound(inbound) + else: + self._commit_reserved_inbound(reservation, inbound) + finally: + if reservation is not None: + reservation.release() async def _attach_connection_identity(self, inbound: InboundMessage) -> InboundMessage: return await attach_connection_identity( @@ -1067,7 +1115,7 @@ class FeishuChannel(Channel): connect_code = self._pending_connect_code(text) if connect_code: if self._main_loop and self._main_loop.is_running(): - fut = asyncio.run_coroutine_threadsafe( + future = self._submit_threadsafe_coroutine( self._bind_connection_from_connect_code( message_id=msg_id, chat_id=chat_id, @@ -1075,8 +1123,11 @@ class FeishuChannel(Channel): code=connect_code, ), self._main_loop, + name="bind_connection", + msg_id=msg_id, ) - fut.add_done_callback(lambda f, mid=msg_id: self._log_future_error(f, "bind_connection", mid)) + if future is None: + logger.info("[Feishu] main loop stopped before channel connection bind could be scheduled") else: logger.warning("[Feishu] main loop not running, cannot bind channel connection") return diff --git a/backend/app/channels/manager.py b/backend/app/channels/manager.py index d7982006f..eb75407fb 100644 --- a/backend/app/channels/manager.py +++ b/backend/app/channels/manager.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio import logging +import math import mimetypes import re import time @@ -54,6 +55,8 @@ logger = logging.getLogger(__name__) DEFAULT_LANGGRAPH_URL = "http://localhost:8001/api" DEFAULT_GATEWAY_URL = "http://localhost:8001" DEFAULT_ASSISTANT_ID = "lead_agent" +DEFAULT_CHANNEL_MAX_CONCURRENCY = 5 +DEFAULT_CHANNEL_SHUTDOWN_GRACE_PERIOD_SECONDS = 3.0 CUSTOM_AGENT_NAME_PATTERN = re.compile(r"^[A-Za-z0-9-]+$") # Lead-agent recursion budget (LangGraph super-steps for the lead graph only). @@ -984,7 +987,8 @@ class ChannelManager: bus: MessageBus, store: ChannelStore, *, - max_concurrency: int = 5, + max_concurrency: int = DEFAULT_CHANNEL_MAX_CONCURRENCY, + shutdown_grace_period_seconds: float = DEFAULT_CHANNEL_SHUTDOWN_GRACE_PERIOD_SECONDS, langgraph_url: str = DEFAULT_LANGGRAPH_URL, gateway_url: str = DEFAULT_GATEWAY_URL, assistant_id: str = DEFAULT_ASSISTANT_ID, @@ -995,9 +999,14 @@ class ChannelManager: inbound_dedupe_store: InboundDedupeStore | None = None, get_stream_bridge: Callable[[], StreamBridge | None] | None = None, ) -> None: + if isinstance(max_concurrency, bool) or not isinstance(max_concurrency, int) or max_concurrency <= 0: + raise ValueError("max_concurrency must be a positive integer") + if isinstance(shutdown_grace_period_seconds, bool) or not isinstance(shutdown_grace_period_seconds, (int, float)) or not math.isfinite(shutdown_grace_period_seconds) or shutdown_grace_period_seconds < 0: + raise ValueError("shutdown_grace_period_seconds must be a non-negative finite number") self.bus = bus self.store = store self._max_concurrency = max_concurrency + self._shutdown_grace_period_seconds = float(shutdown_grace_period_seconds) self._langgraph_url = langgraph_url self._gateway_url = gateway_url self._assistant_id = assistant_id @@ -1023,7 +1032,6 @@ class ChannelManager: self._serialized_thread_runs: dict[tuple[str, str], _SerializedThreadRunState] = {} self._skill_storage: SkillStorage | None = None self._csrf_token = generate_csrf_token() - self._semaphore: asyncio.Semaphore | None = None self._running = False # Distinct from self._running: that flag is also False before the # very first start() (so tests that call internal drain/handler @@ -1031,7 +1039,11 @@ class ChannelManager: # working unchanged). self._stopped tracks specifically whether # stop() has run, for the follow-up drain guard below. self._stopped = False - self._task: asyncio.Task | None = None + # Fixed, long-lived workers own message handling end to end. The pool + # size is the only number of handler coroutines that can exist; inbound + # bursts remain as bounded queue entries instead of semaphore-waiting + # task-per-message fan-out. + self._worker_tasks: set[asyncio.Task[None]] = set() # Inbound webhook dedupe store. Defaults to the in-process Memory store # (pre-#4120 behavior). Multi-pod deployments inject a shared store so # duplicate deliveries landing on different pods are collapsed. @@ -1211,7 +1223,7 @@ class ChannelManager: threaded in — e.g. a ``ChannelManager`` constructed directly without going through ``start_channel_service()`` — so tests and any not-yet-wired deployment never see a dangling background task for - this. When wired, mirrors the existing ``_dispatch_loop`` pattern: + this. When wired, mirrors the worker-task error-reporting pattern: ``asyncio.create_task`` + ``add_done_callback(self._log_task_error)`` so an unexpected watcher failure is surfaced in the logs instead of silently vanishing. The task is also tracked in @@ -1219,7 +1231,7 @@ class ChannelManager: so ``stop()`` can cancel+await any watcher still in flight instead of leaving it to fire a follow-up run after shutdown. """ - if self._get_stream_bridge is None: + if self._get_stream_bridge is None or self._stopped: return run_id = run_result.get("run_id") if isinstance(run_result, dict) else None @@ -1566,77 +1578,174 @@ class ChannelManager: # -- lifecycle --------------------------------------------------------- async def start(self) -> None: - """Start the dispatch loop.""" + """Open inbound admission and start the fixed message-worker pool.""" if self._running: return + unfinished_workers = {task for task in self._worker_tasks if not task.done()} + if unfinished_workers: + raise RuntimeError("cannot restart ChannelManager while workers are still running") + self._worker_tasks.clear() self._running = True self._stopped = False - self._semaphore = asyncio.Semaphore(self._max_concurrency) - self._task = asyncio.create_task(self._dispatch_loop()) - logger.info("ChannelManager started (max_concurrency=%d)", self._max_concurrency) + self.bus.open_inbound() + self._worker_tasks = { + asyncio.create_task( + self._worker_loop(worker_index), + name=f"deerflow-channel-worker-{worker_index}", + ) + for worker_index in range(self._max_concurrency) + } + for task in self._worker_tasks: + task.add_done_callback(self._log_task_error) + task.add_done_callback(self._worker_tasks.discard) + logger.info( + "ChannelManager started (workers=%d, inbound_queue_maxsize=%d, shutdown_grace=%.1fs)", + self._max_concurrency, + self.bus.inbound_queue_maxsize, + self._shutdown_grace_period_seconds, + ) - async def stop(self) -> None: - """Stop the dispatch loop.""" + def begin_shutdown(self) -> int: + """Close admission while leaving workers alive to drain accepted work.""" self._running = False self._stopped = True - if self._task: - self._task.cancel() - try: - await self._task - except asyncio.CancelledError: - pass - self._task = None + return self.bus.close_inbound() - # Follow-up watchers are long-lived background tasks (they await a - # run's full stream, which can take minutes) started outside the - # dispatch loop, so cancelling self._task above does not touch them. - # Left unmanaged, one still subscribed to a run that ends AFTER this - # point would drain its buffer and fire a brand new runs.create() - # into a manager that has already been stopped. + async def stop(self) -> None: + """Drain accepted work, then cancel and await every owned task. + + Admission closes immediately, but workers keep processing the accepted + queue for ``shutdown_grace_period_seconds`` so provider acknowledgments + are normally followed by a final response. Once that grace expires, + workers are cancelled. A successful return means no worker or follow-up + watcher can continue using channel resources. The Gateway lifespan owns + the process-level timeout; if it cancels this coroutine, transports stay + attached to the service and shutdown can be retried. + """ + invalidated_reservations = self.begin_shutdown() + loop = asyncio.get_running_loop() + grace_deadline = loop.time() + self._shutdown_grace_period_seconds + worker_tasks = list(self._worker_tasks) watcher_tasks = list(self._followup_watcher_tasks) + + # Watchers can create follow-up runs and are not acknowledgments for + # this accepted queue. Stop them immediately; active workers are still + # allowed to finish their current/queued messages during the grace. for task in watcher_tasks: task.cancel() - for task in watcher_tasks: - try: - await task - except asyncio.CancelledError: - pass - except Exception: - logger.exception("[Manager] follow-up watcher task raised during stop()") - self._followup_watcher_tasks.clear() + + join_task = asyncio.create_task(self.bus.join_inbound(), name="deerflow-channel-inbound-drain") + drained = False + try: + done, _ = await asyncio.wait({join_task}, timeout=max(0.0, grace_deadline - loop.time())) + drained = join_task in done and not join_task.cancelled() and join_task.exception() is None + except asyncio.CancelledError: + for task in worker_tasks: + task.cancel() + discarded_messages = self.bus.discard_pending_inbound() + if not join_task.done(): + join_task.cancel() + if invalidated_reservations or discarded_messages: + logger.warning( + "[Manager] cancelled shutdown rejected pending inbound work: reservations=%d, queued_messages=%d", + invalidated_reservations, + discarded_messages, + ) + raise + + if not drained: + logger.warning( + "[Manager] graceful shutdown period expired after %.1fs; cancelling active handlers", + self._shutdown_grace_period_seconds, + ) + + for task in worker_tasks: + task.cancel() + + # Anything still queued never started and must not keep Queue.join() + # pending after its workers have been cancelled. + discarded_messages = self.bus.discard_pending_inbound() + if not join_task.done(): + join_task.cancel() + + owned_tasks = {task for task in (*worker_tasks, *watcher_tasks, join_task) if not task.done()} + if owned_tasks: + await asyncio.gather(*owned_tasks, return_exceptions=True) + + for task in tuple(self._worker_tasks): + if task.done(): + self._worker_tasks.discard(task) + for task in tuple(self._followup_watcher_tasks): + if task.done(): + self._followup_watcher_tasks.discard(task) + + if invalidated_reservations or discarded_messages: + logger.warning( + "[Manager] shutdown rejected pending inbound work: reservations=%d, queued_messages=%d", + invalidated_reservations, + discarded_messages, + ) logger.info("ChannelManager stopped") - # -- dispatch loop ----------------------------------------------------- + # -- worker pool ------------------------------------------------------- - async def _dispatch_loop(self) -> None: - logger.info("[Manager] dispatch loop started, waiting for inbound messages") - while self._running: + async def _worker_loop(self, worker_index: int) -> None: + logger.info("[Manager] inbound worker %d started", worker_index) + # Closing admission flips ``_running`` to False, but queued messages + # were already accepted (and providers may already have acknowledged + # them). Keep draining until the queue is empty; stop() then cancels + # workers that are idle in get_inbound(). + while self._running or not self.bus.inbound_queue.empty(): try: - msg = await asyncio.wait_for(self.bus.get_inbound(), timeout=1.0) - except TimeoutError: - continue + msg = await self.bus.get_inbound() except asyncio.CancelledError: - break + raise - # Dedupe before logging "received" so a provider retrying an event N - # times does not log N accepts; duplicates are logged once as ignored. - # Note: this manager-level dedupe only guards the agent run / final - # answer. Provider adapters may emit ack side-effects (a "Working on - # it…" reply, an "eyes" reaction) before publish_inbound, so those are - # intentionally not deduped here. - if await self._is_duplicate_inbound(msg): - continue - logger.info( - "[Manager] received inbound: channel=%s, chat_id=%s, type=%s, text_len=%d, files=%d", - msg.channel_name, - msg.chat_id, - msg.msg_type.value, - len(msg.text or ""), - len(msg.files), - ) - task = asyncio.create_task(self._handle_message(msg)) - task.add_done_callback(self._log_task_error) + dedupe_recorded = False + try: + # Dedupe before logging "received" so a provider retrying an + # event N times does not log N accepts. Provider ack side + # effects may still happen before this manager-level dedupe. + if await self._is_duplicate_inbound(msg): + continue + dedupe_recorded = self._inbound_dedupe_key(msg) is not None + logger.info( + "[Manager] received inbound: channel=%s, chat_id=%s, type=%s, text_len=%d, files=%d", + msg.channel_name, + msg.chat_id, + msg.msg_type.value, + len(msg.text or ""), + len(msg.files), + ) + # Deliberately awaited inline: never create a task per message. + await self._handle_message(msg) + except asyncio.CancelledError: + # A cancellation after dedupe admission must make provider + # redelivery retryable rather than retaining a TTL-long key for + # work that never completed. + if dedupe_recorded: + try: + await self._release_inbound_dedupe_key(msg) + except Exception: + logger.exception("[Manager] failed to release inbound dedupe key during worker cancellation") + raise + except Exception: + logger.exception( + "[Manager] inbound worker %d failed handling channel=%s chat_id=%s", + worker_index, + msg.channel_name, + msg.chat_id, + ) + if dedupe_recorded: + try: + await self._release_inbound_dedupe_key(msg) + except Exception: + # A dedupe backend outage must not shrink the fixed + # worker pool by letting cleanup escape this loop. + logger.exception("[Manager] failed to release inbound dedupe key after worker error") + finally: + self.bus.inbound_task_done() @staticmethod def _inbound_dedupe_key(msg: InboundMessage) -> tuple[str, str, str, str] | None: @@ -1716,8 +1825,7 @@ class ChannelManager: async def _handle_message(self, msg: InboundMessage) -> None: msg = _apply_effective_owner(msg) try: - # Non-command chat can be rejected before it consumes a semaphore - # slot. Commands are handled below because provider adapters consume + # Commands are handled below because provider adapters consume # binding commands before manager dispatch, and _handle_command() # applies its own admission gate for manager-level commands. bound_identity_rejection = None @@ -1727,11 +1835,10 @@ class ChannelManager: await self._reject_unbound_channel_message(msg, bound_identity_rejection=bound_identity_rejection) return - async with self._semaphore: - if msg.msg_type == InboundMessageType.COMMAND: - await self._handle_command(msg) - else: - await self._handle_chat(msg, bound_identity_checked=True) + if msg.msg_type == InboundMessageType.COMMAND: + await self._handle_command(msg) + else: + await self._handle_chat(msg, bound_identity_checked=True) except InvalidChannelSessionConfigError as exc: logger.warning( "Invalid channel session config for %s (chat=%s): %s", diff --git a/backend/app/channels/message_bus.py b/backend/app/channels/message_bus.py index 25fc3f1d8..c36ad5563 100644 --- a/backend/app/channels/message_bus.py +++ b/backend/app/channels/message_bus.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio import logging +import threading import time from collections.abc import Callable, Coroutine from dataclasses import dataclass, field @@ -13,6 +14,8 @@ from typing import Any logger = logging.getLogger(__name__) +DEFAULT_INBOUND_QUEUE_MAXSIZE = 1000 + PENDING_CLARIFICATION_METADATA_KEY = "pending_clarification" RESOLVED_FROM_PENDING_CLARIFICATION_METADATA_KEY = "resolved_from_pending_clarification" # Adapter-owned bytes may use this transient key while crossing the channel @@ -134,6 +137,40 @@ class OutboundMessage: OutboundCallback = Callable[[OutboundMessage], Coroutine[Any, Any, None]] +class InboundQueueFullError(RuntimeError): + """Raised when bounded inbound admission has no capacity.""" + + +class InboundQueueClosedError(RuntimeError): + """Raised when inbound admission has closed during shutdown.""" + + +class InboundReservationExpiredError(RuntimeError): + """Raised when a reservation was already committed or invalidated.""" + + +class InboundReservation: + """One capacity slot reserved before a provider hands work to the bus. + + Some provider SDKs invoke DeerFlow on a foreign thread. Reserving before + scheduling their final identity/ack preparation onto the Gateway loop + bounds both the queue and those scheduled callbacks. A reservation must be + committed exactly once or released in a ``finally`` block. + """ + + def __init__(self, bus: MessageBus, token: object) -> None: + self._bus = bus + self._token = token + + def commit(self, msg: InboundMessage) -> None: + """Commit the reserved message from the MessageBus event loop.""" + self._bus._commit_inbound(self._token, msg) + + def release(self) -> None: + """Release the slot if it has not already been committed or closed.""" + self._bus._release_inbound_reservation(self._token) + + class MessageBus: """Async pub/sub hub connecting channels and the agent dispatcher. @@ -142,15 +179,86 @@ class MessageBus: via registered callbacks. """ - def __init__(self) -> None: - self._inbound_queue: asyncio.Queue[InboundMessage] = asyncio.Queue() + def __init__(self, *, inbound_queue_maxsize: int = DEFAULT_INBOUND_QUEUE_MAXSIZE) -> None: + if isinstance(inbound_queue_maxsize, bool) or not isinstance(inbound_queue_maxsize, int) or inbound_queue_maxsize <= 0: + raise ValueError("inbound_queue_maxsize must be a positive integer") + + self._inbound_queue: asyncio.Queue[InboundMessage] = asyncio.Queue(maxsize=inbound_queue_maxsize) + # Provider callbacks may reserve capacity from SDK-owned threads before + # scheduling async identity/ack preparation on the Gateway loop. + self._inbound_admission_lock = threading.Lock() + self._inbound_queued = 0 + self._inbound_reservations: set[object] = set() + self._accepting_inbound = True + self._full_rejection_count = 0 + self._last_full_warning_at = 0.0 self._outbound_listeners: list[OutboundCallback] = [] # -- inbound ----------------------------------------------------------- async def publish_inbound(self, msg: InboundMessage) -> None: - """Enqueue an inbound message from a channel.""" - await self._inbound_queue.put(msg) + """Admit a message immediately or raise instead of waiting for space. + + This deliberately uses reservation + ``put_nowait`` rather than an + awaited ``Queue.put``. Under overload, producers get an explicit + rejection and cannot accumulate an unbounded set of pending put tasks. + The initial zero-delay sleep is a scheduling handoff, not a capacity + wait: a producer that publishes a batch (notably GitHub webhook + fan-out) gives fixed workers a chance to dequeue between entries so a + batch larger than the queue does not repeatedly fail on the same + prefix during redelivery. It occurs before admission so cancellation + cannot report failure after this call already committed the message. + """ + await asyncio.sleep(0) + reservation = self.reserve_inbound(msg) + try: + reservation.commit(msg) + finally: + reservation.release() + + def reserve_inbound(self, msg: InboundMessage) -> InboundReservation: + """Reserve one bounded intake slot, safely callable from SDK threads.""" + token = object() + should_warn = False + rejection_count = 0 + with self._inbound_admission_lock: + if not self._accepting_inbound: + raise InboundQueueClosedError("channel inbound intake is closed") + admitted = self._inbound_queued + len(self._inbound_reservations) + if admitted >= self._inbound_queue.maxsize: + self._full_rejection_count += 1 + rejection_count = self._full_rejection_count + now = time.monotonic() + if now - self._last_full_warning_at >= 1.0: + self._last_full_warning_at = now + should_warn = True + else: + self._inbound_reservations.add(token) + return InboundReservation(self, token) + + if should_warn: + logger.warning( + "[Bus] inbound capacity exhausted: channel=%s, chat_id=%s, capacity=%d, rejected_total=%d", + msg.channel_name, + msg.chat_id, + self._inbound_queue.maxsize, + rejection_count, + ) + raise InboundQueueFullError(f"channel inbound queue is full (capacity={self._inbound_queue.maxsize})") + + def _commit_inbound(self, token: object, msg: InboundMessage) -> None: + with self._inbound_admission_lock: + if token not in self._inbound_reservations: + raise InboundReservationExpiredError("inbound reservation is no longer active") + self._inbound_reservations.remove(token) + if not self._accepting_inbound: + raise InboundQueueClosedError("channel inbound intake closed before reservation commit") + try: + self._inbound_queue.put_nowait(msg) + except asyncio.QueueFull as exc: # pragma: no cover - reservation accounting invariant + raise RuntimeError("inbound reservation accounting exceeded queue capacity") from exc + self._inbound_queued += 1 + logger.info( "[Bus] inbound enqueued: channel=%s, chat_id=%s, type=%s, queue_size=%d", msg.channel_name, @@ -159,12 +267,70 @@ class MessageBus: self._inbound_queue.qsize(), ) + def _release_inbound_reservation(self, token: object) -> None: + with self._inbound_admission_lock: + if token in self._inbound_reservations: + self._inbound_reservations.remove(token) + async def get_inbound(self) -> InboundMessage: """Block until the next inbound message is available.""" - return await self._inbound_queue.get() + msg = await self._inbound_queue.get() + with self._inbound_admission_lock: + self._inbound_queued -= 1 + return msg + + def get_inbound_nowait(self) -> InboundMessage: + """Return one queued message immediately and release admission capacity.""" + msg = self._inbound_queue.get_nowait() + with self._inbound_admission_lock: + self._inbound_queued -= 1 + return msg + + def inbound_task_done(self) -> None: + """Mark one dequeued inbound message as fully handled.""" + self._inbound_queue.task_done() + + async def join_inbound(self) -> None: + """Wait until every admitted queue item has completed or been dropped.""" + await self._inbound_queue.join() + + def close_inbound(self) -> int: + """Reject new intake and invalidate uncommitted reservations. + + Returns the number of provider-side reservations invalidated. Queued + messages remain until workers finish or ``discard_pending_inbound`` is + called by shutdown. + """ + with self._inbound_admission_lock: + self._accepting_inbound = False + invalidated = len(self._inbound_reservations) + self._inbound_reservations.clear() + return invalidated + + def open_inbound(self) -> None: + """Re-open admission when a stopped manager is explicitly restarted.""" + with self._inbound_admission_lock: + self._accepting_inbound = True + + def discard_pending_inbound(self) -> int: + """Drop queued, not-yet-started messages during shutdown.""" + discarded = 0 + while True: + try: + self.get_inbound_nowait() + except asyncio.QueueEmpty: + break + self._inbound_queue.task_done() + discarded += 1 + return discarded + + @property + def inbound_queue_maxsize(self) -> int: + return self._inbound_queue.maxsize @property def inbound_queue(self) -> asyncio.Queue[InboundMessage]: + """Expose the queue for read-only size/empty inspection.""" return self._inbound_queue # -- outbound ---------------------------------------------------------- diff --git a/backend/app/channels/service.py b/backend/app/channels/service.py index 972bdc0ee..6a27710de 100644 --- a/backend/app/channels/service.py +++ b/backend/app/channels/service.py @@ -4,13 +4,14 @@ from __future__ import annotations import asyncio import logging +import math import os from collections.abc import Callable from typing import TYPE_CHECKING, Any from app.channels.base import Channel -from app.channels.manager import DEFAULT_GATEWAY_URL, DEFAULT_LANGGRAPH_URL, ChannelManager -from app.channels.message_bus import MessageBus +from app.channels.manager import DEFAULT_CHANNEL_MAX_CONCURRENCY, DEFAULT_CHANNEL_SHUTDOWN_GRACE_PERIOD_SECONDS, DEFAULT_GATEWAY_URL, DEFAULT_LANGGRAPH_URL, ChannelManager +from app.channels.message_bus import DEFAULT_INBOUND_QUEUE_MAXSIZE, MessageBus from app.channels.runtime_config_store import merge_runtime_channel_configs from app.channels.store import ChannelStore @@ -65,6 +66,26 @@ def _resolve_service_url(config: dict[str, Any], config_key: str, env_key: str, return default +def _resolve_positive_int(config: dict[str, Any], config_key: str, default: int) -> int: + value = config.pop(config_key, None) + if value is None: + return default + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + logger.warning("Invalid channels.%s=%r; using default %d", config_key, value, default) + return default + return value + + +def _resolve_non_negative_float(config: dict[str, Any], config_key: str, default: float) -> float: + value = config.pop(config_key, None) + if value is None: + return default + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0: + logger.warning("Invalid channels.%s=%r; using default %.1f", config_key, value, default) + return default + return float(value) + + def _merge_channel_connection_runtime_config(channels_config: dict[str, Any], app_config: AppConfig) -> None: connection_config = getattr(app_config, "channel_connections", None) merge_runtime_channel_configs(channels_config, connection_config) @@ -104,11 +125,14 @@ class ChannelService: app_config: AppConfig | None = None, get_stream_bridge: Callable[[], StreamBridge | None] | None = None, ) -> None: - self.bus = MessageBus() + config = dict(channels_config or {}) + inbound_queue_maxsize = _resolve_positive_int(config, "inbound_queue_maxsize", DEFAULT_INBOUND_QUEUE_MAXSIZE) + max_concurrency = _resolve_positive_int(config, "max_concurrency", DEFAULT_CHANNEL_MAX_CONCURRENCY) + shutdown_grace_period_seconds = _resolve_non_negative_float(config, "shutdown_grace_period_seconds", DEFAULT_CHANNEL_SHUTDOWN_GRACE_PERIOD_SECONDS) + self.bus = MessageBus(inbound_queue_maxsize=inbound_queue_maxsize) self.store = ChannelStore() self._connection_repo = connection_repo self._get_stream_bridge = get_stream_bridge - config = dict(channels_config or {}) langgraph_url = _resolve_service_url(config, "langgraph_url", _CHANNELS_LANGGRAPH_URL_ENV, DEFAULT_LANGGRAPH_URL) gateway_url = _resolve_service_url(config, "gateway_url", _CHANNELS_GATEWAY_URL_ENV, DEFAULT_GATEWAY_URL) default_session = config.pop("session", None) @@ -118,6 +142,8 @@ class ChannelService: self.manager = ChannelManager( bus=self.bus, store=self.store, + max_concurrency=max_concurrency, + shutdown_grace_period_seconds=shutdown_grace_period_seconds, langgraph_url=langgraph_url, gateway_url=gateway_url, default_session=default_session if isinstance(default_session, dict) else None, @@ -243,17 +269,33 @@ class ChannelService: return False async def stop(self) -> None: - """Stop all channels and the manager.""" + """Drain accepted messages while channels can still deliver replies.""" + self._running = False + # Reject new provider work first. Existing workers keep draining during + # manager.stop(), and channel transports remain alive until that drain + # completes so an already-sent "Working on it..." can still receive its + # final update. + await self.manager.stop() + stop_errors: list[Exception] = [] for name, channel in list(self._channels.items()): try: await channel.stop() - logger.info("Channel stopped") - except Exception: + except asyncio.CancelledError: + # Keep this and the remaining transports owned by the service. + # The Gateway deadline interrupted shutdown, so detaching them + # would hide resources that may still be in use. + raise + except Exception as exc: logger.exception("Error stopping channel") - self._channels.clear() + stop_errors.append(exc) + else: + if self._channels.get(name) is channel: + self._channels.pop(name, None) + logger.info("Channel stopped") + + if stop_errors: + raise ExceptionGroup("one or more channels failed to stop", stop_errors) - await self.manager.stop() - self._running = False logger.info("ChannelService stopped") def _load_channel_config(self, name: str) -> dict[str, Any] | None: @@ -460,5 +502,7 @@ async def stop_channel_service() -> None: """Stop the global ChannelService.""" global _channel_service if _channel_service is not None: - await _channel_service.stop() - _channel_service = None + service = _channel_service + await service.stop() + if _channel_service is service: + _channel_service = None diff --git a/backend/app/channels/slack.py b/backend/app/channels/slack.py index 6e0470d5c..3e7b17152 100644 --- a/backend/app/channels/slack.py +++ b/backend/app/channels/slack.py @@ -13,7 +13,7 @@ from markdown_to_mrkdwn import SlackMarkdownConverter from app.channels.base import Channel from app.channels.commands import is_known_channel_command from app.channels.connection_identity import attach_connection_identity -from app.channels.message_bus import InboundMessageType, MessageBus, OutboundMessage, ResolvedAttachment +from app.channels.message_bus import InboundMessageType, InboundReservation, MessageBus, OutboundMessage, ResolvedAttachment logger = logging.getLogger(__name__) @@ -137,6 +137,7 @@ class SlackChannel(Channel): self._socket_client.socket_mode_request_listeners.append(self._on_socket_event) + self._open_threadsafe_future_intake() self._running = True self.bus.subscribe_outbound(self._on_outbound) @@ -147,6 +148,7 @@ class SlackChannel(Channel): async def stop(self) -> None: self._running = False self.bus.unsubscribe_outbound(self._on_outbound) + await self._close_and_drain_threadsafe_futures() if self._socket_client: self._socket_client.close() self._socket_client = None @@ -290,6 +292,8 @@ class SlackChannel(Channel): def _on_socket_event(self, client, req) -> None: """Called by slack-sdk for each Socket Mode event.""" + if not self._running: + return try: # Acknowledge the event response = self._SocketModeResponse(envelope_id=req.envelope_id) @@ -334,14 +338,18 @@ class SlackChannel(Channel): connect_code = self._pending_connect_code(text) if connect_code: if self._loop and self._loop.is_running(): - asyncio.run_coroutine_threadsafe( + future = self._submit_threadsafe_coroutine( self._bind_connection_from_connect_code( event=event, team_id=str(team_id or ""), code=connect_code, ), self._loop, + name="bind_connection", + msg_id=event.get("ts"), ) + if future is None: + logger.info("[Slack] main loop stopped before channel connection bind could be scheduled") return # Check allowed users after connect-code handling so browser-initiated @@ -377,18 +385,44 @@ class SlackChannel(Channel): inbound.topic_id = thread_ts if self._loop and self._loop.is_running(): + reservation = self._reserve_inbound(inbound) + if reservation is None: + return # Acknowledge with an eyes reaction self._add_reaction(channel_id, event.get("ts", thread_ts), "eyes") # Send "running" reply first (fire-and-forget from SDK thread) self._send_running_reply(channel_id, thread_ts) - if self._connection_repo is None: - asyncio.run_coroutine_threadsafe(self.bus.publish_inbound(inbound), self._loop) - else: - asyncio.run_coroutine_threadsafe(self._publish_inbound_with_connection(inbound, team_id=team_id), self._loop) + try: + if self._connection_repo is None: + # Reservation bounds callbacks scheduled from the SDK + # thread; no coroutine/Future waits for queue capacity. + self._loop.call_soon_threadsafe(self._commit_reserved_inbound, reservation, inbound) + else: + future = self._submit_threadsafe_coroutine( + self._publish_inbound_with_connection(inbound, reservation=reservation, team_id=team_id), + self._loop, + name="publish_inbound", + msg_id=event.get("ts", thread_ts), + reservation=reservation, + ) + if future is None: + logger.info("[Slack] main loop stopped before reserved inbound could be scheduled") + except RuntimeError: + reservation.release() + logger.info("[Slack] main loop stopped before reserved inbound could be scheduled") - async def _publish_inbound_with_connection(self, inbound, *, team_id: str | None = None) -> None: - inbound = await self._attach_connection_identity(inbound, team_id=team_id) - await self.bus.publish_inbound(inbound) + async def _publish_inbound_with_connection( + self, + inbound, + *, + reservation: InboundReservation, + team_id: str | None = None, + ) -> None: + try: + inbound = await self._attach_connection_identity(inbound, team_id=team_id) + self._commit_reserved_inbound(reservation, inbound) + finally: + reservation.release() async def _attach_connection_identity(self, inbound, *, team_id: str | None = None): workspace_id = str(team_id or inbound.metadata.get("team_id") or "") diff --git a/backend/app/channels/telegram.py b/backend/app/channels/telegram.py index 3b5723292..2c9e2e123 100644 --- a/backend/app/channels/telegram.py +++ b/backend/app/channels/telegram.py @@ -15,6 +15,7 @@ from app.channels.message_bus import ( INBOUND_FILE_CONTENT_KEY, InboundMessage, InboundMessageType, + InboundReservation, MessageBus, OutboundMessage, ResolvedAttachment, @@ -97,6 +98,7 @@ class TelegramChannel(Channel): return self._main_loop = asyncio.get_event_loop() + self._open_threadsafe_future_intake() self._running = True self.bus.subscribe_outbound(self._on_outbound) @@ -134,6 +136,7 @@ class TelegramChannel(Channel): async def stop(self) -> None: self._running = False self.bus.unsubscribe_outbound(self._on_outbound) + await self._close_and_drain_threadsafe_futures() shutdown_loop = asyncio.get_running_loop() deadline = shutdown_loop.time() + TELEGRAM_SHUTDOWN_TIMEOUT_SECONDS telegram_loop = self._tg_loop @@ -799,9 +802,23 @@ class TelegramChannel(Channel): return await update.message.reply_text("Welcome to DeerFlow! Send me a message to start a conversation.\nType /help for available commands.") - async def _process_incoming_with_reply(self, chat_id: str, msg_id: int, inbound: InboundMessage) -> None: - await self._send_running_reply(chat_id, msg_id) - await self.bus.publish_inbound(inbound) + async def _process_incoming_with_reply( + self, + chat_id: str, + msg_id: int, + inbound: InboundMessage, + *, + reservation: InboundReservation | None = None, + ) -> None: + try: + await self._send_running_reply(chat_id, msg_id) + if reservation is None: + await self.bus.publish_inbound(inbound) + else: + self._commit_reserved_inbound(reservation, inbound) + finally: + if reservation is not None: + reservation.release() async def _cmd_generic(self, update, context) -> None: """Forward slash commands to the channel manager.""" @@ -833,11 +850,30 @@ class TelegramChannel(Channel): metadata={"message_id": msg_id}, ) inbound.topic_id = topic_id - inbound = await self._attach_connection_identity(inbound) if self._main_loop and self._main_loop.is_running(): - fut = asyncio.run_coroutine_threadsafe(self._process_incoming_with_reply(chat_id, update.message.message_id, inbound), self._main_loop) - fut.add_done_callback(lambda f: self._log_future_error(f, "process_incoming_with_reply", update.message.message_id)) + reservation = self._reserve_inbound(inbound) + if reservation is None: + return + try: + inbound = await self._attach_connection_identity(inbound) + future = self._submit_threadsafe_coroutine( + self._process_incoming_with_reply( + chat_id, + update.message.message_id, + inbound, + reservation=reservation, + ), + self._main_loop, + name="process_incoming_with_reply", + msg_id=update.message.message_id, + reservation=reservation, + ) + if future is None: + logger.info("[Telegram] main loop stopped before reserved command could be scheduled") + except Exception: + reservation.release() + raise else: logger.warning("[Telegram] Main loop not running. Cannot publish inbound message.") @@ -885,10 +921,29 @@ class TelegramChannel(Channel): metadata={"message_id": msg_id}, ) inbound.topic_id = topic_id - inbound = await self._attach_connection_identity(inbound) if self._main_loop and self._main_loop.is_running(): - fut = asyncio.run_coroutine_threadsafe(self._process_incoming_with_reply(chat_id, update.message.message_id, inbound), self._main_loop) - fut.add_done_callback(lambda f: self._log_future_error(f, "process_incoming_with_reply", update.message.message_id)) + reservation = self._reserve_inbound(inbound) + if reservation is None: + return + try: + inbound = await self._attach_connection_identity(inbound) + future = self._submit_threadsafe_coroutine( + self._process_incoming_with_reply( + chat_id, + update.message.message_id, + inbound, + reservation=reservation, + ), + self._main_loop, + name="process_incoming_with_reply", + msg_id=update.message.message_id, + reservation=reservation, + ) + if future is None: + logger.info("[Telegram] main loop stopped before reserved inbound could be scheduled") + except Exception: + reservation.release() + raise else: logger.warning("[Telegram] Main loop not running. Cannot publish inbound message.") diff --git a/backend/app/channels/wechat.py b/backend/app/channels/wechat.py index 194f7e78e..6a5d4b2e1 100644 --- a/backend/app/channels/wechat.py +++ b/backend/app/channels/wechat.py @@ -668,7 +668,9 @@ class WechatChannel(Channel): ) inbound.topic_id = None inbound = await self._attach_connection_identity(inbound) - await self.bus.publish_inbound(inbound) + # The iLink poll loop processes updates sequentially on the Gateway + # loop, so no provider-side task needs a pre-handoff reservation. + await self._publish_inbound_or_drop(inbound) async def _attach_connection_identity(self, inbound: InboundMessage) -> InboundMessage: return await attach_connection_identity( diff --git a/backend/app/channels/wecom.py b/backend/app/channels/wecom.py index a92d3fe22..10990f58c 100644 --- a/backend/app/channels/wecom.py +++ b/backend/app/channels/wecom.py @@ -385,13 +385,21 @@ class WeComChannel(Channel): self._ws_frames[msg_id] = frame self._ws_stream_ids[msg_id] = stream_id + reservation = self._reserve_inbound(inbound) + if reservation is None: + self._ws_frames.pop(msg_id, None) + self._ws_stream_ids.pop(msg_id, None) + return try: - await self._ws_client.reply_stream(frame, stream_id, self._working_message, False) - except Exception: - pass + try: + await self._ws_client.reply_stream(frame, stream_id, self._working_message, False) + except Exception: + pass - inbound = await self._attach_connection_identity(inbound) - await self.bus.publish_inbound(inbound) + inbound = await self._attach_connection_identity(inbound) + self._commit_reserved_inbound(reservation, inbound) + finally: + reservation.release() async def _attach_connection_identity(self, inbound: InboundMessage) -> InboundMessage: return await attach_connection_identity( diff --git a/backend/docs/IM_CHANNEL_CONNECTIONS.md b/backend/docs/IM_CHANNEL_CONNECTIONS.md index a7c64ad6e..08ab7379a 100644 --- a/backend/docs/IM_CHANNEL_CONNECTIONS.md +++ b/backend/docs/IM_CHANNEL_CONNECTIONS.md @@ -97,14 +97,14 @@ sequenceDiagram autonumber participant Platform as Provider
(Slack/Telegram/...) participant Worker as Provider worker - participant Bus as MessageBus
InboundMessage queue - participant Mgr as ChannelManager + participant Bus as MessageBus
bounded admission + queue + participant Mgr as ChannelManager
fixed worker pool participant Client as langgraph_sdk
async client participant Gateway as Gateway
/api/* routers Platform->>Worker: inbound chat message
(resolved to connection_id + owner_user_id) - Worker->>Bus: publish_inbound(InboundMessage) - Bus->>Mgr: msg = get_inbound() + Worker->>Bus: reserve by Gateway handoff, then commit InboundMessage + Bus->>Mgr: fixed worker gets msg and awaits handler inline Mgr->>Mgr: _channel_storage_user_id(msg)
→ owner-bound user_id Mgr->>Mgr: _get_bound_identity_rejection()
(re-check identity by provider+ext+ws) Mgr->>Client: _get_or_create_thread(thread_id or new) @@ -129,6 +129,18 @@ sequenceDiagram Worker->>Platform: post reply (Telegram editMessageText,
Feishu patch card, etc.) ``` +### Inbound capacity and overload behavior + +Three top-level `channels` settings control the MessageBus/manager lifecycle: `inbound_queue_maxsize` (default `1000`) covers queued messages plus provider-side reservations that may still be doing final identity/ack preparation, `max_concurrency` (default `5`) is the exact number of long-lived `ChannelManager` workers, and `shutdown_grace_period_seconds` (default `3`) bounds graceful draining before active handlers are cancelled. Active handlers run inline in those workers, so a burst cannot create a task per message. The maximum manager-owned live intake is therefore the pending capacity plus the fixed worker count. + +Admission never waits for queue space, because waiting producer coroutines would simply move the unbounded backlog outside the queue. At capacity: + +- Slack, Discord, Feishu/Lark, DingTalk, Telegram, WeChat, and WeCom drop the new message before DeerFlow sends its working acknowledgment. `MessageBus` emits a rate-limited warning with a cumulative rejection count. +- Buzz leaves the per-channel replay watermark unchanged and reconnects, allowing relay history to replay the event. +- GitHub webhook fan-out returns `503`. GitHub records the delivery as failed; an operator or recovery job can retry it through the Recent Deliveries UI or REST redelivery API (GitHub does not retry failed deliveries automatically). + +Shutdown first closes admission and cancels follow-up watchers, but keeps provider transports alive while workers drain accepted messages for up to `shutdown_grace_period_seconds`. Once that grace expires, it cancels active handlers, discards queue entries that never began, and awaits every manager-owned worker and watcher. Provider coroutines submitted from SDK threads are likewise retained, cancelled, and awaited before their channel tears down SDK resources. A successful stop therefore leaves no owned handler able to use a closed transport. The Gateway's outer shutdown timeout remains the process-level bound; if it cancels cleanup, the service retains its transports and singleton instead of reporting a successful stop or hiding unfinished ownership. + ## Sync vs Streaming Channels The two paths split on `ChannelRunPolicy.supports_streaming` (per-channel registration in `CHANNEL_CAPABILITIES`): @@ -262,6 +274,10 @@ Configure the actual IM bots under the existing `channels` block: ```yaml channels: + inbound_queue_maxsize: 1000 + max_concurrency: 5 + shutdown_grace_period_seconds: 3 + telegram: enabled: true bot_token: $TELEGRAM_BOT_TOKEN diff --git a/backend/tests/test_buzz_channel.py b/backend/tests/test_buzz_channel.py index d00f01c58..cbb9f5825 100644 --- a/backend/tests/test_buzz_channel.py +++ b/backend/tests/test_buzz_channel.py @@ -14,7 +14,7 @@ from app.channels import buzz_nostr from app.channels.base import Channel from app.channels.buzz import EDIT_MAX_BYTES, MAX_CACHED_CHANNELS, MAX_CHANNEL_SUBSCRIPTIONS, MAX_RESUBSCRIBE_ATTEMPTS, MEMBERSHIP_LOOKBACK_SECONDS, BuzzChannel, _chunk_text from app.channels.manager import CHANNEL_CAPABILITIES -from app.channels.message_bus import InboundMessageType, MessageBus, OutboundMessage +from app.channels.message_bus import InboundMessage, InboundMessageType, InboundQueueFullError, MessageBus, OutboundMessage from app.channels.run_policy import CHANNEL_RUN_POLICY from app.channels.service import _CHANNEL_CREDENTIAL_KEYS, _CHANNEL_REGISTRY @@ -217,6 +217,23 @@ def test_mentioned_allowed_author_is_published(): assert msg.msg_type == InboundMessageType.CHAT +def test_full_intake_escapes_relay_frame_handler_without_advancing_watermark(): + async def run(): + ch = _channel() + ch.bus = MessageBus(inbound_queue_maxsize=1) + ch._publish = ch.bus.publish_inbound + ch._keys = buzz_nostr.parse_private_key(SK3_HEX) + await ch.bus.publish_inbound(InboundMessage(channel_name="test", chat_id="busy", user_id="busy", text="busy")) + ev = _event(created_at=1700000101) + + with pytest.raises(InboundQueueFullError): + await ch.handle_relay_frame(json.dumps(["EVENT", "sub1", ev])) + + assert CHANNEL not in ch._seen_created_at + + asyncio.run(run()) + + def test_disallowed_author_is_dropped(): ch, captured = _started() _dispatch(ch, _event(sk=SK_OUTSIDER)) diff --git a/backend/tests/test_channel_intake_backpressure.py b/backend/tests/test_channel_intake_backpressure.py new file mode 100644 index 000000000..20dd86c7d --- /dev/null +++ b/backend/tests/test_channel_intake_backpressure.py @@ -0,0 +1,541 @@ +"""Regression tests for bounded IM-channel intake and handler ownership.""" + +from __future__ import annotations + +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from app.channels import service as service_module +from app.channels.manager import ChannelManager +from app.channels.message_bus import ( + InboundMessage, + InboundQueueClosedError, + InboundQueueFullError, + InboundReservationExpiredError, + MessageBus, +) +from app.channels.service import ChannelService +from app.channels.slack import SlackChannel +from app.channels.store import ChannelStore + + +def _message(index: int, *, with_dedupe_identity: bool = False) -> InboundMessage: + metadata = {"team_id": "T1", "message_id": f"m-{index}"} if with_dedupe_identity else {} + return InboundMessage( + channel_name="slack", + chat_id="C1", + user_id="U1", + text=f"message-{index}", + metadata=metadata, + ) + + +@pytest.mark.asyncio +async def test_inbound_queue_rejects_immediately_when_capacity_is_reserved() -> None: + bus = MessageBus(inbound_queue_maxsize=1) + + reservation = bus.reserve_inbound(_message(1)) + with pytest.raises(InboundQueueFullError): + bus.reserve_inbound(_message(2)) + + reservation.commit(_message(1)) + assert bus.inbound_queue.qsize() == 1 + assert bus.get_inbound_nowait().text == "message-1" + bus.inbound_task_done() + await bus.publish_inbound(_message(2)) + assert (await bus.get_inbound()).text == "message-2" + bus.inbound_task_done() + await bus.join_inbound() + + +def test_shutdown_invalidates_provider_side_reservations() -> None: + bus = MessageBus(inbound_queue_maxsize=2) + direct_reservation = bus.reserve_inbound(_message(1)) + adapter_reservation = bus.reserve_inbound(_message(2)) + channel = SlackChannel(bus=bus, config={}) + + assert bus.close_inbound() == 2 + with pytest.raises(InboundReservationExpiredError): + direct_reservation.commit(_message(1)) + direct_reservation.release() + assert channel._commit_reserved_inbound(adapter_reservation, _message(2)) is False + assert bus.inbound_queue.empty() + + +def test_provider_thread_reservations_share_one_hard_capacity_limit() -> None: + capacity = 8 + contenders = 32 + bus = MessageBus(inbound_queue_maxsize=capacity) + barrier = threading.Barrier(contenders) + + def reserve(index: int): + barrier.wait() + try: + return bus.reserve_inbound(_message(index)) + except InboundQueueFullError: + return None + + with ThreadPoolExecutor(max_workers=contenders) as executor: + reservations = list(executor.map(reserve, range(contenders))) + + admitted = [reservation for reservation in reservations if reservation is not None] + assert len(admitted) == capacity + for reservation in admitted: + reservation.release() + + +@pytest.mark.asyncio +async def test_realtime_provider_drops_before_ack_when_queue_is_full() -> None: + bus = MessageBus(inbound_queue_maxsize=1) + await bus.publish_inbound(_message(0)) + channel = SlackChannel(bus=bus, config={}) + channel._loop = MagicMock() + channel._loop.is_running.return_value = True + channel._add_reaction = MagicMock() + channel._send_running_reply = MagicMock() + + channel._handle_message_event( + { + "user": "U1", + "text": "overloaded", + "channel": "C1", + "ts": "1710000000.000100", + } + ) + + channel._add_reaction.assert_not_called() + channel._send_running_reply.assert_not_called() + channel._loop.call_soon_threadsafe.assert_not_called() + assert bus.inbound_queue.qsize() == 1 + + +@pytest.mark.asyncio +async def test_fixed_worker_pool_bounds_handler_and_queue_tasks(tmp_path: Path) -> None: + bus = MessageBus(inbound_queue_maxsize=3) + manager = ChannelManager( + bus=bus, + store=ChannelStore(path=tmp_path / "store.json"), + max_concurrency=2, + ) + release_handlers = asyncio.Event() + started: list[str] = [] + + async def hold_handler(msg: InboundMessage) -> None: + started.append(msg.text) + await release_handlers.wait() + + manager._handle_message = hold_handler # type: ignore[method-assign] + await manager.start() + try: + assert len(manager._worker_tasks) == 2 + + await bus.publish_inbound(_message(0)) + await bus.publish_inbound(_message(1)) + async with asyncio.timeout(1): + while len(started) < 2: + await asyncio.sleep(0) + + for index in range(2, 5): + await bus.publish_inbound(_message(index)) + with pytest.raises(InboundQueueFullError): + await bus.publish_inbound(_message(5)) + + assert bus.inbound_queue.qsize() == 3 + assert len(manager._worker_tasks) == 2 + # The handlers execute inline in the fixed workers. There must not be a + # separate task per admitted message waiting on a semaphore. + assert not any(getattr(task.get_coro(), "__name__", "") == "hold_handler" for task in asyncio.all_tasks()) + + release_handlers.set() + await asyncio.wait_for(bus.join_inbound(), timeout=1) + assert sorted(started) == [f"message-{index}" for index in range(5)] + finally: + release_handlers.set() + await manager.stop() + + +@pytest.mark.asyncio +async def test_stop_gracefully_drains_every_accepted_message_before_cancelling_workers(tmp_path: Path) -> None: + bus = MessageBus(inbound_queue_maxsize=2) + manager = ChannelManager( + bus=bus, + store=ChannelStore(path=tmp_path / "store.json"), + max_concurrency=1, + shutdown_grace_period_seconds=0.5, + ) + first_started = asyncio.Event() + release_first = asyncio.Event() + completed: list[str] = [] + cancelled = False + + async def graceful_handler(msg: InboundMessage) -> None: + nonlocal cancelled + if msg.text == "message-1": + first_started.set() + try: + await release_first.wait() + except asyncio.CancelledError: + cancelled = True + raise + completed.append(msg.text) + + manager._handle_message = graceful_handler # type: ignore[method-assign] + await manager.start() + await bus.publish_inbound(_message(1)) + await asyncio.wait_for(first_started.wait(), timeout=1) + await bus.publish_inbound(_message(2)) + + stop_task = asyncio.create_task(manager.stop()) + await asyncio.sleep(0.02) + + assert not stop_task.done() + assert cancelled is False + + release_first.set() + await asyncio.wait_for(stop_task, timeout=1) + + assert completed == ["message-1", "message-2"] + assert cancelled is False + assert bus.inbound_queue.empty() + await asyncio.wait_for(bus.join_inbound(), timeout=1) + + +@pytest.mark.asyncio +async def test_stop_cancels_after_grace_drops_queue_and_releases_dedupe(tmp_path: Path) -> None: + bus = MessageBus(inbound_queue_maxsize=2) + manager = ChannelManager( + bus=bus, + store=ChannelStore(path=tmp_path / "store.json"), + max_concurrency=1, + shutdown_grace_period_seconds=0, + ) + handler_started = asyncio.Event() + handler_cancelled = asyncio.Event() + + async def blocked_handler(_msg: InboundMessage) -> None: + handler_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + handler_cancelled.set() + raise + + manager._handle_message = blocked_handler # type: ignore[method-assign] + await manager.start() + + active = _message(1, with_dedupe_identity=True) + await bus.publish_inbound(active) + await asyncio.wait_for(handler_started.wait(), timeout=1) + await bus.publish_inbound(_message(2)) + await bus.publish_inbound(_message(3)) + workers = tuple(manager._worker_tasks) + + await manager.stop() + + assert handler_cancelled.is_set() + assert manager._worker_tasks == set() + assert all(task.done() for task in workers) + assert bus.inbound_queue.empty() + await asyncio.wait_for(bus.join_inbound(), timeout=1) + with pytest.raises(InboundQueueClosedError): + await bus.publish_inbound(_message(4)) + + dedupe_key = manager._inbound_dedupe_key(active) + assert dedupe_key is not None + # Cancellation must make the delivery retryable instead of black-holing it + # in the dedupe store until TTL expiry. + assert await manager._inbound_dedupe_store.try_record(dedupe_key) is False + + +@pytest.mark.asyncio +async def test_successful_stop_waits_for_cancel_resistant_workers_and_watchers(tmp_path: Path) -> None: + bus = MessageBus(inbound_queue_maxsize=1) + manager = ChannelManager( + bus=bus, + store=ChannelStore(path=tmp_path / "store.json"), + max_concurrency=1, + shutdown_grace_period_seconds=0.01, + ) + handler_started = asyncio.Event() + handler_cancelled = asyncio.Event() + watcher_started = asyncio.Event() + watcher_cancelled = asyncio.Event() + release_after_cancel = asyncio.Event() + + async def cancellation_resistant_handler(_msg: InboundMessage) -> None: + handler_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + handler_cancelled.set() + await release_after_cancel.wait() + + async def cancellation_resistant_watcher() -> None: + watcher_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + watcher_cancelled.set() + await release_after_cancel.wait() + + manager._handle_message = cancellation_resistant_handler # type: ignore[method-assign] + await manager.start() + watcher = asyncio.create_task(cancellation_resistant_watcher()) + manager._followup_watcher_tasks.add(watcher) + watcher.add_done_callback(manager._followup_watcher_tasks.discard) + await bus.publish_inbound(_message(1)) + await asyncio.wait_for(handler_started.wait(), timeout=1) + await asyncio.wait_for(watcher_started.wait(), timeout=1) + workers = tuple(manager._worker_tasks) + + stop_task = asyncio.create_task(manager.stop()) + await asyncio.wait_for(handler_cancelled.wait(), timeout=1) + await asyncio.wait_for(watcher_cancelled.wait(), timeout=1) + await asyncio.sleep(0.02) + + assert not stop_task.done() + assert any(not worker.done() for worker in workers) + assert not watcher.done() + + release_after_cancel.set() + await asyncio.wait_for(stop_task, timeout=1) + + assert manager._worker_tasks == set() + assert manager._followup_watcher_tasks == set() + assert all(worker.done() for worker in workers) + assert watcher.done() + await asyncio.wait_for(bus.join_inbound(), timeout=1) + + +@pytest.mark.asyncio +async def test_outer_shutdown_cancellation_does_not_start_an_unbounded_second_join(tmp_path: Path) -> None: + bus = MessageBus(inbound_queue_maxsize=1) + manager = ChannelManager( + bus=bus, + store=ChannelStore(path=tmp_path / "store.json"), + max_concurrency=1, + shutdown_grace_period_seconds=60, + ) + handler_started = asyncio.Event() + handler_cancelled = asyncio.Event() + release_after_cancel = asyncio.Event() + + async def cancellation_resistant_handler(_msg: InboundMessage) -> None: + handler_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + handler_cancelled.set() + await release_after_cancel.wait() + + manager._handle_message = cancellation_resistant_handler # type: ignore[method-assign] + await manager.start() + await bus.publish_inbound(_message(1)) + await asyncio.wait_for(handler_started.wait(), timeout=1) + workers = tuple(manager._worker_tasks) + + stop_task = asyncio.create_task(manager.stop()) + await asyncio.sleep(0) + started_at = asyncio.get_running_loop().time() + stop_task.cancel() + with pytest.raises(asyncio.CancelledError): + await stop_task + + assert asyncio.get_running_loop().time() - started_at < 0.2 + assert handler_cancelled.is_set() + assert any(not worker.done() for worker in workers) + assert manager._worker_tasks + + release_after_cancel.set() + await asyncio.wait_for(manager.stop(), timeout=1) + assert manager._worker_tasks == set() + assert all(worker.done() for worker in workers) + await asyncio.wait_for(bus.join_inbound(), timeout=1) + + +@pytest.mark.asyncio +async def test_provider_stop_drains_cross_thread_preparation_futures() -> None: + from app.channels.dingtalk import DingTalkChannel + from app.channels.feishu import FeishuChannel + from app.channels.telegram import TelegramChannel + + loop = asyncio.get_running_loop() + providers = ( + FeishuChannel(MessageBus(), config={}), + DingTalkChannel(MessageBus(), config={}), + TelegramChannel(MessageBus(), config={}), + SlackChannel(MessageBus(), config={}), + ) + + for channel in providers: + started = asyncio.Event() + cancelled = asyncio.Event() + finished = asyncio.Event() + + async def preparation() -> None: + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + finally: + finished.set() + + channel._running = True + channel._main_loop = loop + channel._open_threadsafe_future_intake() + future = await asyncio.to_thread( + channel._submit_threadsafe_coroutine, + preparation(), + loop, + name="test_preparation", + msg_id="message-1", + ) + assert future is not None + await asyncio.wait_for(started.wait(), timeout=1) + + await asyncio.wait_for(channel.stop(), timeout=1) + + assert cancelled.is_set() + assert finished.is_set() + assert future.done() + assert channel._threadsafe_futures == set() + + +def test_channel_service_threads_intake_limits_into_bus_and_worker_pool() -> None: + service = ChannelService( + channels_config={ + "inbound_queue_maxsize": 17, + "max_concurrency": 3, + "shutdown_grace_period_seconds": 2.5, + } + ) + + assert service.bus.inbound_queue_maxsize == 17 + assert service.manager._max_concurrency == 3 + assert service.manager._shutdown_grace_period_seconds == 2.5 + assert "inbound_queue_maxsize" not in service._config + assert "max_concurrency" not in service._config + assert "shutdown_grace_period_seconds" not in service._config + + +@pytest.mark.parametrize( + "invalid_value", + [True, -1, float("nan"), float("inf"), "3"], +) +def test_invalid_shutdown_grace_period_uses_finite_default(invalid_value: object) -> None: + service = ChannelService(channels_config={"shutdown_grace_period_seconds": invalid_value}) + + assert service.manager._shutdown_grace_period_seconds == 3.0 + + +@pytest.mark.asyncio +async def test_cancelled_service_stop_preserves_unfinished_channel_for_retry() -> None: + service = ChannelService(channels_config={"inbound_queue_maxsize": 1, "max_concurrency": 1}) + await service.start() + channel_stop_started = asyncio.Event() + release_channel_stop = asyncio.Event() + + class SlowChannel: + async def stop(self) -> None: + channel_stop_started.set() + await release_channel_stop.wait() + + slow_channel = SlowChannel() + service._channels["slow"] = slow_channel # type: ignore[assignment] + workers = tuple(service.manager._worker_tasks) + stop_task = asyncio.create_task(service.stop()) + await asyncio.wait_for(channel_stop_started.wait(), timeout=1) + stop_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await stop_task + + assert service.manager._worker_tasks == set() + assert all(worker.done() for worker in workers) + assert service._channels == {"slow": slow_channel} + with pytest.raises(InboundQueueClosedError): + await service.bus.publish_inbound(_message(9)) + + release_channel_stop.set() + await asyncio.wait_for(service.stop(), timeout=1) + assert service._channels == {} + + +@pytest.mark.asyncio +async def test_cancelled_manager_drain_preserves_service_channels() -> None: + service = ChannelService(channels_config={}) + manager_stop_started = asyncio.Event() + channel_stop_calls = 0 + + async def blocked_manager_stop() -> None: + manager_stop_started.set() + await asyncio.Event().wait() + + class Transport: + async def stop(self) -> None: + nonlocal channel_stop_calls + channel_stop_calls += 1 + + transport = Transport() + service.manager.stop = blocked_manager_stop # type: ignore[method-assign] + service._channels["transport"] = transport # type: ignore[assignment] + stop_task = asyncio.create_task(service.stop()) + await asyncio.wait_for(manager_stop_started.wait(), timeout=1) + + stop_task.cancel() + with pytest.raises(asyncio.CancelledError): + await stop_task + + assert service._channels == {"transport": transport} + assert channel_stop_calls == 0 + + +@pytest.mark.asyncio +async def test_failed_channel_stop_preserves_service_ownership() -> None: + service = ChannelService(channels_config={}) + + class BrokenTransport: + async def stop(self) -> None: + raise RuntimeError("provider cleanup failed") + + transport = BrokenTransport() + service._channels["transport"] = transport # type: ignore[assignment] + + with pytest.raises(ExceptionGroup, match="one or more channels failed to stop"): + await service.stop() + + assert service._channels == {"transport": transport} + + +@pytest.mark.asyncio +async def test_global_service_is_retained_until_shutdown_succeeds() -> None: + service = ChannelService(channels_config={}) + stop_started = asyncio.Event() + release_stop = asyncio.Event() + + async def blocked_stop() -> None: + stop_started.set() + await release_stop.wait() + + service.stop = blocked_stop # type: ignore[method-assign] + service_module._channel_service = service + stop_task = asyncio.create_task(service_module.stop_channel_service()) + await asyncio.wait_for(stop_started.wait(), timeout=1) + + assert service_module.get_channel_service() is service + + stop_task.cancel() + with pytest.raises(asyncio.CancelledError): + await stop_task + assert service_module.get_channel_service() is service + + release_stop.set() + await service_module.stop_channel_service() + assert service_module.get_channel_service() is None diff --git a/backend/tests/test_channels.py b/backend/tests/test_channels.py index 9d3378fc3..46806644f 100644 --- a/backend/tests/test_channels.py +++ b/backend/tests/test_channels.py @@ -1004,7 +1004,7 @@ class TestChannelManager: _run(go()) - def test_dispatch_loop_dedupes_stable_provider_message_id(self, tmp_path): + def test_worker_pool_dedupes_stable_provider_message_id(self, tmp_path): from app.channels.manager import ChannelManager async def go(): @@ -1186,7 +1186,7 @@ class TestChannelManager: ("feishu", "oc_abc"), ), ) - def test_dispatch_loop_dedupes_unbound_chat_scoped_redelivery(self, tmp_path, monkeypatch, channel, chat_id): + def test_worker_pool_dedupes_unbound_chat_scoped_redelivery(self, tmp_path, monkeypatch, channel, chat_id): """Provider redelivery of an unbound chat-scoped message runs the agent once. Shaped like wechat.py / telegram.py inbound metadata (message_id only, no @@ -1483,7 +1483,7 @@ class TestChannelManager: # silently drop this user's run (willem-bd, PR #4104 review). assert await manager._is_duplicate_inbound(_gh("d1", owner_user_id="bob")) is False - def test_dispatch_loop_releases_dedupe_key_when_handling_fails(self, tmp_path): + def test_worker_pool_releases_dedupe_key_when_handling_fails(self, tmp_path): """A transient handling failure must not black-hole a provider redelivery (ShenAC #1).""" from app.channels.manager import ChannelManager @@ -4826,7 +4826,7 @@ class TestChannelManagerBoundIdentityPolicy: _run(go()) - def test_unbound_auth_enabled_chat_is_rejected_before_semaphore(self, monkeypatch): + def test_unbound_auth_enabled_chat_is_rejected_before_run_creation(self, monkeypatch): from app.channels.manager import BOUND_IDENTITY_REQUIRED_MESSAGE, ChannelManager monkeypatch.delenv("DEER_FLOW_AUTH_DISABLED", raising=False) @@ -4842,8 +4842,6 @@ class TestChannelManagerBoundIdentityPolicy: bus.subscribe_outbound(capture) await manager.start() - assert manager._semaphore is not None - await manager._semaphore.acquire() try: await asyncio.wait_for( manager._handle_message( @@ -4857,7 +4855,6 @@ class TestChannelManagerBoundIdentityPolicy: timeout=0.5, ) finally: - manager._semaphore.release() await manager.stop() assert len(outbound_received) == 1 @@ -6846,7 +6843,6 @@ class TestWeComChannel: async def go(): bus = MessageBus() - bus.publish_inbound = AsyncMock() channel = WeComChannel(bus, config={}) channel._ws_client = SimpleNamespace(reply_stream=AsyncMock()) @@ -6869,9 +6865,8 @@ class TestWeComChannel: await channel._publish_ws_inbound(frame, "hello", files=files) channel._ws_client.reply_stream.assert_awaited_once_with(frame, "stream-1", "Working on it...", False) - bus.publish_inbound.assert_awaited_once() - - inbound = bus.publish_inbound.await_args.args[0] + inbound = await bus.get_inbound() + bus.inbound_task_done() assert inbound.channel_name == "wecom" assert inbound.chat_id == "user-1" assert inbound.user_id == "user-1" @@ -6890,7 +6885,6 @@ class TestWeComChannel: async def go(): bus = MessageBus() - bus.publish_inbound = AsyncMock() channel = WeComChannel(bus, config={"working_message": "Please wait..."}) channel._ws_client = SimpleNamespace(reply_stream=AsyncMock()) channel._working_message = "Please wait..." @@ -6919,7 +6913,6 @@ class TestWeComChannel: async def go(): bus = MessageBus() - bus.publish_inbound = AsyncMock() channel = WeComChannel(bus, config={}) channel._ws_client = SimpleNamespace(reply_stream=AsyncMock()) @@ -6938,7 +6931,8 @@ class TestWeComChannel: await channel._publish_ws_inbound(frame, "/mnt/user-data/uploads/report.pdf") - inbound = bus.publish_inbound.await_args.args[0] + inbound = await bus.get_inbound() + bus.inbound_task_done() assert inbound.text == "/mnt/user-data/uploads/report.pdf" assert inbound.msg_type == InboundMessageType.CHAT @@ -7854,6 +7848,13 @@ class TestSlackAllowedUsers: coro.close() return MagicMock() + @staticmethod + def _immediate_loop(): + loop = MagicMock() + loop.is_running.return_value = True + loop.call_soon_threadsafe.side_effect = lambda callback, *args: callback(*args) + return loop + def test_numeric_allowed_users_match_string_event_user_id(self): from app.channels.slack import SlackChannel @@ -7863,8 +7864,7 @@ class TestSlackAllowedUsers: bus=bus, config={"allowed_users": [123456]}, ) - channel._loop = MagicMock() - channel._loop.is_running.return_value = True + channel._loop = self._immediate_loop() channel._add_reaction = MagicMock() channel._send_running_reply = MagicMock() @@ -7878,13 +7878,13 @@ class TestSlackAllowedUsers: with patch( "app.channels.slack.asyncio.run_coroutine_threadsafe", side_effect=self._submit_coro, - ) as submit: + ): channel._handle_message_event(event) channel._add_reaction.assert_called_once_with("C123", "1710000000.000100", "eyes") channel._send_running_reply.assert_called_once_with("C123", "1710000000.000100") - submit.assert_called_once() - inbound = bus.publish_inbound.call_args.args[0] + channel._loop.call_soon_threadsafe.assert_called_once() + inbound = bus.get_inbound_nowait() assert inbound.user_id == "123456" assert inbound.chat_id == "C123" assert inbound.text == "hello from slack" @@ -7898,8 +7898,7 @@ class TestSlackAllowedUsers: bus=bus, config={"allowed_users": "U123456"}, ) - channel._loop = MagicMock() - channel._loop.is_running.return_value = True + channel._loop = self._immediate_loop() channel._add_reaction = MagicMock() channel._send_running_reply = MagicMock() @@ -7913,13 +7912,13 @@ class TestSlackAllowedUsers: with patch( "app.channels.slack.asyncio.run_coroutine_threadsafe", side_effect=self._submit_coro, - ) as submit: + ): channel._handle_message_event(event) channel._add_reaction.assert_called_once_with("C123", "1710000000.000100", "eyes") channel._send_running_reply.assert_called_once_with("C123", "1710000000.000100") - submit.assert_called_once() - inbound = bus.publish_inbound.call_args.args[0] + channel._loop.call_soon_threadsafe.assert_called_once() + inbound = bus.get_inbound_nowait() assert inbound.user_id == "U123456" assert inbound.chat_id == "C123" assert inbound.text == "hello from slack" @@ -7965,8 +7964,7 @@ class TestSlackAllowedUsers: bus = MessageBus() bus.publish_inbound = AsyncMock() channel = SlackChannel(bus=bus, config={"bot_user_id": "UBOT"}) - channel._loop = MagicMock() - channel._loop.is_running.return_value = True + channel._loop = self._immediate_loop() channel._add_reaction = MagicMock() channel._send_running_reply = MagicMock() @@ -7984,7 +7982,7 @@ class TestSlackAllowedUsers: ): channel._handle_message_event(event) - inbound = bus.publish_inbound.call_args.args[0] + inbound = bus.get_inbound_nowait() assert inbound.text == "/help" assert inbound.msg_type == InboundMessageType.COMMAND @@ -7994,8 +7992,7 @@ class TestSlackAllowedUsers: bus = MessageBus() bus.publish_inbound = AsyncMock() channel = SlackChannel(bus=bus, config={"bot_user_id": "UBOT"}) - channel._loop = MagicMock() - channel._loop.is_running.return_value = True + channel._loop = self._immediate_loop() channel._add_reaction = MagicMock() channel._send_running_reply = MagicMock() @@ -8013,7 +8010,7 @@ class TestSlackAllowedUsers: ): channel._handle_message_event(event) - inbound = bus.publish_inbound.call_args.args[0] + inbound = bus.get_inbound_nowait() assert inbound.text == "/help" assert inbound.msg_type == InboundMessageType.COMMAND @@ -8023,8 +8020,7 @@ class TestSlackAllowedUsers: bus = MessageBus() bus.publish_inbound = AsyncMock() channel = SlackChannel(bus=bus, config={"bot_user_id": "UBOT"}) - channel._loop = MagicMock() - channel._loop.is_running.return_value = True + channel._loop = self._immediate_loop() channel._add_reaction = MagicMock() channel._send_running_reply = MagicMock() @@ -8042,7 +8038,7 @@ class TestSlackAllowedUsers: ): channel._handle_message_event(event) - inbound = bus.publish_inbound.call_args.args[0] + inbound = bus.get_inbound_nowait() assert inbound.text == "/data-analysis analyze uploads/foo.csv" assert inbound.msg_type == InboundMessageType.CHAT @@ -8052,8 +8048,7 @@ class TestSlackAllowedUsers: bus = MessageBus() bus.publish_inbound = AsyncMock() channel = SlackChannel(bus=bus, config={"bot_user_id": "UBOT"}) - channel._loop = MagicMock() - channel._loop.is_running.return_value = True + channel._loop = self._immediate_loop() channel._add_reaction = MagicMock() channel._send_running_reply = MagicMock() @@ -8071,7 +8066,7 @@ class TestSlackAllowedUsers: ): channel._handle_message_event(event) - inbound = bus.publish_inbound.call_args.args[0] + inbound = bus.get_inbound_nowait() assert inbound.text == "<@UASSIGNEE> please review this" assert inbound.msg_type == InboundMessageType.CHAT @@ -8081,8 +8076,7 @@ class TestSlackAllowedUsers: bus = MessageBus() bus.publish_inbound = AsyncMock() channel = SlackChannel(bus=bus, config={"bot_user_id": "UBOT"}) - channel._loop = MagicMock() - channel._loop.is_running.return_value = True + channel._loop = self._immediate_loop() channel._add_reaction = MagicMock() channel._send_running_reply = MagicMock() @@ -8100,7 +8094,7 @@ class TestSlackAllowedUsers: ): channel._handle_message_event(event) - inbound = bus.publish_inbound.call_args.args[0] + inbound = bus.get_inbound_nowait() assert inbound.text == "<@UASSIGNEE> <@UBOT> please review this" assert inbound.msg_type == InboundMessageType.CHAT @@ -8110,8 +8104,7 @@ class TestSlackAllowedUsers: bus = MessageBus() bus.publish_inbound = AsyncMock() channel = SlackChannel(bus=bus, config={}) - channel._loop = MagicMock() - channel._loop.is_running.return_value = True + channel._loop = self._immediate_loop() channel._add_reaction = MagicMock() channel._send_running_reply = MagicMock() @@ -8129,7 +8122,7 @@ class TestSlackAllowedUsers: ): channel._handle_message_event(event) - inbound = bus.publish_inbound.call_args.args[0] + inbound = bus.get_inbound_nowait() assert inbound.text == "<@UASSIGNEE> /help <@UBOT>" assert inbound.msg_type == InboundMessageType.CHAT @@ -8140,8 +8133,8 @@ class TestSlackAllowedUsers: bus.publish_inbound = AsyncMock() channel = SlackChannel(bus=bus, config={}) channel._SocketModeResponse = lambda envelope_id: SimpleNamespace(envelope_id=envelope_id) - channel._loop = MagicMock() - channel._loop.is_running.return_value = True + channel._loop = self._immediate_loop() + channel._running = True channel._add_reaction = MagicMock() channel._send_running_reply = MagicMock() @@ -8167,7 +8160,7 @@ class TestSlackAllowedUsers: ): channel._on_socket_event(client, req) - inbound = bus.publish_inbound.call_args.args[0] + inbound = bus.get_inbound_nowait() assert channel._bot_user_id == "UBOT" assert inbound.text == "/help" assert inbound.msg_type == InboundMessageType.COMMAND @@ -8182,8 +8175,7 @@ class TestSlackAllowedUsers: bus=bus, config={"allowed_users": 123456}, ) - channel._loop = MagicMock() - channel._loop.is_running.return_value = True + channel._loop = self._immediate_loop() channel._add_reaction = MagicMock() channel._send_running_reply = MagicMock() @@ -8197,12 +8189,12 @@ class TestSlackAllowedUsers: with patch( "app.channels.slack.asyncio.run_coroutine_threadsafe", side_effect=self._submit_coro, - ) as submit: + ): channel._handle_message_event(event) assert "Slack allowed_users should be a list" in caplog.text - submit.assert_called_once() - inbound = bus.publish_inbound.call_args.args[0] + channel._loop.call_soon_threadsafe.assert_called_once() + inbound = bus.get_inbound_nowait() assert inbound.user_id == "123456" def test_raises_after_all_retries_exhausted(self): diff --git a/backend/tests/test_dingtalk_channel.py b/backend/tests/test_dingtalk_channel.py index f9776830a..e0f27e258 100644 --- a/backend/tests/test_dingtalk_channel.py +++ b/backend/tests/test_dingtalk_channel.py @@ -34,6 +34,12 @@ def _run(coro): loop.close() +async def _take_inbound(bus: MessageBus): + inbound = await asyncio.wait_for(bus.get_inbound(), timeout=1) + bus.inbound_task_done() + return inbound + + # --------------------------------------------------------------------------- # Helper: build mock ChatbotMessage # --------------------------------------------------------------------------- @@ -258,8 +264,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.channel_name == "dingtalk" assert inbound.chat_id == "user_001" assert inbound.user_id == "user_001" @@ -292,8 +297,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.channel_name == "dingtalk" assert inbound.chat_id == "conv_group_001" assert inbound.user_id == "user_002" @@ -330,8 +334,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.msg_type == InboundMessageType.COMMAND, f"{text!r} should be COMMAND" assert not inbound.text.startswith("@"), "leading mention must be stripped for dispatch" assert inbound.text.split(maxsplit=1)[0] in KNOWN_CHANNEL_COMMANDS @@ -362,8 +365,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.msg_type == InboundMessageType.CHAT assert inbound.text == "@bot please summarise this" @@ -393,8 +395,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.chat_id == "conv_group_001" assert inbound.topic_id == "msg_group_002" assert inbound.metadata["conversation_type"] == _CONVERSATION_TYPE_GROUP @@ -431,6 +432,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) bus.publish_inbound.assert_not_awaited() + assert bus.inbound_queue.empty() _run(go()) @@ -459,6 +461,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) bus.publish_inbound.assert_not_awaited() + assert bus.inbound_queue.empty() _run(go()) @@ -490,8 +493,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.chat_id == "conv_group_003" assert inbound.topic_id == "msg_group_003" @@ -512,8 +514,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.msg_type == InboundMessageType.COMMAND _run(go()) @@ -533,8 +534,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.msg_type == InboundMessageType.CHAT _run(go()) @@ -553,6 +553,7 @@ class TestOnChatbotMessage: await asyncio.sleep(0.1) bus.publish_inbound.assert_not_awaited() + assert bus.inbound_queue.empty() _run(go()) @@ -577,7 +578,7 @@ class TestAllowedUsersFiltering: channel._on_chatbot_message(msg) await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() + await _take_inbound(bus) _run(go()) @@ -595,6 +596,7 @@ class TestAllowedUsersFiltering: await asyncio.sleep(0.1) bus.publish_inbound.assert_not_awaited() + assert bus.inbound_queue.empty() _run(go()) @@ -615,6 +617,7 @@ class TestAllowedUsersFiltering: await asyncio.sleep(0.1) bus.publish_inbound.assert_not_awaited() + assert bus.inbound_queue.empty() # The parsed-message INFO log (with message content) must not fire for # a blocked sender — allowed_users still acts as a privacy/noise filter. assert "parsed message" not in caplog.text @@ -638,6 +641,7 @@ class TestAllowedUsersFiltering: await asyncio.sleep(0.1) channel._bind_connection_from_connect_code.assert_awaited_once() bus.publish_inbound.assert_not_awaited() + assert bus.inbound_queue.empty() _run(go()) @@ -655,7 +659,7 @@ class TestAllowedUsersFiltering: channel._on_chatbot_message(msg) await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() + await _take_inbound(bus) _run(go()) @@ -871,7 +875,7 @@ class TestTopicIdMapping: channel._on_chatbot_message(msg) await asyncio.sleep(0.1) - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.topic_id is None _run(go()) @@ -894,7 +898,7 @@ class TestTopicIdMapping: channel._on_chatbot_message(msg) await asyncio.sleep(0.1) - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.topic_id == "msg_group_001" _run(go()) @@ -1645,29 +1649,34 @@ class TestCardMode: assert channel._dingtalk_client is mock_client def test_chatbot_message_stored_for_card_mode(self): - bus = MessageBus() - channel = DingTalkChannel(bus, config={"card_template_id": "tpl_123"}) + async def go(): + bus = MessageBus() + channel = DingTalkChannel(bus, config={"card_template_id": "tpl_123"}) - mock_message = MagicMock() - mock_message.sender_staff_id = "user_001" - mock_message.conversation_type = "1" - mock_message.conversation_id = "" - mock_message.message_id = "msg_001" - mock_message.sender_nick = "TestUser" - mock_message.message_type = "text" - mock_message.text = MagicMock(content="hello") - mock_message.rich_text_content = None + mock_message = MagicMock() + mock_message.sender_staff_id = "user_001" + mock_message.conversation_type = "1" + mock_message.conversation_id = "" + mock_message.message_id = "msg_001" + mock_message.sender_nick = "TestUser" + mock_message.message_type = "text" + mock_message.text = MagicMock(content="hello") + mock_message.rich_text_content = None - channel._main_loop = MagicMock() - channel._main_loop.is_running.return_value = False - channel._allowed_users = set() - channel._running = True + channel._main_loop = asyncio.get_running_loop() + channel._allowed_users = set() + channel._running = True + channel._send_running_reply = AsyncMock() - channel._on_chatbot_message(mock_message) + channel._on_chatbot_message(mock_message) + await asyncio.sleep(0.1) - assert len(channel._incoming_messages) == 1 - stored_msg = list(channel._incoming_messages.values())[0] - assert stored_msg is mock_message + assert len(channel._incoming_messages) == 1 + stored_msg = list(channel._incoming_messages.values())[0] + assert stored_msg is mock_message + await _take_inbound(bus) + + _run(go()) def test_card_replier_cleanup_on_final(self): async def go(): @@ -1867,8 +1876,7 @@ class TestOnChatbotMessageFiles: channel._on_chatbot_message(_make_picture_message(download_code="dc_pic", sender_staff_id="u1", message_id="m1")) await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.msg_type == InboundMessageType.CHAT assert inbound.files == [{"type": "image", "download_code": "dc_pic", "filename": "image.png"}] @@ -1887,8 +1895,7 @@ class TestOnChatbotMessageFiles: channel._on_chatbot_message(_make_file_message(download_code="dc_doc", file_name="data.csv", sender_staff_id="u1", message_id="m1")) await asyncio.sleep(0.1) - bus.publish_inbound.assert_awaited_once() - inbound = bus.publish_inbound.await_args.args[0] + inbound = await _take_inbound(bus) assert inbound.files == [{"type": "file", "download_code": "dc_doc", "filename": "data.csv"}] _run(go()) @@ -1909,6 +1916,7 @@ class TestOnChatbotMessageFiles: await asyncio.sleep(0.1) bus.publish_inbound.assert_not_awaited() + assert bus.inbound_queue.empty() _run(go()) diff --git a/backend/tests/test_discord_channel.py b/backend/tests/test_discord_channel.py index ecbad5de9..16fd9bfc0 100644 --- a/backend/tests/test_discord_channel.py +++ b/backend/tests/test_discord_channel.py @@ -6,13 +6,13 @@ import asyncio import builtins import threading from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from app.channels.discord import DiscordChannel from app.channels.manager import CHANNEL_CAPABILITIES -from app.channels.message_bus import InboundMessageType, MessageBus, OutboundMessage, ResolvedAttachment +from app.channels.message_bus import InboundMessage, InboundMessageType, MessageBus, OutboundMessage, ResolvedAttachment from app.channels.service import _CHANNEL_REGISTRY @@ -46,11 +46,10 @@ def _make_discord_message(text: str): async def test_discord_bot_mention_slash_skill_routes_as_chat() -> None: bus = MessageBus() channel = DiscordChannel(bus=bus, config={"bot_token": "token"}) - captured = [] channel._running = True channel._client = SimpleNamespace(user=SimpleNamespace(id=999, mention="<@999>")) channel._discord_module = SimpleNamespace(Thread=type("FakeThread", (), {})) - channel._publish = captured.append + channel._main_loop = asyncio.get_running_loop() async def noop(*_args, **_kwargs): return None @@ -59,9 +58,10 @@ async def test_discord_bot_mention_slash_skill_routes_as_chat() -> None: channel._add_reaction = noop await channel._on_message(_make_discord_message("<@999> /data-analysis analyze uploads/foo.csv")) + await asyncio.sleep(0) - assert len(captured) == 1 - inbound = captured[0] + inbound = bus.get_inbound_nowait() + bus.inbound_task_done() assert inbound.text == "/data-analysis analyze uploads/foo.csv" assert inbound.msg_type == InboundMessageType.CHAT assert inbound.topic_id == "456" @@ -71,11 +71,10 @@ async def test_discord_bot_mention_slash_skill_routes_as_chat() -> None: async def test_discord_bot_mention_known_command_routes_as_command() -> None: bus = MessageBus() channel = DiscordChannel(bus=bus, config={"bot_token": "token"}) - captured = [] channel._running = True channel._client = SimpleNamespace(user=SimpleNamespace(id=999, mention="<@999>")) channel._discord_module = SimpleNamespace(Thread=type("FakeThread", (), {})) - channel._publish = captured.append + channel._main_loop = asyncio.get_running_loop() async def noop(*_args, **_kwargs): return None @@ -84,14 +83,67 @@ async def test_discord_bot_mention_known_command_routes_as_command() -> None: channel._add_reaction = noop await channel._on_message(_make_discord_message("<@999> /help")) + await asyncio.sleep(0) - assert len(captured) == 1 - inbound = captured[0] + inbound = bus.get_inbound_nowait() + bus.inbound_task_done() assert inbound.text == "/help" assert inbound.msg_type == InboundMessageType.COMMAND assert inbound.topic_id == "456" +@pytest.mark.asyncio +async def test_discord_full_queue_rejects_before_thread_or_identity_side_effects() -> None: + bus = MessageBus(inbound_queue_maxsize=1) + await bus.publish_inbound( + InboundMessage( + channel_name="slack", + chat_id="C1", + user_id="U1", + text="already queued", + ) + ) + channel = DiscordChannel(bus=bus, config={"bot_token": "token", "thread_mode": True}) + channel._running = True + channel._client = SimpleNamespace(user=SimpleNamespace(id=999, mention="<@999>")) + channel._discord_module = SimpleNamespace(Thread=type("FakeThread", (), {})) + channel._main_loop = asyncio.get_running_loop() + channel._create_thread = AsyncMock() + channel._attach_connection_identity = AsyncMock() + + await channel._on_message(_make_discord_message("hello")) + + channel._create_thread.assert_not_awaited() + channel._attach_connection_identity.assert_not_awaited() + assert channel._active_threads == {} + assert bus.inbound_queue.qsize() == 1 + + +@pytest.mark.asyncio +async def test_discord_releases_early_reservation_when_thread_creation_raises() -> None: + bus = MessageBus(inbound_queue_maxsize=1) + channel = DiscordChannel(bus=bus, config={"bot_token": "token", "thread_mode": True}) + channel._running = True + channel._client = SimpleNamespace(user=SimpleNamespace(id=999, mention="<@999>")) + channel._discord_module = SimpleNamespace(Thread=type("FakeThread", (), {})) + channel._main_loop = asyncio.get_running_loop() + channel._create_thread = AsyncMock(side_effect=RuntimeError("thread create failed")) + + with pytest.raises(RuntimeError, match="thread create failed"): + await channel._on_message(_make_discord_message("hello")) + + # The exception path must return the capacity slot to the shared bus. + await bus.publish_inbound( + InboundMessage( + channel_name="slack", + chat_id="C1", + user_id="U1", + text="capacity was released", + ) + ) + assert bus.inbound_queue.qsize() == 1 + + # --------------------------------------------------------------------------- # send_file file-handle lifecycle # --------------------------------------------------------------------------- diff --git a/backend/tests/test_github_dispatcher.py b/backend/tests/test_github_dispatcher.py index d39097add..9a361114c 100644 --- a/backend/tests/test_github_dispatcher.py +++ b/backend/tests/test_github_dispatcher.py @@ -9,12 +9,13 @@ shape lands on the bus per matching binding. from __future__ import annotations +import asyncio from pathlib import Path import pytest import yaml -from app.channels.message_bus import InboundMessage, MessageBus +from app.channels.message_bus import InboundMessage, InboundQueueFullError, MessageBus from app.gateway.github.dispatcher import fanout_event @@ -46,6 +47,77 @@ async def _drain(bus: MessageBus) -> list[InboundMessage]: return out +@pytest.mark.asyncio +async def test_fanout_redelivery_progresses_beyond_queue_sized_prefix(base_dir: Path) -> None: + """A redelivery must not keep failing on the same queue-sized prefix.""" + from app.channels.manager import ChannelManager + from app.channels.store import ChannelStore + + agent_names = {"alpha", "bravo", "charlie"} + for name in agent_names: + _write_agent( + base_dir, + "default", + name, + { + "name": name, + "github": { + "bindings": [ + { + "repo": "a/b", + "triggers": {"pull_request": {"actions": ["opened"]}}, + } + ] + }, + }, + ) + + bus = MessageBus(inbound_queue_maxsize=1) + manager = ChannelManager( + bus=bus, + store=ChannelStore(path=base_dir / "fanout-store.json"), + max_concurrency=1, + ) + handled: list[str] = [] + first_started = asyncio.Event() + release_first = asyncio.Event() + + async def handle(msg: InboundMessage) -> None: + if not handled: + first_started.set() + await release_first.wait() + handled.append(msg.metadata["agent_name"]) + + manager._handle_message = handle # type: ignore[method-assign] + await manager.start() + try: + payload = { + "action": "opened", + "pull_request": {"number": 1, "title": "x", "user": {"login": "u"}, "body": ""}, + "repository": {"full_name": "a/b"}, + "sender": {"login": "u"}, + } + with pytest.raises(InboundQueueFullError): + await fanout_event(bus, "pull_request", "delivery-over-capacity", payload) + + await asyncio.wait_for(first_started.wait(), timeout=1) + release_first.set() + await asyncio.wait_for(bus.join_inbound(), timeout=1) + assert len(handled) == 2 + + # The same delivery id republishes the already-admitted prefix, which + # the manager dedupe consumes quickly. The scheduling handoff in + # publish_inbound lets the worker do that before the next admission, + # so the previously starved suffix reaches the queue on this retry. + result = await fanout_event(bus, "pull_request", "delivery-over-capacity", payload) + await asyncio.wait_for(bus.join_inbound(), timeout=1) + + assert set(result["fired_agents"]) == agent_names + assert set(handled) == agent_names + finally: + await manager.stop() + + # --------------------------------------------------------------------------- # Bot-loop prevention — per-agent self-event gate # --------------------------------------------------------------------------- diff --git a/config.example.yaml b/config.example.yaml index 62e5c7865..5ce5cf143 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -2216,6 +2216,14 @@ run_ownership: # langgraph_url: http://localhost:8001/api # # Gateway API URL for auxiliary queries like /models, /memory (default: http://localhost:8001) # gateway_url: http://localhost:8001 +# # Maximum queued or provider-reserved inbound messages. Must be a positive integer. +# inbound_queue_maxsize: 1000 +# # Fixed number of long-lived inbound handler workers. Must be a positive integer. +# max_concurrency: 5 +# # Seconds to drain accepted inbound work before cancelling active handlers. +# # Must be a non-negative finite number. Cancelled handlers are awaited; the Gateway's +# # outer shutdown timeout remains the process-level bound for incomplete cleanup. +# shutdown_grace_period_seconds: 3 # # # # Docker Compose note: # # If channels run inside the gateway container, use container DNS names instead