diff --git a/docs/architecture.md b/docs/architecture.md index 91506131..02ba6a11 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -1200,8 +1200,13 @@ The bridge dispatches it to `POST /v1/api/cancel` on the server owning the works which sets the cooperative cancel flag and unblocks any pending approval/plan waits. **Completion detection:** The bridge tracks which `correlation_id` maps to which -`ws_id` for active sends. When the global SSE reports `ws_state → idle` for a tracked -workstream, the bridge emits a synthetic `TurnCompleteEvent` with the correlation ID. +`ws_id` for active sends. Content tokens from the per-ws SSE stream are accumulated +in a per-workstream buffer (`_ws_content_buffer`, capped at 256 KB). When the global +SSE reports `ws_state → idle` for a tracked workstream, the bridge emits a synthetic +`TurnCompleteEvent` carrying the correlation ID and the accumulated response text in +`content`. This catch-up mechanism lets downstream consumers (e.g. the Discord bot) +recover the full response even when individual `ContentEvent`s were missed due to the +race between the two independent SSE connections. **Multi-node routing:** Each bridge retrieves its `node_id` from the server's `/health` endpoint on startup (with exponential backoff retry). The server @@ -1365,11 +1370,23 @@ directly over HTTP for lower latency: `_exec_notify()` queries the `services` database table for healthy channel gateways (heartbeat within 120 seconds), authenticates with a service JWT (`aud: turnstone-channel`), and POSTs to `POST /v1/api/notify` on the first healthy gateway. The -gateway validates the JWT, resolves the target (username lookup via +payload includes the originating `ws_id` for reply routing. The gateway +validates the JWT, resolves the target (username lookup via `channel_users` or direct `channel_type`+`channel_id`), and delegates to -the appropriate `ChannelAdapter.send()`. Delivery retries up to 3 times -with backoff, re-querying the service registry on each attempt. See -[Notification Flow diagram](diagrams/png/17-notify-flow.png). +`ChannelAdapter.send_notification()` which sends the message and tracks +the outgoing message ID → `(ws_id, target_user_id)` mapping. Delivery +retries up to 3 times with backoff, re-querying the service registry on +each attempt. See [Notification Flow diagram](diagrams/png/17-notify-flow.png). + +**Bidirectional replies:** When a user replies to a notification DM, the +Discord bot looks up the originating `ws_id` from the tracked message ID, +verifies the replying user matches the notification recipient, and routes +the reply to the workstream via `router.send_message()`. The workstream's +response is forwarded back to the DM via a temporary entry in +`_notify_reply_channels`. On `TurnCompleteEvent`, the response message is +itself tracked for further replies, enabling multi-turn DM conversations +without requiring the user to open the web UI. Tracking entries are capped +at 100 (FIFO eviction) and cleaned up on workstream close. --- diff --git a/docs/channels.md b/docs/channels.md index b4c869f4..e51f1688 100644 --- a/docs/channels.md +++ b/docs/channels.md @@ -30,8 +30,8 @@ Key components: - **ChannelAdapter protocol** (`turnstone/channels/_protocol.py`) — generic interface for any messaging platform. Defines `start()`, `stop()`, - `send()`, `edit_message()`, `send_approval_request()`, - `send_plan_review()`, and `create_thread()`. + `send()`, `send_notification()`, `edit_message()`, + `send_approval_request()`, `send_plan_review()`, and `create_thread()`. - **ChannelRouter** (`turnstone/channels/_routing.py`) — maps channel/thread IDs to turnstone workstream IDs. Handles workstream creation via MQ, stale route detection, and user identity resolution. @@ -271,13 +271,43 @@ gateway directly over HTTP: 2. `_exec_notify()` queries the `services` table for healthy channel gateways (heartbeat within the last 120 seconds) 3. The server mints a service JWT (`aud: turnstone-channel`) via - `ServiceTokenManager` and POSTs to the first healthy gateway + `ServiceTokenManager` and POSTs to the first healthy gateway. The + payload includes the originating `ws_id` for reply routing. 4. The gateway validates the JWT, resolves the target, and calls - `adapter.send()` on the appropriate platform adapter + `adapter.send_notification()` which sends the message and tracks + the outgoing message ID for reply routing 5. On failure, the server tries the next gateway. If all fail, it retries up to 2 more times (delays: 1s, 3s), re-querying the service registry on each attempt +### Bidirectional Replies + +Notifications support multi-turn DM conversations. When a user replies +to a notification DM: + +1. The bot looks up the originating `ws_id` from the tracked message ID + (`_notify_ws_map`) +2. Verifies the replying user matches the original notification + recipient (defence in depth — Discord DMs are already private) +3. Routes the reply to the workstream via `router.send_message()` +4. Registers the DM channel for response forwarding + (`_notify_reply_channels`) +5. When the workstream responds (`TurnCompleteEvent`), the response is + forwarded to the DM +6. The response message is itself tracked, so the user can reply again + for another turn + +This enables scenarios like an oncall engineer responding to a CI/CD +failure notification from their phone before opening a laptop. + +**Limits:** + +- Tracking map capped at 100 entries (FIFO eviction of oldest) +- Entries cleaned up on workstream close/unsubscribe +- Replying to an expired notification sends + *"This notification is no longer active."* +- DM reply content capped at 4096 characters + ### Service Registry The channel gateway registers itself in the `services` database table @@ -328,12 +358,18 @@ class ChannelAdapter(Protocol): async def start(self) -> None: ... async def stop(self) -> None: ... async def send(self, channel_id: str, content: str) -> str: ... + async def send_notification(self, channel_id: str, content: str, ws_id: str) -> str: ... async def edit_message(self, channel_id: str, message_id: str, content: str) -> None: ... async def send_approval_request(self, channel_id: str, ws_id: str, correlation_id: str, items: list[dict]) -> None: ... async def send_plan_review(self, channel_id: str, ws_id: str, correlation_id: str, content: str) -> None: ... async def create_thread(self, parent_channel_id: str, name: str, message_id: str = "") -> str: ... ``` +`send_notification()` is like `send()` but associates the outgoing +message with a `ws_id` so that user replies can be routed back to the +originating workstream. Adapters must track the mapping from outgoing +message ID to `(ws_id, target_user_id)` and handle DM replies. + To add a new platform: 1. Create `turnstone/channels//` package diff --git a/docs/diagrams/06-mq-protocol.puml b/docs/diagrams/06-mq-protocol.puml index 3c896582..854e172f 100644 --- a/docs/diagrams/06-mq-protocol.puml +++ b/docs/diagrams/06-mq-protocol.puml @@ -174,6 +174,7 @@ package "Outbound Events (Bridge → Client)" #E3F2FD { } class TurnCompleteEvent { type = "turn_complete" + + content: str } } diff --git a/docs/diagrams/16-channel-architecture.puml b/docs/diagrams/16-channel-architecture.puml index f92439d8..4f82db25 100644 --- a/docs/diagrams/16-channel-architecture.puml +++ b/docs/diagrams/16-channel-architecture.puml @@ -53,6 +53,7 @@ class "DiscordBot" as Bot <> { +on_message(msg) +on_interaction(interaction) +send(channel_id, content) + +send_notification(channel_id, content, ws_id) +run(token) -- discord.py Client @@ -61,6 +62,9 @@ class "DiscordBot" as Bot <> { Creates threads for workstreams Renders approval buttons escape_mentions() on send + -- + _notify_ws_map: msg_id → (ws_id, user_id) + _notify_reply_channels: ws_id → (dm, user_id) } class "ChannelRouter" as Router <> { @@ -240,11 +244,20 @@ note bottom of SVC 3. Queries services table for healthy gateways 4. Mints JWT (aud: turnstone-channel) via ServiceTokenManager - 5. POSTs to first healthy gateway + 5. POSTs to first healthy gateway (incl. ws_id) 6. Gateway validates JWT, resolves target - 7. adapter.send() → Discord API + 7. adapter.send_notification() → Discord API + (tracks msg_id → ws_id for reply routing) 8. On failure: retry up to 3× (1s, 3s backoff) 9. SSRF: only http(s) URLs allowed + + **Bidirectional DM Replies** + 1. User replies to notification DM + 2. Bot looks up ws_id from _notify_ws_map + 3. Verifies author == notification recipient + 4. Routes reply via router.send_message() + 5. Response forwarded to DM on TurnCompleteEvent + 6. Response tracked for multi-turn conversation end note @enduml diff --git a/docs/diagrams/17-notify-flow.puml b/docs/diagrams/17-notify-flow.puml index 70dd8a28..ed2150bf 100644 --- a/docs/diagrams/17-notify-flow.puml +++ b/docs/diagrams/17-notify-flow.puml @@ -103,6 +103,41 @@ alt all retries exhausted Session --> Session : "Error: notification delivery failed" end +== Bidirectional Reply (User responds to notification DM) == + +Discord -> Adapter : user replies to\nnotification message +Adapter -> Adapter : lookup message_id\nin _notify_ws_map +note right + Maps message_id → + (ws_id, target_user_id) + Atomic pop prevents TOCTOU +end note + +alt message not tracked + Adapter -> Discord : "This notification\nis no longer active." +else tracked + Adapter -> Adapter : verify author ==\ntarget_user_id + Adapter -> Adapter : resolve_user()\n(unlinked → drop) + Adapter -> Adapter : router.send_message(ws_id, content) + note right + Routes reply via MQ to + the originating workstream. + Registers DM channel in + _notify_reply_channels[ws_id] + end note + + ... workstream processes reply ... + + Adapter <- Adapter : TurnCompleteEvent\n(with content) + Adapter -> Discord : forward response to DM + Adapter -> Adapter : track response message\nfor multi-turn replies + note right + Response message_id added + to _notify_ws_map — user can + reply again indefinitely + end note +end + == Service Registry (Background) == note over Gateway, Storage diff --git a/docs/diagrams/png/06-mq-protocol.png b/docs/diagrams/png/06-mq-protocol.png index 6d294cde..e810b5f4 100644 --- a/docs/diagrams/png/06-mq-protocol.png +++ b/docs/diagrams/png/06-mq-protocol.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:2229801220548e4794baa67e27a0a39dc7968c826a28fa8144763e678c8ed733 -size 192556 +oid sha256:6e94a10f039a7f69517e84d0946e0c649035c15b38ebc2314e7b9cd501eb244d +size 192559 diff --git a/docs/diagrams/png/16-channel-architecture.png b/docs/diagrams/png/16-channel-architecture.png index 16e5b118..26e96d0d 100644 --- a/docs/diagrams/png/16-channel-architecture.png +++ b/docs/diagrams/png/16-channel-architecture.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:1380065cbb5f95b5ea7dc6b2a00986c455b82888af60784980dffbd936460dcf -size 431129 +oid sha256:6fc99bb8d84d6e9f3dac9d5c12ac7f569a041b29431c57612c24b50f332982ed +size 462992 diff --git a/docs/diagrams/png/17-notify-flow.png b/docs/diagrams/png/17-notify-flow.png index 8208bed1..b9dde969 100644 --- a/docs/diagrams/png/17-notify-flow.png +++ b/docs/diagrams/png/17-notify-flow.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:f0f6097840fccdbfe16cd5e4c9f5d063b2c36942460a944df68b8ec947e63ea3 -size 221452 +oid sha256:cc4c511c34a2e5d286fd128c3509405a5b240ca02a4bafb395d2e94d002a5b8b +size 293203 diff --git a/tests/test_bridge_events.py b/tests/test_bridge_events.py index 550a0209..e3d805b7 100644 --- a/tests/test_bridge_events.py +++ b/tests/test_bridge_events.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock, patch from turnstone.mq.bridge import Bridge -from turnstone.mq.protocol import StateChangeEvent, TurnCompleteEvent +from turnstone.mq.protocol import ContentEvent, StateChangeEvent, TurnCompleteEvent def _make_bridge(): @@ -67,3 +67,90 @@ class TestIdleTurnComplete: assert len(state_changes) == 1 assert state_changes[0].state == "thinking" assert len(turn_completes) == 0 + + +class TestContentBuffer: + """Bridge should accumulate content tokens and attach to TurnCompleteEvent.""" + + def test_content_buffer_accumulated_in_turn_complete(self): + """Content events should be accumulated and included in TurnCompleteEvent.""" + bridge = _make_bridge() + + # Simulate content events from per-ws SSE + bridge._handle_ws_event("ws-1", {"type": "content", "text": "Hello "}) + bridge._handle_ws_event("ws-1", {"type": "content", "text": "world"}) + + published = [] + with patch.object( + bridge, "_publish_ws", side_effect=lambda ws, ev: published.append((ws, ev)) + ): + bridge._handle_global_event({"type": "ws_state", "ws_id": "ws-1", "state": "idle"}) + + turn_completes = [(ws, ev) for ws, ev in published if isinstance(ev, TurnCompleteEvent)] + assert len(turn_completes) == 1 + _, ev = turn_completes[0] + assert ev.content == "Hello world" + # Buffer should be cleared + assert "ws-1" not in bridge._ws_content_buffer + + def test_content_buffer_empty_for_no_content_turn(self): + """TurnCompleteEvent.content should be empty when no content events fired.""" + bridge = _make_bridge() + + published = [] + with patch.object( + bridge, "_publish_ws", side_effect=lambda ws, ev: published.append((ws, ev)) + ): + bridge._handle_global_event({"type": "ws_state", "ws_id": "ws-1", "state": "idle"}) + + turn_completes = [(ws, ev) for ws, ev in published if isinstance(ev, TurnCompleteEvent)] + assert len(turn_completes) == 1 + _, ev = turn_completes[0] + assert ev.content == "" + + def test_content_buffer_cleared_on_ws_closed(self): + """ws_closed should clean up the content buffer.""" + bridge = _make_bridge() + + bridge._handle_ws_event("ws-1", {"type": "content", "text": "orphan"}) + assert "ws-1" in bridge._ws_content_buffer + + bridge._handle_global_event({"type": "ws_closed", "ws_id": "ws-1"}) + assert "ws-1" not in bridge._ws_content_buffer + + def test_content_buffer_publishes_content_event(self): + """Content events should still be published to per-ws channel.""" + bridge = _make_bridge() + + published = [] + with patch.object( + bridge, "_publish_ws", side_effect=lambda ws, ev: published.append((ws, ev)) + ): + bridge._handle_ws_event("ws-1", {"type": "content", "text": "hello"}) + + content_events = [(ws, ev) for ws, ev in published if isinstance(ev, ContentEvent)] + assert len(content_events) == 1 + _, ev = content_events[0] + assert ev.text == "hello" + + def test_multi_round_content_accumulates(self): + """Content from multiple tool-use rounds accumulates in a single turn.""" + bridge = _make_bridge() + + # Round 1 + bridge._handle_ws_event("ws-1", {"type": "content", "text": "I'll run "}) + bridge._handle_ws_event("ws-1", {"type": "stream_end"}) + # Round 2 (after tool execution) + bridge._handle_ws_event("ws-1", {"type": "content", "text": "the command."}) + bridge._handle_ws_event("ws-1", {"type": "stream_end"}) + + published = [] + with patch.object( + bridge, "_publish_ws", side_effect=lambda ws, ev: published.append((ws, ev)) + ): + bridge._handle_global_event({"type": "ws_state", "ws_id": "ws-1", "state": "idle"}) + + turn_completes = [(ws, ev) for ws, ev in published if isinstance(ev, TurnCompleteEvent)] + assert len(turn_completes) == 1 + _, ev = turn_completes[0] + assert ev.content == "I'll run the command." diff --git a/tests/test_channel_discord.py b/tests/test_channel_discord.py index 90b1224f..32186c1a 100644 --- a/tests/test_channel_discord.py +++ b/tests/test_channel_discord.py @@ -21,7 +21,7 @@ def _run(coro): return asyncio.run(coro) -def _make_message(*, bot=False, guild=True, content="hello", channel=None): +def _make_message(*, bot=False, guild=True, content="hello", channel=None, reference=None): """Build a mock ``discord.Message``.""" msg = MagicMock(spec=discord.Message) msg.author = MagicMock() @@ -31,6 +31,7 @@ def _make_message(*, bot=False, guild=True, content="hello", channel=None): msg.guild = MagicMock() if guild else None msg.channel = channel or MagicMock() msg.mentions = [] + msg.reference = reference return msg @@ -204,6 +205,8 @@ class TestMessageCog: ts.router.send_message = AsyncMock() ts.config = MagicMock() ts._ws_tasks = {} + ts._notify_ws_map = {} + ts._notify_reply_channels = {} bot.turnstone = ts cog = MessageCog(bot) @@ -328,6 +331,7 @@ class TestWsEventFinalization: bot.config.auto_approve_tools = [] bot._streaming = {} bot._pending_approval_msgs = {} + bot._notify_reply_channels = {} # Use the real _on_ws_event method bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) @@ -356,6 +360,7 @@ class TestWsEventFinalization: bot = MagicMock(spec=TurnstoneBot) bot._streaming = {} bot._pending_approval_msgs = {} + bot._notify_reply_channels = {} bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) thread = AsyncMock() @@ -387,6 +392,7 @@ class TestApprovalVerdictDisplay: bot.config.auto_approve_tools = [] bot._streaming = {} bot._pending_approval_msgs = {} + bot._notify_reply_channels = {} bot._should_auto_approve = MagicMock(return_value=False) bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) return bot @@ -504,6 +510,7 @@ class TestApprovalVerdictDisplay: bot = MagicMock(spec=TurnstoneBot) bot._streaming = {} bot._pending_approval_msgs = {"ws-1": MagicMock()} + bot._notify_reply_channels = {} bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) thread = AsyncMock() @@ -513,6 +520,315 @@ class TestApprovalVerdictDisplay: assert "ws-1" not in bot._pending_approval_msgs +class TestContentCatchup: + """TurnCompleteEvent with content field provides catch-up for missed ContentEvents.""" + + def _make_bot(self): + from turnstone.channels.discord.bot import TurnstoneBot + + bot = MagicMock(spec=TurnstoneBot) + bot.config = MagicMock() + bot.config.max_message_length = 2000 + bot.config.streaming_edit_interval = 1.5 + bot.config.auto_approve = False + bot.config.auto_approve_tools = [] + bot._streaming = {} + bot._pending_approval_msgs = {} + bot._notify_reply_channels = {} + bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) + return bot + + def test_catchup_sends_content_when_no_streaming(self): + """TurnCompleteEvent with content but no SM sends catch-up message.""" + from turnstone.mq.protocol import TurnCompleteEvent + + bot = self._make_bot() + thread = AsyncMock() + + raw = TurnCompleteEvent( + ws_id="ws-1", correlation_id="", content="Caught up response" + ).to_json() + _run(bot._on_ws_event("ws-1", thread, raw)) + + thread.send.assert_awaited_once_with("Caught up response") + + def test_catchup_skipped_when_streaming_exists(self): + """TurnCompleteEvent with content and existing SM uses SM finalize, not catch-up.""" + from turnstone.mq.protocol import ContentEvent, TurnCompleteEvent + + bot = self._make_bot() + thread = AsyncMock() + + # Feed content event to create SM + content_raw = ContentEvent(ws_id="ws-1", text="Streamed").to_json() + _run(bot._on_ws_event("ws-1", thread, content_raw)) + assert "ws-1" in bot._streaming + + # Now TurnCompleteEvent with content — SM should be finalized, not catch-up + complete_raw = TurnCompleteEvent( + ws_id="ws-1", correlation_id="", content="Streamed" + ).to_json() + _run(bot._on_ws_event("ws-1", thread, complete_raw)) + assert "ws-1" not in bot._streaming + + def test_catchup_empty_content_no_message(self): + """TurnCompleteEvent with empty content and no SM sends nothing.""" + from turnstone.mq.protocol import TurnCompleteEvent + + bot = self._make_bot() + thread = AsyncMock() + + raw = TurnCompleteEvent(ws_id="ws-1", correlation_id="", content="").to_json() + _run(bot._on_ws_event("ws-1", thread, raw)) + + thread.send.assert_not_awaited() + + +class TestNotificationTracking: + """Tests for notification message tracking and DM reply routing.""" + + def test_send_notification_tracks_message(self): + """send_notification should store message_id -> (ws_id, target_user) mapping.""" + from turnstone.channels.discord.bot import TurnstoneBot + + bot = MagicMock(spec=TurnstoneBot) + bot._notify_ws_map = {} + bot._MAX_NOTIFY_TRACKING = 100 + bot.send = AsyncMock(return_value="12345") + bot.send_notification = TurnstoneBot.send_notification.__get__(bot, TurnstoneBot) + bot._track_notification = TurnstoneBot._track_notification.__get__(bot, TurnstoneBot) + + _run(bot.send_notification("chan-1", "Hello", "ws-abc")) + + assert 12345 in bot._notify_ws_map + assert bot._notify_ws_map[12345] == ("ws-abc", "chan-1") + + def test_send_notification_evicts_old_entries(self): + """Oldest notification tracking entries are evicted when cap is reached.""" + from turnstone.channels.discord.bot import TurnstoneBot + + bot = MagicMock(spec=TurnstoneBot) + bot._MAX_NOTIFY_TRACKING = 3 + bot._notify_ws_map = { + 1: ("ws-1", "u1"), + 2: ("ws-2", "u2"), + 3: ("ws-3", "u3"), + } + bot.send = AsyncMock(return_value="4") + bot.send_notification = TurnstoneBot.send_notification.__get__(bot, TurnstoneBot) + bot._track_notification = TurnstoneBot._track_notification.__get__(bot, TurnstoneBot) + + _run(bot.send_notification("chan-1", "Hello", "ws-4")) + + assert 4 in bot._notify_ws_map + assert 1 not in bot._notify_ws_map # oldest evicted + assert len(bot._notify_ws_map) <= 3 + + def test_dm_reply_routes_to_workstream(self): + """DM reply to a tracked notification routes the message to the workstream.""" + from turnstone.channels.discord.cog import MessageCog + + bot = MagicMock() + bot.user = MagicMock() + bot.user.id = 99999 + + ts = MagicMock() + ts._is_allowed_channel = MagicMock(return_value=True) + ts.storage = MagicMock() + ts.router = MagicMock() + ts.router.resolve_user = AsyncMock(return_value="u_abc") + ts.router.send_message = AsyncMock() + ts.config = MagicMock() + # Maps message_id -> (ws_id, target_discord_user_id) + ts._notify_ws_map = {77777: ("ws-target", "12345")} + ts._notify_reply_channels = {} + bot.turnstone = ts + + cog = MessageCog(bot) + + # Build a DM reply to the tracked notification message + ref = MagicMock() + ref.message_id = 77777 + msg = _make_message(guild=False, content="additional context", reference=ref) + # msg.author.id defaults to 12345 from _make_message + + _run(cog._on_message(msg)) + + ts.router.send_message.assert_awaited_once_with("ws-target", "additional context") + assert "ws-target" in ts._notify_reply_channels + dm_chan, target_uid = ts._notify_reply_channels["ws-target"] + assert target_uid == "12345" + assert 77777 not in ts._notify_ws_map # cleaned up + + def test_dm_reply_user_mismatch_rejected_and_preserved(self): + """DM reply from wrong user is rejected; entry re-inserted for legitimate user.""" + from turnstone.channels.discord.cog import MessageCog + + bot = MagicMock() + bot.user = MagicMock() + bot.user.id = 99999 + + ts = MagicMock() + ts.router = MagicMock() + ts.router.resolve_user = AsyncMock(return_value="u_abc") + ts.router.send_message = AsyncMock() + # Target user is "99999" but replying user has author.id = 12345 + ts._notify_ws_map = {77777: ("ws-target", "99999")} + ts._notify_reply_channels = {} + bot.turnstone = ts + + cog = MessageCog(bot) + ref = MagicMock() + ref.message_id = 77777 + msg = _make_message(guild=False, content="impostor", reference=ref) + + _run(cog._on_message(msg)) + + ts.router.send_message.assert_not_awaited() + # Entry should be re-inserted so the legitimate user can still reply. + assert 77777 in ts._notify_ws_map + assert ts._notify_ws_map[77777] == ("ws-target", "99999") + + def test_dm_reply_stale_notification_feedback(self): + """DM reply to an expired/unknown notification should inform the user.""" + from turnstone.channels.discord.cog import MessageCog + + bot = MagicMock() + bot.user = MagicMock() + bot.user.id = 99999 + + ts = MagicMock() + ts.router = MagicMock() + ts.router.send_message = AsyncMock() + ts._notify_ws_map = {} # empty — no tracked notifications + ts._notify_reply_channels = {} + bot.turnstone = ts + + cog = MessageCog(bot) + + ref = MagicMock() + ref.message_id = 99999 # not in map + dm_channel = AsyncMock() + msg = _make_message(guild=False, content="reply", reference=ref, channel=dm_channel) + + _run(cog._on_message(msg)) + + # Should NOT route to any workstream + ts.router.send_message.assert_not_awaited() + # Should send feedback to the DM channel + dm_channel.send.assert_awaited_once_with("*This notification is no longer active.*") + + def test_dm_without_reference_ignored(self): + """DM without a message reference should be ignored.""" + from turnstone.channels.discord.cog import MessageCog + + bot = MagicMock() + bot.user = MagicMock() + bot.user.id = 99999 + + ts = MagicMock() + ts.router = MagicMock() + ts.router.send_message = AsyncMock() + ts._notify_ws_map = {77777: ("ws-target", "12345")} + ts._notify_reply_channels = {} + bot.turnstone = ts + + cog = MessageCog(bot) + msg = _make_message(guild=False) # reference=None + + _run(cog._on_message(msg)) + + ts.router.send_message.assert_not_awaited() + + def test_dm_reply_unlinked_user_ignored(self): + """DM reply from an unlinked user should be ignored.""" + from turnstone.channels.discord.cog import MessageCog + + bot = MagicMock() + bot.user = MagicMock() + bot.user.id = 99999 + + ts = MagicMock() + ts.router = MagicMock() + ts.router.resolve_user = AsyncMock(return_value=None) + ts.router.send_message = AsyncMock() + ts._notify_ws_map = {77777: ("ws-target", "12345")} + ts._notify_reply_channels = {} + bot.turnstone = ts + + cog = MessageCog(bot) + + ref = MagicMock() + ref.message_id = 77777 + msg = _make_message(guild=False, content="reply", reference=ref) + + _run(cog._on_message(msg)) + + ts.router.send_message.assert_not_awaited() + + def test_turn_complete_forwards_to_dm(self): + """TurnCompleteEvent should forward content to notification reply DM.""" + from turnstone.channels.discord.bot import TurnstoneBot + from turnstone.mq.protocol import TurnCompleteEvent + + bot = MagicMock(spec=TurnstoneBot) + bot.config = MagicMock() + bot.config.max_message_length = 2000 + bot._streaming = {} + bot._pending_approval_msgs = {} + bot._notify_ws_map = {} + bot._MAX_NOTIFY_TRACKING = 100 + + dm_channel = AsyncMock() + sent_msg = MagicMock() + sent_msg.id = 88888 + dm_channel.send = AsyncMock(return_value=sent_msg) + bot._notify_reply_channels = {"ws-1": (dm_channel, "u123")} + bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) + bot._track_notification = TurnstoneBot._track_notification.__get__(bot, TurnstoneBot) + + thread = AsyncMock() + + raw = TurnCompleteEvent( + ws_id="ws-1", correlation_id="", content="Here's the response" + ).to_json() + _run(bot._on_ws_event("ws-1", thread, raw)) + + # Should send to DM channel + dm_channel.send.assert_awaited_once_with("Here's the response") + # Should clean up forwarding + assert "ws-1" not in bot._notify_reply_channels + # Response message should be tracked for multi-turn replies + assert 88888 in bot._notify_ws_map + assert bot._notify_ws_map[88888] == ("ws-1", "u123") + + def test_turn_complete_cleans_up_dm_even_without_content(self): + """TurnCompleteEvent without content should still clean up DM tracking.""" + from turnstone.channels.discord.bot import TurnstoneBot + from turnstone.mq.protocol import TurnCompleteEvent + + bot = MagicMock(spec=TurnstoneBot) + bot._streaming = {} + bot._pending_approval_msgs = {} + bot._notify_ws_map = {} + + dm_channel = AsyncMock() + bot._notify_reply_channels = {"ws-1": (dm_channel, "u123")} + bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) + + thread = AsyncMock() + + raw = TurnCompleteEvent(ws_id="ws-1", correlation_id="", content="").to_json() + _run(bot._on_ws_event("ws-1", thread, raw)) + + # DM should not be sent to (no content) + dm_channel.send.assert_not_awaited() + # But should still be cleaned up + assert "ws-1" not in bot._notify_reply_channels + # No response tracked (nothing was sent) + assert len(bot._notify_ws_map) == 0 + + class TestChannelCLI: """Tests for the channel CLI entry point.""" diff --git a/turnstone/channels/_http.py b/turnstone/channels/_http.py index 6b93bc63..a62e8dbd 100644 --- a/turnstone/channels/_http.py +++ b/turnstone/channels/_http.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio import json +import re import socket import uuid from typing import TYPE_CHECKING, Any @@ -26,6 +27,9 @@ if TYPE_CHECKING: log = get_logger(__name__) +# ws_id is a hex string (8–32 chars depending on entry point). +_WS_ID_RE = re.compile(r"^[0-9a-f]{8,32}$") + async def _handle_health(request: Request) -> JSONResponse: return JSONResponse({"status": "ok", "service": "channel"}) @@ -81,6 +85,9 @@ async def _handle_notify(request: Request) -> JSONResponse: target = body.get("target") message = body.get("message", "").strip() if isinstance(body.get("message"), str) else "" title = body.get("title", "").strip() if isinstance(body.get("title"), str) else "" + ws_id = body.get("ws_id", "").strip() if isinstance(body.get("ws_id"), str) else "" + if ws_id and not _WS_ID_RE.match(ws_id): + return JSONResponse({"error": "invalid ws_id format"}, status_code=400) if not target or not message: return JSONResponse({"error": "target and message are required"}, status_code=400) @@ -132,7 +139,10 @@ async def _handle_notify(request: Request) -> JSONResponse: ) continue try: - msg_id = await adapter.send(channel_id, content) + if ws_id: + msg_id = await adapter.send_notification(channel_id, content, ws_id) + else: + msg_id = await adapter.send(channel_id, content) results.append( { "channel_type": channel_type, diff --git a/turnstone/channels/_protocol.py b/turnstone/channels/_protocol.py index bf7cfa3a..9d8a6361 100644 --- a/turnstone/channels/_protocol.py +++ b/turnstone/channels/_protocol.py @@ -41,6 +41,14 @@ class ChannelAdapter(Protocol): """Send a message to a channel. Returns the platform message ID.""" ... + async def send_notification(self, channel_id: str, content: str, ws_id: str) -> str: + """Send a notification and track the reply mapping. Returns message ID. + + Like :meth:`send` but associates the outgoing message with *ws_id* + so that replies can be routed back to the originating workstream. + """ + ... + async def edit_message(self, channel_id: str, message_id: str, content: str) -> None: """Edit an existing message in a channel.""" ... diff --git a/turnstone/channels/discord/bot.py b/turnstone/channels/discord/bot.py index 2d704156..826b2fb5 100644 --- a/turnstone/channels/discord/bot.py +++ b/turnstone/channels/discord/bot.py @@ -124,6 +124,7 @@ class TurnstoneBot: """ channel_type: str = "discord" + _MAX_NOTIFY_TRACKING: int = 100 def __init__( self, @@ -151,6 +152,16 @@ class TurnstoneBot: # workstream so that IntentVerdictEvent can update it with LLM judge # results. self._pending_approval_msgs: dict[str, discord.Message] = {} + # Notification reply tracking: maps Discord message ID → + # (ws_id, target_discord_user_id) so that DM replies can be routed + # back to the originating workstream. The target user ID is checked + # on reply to prevent cross-user message injection. + self._notify_ws_map: dict[int, tuple[str, str]] = {} + # Temporary DM forwarding: maps ws_id → (DM channel, target_user_id) + # for forwarding the workstream's next response back to the + # notification reply DM. The target_user_id is carried so the + # response message can be re-tracked for multi-turn DM conversations. + self._notify_reply_channels: dict[str, tuple[discord.abc.Messageable, str]] = {} intents = discord.Intents.default() intents.message_content = True @@ -256,6 +267,11 @@ class TurnstoneBot: self._subscribed_ws.discard(ws_id) self._streaming.pop(ws_id, None) self._pending_approval_msgs.pop(ws_id, None) + self._notify_reply_channels.pop(ws_id, None) + # Purge stale notification tracking entries for this workstream. + stale = [mid for mid, entry in self._notify_ws_map.items() if entry[0] == ws_id] + for mid in stale: + del self._notify_ws_map[mid] log.info("discord.unsubscribed", ws_id=ws_id) # -- event dispatch ------------------------------------------------------ @@ -360,6 +376,26 @@ class TurnstoneBot: sm = self._streaming.pop(ws_id, None) if sm is not None: await sm.finalize() + elif event.content: + # Catch-up: content events were missed (race between global + # SSE and per-ws SSE) — send the full response directly. + for chunk in chunk_message(event.content, self.config.max_message_length): + await thread.send(chunk) + # Forward response to notification reply DM if active. + dm_entry = self._notify_reply_channels.pop(ws_id, None) + if dm_entry is not None and event.content: + dm_channel, target_user_id = dm_entry + last_msg: discord.Message | None = None + for chunk in chunk_message(event.content, self.config.max_message_length): + try: + last_msg = await dm_channel.send(chunk) + except Exception: + log.debug("discord.notify_reply_dm_failed", ws_id=ws_id) + break + # Track the response message so the user can reply again + # for multi-turn DM conversations. + if last_msg is not None: + self._track_notification(last_msg.id, ws_id, target_user_id) # Clean up pending approval message tracking. self._pending_approval_msgs.pop(ws_id, None) @@ -390,6 +426,18 @@ class TurnstoneBot: return False return True + def _track_notification(self, message_id: int, ws_id: str, target_user_id: str) -> None: + """Record a notification message for reply routing. + + Evicts the oldest entry when the map exceeds + ``_MAX_NOTIFY_TRACKING``. Relies on dict insertion order + (Python 3.7+). + """ + while len(self._notify_ws_map) >= self._MAX_NOTIFY_TRACKING: + oldest = next(iter(self._notify_ws_map)) + del self._notify_ws_map[oldest] + self._notify_ws_map[message_id] = (ws_id, target_user_id) + def _is_allowed_channel(self, channel_id: int) -> bool: """Return True if *channel_id* is in the allowed list (or list is empty).""" if not self.config.allowed_channels: @@ -430,6 +478,26 @@ class TurnstoneBot: return str(msg.id) if msg else "" + async def send_notification(self, channel_id: str, content: str, ws_id: str) -> str: + """Send a notification DM and track the message for reply routing. + + Like :meth:`send` but records a mapping from the outgoing Discord + message ID to ``(ws_id, channel_id)`` so that a user reply can be + routed back to the originating workstream. The *channel_id* is the + Discord user ID the notification was sent to — verified on reply to + prevent cross-user message injection. + """ + msg_id_str = await self.send(channel_id, content) + if msg_id_str and ws_id: + self._track_notification(int(msg_id_str), ws_id, channel_id) + log.debug( + "discord.notification_tracked", + message_id=msg_id_str, + ws_id=ws_id, + target_user=channel_id, + ) + return msg_id_str + async def stop(self) -> None: """Disconnect the bot and clean up subscriptions.""" for ws_id in list(self._subscribed_ws): diff --git a/turnstone/channels/discord/cog.py b/turnstone/channels/discord/cog.py index ff488fcc..f35b2fd6 100644 --- a/turnstone/channels/discord/cog.py +++ b/turnstone/channels/discord/cog.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: log = get_logger(__name__) _THREAD_NAME_MAX = 100 +_DM_REPLY_MAX_LENGTH = 4096 # Discord's own message limit class MessageCog: @@ -102,8 +103,9 @@ class MessageCog: if message.author == self.bot.user or message.author.bot: return - # Ignore DMs. + # DM handling — route replies to tracked notifications. if message.guild is None: + await self._handle_dm(message) return channel = message.channel @@ -198,6 +200,63 @@ class MessageCog: author=str(message.author), ) + # -- DM reply handling --------------------------------------------------- + + async def _handle_dm(self, message: discord.Message) -> None: + """Route DM replies to tracked notification workstreams.""" + # Only handle explicit replies to a tracked notification message. + ref = message.reference + if ref is None or ref.message_id is None: + return + + # Atomic pop prevents TOCTOU race across await points. + entry = self.ts._notify_ws_map.pop(ref.message_id, None) + if entry is None: + # NOTE: This also fires for replies to non-notification bot + # messages in DMs (false positive). Acceptable because DM + # interactions are almost exclusively notification-driven. + await message.channel.send("*This notification is no longer active.*") + return + + ws_id, target_user_id = entry + + # Defence in depth: verify the replying user is the notification + # recipient. Discord enforces this (DMs are private), but a + # server-side check prevents cross-user injection via compromised + # accounts or API-level forgery. + if str(message.author.id) != target_user_id: + # Re-insert so the legitimate user can still reply. + self.ts._notify_ws_map[ref.message_id] = entry + log.warning( + "discord.notification_reply_user_mismatch", + expected=target_user_id, + actual=str(message.author.id), + ) + return + + # Resolve user identity — unlinked users are silently ignored. + # Re-insert the tracking entry so the user can retry after linking. + user_id = await self.ts.router.resolve_user("discord", str(message.author.id)) + if user_id is None: + self.ts._notify_ws_map[ref.message_id] = entry + return + + # Route the reply to the originating workstream. + content = message.content[:_DM_REPLY_MAX_LENGTH] + await self.ts.router.send_message(ws_id, content) + + # Register the DM channel for response forwarding. The bot's + # _on_ws_event handler will send the next turn's response here, + # track the response for further replies, and clean up on + # TurnCompleteEvent. + self.ts._notify_reply_channels[ws_id] = (message.channel, target_user_id) + + log.info( + "discord.notification_reply_routed", + ws_id=ws_id, + author=str(message.author), + ) + # -- slash commands ------------------------------------------------------ async def _cmd_link(self, interaction: discord.Interaction, token: str) -> None: diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 384c3774..e263e277 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -3963,6 +3963,7 @@ class ChatSession: "target": target, "message": item["message"], "title": item.get("title", ""), + "ws_id": self._ws_id, } # Build auth headers for service-to-service call diff --git a/turnstone/mq/bridge.py b/turnstone/mq/bridge.py index bd3278a3..f291ced6 100644 --- a/turnstone/mq/bridge.py +++ b/turnstone/mq/bridge.py @@ -16,6 +16,7 @@ import os import threading import time import uuid +from collections import deque from typing import TYPE_CHECKING, Any import httpx @@ -57,6 +58,11 @@ log = logging.getLogger("turnstone.mq.bridge") # Server's default safe tools (auto-approved without user confirmation) DEFAULT_SAFE_TOOLS = frozenset(["read_file", "search", "man", "memory", "recall"]) +# Maximum total character count of the per-ws content buffer. Prevents +# unbounded memory growth if a workstream produces very long responses or +# the idle event never fires (e.g. bug / disconnect). +_MAX_CONTENT_BUFFER_CHARS = 256 * 1024 + class Bridge: """Connects a message broker to turnstone-server's HTTP API. @@ -106,6 +112,12 @@ class Bridge: self._active_sends: dict[str, str] = {} # ws_id → correlation_id self._pending_approvals: dict[str, str] = {} # ws_id → request_id self._pending_plan_reviews: dict[str, str] = {} # ws_id → request_id + # Content buffer: accumulates assistant text per workstream from + # per-ws SSE. Attached to TurnCompleteEvent when idle is detected + # via the global SSE so downstream consumers can catch up if the + # streaming path missed events (race between the two SSE connections). + self._ws_content_buffer: dict[str, deque[str]] = {} + self._ws_content_buffer_size: dict[str, int] = {} # running char total self._running = True @property @@ -587,7 +599,23 @@ class Bridge: etype = data.get("type", "") if etype == "content": - self._publish_ws(ws_id, ContentEvent(ws_id=ws_id, text=data.get("text", ""))) + text = data.get("text", "") + if text: + with self._lock: + if ws_id not in self._ws_content_buffer: + self._ws_content_buffer[ws_id] = deque() + self._ws_content_buffer_size[ws_id] = 0 + buf = self._ws_content_buffer[ws_id] + buf.append(text) + self._ws_content_buffer_size[ws_id] += len(text) + # Cap buffer per workstream to prevent DoS from + # extremely long responses or missing idle events. + while ( + self._ws_content_buffer_size[ws_id] > _MAX_CONTENT_BUFFER_CHARS + and len(buf) > 1 + ): + self._ws_content_buffer_size[ws_id] -= len(buf.popleft()) + self._publish_ws(ws_id, ContentEvent(ws_id=ws_id, text=text)) elif etype == "reasoning": self._publish_ws(ws_id, ReasoningEvent(ws_id=ws_id, text=data.get("text", ""))) elif etype == "tool_info": @@ -825,7 +853,16 @@ class Bridge: if state == "idle": with self._lock: cid = self._active_sends.pop(ws_id, None) - self._publish_ws(ws_id, TurnCompleteEvent(ws_id=ws_id, correlation_id=cid or "")) + content_parts = self._ws_content_buffer.pop(ws_id, deque()) + self._ws_content_buffer_size.pop(ws_id, None) + self._publish_ws( + ws_id, + TurnCompleteEvent( + ws_id=ws_id, + correlation_id=cid or "", + content="".join(content_parts), + ), + ) elif etype == "ws_rename": self._publish_global(WorkstreamRenameEvent(ws_id=ws_id, name=data.get("name", ""))) @@ -840,6 +877,8 @@ class Bridge: self._ws_auto_approve.pop(ws_id, None) self._ws_approve_tools.pop(ws_id, None) self._active_sends.pop(ws_id, None) + self._ws_content_buffer.pop(ws_id, None) + self._ws_content_buffer_size.pop(ws_id, None) # -- heartbeat ----------------------------------------------------------- diff --git a/turnstone/mq/protocol.py b/turnstone/mq/protocol.py index d47aadd8..8f96d26c 100644 --- a/turnstone/mq/protocol.py +++ b/turnstone/mq/protocol.py @@ -268,9 +268,15 @@ class TurnCompleteEvent(OutboundEvent): This is a synthetic event produced by the bridge when it detects the ws_state transition to 'idle'. ``correlation_id`` is set for MQ-initiated turns and empty for turns initiated from the server UI. + + ``content`` carries the full assistant response text accumulated from + per-ws SSE content tokens. Downstream consumers (e.g. Discord bot) can + use it as a catch-up when the streaming path missed events due to the + race between the global SSE and per-ws SSE connections. """ type: str = "turn_complete" + content: str = "" @dataclass