diff --git a/docs/architecture.md b/docs/architecture.md index dcb3c26d..db0209df 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -85,7 +85,7 @@ turnstone/ static/ Cluster dashboard web UI (page-specific HTML, CSS, JS) channels/ cli.py Unified channel gateway entry point (turnstone-channel) - _protocol.py ChannelAdapter protocol, ChannelEvent dataclass + _protocol.py ChannelAdapter protocol _routing.py ChannelRouter — channel/thread ↔ workstream mapping via HTTP _config.py Base ChannelConfig dataclass discord/ Discord adapter (bot, cog, views, streaming, config) diff --git a/docs/channels.md b/docs/channels.md index 05dd1c64..e968383a 100644 --- a/docs/channels.md +++ b/docs/channels.md @@ -35,8 +35,7 @@ Key components: - **ChannelAdapter protocol** (`turnstone/channels/_protocol.py`) — generic interface for any messaging platform. Defines `start()`, `stop()`, - `send()`, `send_notification()`, `edit_message()`, - `send_approval_request()`, `send_plan_review()`, and `create_thread()`. + `send()`, and `send_notification()`. - **ChannelRouter** (`turnstone/channels/_routing.py`) — maps channel/thread IDs to turnstone workstream IDs. Handles workstream creation via HTTP, stale route detection, and user identity resolution. @@ -425,10 +424,6 @@ class ChannelAdapter(Protocol): 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 @@ -436,6 +431,11 @@ 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. +Platform-specific concerns — approval prompts, plan reviews, message +edits, thread creation — live inside the adapter implementation and are +not part of the protocol surface. Each adapter drives those via its +own `_on_ws_event` dispatcher using SDK-native APIs. + To add a new platform: 1. Create `turnstone/channels//` package diff --git a/tests/test_channel_discord.py b/tests/test_channel_discord.py index 48996a6a..aad4f3a0 100644 --- a/tests/test_channel_discord.py +++ b/tests/test_channel_discord.py @@ -28,6 +28,21 @@ def _run(coro): return asyncio.run(coro) +def _bind_ws_event_handlers(bot, cls): + """Bind ``_on_ws_event`` + every ``_handle_*`` method from *cls* to *bot*. + + ``MagicMock(spec=cls)`` stubs async methods as ``AsyncMock`` no-ops, + so dispatcher tests that invoke the real ``_on_ws_event`` must also + bind the per-event handlers it delegates to. + """ + bot._on_ws_event = cls._on_ws_event.__get__(bot, cls) + for name in dir(cls): + if name.startswith("_handle_"): + attr = getattr(cls, name) + if callable(attr): + setattr(bot, name, attr.__get__(bot, cls)) + + def _make_message(*, bot=False, guild=True, content="hello", channel=None, reference=None): """Build a mock ``discord.Message``.""" msg = MagicMock(spec=discord.Message) @@ -128,7 +143,7 @@ class TestStreamingMessage: _run(sm.append("hello ")) _run(sm.append("world")) - assert "".join(sm._buffer) == "hello world" + assert sm.accumulated_text == "hello world" def test_finalize_sends_when_no_prior_message(self): from turnstone.channels.discord.bot import StreamingMessage @@ -153,7 +168,7 @@ class TestStreamingMessage: # First append triggers flush (interval=0) which creates the message. _run(sm.append("hi")) - assert sm._message is sent_msg + assert sm.message is sent_msg _run(sm.append(" there")) _run(sm.finalize()) @@ -352,20 +367,20 @@ class TestAskModelSelection: class TestParseFooter: """Tests for _parse_footer in views.py.""" - def test_valid_footer(self): + def test_valid_footer_with_owner(self): from turnstone.channels.discord.views import _parse_footer + interaction = _make_interaction(footer_text="ws_abc|corr_123|12345") + result = _parse_footer(interaction) + assert result == ("ws_abc", "corr_123", "12345") + + def test_footer_without_owner_returns_empty_owner(self): + from turnstone.channels.discord.views import _parse_footer + + # Legacy footer without an owner field (pre-upgrade posts). interaction = _make_interaction(footer_text="ws_abc|corr_123") result = _parse_footer(interaction) - assert result == ("ws_abc", "corr_123") - - def test_footer_with_pipe_in_correlation(self): - from turnstone.channels.discord.views import _parse_footer - - interaction = _make_interaction(footer_text="ws_abc|corr|extra") - result = _parse_footer(interaction) - # split("|", 1) means the second part includes everything after first pipe. - assert result == ("ws_abc", "corr|extra") + assert result == ("ws_abc", "corr_123", "") def test_no_message_returns_none(self): from turnstone.channels.discord.views import _parse_footer @@ -428,7 +443,7 @@ class TestWsEventFinalization: bot._notify_reply_channels = {} # Use the real _on_ws_event method - bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) + _bind_ws_event_handlers(bot, TurnstoneBot) thread = AsyncMock() @@ -457,7 +472,7 @@ class TestWsEventFinalization: bot._tool_info_msgs = {} bot._pending_approval_msgs = {} bot._notify_reply_channels = {} - bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) + _bind_ws_event_handlers(bot, TurnstoneBot) thread = AsyncMock() @@ -478,6 +493,7 @@ class TestApprovalVerdictDisplay: def _make_bot(self): """Build a mock TurnstoneBot with _on_ws_event bound.""" + from turnstone.channels._routing import PolicyVerdict from turnstone.channels.discord.bot import TurnstoneBot bot = MagicMock(spec=TurnstoneBot) @@ -493,7 +509,9 @@ class TestApprovalVerdictDisplay: 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) + bot.router = MagicMock() + bot.router.evaluate_tool_policies = AsyncMock(return_value=PolicyVerdict(kind="none")) + _bind_ws_event_handlers(bot, TurnstoneBot) return bot def test_approval_with_heuristic_verdict(self): @@ -612,7 +630,7 @@ class TestApprovalVerdictDisplay: bot._tool_info_msgs = {} bot._pending_approval_msgs = {"ws-1": MagicMock()} bot._notify_reply_channels = {} - bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) + _bind_ws_event_handlers(bot, TurnstoneBot) thread = AsyncMock() event = StreamEndEvent(ws_id="ws-1") @@ -638,7 +656,7 @@ class TestStreamEndBehavior: bot._tool_info_msgs = {} bot._pending_approval_msgs = {} bot._notify_reply_channels = {} - bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) + _bind_ws_event_handlers(bot, TurnstoneBot) return bot def test_stream_end_no_streaming_no_send(self): @@ -674,38 +692,88 @@ class TestStreamEndBehavior: 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.""" + def _make_dm_bot(self, *, sent_message_id: int): + """Build a MagicMock bot whose notification target resolves to a DM.""" from turnstone.channels.discord.bot import TurnstoneBot bot = MagicMock(spec=TurnstoneBot) + bot.config = MagicMock() + bot.config.max_message_length = 2000 + + sent_msg = MagicMock() + sent_msg.id = sent_message_id + + dm_channel = MagicMock() + dm_channel.send = AsyncMock(return_value=sent_msg) + + user = MagicMock() + user.id = 7777 + user.create_dm = AsyncMock(return_value=dm_channel) + + inner_bot = MagicMock() + inner_bot.get_channel = MagicMock(return_value=None) + inner_bot.fetch_user = AsyncMock(return_value=user) + bot._bot = inner_bot + + bot.send_notification = TurnstoneBot.send_notification.__get__(bot, TurnstoneBot) + bot._track_notification = TurnstoneBot._track_notification.__get__(bot, TurnstoneBot) + return bot + + def test_send_notification_tracks_dm_with_user_id(self): + """send_notification for a DM records (ws_id, resolved_user_id).""" + bot = self._make_dm_bot(sent_message_id=12345) bot._notify_ws_map = {} bot._MAX_NOTIFY_TRACKING = 100 - bot.send = AsyncMock(return_value="12345") + + _run(bot.send_notification("7777", "Hello", "ws-abc")) + + # Tracked under the resolved Discord user ID, not the raw argument. + assert 12345 in bot._notify_ws_map + assert bot._notify_ws_map[12345] == ("ws-abc", "7777") + + def test_send_notification_to_guild_channel_is_not_tracked(self): + """Notifications delivered to a guild channel must not register reply tracking. + + The reply-channel_id check treats the stored value as a Discord + user ID, so storing a channel ID would reject every legitimate + reply. + """ + from turnstone.channels.discord.bot import TurnstoneBot + + bot = MagicMock(spec=TurnstoneBot) + bot.config = MagicMock() + bot.config.max_message_length = 2000 + bot._notify_ws_map = {} + bot._MAX_NOTIFY_TRACKING = 100 + + sent_msg = MagicMock() + sent_msg.id = 99999 + + channel = MagicMock() + channel.send = AsyncMock(return_value=sent_msg) + + inner_bot = MagicMock() + inner_bot.get_channel = MagicMock(return_value=channel) + bot._bot = inner_bot + 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")) + _run(bot.send_notification("888888", "Hello", "ws-abc")) - assert 12345 in bot._notify_ws_map - assert bot._notify_ws_map[12345] == ("ws-abc", "chan-1") + assert bot._notify_ws_map == {} 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 = self._make_dm_bot(sent_message_id=4) 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")) + _run(bot.send_notification("7777", "Hello", "ws-4")) assert 4 in bot._notify_ws_map assert 1 not in bot._notify_ws_map # oldest evicted @@ -878,7 +946,7 @@ class TestNotificationTracking: 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) + _bind_ws_event_handlers(bot, TurnstoneBot) bot._track_notification = TurnstoneBot._track_notification.__get__(bot, TurnstoneBot) thread = AsyncMock() @@ -913,7 +981,7 @@ class TestNotificationTracking: dm_channel = AsyncMock() bot._notify_reply_channels = {"ws-1": (dm_channel, "u123")} - bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) + _bind_ws_event_handlers(bot, TurnstoneBot) thread = AsyncMock() @@ -1052,40 +1120,83 @@ class TestTryParseMedia: class TestIsSafeImageUrl: """Tests for _is_safe_image_url in _formatter.py.""" - def test_http_url(self): + @staticmethod + def _patch_resolver(monkeypatch, ips): + """Replace socket.getaddrinfo with a stub returning *ips*.""" + import socket + + def fake(host, port, family=0, *args, **kwargs): # noqa: ARG001 + return [(family, 0, 0, "", (ip, 0)) for ip in ips] + + monkeypatch.setattr(socket, "getaddrinfo", fake) + + def test_http_url(self, monkeypatch): from turnstone.channels._formatter import _is_safe_image_url - assert _is_safe_image_url("http://jellyfin:8096/Items/abc/Images/Primary") is True + self._patch_resolver(monkeypatch, ["203.0.113.5"]) + assert _run(_is_safe_image_url("http://jellyfin:8096/Items/abc/Images/Primary")) is True - def test_https_url(self): + def test_https_url(self, monkeypatch): from turnstone.channels._formatter import _is_safe_image_url - assert _is_safe_image_url("https://jellyfin.example.com/Items/abc/Images/Primary") is True + self._patch_resolver(monkeypatch, ["203.0.113.5"]) + assert ( + _run(_is_safe_image_url("https://jellyfin.example.com/Items/abc/Images/Primary")) + is True + ) def test_ftp_rejected(self): from turnstone.channels._formatter import _is_safe_image_url - assert _is_safe_image_url("ftp://evil.com/image.jpg") is False + assert _run(_is_safe_image_url("ftp://evil.com/image.jpg")) is False def test_file_rejected(self): from turnstone.channels._formatter import _is_safe_image_url - assert _is_safe_image_url("file:///etc/passwd") is False + assert _run(_is_safe_image_url("file:///etc/passwd")) is False def test_userinfo_rejected(self): from turnstone.channels._formatter import _is_safe_image_url - assert _is_safe_image_url("http://user:pass@jellyfin:8096/image") is False + assert _run(_is_safe_image_url("http://user:pass@jellyfin:8096/image")) is False def test_empty_rejected(self): from turnstone.channels._formatter import _is_safe_image_url - assert _is_safe_image_url("") is False + assert _run(_is_safe_image_url("")) is False def test_private_ip_allowed(self): from turnstone.channels._formatter import _is_safe_image_url - assert _is_safe_image_url("http://192.168.0.6:8096/Items/abc/Images/Primary") is True + assert _run(_is_safe_image_url("http://192.168.0.6:8096/Items/abc/Images/Primary")) is True + + def test_dns_rebinding_rejected(self, monkeypatch): + """Hostname that resolves to a loopback IP must be rejected.""" + from turnstone.channels._formatter import _is_safe_image_url + + self._patch_resolver(monkeypatch, ["127.0.0.1"]) + assert _run(_is_safe_image_url("http://rebind.example.com/image")) is False + + def test_metadata_endpoint_rejected(self): + """AWS/GCP metadata IP is link-local → rejected.""" + from turnstone.channels._formatter import _is_safe_image_url + + assert _run(_is_safe_image_url("http://169.254.169.254/latest/meta-data/")) is False + + def test_ipv6_aws_nitro_metadata_rejected(self, monkeypatch): + """fd00:ec2::254 is IPv6 ULA (is_private) but must be blocked — + the IPv4 169.254.169.254 check left this analogue open.""" + from turnstone.channels._formatter import _is_safe_image_url + + self._patch_resolver(monkeypatch, ["fd00:ec2::254"]) + assert _run(_is_safe_image_url("http://nitro.example.com/")) is False + + def test_ipv6_ecs_task_metadata_rejected(self, monkeypatch): + """ECS Task Metadata lives in the same fd00:ec2::/32 prefix.""" + from turnstone.channels._formatter import _is_safe_image_url + + self._patch_resolver(monkeypatch, ["fd00:ec2::23"]) + assert _run(_is_safe_image_url("http://ecs-meta.example.com/")) is False class TestBuildMediaEmbed: @@ -1187,7 +1298,7 @@ class TestThinkingIndicator: 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) + _bind_ws_event_handlers(bot, TurnstoneBot) return bot def test_thinking_start_sends_message(self): @@ -1244,7 +1355,7 @@ class TestThinkingIndicator: # Thinking message becomes the StreamingMessage base — no delete. assert "ws-1" not in bot._thinking_msgs sm = bot._streaming["ws-1"] - assert sm._message is thinking_msg + assert sm.message is thinking_msg def test_stream_end_clears_thinking_message(self): from turnstone.sdk.events import StreamEndEvent @@ -1287,7 +1398,7 @@ class TestToolInfoEvent: 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) + _bind_ws_event_handlers(bot, TurnstoneBot) return bot def test_sends_per_item_embed(self): @@ -1382,7 +1493,7 @@ class TestToolResultEvent: bot._notify_reply_channels = {} bot._http_client = MagicMock() bot._should_auto_approve = MagicMock(return_value=False) - bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot) + _bind_ws_event_handlers(bot, TurnstoneBot) return bot def test_marks_info_done_and_sends_result(self): @@ -1537,7 +1648,7 @@ class TestApprovalResolved: 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) + _bind_ws_event_handlers(bot, TurnstoneBot) return bot def test_disables_buttons_on_timeout(self): @@ -1605,3 +1716,226 @@ class TestChannelCLI: main() assert exc_info.value.code == 1 + + +# --------------------------------------------------------------------------- +# Approval / plan-review interaction views — owner-check regression tests +# --------------------------------------------------------------------------- + + +def _make_view_interaction(user_id: int, footer: str | None) -> MagicMock: + """Build a minimal interaction for ApprovalView / PlanReviewView tests.""" + interaction = MagicMock(spec=discord.Interaction) + interaction.user = MagicMock() + interaction.user.id = user_id + interaction.response = MagicMock() + interaction.response.send_message = AsyncMock() + interaction.response.defer = AsyncMock() + interaction.response.send_modal = AsyncMock() + interaction.followup = MagicMock() + interaction.followup.send = AsyncMock() + interaction.message = MagicMock() + if footer is None: + interaction.message.embeds = [] + else: + embed = MagicMock() + embed.footer.text = footer + interaction.message.embeds = [embed] + return interaction + + +def _make_view_bot() -> MagicMock: + """Build a TurnstoneBot double with just the surface the views read.""" + from turnstone.channels.discord.bot import TurnstoneBot + + bot = MagicMock(spec=TurnstoneBot) + bot.router = MagicMock() + bot.router.resolve_user = AsyncMock(return_value="turnstone-user-1") + bot.router.send_approval = AsyncMock() + bot.router.send_plan_feedback = AsyncMock() + bot._pending_approval_msgs = {} + return bot + + +class TestApprovalViewOwnerCheck: + """ApprovalView rejects clicks from anyone other than the session owner.""" + + def test_owner_approve_allowed(self, monkeypatch): + from turnstone.channels.discord.views import ApprovalView + + # Avoid real disable_message_buttons (touches discord.ui internals). + monkeypatch.setattr( + "turnstone.channels.discord.views._disable_buttons", + AsyncMock(), + ) + view = ApprovalView(_make_view_bot()) + interaction = _make_view_interaction(user_id=42, footer="ws-1|corr-1|42") + + _run(view._handle(interaction, approved=True, always=False)) + + view.bot.router.send_approval.assert_awaited_once_with( + ws_id="ws-1", + correlation_id="corr-1", + approved=True, + always=False, + ) + + def test_non_owner_rejected(self): + from turnstone.channels.discord.views import ApprovalView + + view = ApprovalView(_make_view_bot()) + interaction = _make_view_interaction(user_id=999, footer="ws-1|corr-1|42") + + _run(view._handle(interaction, approved=True, always=False)) + + view.bot.router.send_approval.assert_not_awaited() + interaction.response.send_message.assert_awaited_once() + msg_kwargs = interaction.response.send_message.call_args + assert "Only the session owner" in msg_kwargs.args[0] + assert msg_kwargs.kwargs.get("ephemeral") is True + + def test_legacy_footer_without_owner_rejected(self): + from turnstone.channels.discord.views import ApprovalView + + view = ApprovalView(_make_view_bot()) + # Pre-upgrade footer with only ws_id|correlation_id — fail closed. + interaction = _make_view_interaction(user_id=42, footer="ws-1|corr-1") + + _run(view._handle(interaction, approved=True, always=False)) + + view.bot.router.send_approval.assert_not_awaited() + + +class TestPlanReviewViewOwnerCheck: + """PlanReviewView rejects clicks from anyone other than the session owner.""" + + def test_owner_approve_allowed(self, monkeypatch): + from turnstone.channels.discord.views import PlanReviewView + + monkeypatch.setattr( + "turnstone.channels.discord.views._disable_buttons", + AsyncMock(), + ) + view = PlanReviewView(_make_view_bot()) + interaction = _make_view_interaction(user_id=42, footer="ws-1|corr-1|42") + + _run(view._handle_approve(interaction)) + + view.bot.router.send_plan_feedback.assert_awaited_once_with( + ws_id="ws-1", + correlation_id="corr-1", + feedback="", + ) + + def test_non_owner_approve_rejected(self): + from turnstone.channels.discord.views import PlanReviewView + + view = PlanReviewView(_make_view_bot()) + interaction = _make_view_interaction(user_id=999, footer="ws-1|corr-1|42") + + _run(view._handle_approve(interaction)) + + view.bot.router.send_plan_feedback.assert_not_awaited() + interaction.response.send_message.assert_awaited_once() + + def test_non_owner_changes_modal_rejected(self): + from turnstone.channels.discord.views import PlanReviewView + + view = PlanReviewView(_make_view_bot()) + interaction = _make_view_interaction(user_id=999, footer="ws-1|corr-1|42") + + _run(view._handle_changes(interaction)) + + interaction.response.send_modal.assert_not_awaited() + interaction.response.send_message.assert_awaited_once() + + +class TestDiscordThreadOwnerCheck: + """Sec-3 gate: only the thread creator can send messages into the workstream.""" + + @staticmethod + def _make_cog_and_ts(): + """Build a MessageCog wired to a minimal TurnstoneBot double.""" + from turnstone.channels.discord.cog import MessageCog + + bot = MagicMock() + bot.user = MagicMock() + bot.user.id = 99999 + bot.user.mentioned_in = MagicMock(return_value=False) + + ts = MagicMock() + ts._is_allowed_channel = MagicMock(return_value=True) + ts.storage = MagicMock() + ts.router = MagicMock() + ts.router.lookup_ws_id = AsyncMock(return_value="ws-1") + ts.router.resolve_user = AsyncMock(return_value="turnstone-user-1") + ts.router.send_message = AsyncMock() + ts.router.get_or_create_workstream = AsyncMock(return_value=("ws-1", False)) + ts.config = MagicMock() + ts._ws_tasks = {} + ts._subscribed_ws = {"ws-1"} + ts._notify_ws_map = {} + ts._notify_reply_channels = {} + ts.get_thread_invoker = MagicMock(return_value=None) + ts.subscribe_ws = AsyncMock() + bot.turnstone = ts + + return MessageCog(bot), ts + + def test_non_owner_thread_message_dropped(self): + """A linked user who is NOT the thread creator gets their message + silently dropped — router.send_message must not fire.""" + cog, ts = self._make_cog_and_ts() + + # Build a thread whose owner_id is different from the message author. + thread = MagicMock(spec=discord.Thread) + thread.id = 555 + thread.parent_id = 111 + thread.owner_id = 42 # thread creator + thread.name = "some-thread" + + msg = _make_message(guild=True, channel=thread) + msg.author.id = 999 # non-owner trying to inject + + _run(cog._on_message(msg)) + + ts.router.send_message.assert_not_awaited() + ts.router.get_or_create_workstream.assert_not_awaited() + + def test_ask_thread_followup_allowed_when_invoker_registered(self): + """/ask creates threads with owner_id=bot; follow-ups from the + registered invoker must still reach the workstream.""" + cog, ts = self._make_cog_and_ts() + # Simulate what _cmd_ask does after channel.create_thread(). + ts.get_thread_invoker = MagicMock(return_value=111) + + thread = MagicMock(spec=discord.Thread) + thread.id = 555 + thread.parent_id = 222 + thread.owner_id = 99999 # bot owns the thread after channel.create_thread + thread.name = "ask-thread" + + msg = _make_message(guild=True, channel=thread) + msg.author.id = 111 # the human who ran /ask + + _run(cog._on_message(msg)) + + ts.router.send_message.assert_awaited_once_with("ws-1", msg.content) + + def test_ask_thread_rejects_other_user_even_when_invoker_registered(self): + """Registered invoker lock: only that user's follow-ups pass.""" + cog, ts = self._make_cog_and_ts() + ts.get_thread_invoker = MagicMock(return_value=111) + + thread = MagicMock(spec=discord.Thread) + thread.id = 555 + thread.parent_id = 222 + thread.owner_id = 99999 # bot-owned + thread.name = "ask-thread" + + msg = _make_message(guild=True, channel=thread) + msg.author.id = 222 # someone other than the recorded invoker + + _run(cog._on_message(msg)) + + ts.router.send_message.assert_not_awaited() diff --git a/tests/test_channel_protocol.py b/tests/test_channel_protocol.py index 3a7568eb..c843dd70 100644 --- a/tests/test_channel_protocol.py +++ b/tests/test_channel_protocol.py @@ -1,55 +1,13 @@ -"""Tests for turnstone.channels._protocol and turnstone.channels._formatter.""" +"""Tests for turnstone.channels._formatter.""" from __future__ import annotations from turnstone.channels._formatter import ( chunk_message, format_approval_request, - format_plan_review, format_verdict, truncate, ) -from turnstone.channels._protocol import ChannelEvent - -# --------------------------------------------------------------------------- -# ChannelEvent -# --------------------------------------------------------------------------- - - -class TestChannelEvent: - def test_construction(self) -> None: - evt = ChannelEvent( - channel_type="discord", - channel_id="ch-1", - channel_user_id="u-42", - message="hello", - parent_channel_id="parent", - metadata={"key": "val"}, - ) - assert evt.channel_type == "discord" - assert evt.channel_id == "ch-1" - assert evt.channel_user_id == "u-42" - assert evt.message == "hello" - assert evt.parent_channel_id == "parent" - assert evt.metadata == {"key": "val"} - - def test_defaults(self) -> None: - evt = ChannelEvent( - channel_type="slack", - channel_id="ch-2", - channel_user_id="u-7", - message="hi", - ) - assert evt.parent_channel_id == "" - assert evt.metadata == {} - - def test_metadata_independence(self) -> None: - """Default metadata dicts are independent across instances.""" - a = ChannelEvent(channel_type="x", channel_id="1", channel_user_id="u", message="m") - b = ChannelEvent(channel_type="x", channel_id="2", channel_user_id="u", message="m") - a.metadata["key"] = "val" - assert "key" not in b.metadata - # --------------------------------------------------------------------------- # chunk_message @@ -172,18 +130,6 @@ class TestFormatApprovalRequest: assert "/etc/hosts" in result -# --------------------------------------------------------------------------- -# format_plan_review -# --------------------------------------------------------------------------- - - -class TestFormatPlanReview: - def test_format(self) -> None: - result = format_plan_review("Step 1: do stuff") - assert result.startswith("**Plan review requested:**") - assert "Step 1: do stuff" in result - - # --------------------------------------------------------------------------- # format_verdict # --------------------------------------------------------------------------- diff --git a/tests/test_channel_slack.py b/tests/test_channel_slack.py index a2179071..c0921c5d 100644 --- a/tests/test_channel_slack.py +++ b/tests/test_channel_slack.py @@ -62,12 +62,20 @@ def _make_bot() -> tuple[object, MagicMock, MagicMock]: storage = MagicMock() storage.list_channel_routes_by_type = MagicMock(return_value=[]) + from turnstone.channels._routing import PolicyVerdict + router = MagicMock() router.get_or_create_workstream = AsyncMock(return_value=("ws-1", True)) router.send_message = AsyncMock() router.send_approval = AsyncMock() router.send_plan_feedback = AsyncMock() router.get_node_url = AsyncMock(return_value="http://localhost:8080") + router.evaluate_tool_policies = AsyncMock(return_value=PolicyVerdict(kind="none")) + router.delete_route = AsyncMock() + router.close_workstream = AsyncMock() + # Default: every test Slack user is already linked. Tests that + # exercise the unlinked path override this per-instance. + router.resolve_user = AsyncMock(return_value="turnstone-user-1") router.aclose = AsyncMock() client = AsyncMock() @@ -185,6 +193,98 @@ class TestSlackRoute: == "C123:U456:111.222" ) + def test_round_trip(self) -> None: + """Every shape emitted by to_channel_id must round-trip through parse.""" + from turnstone.channels.slack.routes import SlackRoute + + shapes = [ + SlackRoute(channel="C123"), + SlackRoute(channel="C123", user_id="U456"), + SlackRoute(channel="C123", user_id="U456", thread_ts="111.222"), + ] + for route in shapes: + assert SlackRoute.parse(route.to_channel_id()) == route + + def test_parse_trailing_colon_normalizes(self) -> None: + """``"C123:"`` should normalize to ``SlackRoute("C123")``.""" + from turnstone.channels.slack.routes import SlackRoute + + assert SlackRoute.parse("C123:") == SlackRoute(channel="C123") + assert SlackRoute.parse("C123:U456:") == SlackRoute(channel="C123", user_id="U456") + + def test_parse_extra_colons_folded_into_thread_ts(self) -> None: + """Extra ``:`` past the third field fold into ``thread_ts`` verbatim. + + Slack IDs and timestamps never contain ``:`` so this is safe in + practice; the test locks the documented behaviour. + """ + from turnstone.channels.slack.routes import SlackRoute + + route = SlackRoute.parse("C1:U1:ts:extra") + assert route.channel == "C1" + assert route.user_id == "U1" + assert route.thread_ts == "ts:extra" + + +# --------------------------------------------------------------------------- +# Route recovery + session archival +# --------------------------------------------------------------------------- + + +class TestRecoverRoutes: + """Tests for TurnstoneSlackBot._recover_routes on bot startup.""" + + def test_latest_ts_per_user_wins(self) -> None: + """When multiple routes exist for the same (channel, user), the + newest thread_ts populates _channel_sessions.""" + bot, _router, _client = _make_bot() + bot.storage.list_channel_routes_by_type = MagicMock( # type: ignore[attr-defined] + return_value=[ + {"ws_id": "ws-old", "channel_id": "C1:U1:1000000.000001"}, + {"ws_id": "ws-new", "channel_id": "C1:U1:2000000.000001"}, + ] + ) + bot.subscribe_ws = AsyncMock() # type: ignore[attr-defined] + + _run(bot._recover_routes()) # type: ignore[attr-defined] + + assert bot._channel_sessions == {("C1", "U1"): ("ws-new", "2000000.000001")} # type: ignore[attr-defined] + # Both routes get resubscribed so their SSE streams stay active. + assert bot.subscribe_ws.await_count == 2 # type: ignore[attr-defined] + + def test_non_threaded_routes_skip_session_table(self) -> None: + """A route without a thread_ts still gets subscribed but never + populates _channel_sessions (DMs fall into this shape).""" + bot, _router, _client = _make_bot() + bot.storage.list_channel_routes_by_type = MagicMock( # type: ignore[attr-defined] + return_value=[{"ws_id": "ws-dm", "channel_id": "D1:U9"}] + ) + bot.subscribe_ws = AsyncMock() # type: ignore[attr-defined] + + _run(bot._recover_routes()) # type: ignore[attr-defined] + + assert bot._channel_sessions == {} # type: ignore[attr-defined] + bot.subscribe_ws.assert_awaited_once_with("ws-dm", "D1:U9") # type: ignore[attr-defined] + + +class TestArchiveSession: + """Tests for TurnstoneSlackBot._archive_session cleanup.""" + + def test_archive_drops_route_and_closes_workstream(self) -> None: + bot, router, client = _make_bot() + bot._channel_sessions[("C1", "U1")] = ("ws-old", "1000000.000001") # type: ignore[attr-defined] + bot._subscribed_ws.add("ws-old") # type: ignore[attr-defined] + + _run(bot._archive_session("C1", "U1", "ws-old", "1000000.000001")) # type: ignore[attr-defined] + + router.delete_route.assert_awaited_once_with( # type: ignore[attr-defined] + "slack", "C1:U1:1000000.000001" + ) + router.close_workstream.assert_awaited_once_with("ws-old") # type: ignore[attr-defined] + assert ("C1", "U1") not in bot._channel_sessions # type: ignore[attr-defined] + # archive notice posted in the old thread + client.chat_postMessage.assert_awaited() # type: ignore[attr-defined] + # --------------------------------------------------------------------------- # Preview sanitization @@ -241,7 +341,7 @@ class TestStreamingMessage: _run(sm.append("hello ")) _run(sm.append("world")) - assert "".join(sm._buffer) == "hello world" + assert sm.accumulated_text == "hello world" def test_finalize_sends_when_no_prior_message(self) -> None: from turnstone.channels.slack.bot import StreamingMessage @@ -264,7 +364,7 @@ class TestStreamingMessage: sm = StreamingMessage(client=client, channel="C1", edit_interval=0.0) _run(sm.append("hi")) - assert sm._ts == "123" + assert sm.message_ts == "123" _run(sm.finalize()) client.chat_update.assert_awaited() @@ -489,7 +589,7 @@ class TestApprovalOwnership: "container": {"channel_id": "C01SAPU5414", "message_ts": "111.222"}, } - _run(bot._on_approve(AsyncMock(), body)) # type: ignore[attr-defined] + _run(bot._resolve_approval(AsyncMock(), body, approved=True)) # type: ignore[attr-defined] client.chat_postEphemeral.assert_awaited_once() router.send_approval.assert_not_awaited() @@ -511,7 +611,7 @@ class TestApprovalOwnership: "container": {"channel_id": "C01SAPU5414", "message_ts": "111.222"}, } - _run(bot._on_deny(AsyncMock(), body)) # type: ignore[attr-defined] + _run(bot._resolve_approval(AsyncMock(), body, approved=False)) # type: ignore[attr-defined] client.chat_postEphemeral.assert_awaited_once() router.send_approval.assert_not_awaited() @@ -533,7 +633,7 @@ class TestApprovalOwnership: "container": {"channel_id": "C01SAPU5414", "message_ts": "111.222"}, } - _run(bot._on_approve(AsyncMock(), body)) # type: ignore[attr-defined] + _run(bot._resolve_approval(AsyncMock(), body, approved=True)) # type: ignore[attr-defined] router.send_approval.assert_awaited_once_with(ws_id, "corr-1", approved=True) client.chat_update.assert_awaited_once() @@ -548,6 +648,7 @@ class TestWsEventDispatch: """Tests for SSE event handling in the Slack bot.""" def _make_ws_bot(self) -> tuple[object, MagicMock]: + from turnstone.channels._routing import PolicyVerdict from turnstone.channels.slack.bot import TurnstoneSlackBot from turnstone.channels.slack.config import SlackConfig @@ -560,6 +661,8 @@ class TestWsEventDispatch: router = MagicMock() router.send_approval = AsyncMock() router.send_plan_feedback = AsyncMock() + router.evaluate_tool_policies = AsyncMock(return_value=PolicyVerdict(kind="none")) + router.resolve_user = AsyncMock(return_value="turnstone-user-1") client = AsyncMock() client.chat_postMessage = AsyncMock(return_value={"ok": True, "ts": "123"}) client.chat_update = AsyncMock(return_value={"ok": True}) @@ -633,6 +736,7 @@ class TestWsEventDispatch: assert "Something went wrong" in text def test_approve_request_auto_approve(self) -> None: + from turnstone.channels._routing import PolicyVerdict from turnstone.channels.slack.bot import TurnstoneSlackBot from turnstone.channels.slack.config import SlackConfig from turnstone.channels.slack.routes import SlackRoute @@ -642,6 +746,7 @@ class TestWsEventDispatch: storage = MagicMock() router = MagicMock() router.send_approval = AsyncMock() + router.evaluate_tool_policies = AsyncMock(return_value=PolicyVerdict(kind="none")) client = AsyncMock() client.chat_postMessage = AsyncMock(return_value={"ok": True, "ts": "123"}) @@ -749,8 +854,11 @@ class TestWsEventDispatch: def test_plan_approve_sends_feedback_and_updates_message(self) -> None: bot, client = self._make_ws_bot() + # Register pending review with an owner so the new sec-2 gate passes. + bot._pending_plan_review_ts["ws-1"] = ("C1", "111.222", "U_OWNER") # type: ignore[attr-defined] body = { "actions": [{"value": "ws-1"}], + "user": {"id": "U_OWNER"}, "container": {"channel_id": "C1", "message_ts": "111.222"}, } @@ -759,9 +867,23 @@ class TestWsEventDispatch: bot.router.send_plan_feedback.assert_awaited_once_with("ws-1", "", "") # type: ignore[attr-defined] client.chat_update.assert_awaited_once() + def test_plan_approve_rejects_non_owner(self) -> None: + bot, client = self._make_ws_bot() + bot._pending_plan_review_ts["ws-1"] = ("C1", "111.222", "U_OWNER") # type: ignore[attr-defined] + body = { + "actions": [{"value": "ws-1"}], + "user": {"id": "U_OTHER"}, + "container": {"channel_id": "C1", "message_ts": "111.222"}, + } + + _run(bot._on_plan_approve(AsyncMock(), body)) # type: ignore[attr-defined] + + bot.router.send_plan_feedback.assert_not_awaited() # type: ignore[attr-defined] + client.chat_postEphemeral.assert_awaited_once() + def test_plan_feedback_modal_sends_feedback_and_updates_message(self) -> None: bot, client = self._make_ws_bot() - bot._pending_plan_review_ts["ws-1"] = ("C1", "111.222") # type: ignore[attr-defined] + bot._pending_plan_review_ts["ws-1"] = ("C1", "111.222", "U_OWNER") # type: ignore[attr-defined] view = { "private_metadata": "ws-1", @@ -769,8 +891,9 @@ class TestWsEventDispatch: "values": {"feedback_block": {"feedback_input": {"value": "please revise step 2"}}} }, } + body = {"user": {"id": "U_OWNER"}} - _run(bot._on_plan_feedback_modal(AsyncMock(), {}, view)) # type: ignore[attr-defined] + _run(bot._on_plan_feedback_modal(AsyncMock(), body, view)) # type: ignore[attr-defined] bot.router.send_plan_feedback.assert_awaited_once_with( # type: ignore[attr-defined] "ws-1", @@ -779,6 +902,49 @@ class TestWsEventDispatch: ) client.chat_update.assert_awaited_once() + def test_link_prefix_does_not_hijack_regular_prompt(self) -> None: + """`/turnstone linking up the docs` must not misroute into + _handle_link with `"ing up the docs"` as the token.""" + bot, router, client = _make_bot() + bot._handle_link = AsyncMock() # type: ignore[attr-defined] + # Force the linked-user gate to pass so the natural-language + # prompt can flow through to the session-start branch. + router.get_or_create_workstream = AsyncMock(return_value=("ws-new", True)) + body = { + "channel_id": "C01SAPU5414", + "user_id": "U111", + "text": "linking up the docs", + } + _run(bot._on_slash_command(AsyncMock(), body)) # type: ignore[attr-defined] + bot._handle_link.assert_not_awaited() # type: ignore[attr-defined] + + def test_link_rate_limit_blocks_after_cap(self) -> None: + """Sec-3: /turnstone link must throttle at _LINK_RATE_LIMIT/hour.""" + from turnstone.channels.slack.bot import _LINK_RATE_LIMIT + + bot, _router, client = _make_bot() + # Make the user already linked so _handle_link skips past the + # rate limit check would otherwise take a slot on a successful + # storage hit; we still want to exercise the throttle directly. + for _ in range(_LINK_RATE_LIMIT): + assert bot._allow_link_attempt("U111") # type: ignore[attr-defined] + # Next attempt is blocked. + assert not bot._allow_link_attempt("U111") # type: ignore[attr-defined] + + def test_plan_feedback_modal_rejects_non_owner(self) -> None: + bot, _client = self._make_ws_bot() + bot._pending_plan_review_ts["ws-1"] = ("C1", "111.222", "U_OWNER") # type: ignore[attr-defined] + + view = { + "private_metadata": "ws-1", + "state": {"values": {"feedback_block": {"feedback_input": {"value": "please revise"}}}}, + } + body = {"user": {"id": "U_OTHER"}} + + _run(bot._on_plan_feedback_modal(AsyncMock(), body, view)) # type: ignore[attr-defined] + + bot.router.send_plan_feedback.assert_not_awaited() # type: ignore[attr-defined] + # --------------------------------------------------------------------------- # Notification tracking @@ -969,6 +1135,9 @@ class TestChannelCLI: async def start(self) -> None: return None + async def stop(self) -> None: + return None + class FakeServer: def __init__(self, _config) -> None: pass @@ -980,7 +1149,7 @@ class TestChannelCLI: created_adapters.update(adapters) return MagicMock() - async def _fake_gather(*aws): # type: ignore[no-untyped-def] + async def _fake_gather(*aws, return_exceptions=False): # type: ignore[no-untyped-def] for aw in aws: await aw return [] @@ -1035,6 +1204,9 @@ class TestChannelCLI: async def start(self) -> None: return None + async def stop(self) -> None: + return None + class FakeDiscordBot: channel_type = "discord" @@ -1044,6 +1216,9 @@ class TestChannelCLI: async def start(self) -> None: return None + async def stop(self) -> None: + return None + class FakeServer: def __init__(self, _config) -> None: pass @@ -1055,7 +1230,7 @@ class TestChannelCLI: created_adapters.update(adapters) return MagicMock() - async def _fake_gather(*aws): # type: ignore[no-untyped-def] + async def _fake_gather(*aws, return_exceptions=False): # type: ignore[no-untyped-def] for aw in aws: await aw return [] diff --git a/tests/test_channel_sse.py b/tests/test_channel_sse.py new file mode 100644 index 00000000..13bb62a5 --- /dev/null +++ b/tests/test_channel_sse.py @@ -0,0 +1,397 @@ +"""Tests for the shared SSE reconnect helper in turnstone.channels._sse.""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + + +def _run(coro): # type: ignore[no-untyped-def] + return asyncio.run(coro) + + +class _FakeSSEEvent: + """A fake ``httpx_sse.ServerSentEvent`` with the subset we read.""" + + def __init__(self, event: str, data: str) -> None: + self.event = event + self.data = data + + +class _FakeEventSource: + """Context manager returned by our fake ``aconnect_sse``. + + Captures the (status_code, events) the test wants to deliver. + ``aiter_sse`` yields the events then returns; the caller then hits + the outer ``while True`` loop again, which will pick up the next + queued response via the shared iterator state on _FakeConnect. + """ + + def __init__(self, *, status_code: int, events: list[_FakeSSEEvent]) -> None: + self.response = SimpleNamespace( + status_code=status_code, + request=MagicMock(), + ) + self._events = events + + async def __aenter__(self) -> _FakeEventSource: + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: # noqa: ANN001 + return None + + async def aiter_sse(self): # type: ignore[no-untyped-def] + for event in self._events: + yield event + + +class _FakeConnect: + """Drop-in replacement for ``httpx_sse.aconnect_sse``. + + On each call, pops the next ``_FakeEventSource`` from *queue*. When + the queue is empty, raises ``asyncio.CancelledError`` so the loop + terminates cleanly in tests. + """ + + def __init__(self, queue: list[_FakeEventSource]) -> None: + self._queue = queue + self.call_count = 0 + + def __call__(self, *args, **kwargs): # noqa: ANN001, ANN204 + self.call_count += 1 + if not self._queue: + raise asyncio.CancelledError + return self._queue.pop(0) + + +@pytest.fixture +def _fast_sleep(monkeypatch): + """Patch asyncio.sleep so backoff doesn't actually wait; record calls.""" + sleeps: list[float] = [] + + async def fake_sleep(delay: float) -> None: + sleeps.append(delay) + + monkeypatch.setattr("turnstone.channels._sse.asyncio.sleep", fake_sleep) + return sleeps + + +def _valid_event_data(ws_id: str = "ws-1") -> str: + """A payload ``ServerEvent.from_dict`` will accept (a ContentEvent).""" + return json.dumps( + { + "type": "content", + "ws_id": ws_id, + "text": "hello", + } + ) + + +# --------------------------------------------------------------------------- +# 404 → on_stale + exit +# --------------------------------------------------------------------------- + + +class TestStaleRoute: + def test_404_calls_on_stale_and_returns(self, monkeypatch, _fast_sleep): + from turnstone.channels import _sse + + queue = [_FakeEventSource(status_code=404, events=[])] + fake_connect = _FakeConnect(queue) + monkeypatch.setattr(_sse.httpx_sse, "aconnect_sse", fake_connect) + + on_stale = AsyncMock() + on_event = AsyncMock() + + async def node_url_fn(ws_id: str) -> str: + return "http://node" + + _run( + _sse.run_sse_stream( + http_client=MagicMock(), + log_prefix="test", + ws_id="ws-1", + node_url_fn=node_url_fn, + token_factory=None, + on_event=on_event, + on_stale=on_stale, + ) + ) + + on_stale.assert_awaited_once() + on_event.assert_not_awaited() + # No reconnect after 404. + assert fake_connect.call_count == 1 + assert _fast_sleep == [] + + def test_on_stale_exception_still_exits(self, monkeypatch, _fast_sleep): + """If on_stale raises, the loop must not reconnect.""" + from turnstone.channels import _sse + + queue = [_FakeEventSource(status_code=404, events=[])] + fake_connect = _FakeConnect(queue) + monkeypatch.setattr(_sse.httpx_sse, "aconnect_sse", fake_connect) + + on_stale = AsyncMock(side_effect=RuntimeError("storage down")) + + async def node_url_fn(ws_id: str) -> str: + return "http://node" + + _run( + _sse.run_sse_stream( + http_client=MagicMock(), + log_prefix="test", + ws_id="ws-1", + node_url_fn=node_url_fn, + token_factory=None, + on_event=AsyncMock(), + on_stale=on_stale, + ) + ) + + on_stale.assert_awaited_once() + # Still a single connect — no livelock. + assert fake_connect.call_count == 1 + + +# --------------------------------------------------------------------------- +# 500+ → exponential backoff +# --------------------------------------------------------------------------- + + +class TestBackoff: + def test_500_triggers_backoff_and_retries(self, monkeypatch, _fast_sleep): + from turnstone.channels import _sse + + queue = [ + _FakeEventSource(status_code=503, events=[]), + _FakeEventSource(status_code=503, events=[]), + _FakeEventSource(status_code=503, events=[]), + ] + fake_connect = _FakeConnect(queue) + monkeypatch.setattr(_sse.httpx_sse, "aconnect_sse", fake_connect) + + async def node_url_fn(ws_id: str) -> str: + return "http://node" + + with contextlib.suppress(asyncio.CancelledError): + _run( + _sse.run_sse_stream( + http_client=MagicMock(), + log_prefix="test", + ws_id="ws-1", + node_url_fn=node_url_fn, + token_factory=None, + on_event=AsyncMock(), + on_stale=AsyncMock(), + ) + ) + + assert fake_connect.call_count >= 3 + # First three recorded sleeps are 2s, 4s, 8s (starts at + # SSE_RECONNECT_DELAY, doubles each time, capped at + # SSE_MAX_RECONNECT_DELAY). + assert _fast_sleep[0] == _sse.SSE_RECONNECT_DELAY + assert _fast_sleep[1] == _sse.SSE_RECONNECT_DELAY * 2 + assert _fast_sleep[2] == _sse.SSE_RECONNECT_DELAY * 4 + + def test_backoff_resets_after_successful_dispatch(self, monkeypatch, _fast_sleep): + """After a 200 + successful event dispatch, the next error + restarts backoff at the initial delay.""" + from turnstone.channels import _sse + + good_event = _FakeSSEEvent(event="message", data=_valid_event_data()) + queue = [ + _FakeEventSource(status_code=503, events=[]), + _FakeEventSource(status_code=200, events=[good_event]), + _FakeEventSource(status_code=503, events=[]), + ] + fake_connect = _FakeConnect(queue) + monkeypatch.setattr(_sse.httpx_sse, "aconnect_sse", fake_connect) + + on_event = AsyncMock() + + async def node_url_fn(ws_id: str) -> str: + return "http://node" + + with contextlib.suppress(asyncio.CancelledError): + _run( + _sse.run_sse_stream( + http_client=MagicMock(), + log_prefix="test", + ws_id="ws-1", + node_url_fn=node_url_fn, + token_factory=None, + on_event=on_event, + on_stale=AsyncMock(), + ) + ) + + on_event.assert_awaited() + # Sleep sequence: 2 (after first 503), 2 (reset after 200/event), + # then CancelledError exits. First two sleeps are both the base + # delay — the reset did its job. + assert len(_fast_sleep) >= 2 + assert _fast_sleep[0] == _sse.SSE_RECONNECT_DELAY + assert _fast_sleep[1] == _sse.SSE_RECONNECT_DELAY + + +# --------------------------------------------------------------------------- +# Event dispatch +# --------------------------------------------------------------------------- + + +class TestEventDispatch: + def test_invalid_json_is_skipped(self, monkeypatch, _fast_sleep): + from turnstone.channels import _sse + + bad = _FakeSSEEvent(event="message", data="{not json") + good = _FakeSSEEvent(event="message", data=_valid_event_data()) + queue = [_FakeEventSource(status_code=200, events=[bad, good])] + fake_connect = _FakeConnect(queue) + monkeypatch.setattr(_sse.httpx_sse, "aconnect_sse", fake_connect) + + on_event = AsyncMock() + + async def node_url_fn(ws_id: str) -> str: + return "http://node" + + with contextlib.suppress(asyncio.CancelledError): + _run( + _sse.run_sse_stream( + http_client=MagicMock(), + log_prefix="test", + ws_id="ws-1", + node_url_fn=node_url_fn, + token_factory=None, + on_event=on_event, + on_stale=AsyncMock(), + ) + ) + + # Good event delivered, bad one silently dropped. + assert on_event.await_count == 1 + + def test_on_event_exception_does_not_kill_stream(self, monkeypatch, _fast_sleep): + from turnstone.channels import _sse + + e1 = _FakeSSEEvent(event="message", data=_valid_event_data()) + e2 = _FakeSSEEvent(event="message", data=_valid_event_data()) + queue = [_FakeEventSource(status_code=200, events=[e1, e2])] + fake_connect = _FakeConnect(queue) + monkeypatch.setattr(_sse.httpx_sse, "aconnect_sse", fake_connect) + + on_event = AsyncMock(side_effect=[RuntimeError("boom"), None]) + + async def node_url_fn(ws_id: str) -> str: + return "http://node" + + with contextlib.suppress(asyncio.CancelledError): + _run( + _sse.run_sse_stream( + http_client=MagicMock(), + log_prefix="test", + ws_id="ws-1", + node_url_fn=node_url_fn, + token_factory=None, + on_event=on_event, + on_stale=AsyncMock(), + ) + ) + + # Both events attempted — first raised but second still delivered. + assert on_event.await_count == 2 + + +# --------------------------------------------------------------------------- +# Token factory +# --------------------------------------------------------------------------- + + +class TestTokenFactory: + def test_header_refreshed_per_connection(self, monkeypatch, _fast_sleep): + """token_factory is called once per reconnect so rotating service + JWTs stay fresh.""" + from turnstone.channels import _sse + + # Two reconnects followed by CancelledError to exit. + queue = [ + _FakeEventSource(status_code=503, events=[]), + _FakeEventSource(status_code=503, events=[]), + ] + fake_connect = _FakeConnect(queue) + monkeypatch.setattr(_sse.httpx_sse, "aconnect_sse", fake_connect) + + tokens: list[str] = [] + + def factory() -> str: + tok = f"tok-{len(tokens)}" + tokens.append(tok) + return tok + + async def node_url_fn(ws_id: str) -> str: + return "http://node" + + with contextlib.suppress(asyncio.CancelledError): + _run( + _sse.run_sse_stream( + http_client=MagicMock(), + log_prefix="test", + ws_id="ws-1", + node_url_fn=node_url_fn, + token_factory=factory, + on_event=AsyncMock(), + on_stale=AsyncMock(), + ) + ) + + assert len(tokens) >= 2 + assert tokens[0] != tokens[1] + + +# --------------------------------------------------------------------------- +# httpx errors +# --------------------------------------------------------------------------- + + +class TestTransportErrors: + def test_connect_error_falls_through_to_backoff(self, monkeypatch, _fast_sleep): + """ConnectError is caught and treated as retryable.""" + from turnstone.channels import _sse + + call_order = {"n": 0} + + def fake_connect(*args, **kwargs): # noqa: ANN001, ANN003 + call_order["n"] += 1 + if call_order["n"] == 1: + raise httpx.ConnectError("boom") + # Second attempt: signal the loop to exit. + raise asyncio.CancelledError + + monkeypatch.setattr(_sse.httpx_sse, "aconnect_sse", fake_connect) + + async def node_url_fn(ws_id: str) -> str: + return "http://node" + + with contextlib.suppress(asyncio.CancelledError): + _run( + _sse.run_sse_stream( + http_client=MagicMock(), + log_prefix="test", + ws_id="ws-1", + node_url_fn=node_url_fn, + token_factory=None, + on_event=AsyncMock(), + on_stale=AsyncMock(), + ) + ) + + assert call_order["n"] == 2 + # Backoff ran once after the ConnectError. + assert _fast_sleep == [_sse.SSE_RECONNECT_DELAY] diff --git a/turnstone/channels/__init__.py b/turnstone/channels/__init__.py index d27aef13..2db9ee8c 100644 --- a/turnstone/channels/__init__.py +++ b/turnstone/channels/__init__.py @@ -1,15 +1,13 @@ """Shared channel infrastructure for turnstone communication integrations. -Provides the :class:`ChannelAdapter` protocol, the :class:`ChannelEvent` -normalized event type, the :class:`ChannelRouter` for workstream mapping, -and shared formatting / configuration utilities. +Provides the :class:`ChannelAdapter` protocol, the :class:`ChannelRouter` +for workstream mapping, and shared formatting / configuration utilities. """ -from turnstone.channels._protocol import ChannelAdapter, ChannelEvent +from turnstone.channels._protocol import ChannelAdapter from turnstone.channels._routing import ChannelRouter __all__ = [ "ChannelAdapter", - "ChannelEvent", "ChannelRouter", ] diff --git a/turnstone/channels/_config.py b/turnstone/channels/_config.py index ff3f0884..7db88fb1 100644 --- a/turnstone/channels/_config.py +++ b/turnstone/channels/_config.py @@ -4,6 +4,12 @@ from __future__ import annotations from dataclasses import dataclass, field +# Shared adapter constants. +SSE_RECONNECT_DELAY: float = 2.0 +SSE_MAX_RECONNECT_DELAY: float = 30.0 +MAX_NOTIFY_TRACKING: int = 100 +CREATE_LOCK_CAP: int = 1024 # LRU bound on ChannelRouter per-channel creation locks + @dataclass class ChannelConfig: diff --git a/turnstone/channels/_formatter.py b/turnstone/channels/_formatter.py index 999214cb..8511e8d6 100644 --- a/turnstone/channels/_formatter.py +++ b/turnstone/channels/_formatter.py @@ -25,7 +25,26 @@ def chunk_message(text: str, max_length: int = 2000) -> list[str]: if len(text) <= max_length: return [text] - chunks: list[str] = [] + # Fast path: plain text with no code fences. Skips per-iteration + # fence bookkeeping for the common streaming-response case. + if "```" not in text: + chunks: list[str] = [] + remaining = text + while remaining: + if len(remaining) <= max_length: + chunks.append(remaining) + break + candidate = remaining[:max_length] + split_idx = candidate.rfind("\n") + if split_idx <= 0: + split_idx = candidate.rfind(" ") + if split_idx <= 0: + split_idx = max_length + chunks.append(remaining[:split_idx]) + remaining = remaining[split_idx:].lstrip("\n") + return chunks + + chunks = [] remaining = text in_code_block = False @@ -94,8 +113,6 @@ def format_approval_request(items: list[dict[str, Any]]) -> str: if not preview: args = item.get("function", {}).get("arguments", "") if isinstance(args, dict): - import json - args = json.dumps(args, ensure_ascii=False) preview = str(args) preview = truncate(preview) @@ -139,11 +156,6 @@ def format_verdict(verdict: dict[str, Any]) -> str: return "\n".join(parts) -def format_plan_review(content: str) -> str: - """Format a plan-review prompt with a header.""" - return f"**Plan review requested:**\n\n{content}" - - def format_tool_result(output: str) -> str: """Format a tool result into a compact code-block summary. @@ -202,15 +214,38 @@ def try_parse_media(output: str) -> dict[str, Any] | None: _BLOCKED_HOSTNAMES = frozenset({"localhost", "metadata.google.internal"}) +# Cloud-metadata deny-list applied *before* the `is_private` allowance so +# ULA-hosted vendor metadata endpoints don't slip through the "private IPs +# are fine, we trust the LAN" exception. IPv4 169.254.169.254 is caught +# by `is_link_local`; IPv6 ULA metadata (AWS Nitro IMDS at fd00:ec2::254, +# ECS task metadata at fd00:ec2::23) is `is_private` and needs explicit +# blocking. Add new vendor prefixes here as they're published. +_BLOCKED_IP_NETWORKS: tuple[str, ...] = ( + "fd00:ec2::/32", # AWS Nitro IMDS / ECS task metadata over IPv6 +) -def _is_safe_image_url(url: str) -> bool: + +async def _is_safe_image_url(url: str) -> bool: """Validate that *url* uses http(s), has no embedded credentials, and does - not target loopback or cloud metadata endpoints. + not target loopback, link-local (incl. cloud metadata 169.254.169.254), + or reserved ranges — even after DNS resolution. - Private/LAN IPs are intentionally allowed (media servers are typically - on the local network). + Resolves the hostname and checks every returned address so a DNS + rebinding attack cannot swap a safe-looking public IP for an + internal one between validation and fetch. Private/LAN IPs are + still allowed (media servers typically live on the local network), + so only loopback + link-local + multicast + reserved are rejected. + + NOTE: there is a residual TOCTOU gap because httpx resolves the + hostname again when it actually issues the GET. A 0-TTL rebinding + resolver could still slip an internal IP in between validation and + fetch. Fully closing the gap requires pinning the validated IP on + the connection (a custom httpx transport) — out of scope for this + backfill pass. """ + import asyncio import ipaddress + import socket from urllib.parse import urlparse try: @@ -226,12 +261,41 @@ def _is_safe_image_url(url: str) -> bool: return False if hostname in _BLOCKED_HOSTNAMES: return False + + # Collect candidate IPs: either an IP literal in the URL, or every + # A/AAAA record the resolver returns for a hostname. + candidates: list[str] = [] try: - ip = ipaddress.ip_address(hostname) - if ip.is_loopback or ip.is_link_local: - return False + ipaddress.ip_address(hostname) + candidates.append(hostname) except ValueError: - pass # Not an IP literal — hostname is fine + try: + infos = await asyncio.to_thread(socket.getaddrinfo, hostname, None, socket.AF_UNSPEC) + except socket.gaierror: + return False + # Strip IPv6 zone IDs (e.g. ``fe80::1%eth0``) before parsing — + # ipaddress.ip_address would raise on them and we'd drop the host + # on unrelated metadata. + candidates = [str(info[4][0]).partition("%")[0] for info in infos] + if not candidates: + return False + + blocked_networks = [ipaddress.ip_network(cidr) for cidr in _BLOCKED_IP_NETWORKS] + for raw in candidates: + try: + ip = ipaddress.ip_address(raw) + except ValueError: + return False + if ( + ip.is_loopback + or ip.is_link_local + or ip.is_multicast + or ip.is_reserved + or ip.is_unspecified + ): + return False + if any(ip in net for net in blocked_networks): + return False return True @@ -249,7 +313,7 @@ async def _fetch_thumbnail( are typically on the local network), but scheme is restricted to http(s) and userinfo is rejected. """ - if not _is_safe_image_url(url): + if not await _is_safe_image_url(url): return None try: async with http.stream("GET", url, timeout=timeout) as resp: diff --git a/turnstone/channels/_http.py b/turnstone/channels/_http.py index a0eaf47c..3412d34d 100644 --- a/turnstone/channels/_http.py +++ b/turnstone/channels/_http.py @@ -59,6 +59,16 @@ def _check_auth(request: Request) -> JSONResponse | None: result = validate_jwt(token, jwt_secret, audience=JWT_AUD_CHANNEL) if result is not None: + # Scope check: a valid ``turnstone-channel``-audience token is + # not sufficient on its own — require ``write`` so a low-scope + # service token can't drive notification delivery. + if "write" not in result.scopes: + log.warning( + "notify.auth_insufficient_scope", + user_id=result.user_id, + scopes=sorted(result.scopes), + ) + return JSONResponse({"error": "insufficient scope"}, status_code=403) return None return JSONResponse({"error": "Unauthorized"}, status_code=401) diff --git a/turnstone/channels/_protocol.py b/turnstone/channels/_protocol.py index 9d8a6361..c15df593 100644 --- a/turnstone/channels/_protocol.py +++ b/turnstone/channels/_protocol.py @@ -1,26 +1,12 @@ -"""Channel adapter protocol and normalized event type. +"""Channel adapter protocol. -Defines the :class:`ChannelEvent` data class for inbound events and the -:class:`ChannelAdapter` structural protocol that all bidirectional channel -adapters (Discord, Slack, etc.) must satisfy. +Defines the :class:`ChannelAdapter` structural protocol that bidirectional +channel adapters (Discord, Slack, etc.) must satisfy. """ from __future__ import annotations -from dataclasses import dataclass, field -from typing import Any, Protocol, runtime_checkable - - -@dataclass -class ChannelEvent: - """Normalized inbound event from any channel.""" - - channel_type: str # "discord", "slack" - channel_id: str # thread/channel ID - channel_user_id: str # platform user ID - message: str - parent_channel_id: str = "" # main channel (for thread creation) - metadata: dict[str, Any] = field(default_factory=dict) +from typing import Protocol, runtime_checkable @runtime_checkable @@ -48,31 +34,3 @@ class ChannelAdapter(Protocol): 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.""" - ... - - async def send_approval_request( - self, - channel_id: str, - ws_id: str, - correlation_id: str, - items: list[dict[str, Any]], - ) -> None: - """Send an interactive tool-approval prompt to a channel.""" - ... - - async def send_plan_review( - self, - channel_id: str, - ws_id: str, - correlation_id: str, - content: str, - ) -> None: - """Send a plan-review prompt to a channel.""" - ... - - async def create_thread(self, parent_channel_id: str, name: str, message_id: str = "") -> str: - """Create a thread under a parent channel. Returns the new thread ID.""" - ... diff --git a/turnstone/channels/_routing.py b/turnstone/channels/_routing.py index 976c2454..622162a0 100644 --- a/turnstone/channels/_routing.py +++ b/turnstone/channels/_routing.py @@ -9,13 +9,39 @@ from __future__ import annotations import asyncio import time -from typing import TYPE_CHECKING, Any +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal +from turnstone.channels._config import CREATE_LOCK_CAP from turnstone.core.log import get_logger from turnstone.sdk._types import TurnstoneAPIError from turnstone.sdk.console import AsyncTurnstoneConsole from turnstone.sdk.server import AsyncTurnstoneServer + +@dataclass +class PolicyVerdict: + """Outcome of evaluating admin tool policies for an approval request. + + ``kind`` is one of: + + - ``"none"``: no tool needed approval evaluation (e.g. all items are + errors or already resolved). Adapter should fall through to the + auto-approve branch. + - ``"deny"``: at least one tool was denied by policy. Adapter should + notify the user and forward ``approved=False`` with the feedback. + - ``"allow"``: every tool was allowed by policy. Adapter should + notify the user and forward ``approved=True``. + - ``"defer"``: mixed or unknown verdict. Adapter should fall through + to interactive approval. + """ + + kind: Literal["none", "deny", "allow", "defer"] + denied_tools: list[str] = field(default_factory=list) + tool_names: list[str] = field(default_factory=list) + + if TYPE_CHECKING: from collections.abc import Callable @@ -26,6 +52,8 @@ log = get_logger(__name__) _WS_CREATE_TIMEOUT = 30.0 # seconds _CHANNEL_DEFAULT_TTL = 300.0 # cache channel default alias for 5 minutes _MODELS_CACHE_TTL = 30.0 # cache model list for autocomplete +_ROUTE_CACHE_TTL = 30.0 # cache (channel_type, channel_id) → ws_id lookups +_ROUTE_CACHE_CAP = 4096 # LRU bound on the lookup cache class ChannelRouter: @@ -62,7 +90,7 @@ class ChannelRouter: self._auto_approve = auto_approve self._auto_approve_tools: list[str] = auto_approve_tools or [] self._skill = skill - self._create_locks: dict[str, asyncio.Lock] = {} + self._create_locks: OrderedDict[str, asyncio.Lock] = OrderedDict() # Per-workstream node URLs from console routing responses. # Populated when console_url is set and the create response # includes node_url. @@ -92,6 +120,9 @@ class ChannelRouter: # Cached model list for autocomplete (shorter TTL). self._models_cache: dict[str, Any] = {} self._models_cache_ts: float = 0.0 + # TTL cache for (channel_type, channel_id) → ws_id so hot inbound + # paths don't hit storage on every message. Bounded LRU. + self._route_cache: OrderedDict[tuple[str, str], tuple[str, float]] = OrderedDict() # -- lifecycle ----------------------------------------------------------- @@ -137,11 +168,15 @@ class ChannelRouter: return self._channel_default_alias # Mark refresh window before awaiting so concurrent callers # reuse the cached value instead of triggering duplicate fetches. + prev_ts = self._channel_default_ts self._channel_default_ts = now try: data = await self.list_models() self._channel_default_alias = data.get("channel_default_alias", "") except Exception: + # Roll the timestamp back so the next caller retries instead of + # serving a stale/empty alias for the full TTL window. + self._channel_default_ts = prev_ts log.debug("channel_router.channel_default_fetch_failed", exc_info=True) return self._channel_default_alias @@ -184,9 +219,29 @@ class ChannelRouter: completes. """ key = f"{channel_type}:{channel_id}" - lock = self._create_locks.setdefault(key, asyncio.Lock()) - - old_ws_id: str | None = None + lock = self._create_locks.get(key) + if lock is None: + lock = asyncio.Lock() + self._create_locks[key] = lock + # Bound the map: once a route is persisted, the lock is no longer + # needed on future requests, so evicting the LRU entry is safe — + # UNLESS that entry is currently held by a task awaiting I/O + # inside the critical section. Evicting a held lock breaks + # mutual exclusion because a subsequent cache miss for the + # same key would create a fresh lock and run the create path + # concurrently (→ duplicate server-side workstreams). Scan + # from oldest to newest and pop the first unheld entry; if + # every entry is held we leave the map slightly over-cap + # rather than corrupt ordering. + if len(self._create_locks) > CREATE_LOCK_CAP: + for candidate_key, candidate_lock in list(self._create_locks.items()): + if candidate_key == key: + continue + if not candidate_lock.locked(): + del self._create_locks[candidate_key] + break + else: + self._create_locks.move_to_end(key) async with lock: # 1. Check for existing route. @@ -324,6 +379,47 @@ class ChannelRouter: await self._server.send(message, ws_id) log.debug("channel_router.send_message", ws_id=ws_id) + async def evaluate_tool_policies( + self, + items: list[dict[str, Any]], + ) -> PolicyVerdict: + """Evaluate admin tool policies for an ApproveRequestEvent batch. + + Returns a :class:`PolicyVerdict` summarising the outcome so each + adapter only has to translate the verdict into platform-specific + chat messages. + """ + tool_names = [ + it.get("approval_label", "") or it.get("func_name", "") + for it in items + if it.get("needs_approval") and it.get("func_name") and not it.get("error") + ] + tool_names = [n for n in tool_names if n] + if not tool_names: + return PolicyVerdict(kind="none") + + try: + from turnstone.core.policy import evaluate_tool_policies_batch + + verdicts = await asyncio.to_thread( + evaluate_tool_policies_batch, + self._storage, + tool_names, + ) + except Exception: + # Fail-open: freezing every workstream on a storage hiccup is worse + # than letting the approval fall through to interactive review. + # Log at WARNING so the policy-DB outage is still auditable. + log.warning("channel_router.policy_evaluation_failed", exc_info=True) + return PolicyVerdict(kind="defer", tool_names=tool_names) + + denied = [n for n, v in verdicts.items() if v == "deny"] + if denied: + return PolicyVerdict(kind="deny", denied_tools=denied, tool_names=tool_names) + if all(verdicts.get(n) == "allow" for n in tool_names): + return PolicyVerdict(kind="allow", tool_names=tool_names) + return PolicyVerdict(kind="defer", tool_names=tool_names) + async def send_approval( self, ws_id: str, @@ -364,8 +460,36 @@ class ChannelRouter: # -- route management ---------------------------------------------------- + async def lookup_ws_id(self, channel_type: str, channel_id: str) -> str | None: + """Return the ws_id bound to (channel_type, channel_id), or None. + + TTL-cached so the hot inbound-message path (thread replies, + DM replies) doesn't hit storage on every token. + """ + key = (channel_type, channel_id) + now = time.monotonic() + cached = self._route_cache.get(key) + if cached is not None: + ws_id, expires_at = cached + if now < expires_at: + self._route_cache.move_to_end(key) + return ws_id + # Expired — fall through to a fresh lookup. + del self._route_cache[key] + + route = await asyncio.to_thread(self._storage.get_channel_route, channel_type, channel_id) + if route is None: + return None + + ws_id = route["ws_id"] + self._route_cache[key] = (ws_id, now + _ROUTE_CACHE_TTL) + if len(self._route_cache) > _ROUTE_CACHE_CAP: + self._route_cache.popitem(last=False) + return ws_id + async def delete_route(self, channel_type: str, channel_id: str) -> None: """Remove a channel-to-workstream mapping.""" + self._route_cache.pop((channel_type, channel_id), None) deleted = await asyncio.to_thread( self._storage.delete_channel_route, channel_type, channel_id ) diff --git a/turnstone/channels/_sse.py b/turnstone/channels/_sse.py new file mode 100644 index 00000000..6a7440b8 --- /dev/null +++ b/turnstone/channels/_sse.py @@ -0,0 +1,158 @@ +"""Shared SSE listener loop for channel adapters. + +Both the Discord and Slack adapters subscribe to per-workstream SSE event +streams with identical reconnect / 404-stale-route / backoff behaviour. +:func:`run_sse_stream` extracts that loop so each adapter supplies only +its platform-specific ``on_event`` and ``on_stale`` callbacks. +""" + +from __future__ import annotations + +import asyncio +import json +from typing import TYPE_CHECKING + +import httpx +import httpx_sse + +from turnstone.channels._config import SSE_MAX_RECONNECT_DELAY, SSE_RECONNECT_DELAY +from turnstone.core.log import get_logger +from turnstone.sdk.events import ServerEvent + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + +log = get_logger(__name__) + + +async def run_sse_stream( + *, + http_client: httpx.AsyncClient, + log_prefix: str, + ws_id: str, + node_url_fn: Callable[[str], Awaitable[str]], + token_factory: Callable[[], str] | None, + on_event: Callable[[ServerEvent], Awaitable[None]], + on_stale: Callable[[], Awaitable[None]], +) -> None: + """Run an SSE subscription loop with reconnect/backoff for one workstream. + + Parameters + ---------- + http_client: + Shared ``httpx.AsyncClient`` for all SSE connections. + log_prefix: + Platform tag used in log events (e.g. ``"discord"`` / ``"slack"``). + ws_id: + Workstream identifier, passed as a query parameter. + node_url_fn: + Async callable returning the base server URL for *ws_id* on each + connection attempt (so reconnects pick up router cache refreshes). + token_factory: + Optional callable returning an ``Authorization: Bearer ...`` token + per connection (supports auto-rotating service JWTs). + on_event: + Async callback invoked once per parsed :class:`ServerEvent`. + Exceptions are logged and do not kill the stream. + on_stale: + Async callback invoked when the server returns 404 for *ws_id*, + indicating the workstream was evicted/closed. After ``on_stale`` + returns, the loop exits (does not reconnect). + """ + delay = SSE_RECONNECT_DELAY + url = "" + + while True: + try: + node_base = await node_url_fn(ws_id) + url = f"{node_base}/v1/api/events" + + sse_headers: dict[str, str] | None = None + if token_factory is not None: + sse_headers = {"Authorization": f"Bearer {token_factory()}"} + + async with httpx_sse.aconnect_sse( + http_client, + "GET", + url, + params={"ws_id": ws_id}, + headers=sse_headers, + ) as event_source: + status = event_source.response.status_code + if status == 404: + log.info(f"{log_prefix}.sse_ws_gone", ws_id=ws_id) + # The 404-stops-reconnect invariant belongs to this loop, + # not to the caller — if on_stale raises we still exit. + try: + await on_stale() + except Exception: + log.warning( + f"{log_prefix}.sse_on_stale_failed", + ws_id=ws_id, + exc_info=True, + ) + return + + if status >= 400: + log.warning( + f"{log_prefix}.sse_upstream_error", + ws_id=ws_id, + status=status, + ) + raise httpx.HTTPStatusError( + f"SSE upstream {status}", + request=event_source.response.request, + response=event_source.response, + ) + + delay = SSE_RECONNECT_DELAY # reset on successful connect + async for sse in event_source.aiter_sse(): + if sse.event != "message" and sse.event: + continue + try: + data = json.loads(sse.data) + except json.JSONDecodeError: + log.debug( + f"{log_prefix}.sse_invalid_json", + ws_id=ws_id, + data=sse.data[:200], + ) + continue + + event = ServerEvent.from_dict(data) + try: + await on_event(event) + except Exception: + log.warning( + f"{log_prefix}.event_dispatch_failed", + ws_id=ws_id, + exc_info=True, + ) + + except httpx.HTTPStatusError as exc: + # Already logged at WARNING inside the try block (the raise + # was our own — status was captured there). Caught here to + # fall through to backoff + retry. + log.debug( + f"{log_prefix}.sse_http_status_error", + ws_id=ws_id, + error=str(exc), + ) + except httpx.RemoteProtocolError: + log.debug(f"{log_prefix}.sse_remote_closed", ws_id=ws_id) + except asyncio.CancelledError: + return + except httpx.ReadTimeout: + log.info(f"{log_prefix}.sse_read_timeout", ws_id=ws_id) + except (httpx.ConnectError, httpx.ConnectTimeout) as exc: + log.warning( + f"{log_prefix}.sse_connect_failed", + ws_id=ws_id, + url=url, + error=str(exc), + ) + except Exception: + log.warning(f"{log_prefix}.sse_error", ws_id=ws_id, exc_info=True) + + await asyncio.sleep(delay) + delay = min(delay * 2, SSE_MAX_RECONNECT_DELAY) diff --git a/turnstone/channels/cli.py b/turnstone/channels/cli.py index e9587b02..77cf61a3 100644 --- a/turnstone/channels/cli.py +++ b/turnstone/channels/cli.py @@ -9,15 +9,35 @@ Run as: ``turnstone-channel --discord-token $TURNSTONE_DISCORD_TOKEN`` from __future__ import annotations +import argparse +import asyncio +import contextlib import os import socket import sys +import time +from typing import TYPE_CHECKING, cast + +from turnstone.core.log import add_log_args, configure_logging_from_args, get_logger + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from turnstone.channels._protocol import ChannelAdapter + from turnstone.core.storage import StorageBackend + + # uvicorn ASGIApp is unions of several protocols; use a loose alias here. + _ASGIApp = Callable[..., Awaitable[None]] + +log = get_logger(__name__) + +_DISCOVERY_BUDGET_S = 30.0 # cap total wall-clock wait on startup discovery +_DISCOVERY_INITIAL_DELAY_S = 1.0 # first retry delay +_DISCOVERY_MAX_DELAY_S = 8.0 # cap per-attempt sleep -def main() -> None: - """Parse arguments, initialize storage, and run adapters.""" - import argparse - +def _build_parser() -> argparse.ArgumentParser: + """Construct the CLI argument parser.""" parser = argparse.ArgumentParser( description="turnstone channel gateway — bridges messaging platforms to the turnstone cluster" ) @@ -107,136 +127,106 @@ def main() -> None: ) # -- Logging ------------------------------------------------------------- - from turnstone.core.log import add_log_args - add_log_args(parser) - args = parser.parse_args() + return parser - # -- Logging setup ------------------------------------------------------- - from turnstone.core.log import configure_logging_from_args - configure_logging_from_args(args, "channel") +def _build_token_factories( + jwt_secret: str, +) -> tuple[Callable[[], str] | None, Callable[[], str] | None]: + """Return ``(console_factory, server_factory)`` when a JWT secret is set.""" + if not jwt_secret: + return None, None - from turnstone.core.log import get_logger + from turnstone.core.auth import JWT_AUD_CONSOLE, JWT_AUD_SERVER, ServiceTokenManager - log = get_logger(__name__) - - # -- Storage ------------------------------------------------------------- - from turnstone.core.storage._registry import get_storage, init_storage - - db_backend = os.environ.get("TURNSTONE_DB_BACKEND", "sqlite") - db_url = os.environ.get("TURNSTONE_DB_URL", "") - db_path = os.environ.get("TURNSTONE_DB_PATH", "") - - init_storage( - backend=db_backend, - url=db_url, - path=db_path, + scopes = frozenset({"read", "write", "approve", "service"}) + console_mgr = ServiceTokenManager( + user_id="channel-gateway", + scopes=scopes, + source="channel", + secret=jwt_secret, + audience=JWT_AUD_CONSOLE, + expiry_hours=1, + ) + server_mgr = ServiceTokenManager( + user_id="channel-gateway", + scopes=scopes, + source="channel", + secret=jwt_secret, + audience=JWT_AUD_SERVER, + expiry_hours=1, ) - # -- Auth config --------------------------------------------------------- - jwt_secret = os.environ.get("TURNSTONE_JWT_SECRET", "").strip() + def console_factory() -> str: + return console_mgr.token - # Prefer auto-rotating service JWTs when jwt_secret is available. - # Two separate token factories: one for console (aud=turnstone-console) - # and one for server nodes (aud=turnstone-server, used for SSE). - _console_token_factory = None - _server_token_factory = None - if jwt_secret: - from turnstone.core.auth import JWT_AUD_CONSOLE, JWT_AUD_SERVER, ServiceTokenManager + def server_factory() -> str: + return server_mgr.token - _scopes = frozenset({"read", "write", "approve", "service"}) - _console_mgr = ServiceTokenManager( - user_id="channel-gateway", - scopes=_scopes, - source="channel", - secret=jwt_secret, - audience=JWT_AUD_CONSOLE, - expiry_hours=1, - ) - _server_mgr = ServiceTokenManager( - user_id="channel-gateway", - scopes=_scopes, - source="channel", - secret=jwt_secret, - audience=JWT_AUD_SERVER, - expiry_hours=1, - ) - _console_token_factory = lambda: _console_mgr.token # noqa: E731 - _server_token_factory = lambda: _server_mgr.token # noqa: E731 + return console_factory, server_factory - server_url: str = args.server_url - console_url: str = args.console_url - # Auto-discover console and server from services table. - # Retry until at least one is found — the console/servers may still be - # starting up. If the DB is down the cluster isn't functional anyway. - if not console_url or not server_url: - import time as _time +def _resolve_service_urls( + storage: StorageBackend, + console_url: str, + server_url: str, +) -> tuple[str, str]: + """Fill in missing console / server URLs from the service registry. - try: - from turnstone.core.storage._registry import get_storage as _get_st + Retries with exponential backoff up to ``_DISCOVERY_BUDGET_S`` seconds + since the console / servers may still be starting up. Returns + ``(console_url, server_url)``. + """ + if console_url and server_url: + return console_url, server_url - _disc_storage = _get_st() - log.info("channel.discovering_services") - for _attempt in range(30): # up to 30s - if not console_url: - consoles = _disc_storage.list_services("console", max_age_seconds=3600) - if consoles: - console_url = consoles[0]["url"] - log.info("channel.discovered_console", url=console_url) - if not server_url: - servers = _disc_storage.list_services("server", max_age_seconds=120) - if servers: - server_url = servers[0]["url"] - log.info("channel.discovered_server", url=server_url) - if console_url or server_url: - break - _time.sleep(1) - else: + try: + log.info("channel.discovering_services") + deadline = time.monotonic() + _DISCOVERY_BUDGET_S + delay = _DISCOVERY_INITIAL_DELAY_S + while True: + if not console_url: + consoles = storage.list_services("console", max_age_seconds=3600) + if consoles: + console_url = consoles[0]["url"] + log.info("channel.discovered_console", url=console_url) + if not server_url: + servers = storage.list_services("server", max_age_seconds=120) + if servers: + server_url = servers[0]["url"] + log.info("channel.discovered_server", url=server_url) + if console_url or server_url: + break + + remaining = deadline - time.monotonic() + if remaining <= 0: log.warning( "channel.discovery_timeout", console_url=console_url, server_url=server_url, ) - except Exception: - log.warning("channel.discovery_failed", exc_info=True) + break - if not console_url and not server_url: - print( - "Error: no console or server URL available. Set --server-url, " - "--console-url, or ensure the database is reachable and services " - "are registered.", - file=sys.stderr, - ) - sys.exit(1) + time.sleep(min(delay, remaining)) + delay = min(delay * 2, _DISCOVERY_MAX_DELAY_S) + except Exception: + log.warning("channel.discovery_failed", exc_info=True) - # -- Adapter selection --------------------------------------------------- - if not args.discord_token and not args.slack_token: - print( - "Error: no channel adapters configured. " - "Set --discord-token / $TURNSTONE_DISCORD_TOKEN " - "or --slack-token / $TURNSTONE_SLACK_TOKEN.", - file=sys.stderr, - ) - sys.exit(1) + return console_url, server_url - # Slack config validation (fail fast) - if bool(args.slack_token) != bool(args.slack_app_token): - raise SystemExit("--slack-token and --slack-app-token must be provided together") - # -- Run ----------------------------------------------------------------- - import asyncio - import contextlib - from typing import TYPE_CHECKING, cast - - from turnstone.channels._http import _get_service_id, create_channel_app - - if TYPE_CHECKING: - from turnstone.channels._protocol import ChannelAdapter - - storage = get_storage() +def _build_adapters( + args: argparse.Namespace, + storage: StorageBackend, + *, + server_url: str, + console_url: str, + console_token_factory: Callable[[], str] | None, + server_token_factory: Callable[[], str] | None, +) -> dict[str, ChannelAdapter]: + """Instantiate the channel adapters selected by the provided args.""" adapters: dict[str, ChannelAdapter] = {} if args.discord_token: @@ -262,8 +252,8 @@ def main() -> None: server_url, storage, console_url=console_url, - console_token_factory=_console_token_factory, - server_token_factory=_server_token_factory, + console_token_factory=console_token_factory, + server_token_factory=server_token_factory, ) adapters[discord_bot.channel_type] = cast("ChannelAdapter", discord_bot) @@ -284,16 +274,157 @@ def main() -> None: server_url=server_url, storage=storage, console_url=console_url, - console_token_factory=_console_token_factory, - server_token_factory=_server_token_factory, + console_token_factory=console_token_factory, + server_token_factory=server_token_factory, ) adapters[slack_bot.channel_type] = cast("ChannelAdapter", slack_bot) - channel_app = create_channel_app( - adapters, - storage, - jwt_secret=jwt_secret, + return adapters + + +def _resolve_advertise_url(args: argparse.Namespace) -> str: + """Compute the URL the gateway should advertise in the service registry.""" + override = os.environ.get("TURNSTONE_CHANNEL_ADVERTISE_URL", "").strip() + if override: + return override + + advertise_host = socket.gethostname() if args.http_host in ("0.0.0.0", "::") else args.http_host + scheme = "https" if args.ssl_certfile else "http" + return f"{scheme}://{advertise_host}:{args.http_port}" + + +async def _heartbeat_loop(storage: StorageBackend, service_id: str) -> None: + """Periodically update the channel service heartbeat.""" + from turnstone.core.storage._registry import StorageUnavailableError + + while True: + await asyncio.sleep(30) + try: + await asyncio.to_thread(storage.heartbeat_service, "channel", service_id) + except StorageUnavailableError: + pass # already logged by storage layer + except Exception: + log.exception("channel.heartbeat_failed") + + +async def _run_gateway( + adapters: dict[str, ChannelAdapter], + channel_app: _ASGIApp, + storage: StorageBackend, + args: argparse.Namespace, +) -> None: + """Run all adapters + HTTP server + service heartbeat concurrently.""" + import uvicorn + + from turnstone.channels._http import _get_service_id + + service_id = _get_service_id() + service_url = _resolve_advertise_url(args) + + storage.register_service("channel", service_id, service_url) + log.info("channel.service_registered", service_id=service_id, url=service_url) + + if bool(args.ssl_certfile) != bool(args.ssl_keyfile): + print( + "Both --ssl-certfile and --ssl-keyfile are required for TLS", + file=sys.stderr, + ) + sys.exit(1) + + uv_config = uvicorn.Config( + channel_app, + host=args.http_host, + port=args.http_port, + log_level="warning", + ssl_certfile=args.ssl_certfile, + ssl_keyfile=args.ssl_keyfile, + ssl_ca_certs=args.ssl_ca_certs, ) + server = uvicorn.Server(uv_config) + + heartbeat_task = asyncio.create_task(_heartbeat_loop(storage, service_id)) + try: + await asyncio.gather( + *(adapter.start() for adapter in adapters.values()), + server.serve(), + ) + finally: + heartbeat_task.cancel() + # Await the task so its CancelledError propagates before we tear + # down the adapters and the service registry below. CancelledError + # is the expected outcome after task.cancel(); suppress it. + with contextlib.suppress(asyncio.CancelledError): + await heartbeat_task + + # Stop adapters so SSE tasks, httpx clients, and Slack socket + # handlers close cleanly before we deregister from the service + # registry. + await asyncio.gather( + *(adapter.stop() for adapter in adapters.values()), + return_exceptions=True, + ) + + await asyncio.to_thread(storage.deregister_service, "channel", service_id) + log.info("channel.service_deregistered", service_id=service_id) + + +def main() -> None: + """Parse arguments, initialize storage, and run adapters.""" + from turnstone.channels._http import create_channel_app + from turnstone.core.storage._registry import get_storage, init_storage + + parser = _build_parser() + args = parser.parse_args() + + configure_logging_from_args(args, "channel") + + init_storage( + backend=os.environ.get("TURNSTONE_DB_BACKEND", "sqlite"), + url=os.environ.get("TURNSTONE_DB_URL", ""), + path=os.environ.get("TURNSTONE_DB_PATH", ""), + ) + storage = get_storage() + + jwt_secret = os.environ.get("TURNSTONE_JWT_SECRET", "").strip() + console_token_factory, server_token_factory = _build_token_factories(jwt_secret) + + console_url, server_url = _resolve_service_urls( + storage, + args.console_url, + args.server_url, + ) + + if not console_url and not server_url: + print( + "Error: no console or server URL available. Set --server-url, " + "--console-url, or ensure the database is reachable and services " + "are registered.", + file=sys.stderr, + ) + sys.exit(1) + + if not args.discord_token and not args.slack_token: + print( + "Error: no channel adapters configured. " + "Set --discord-token / $TURNSTONE_DISCORD_TOKEN " + "or --slack-token / $TURNSTONE_SLACK_TOKEN.", + file=sys.stderr, + ) + sys.exit(1) + + if bool(args.slack_token) != bool(args.slack_app_token): + raise SystemExit("--slack-token and --slack-app-token must be provided together") + + adapters = _build_adapters( + args, + storage, + server_url=server_url, + console_url=console_url, + console_token_factory=console_token_factory, + server_token_factory=server_token_factory, + ) + + channel_app = create_channel_app(adapters, storage, jwt_secret=jwt_secret) log.info( "channel.starting", @@ -302,83 +433,8 @@ def main() -> None: server_url=server_url, ) - async def _run_all() -> None: - """Run all adapters + HTTP server + service heartbeat concurrently.""" - import uvicorn - - service_id = _get_service_id() - - # Resolve advertise URL — env override for Docker/K8s, - # otherwise derive from bind address. - advertise_url = os.environ.get("TURNSTONE_CHANNEL_ADVERTISE_URL", "").strip() - if not advertise_url: - if args.http_host in ("0.0.0.0", "::"): - advertise_host = socket.gethostname() - else: - advertise_host = args.http_host - scheme = "https" if args.ssl_certfile else "http" - advertise_url = f"{scheme}://{advertise_host}:{args.http_port}" - service_url = advertise_url - - # Register in service registry - storage.register_service("channel", service_id, service_url) - log.info( - "channel.service_registered", - service_id=service_id, - url=service_url, - ) - - async def _heartbeat_loop() -> None: - """Periodically update service heartbeat.""" - from turnstone.core.storage._registry import StorageUnavailableError - - while True: - await asyncio.sleep(30) - try: - await asyncio.to_thread(storage.heartbeat_service, "channel", service_id) - except StorageUnavailableError: - pass # already logged by storage layer - except Exception: - log.exception("channel.heartbeat_failed") - - # TLS: use cert files if available (from bootstrap or TLSClient) - ssl_certfile = getattr(args, "ssl_certfile", None) - ssl_keyfile = getattr(args, "ssl_keyfile", None) - ssl_ca_certs = getattr(args, "ssl_ca_certs", None) - if bool(ssl_certfile) != bool(ssl_keyfile): - print( - "Both --ssl-certfile and --ssl-keyfile are required for TLS", - file=sys.stderr, - ) - sys.exit(1) - - uv_config = uvicorn.Config( - channel_app, - host=args.http_host, - port=args.http_port, - log_level="warning", - ssl_certfile=ssl_certfile, - ssl_keyfile=ssl_keyfile, - ssl_ca_certs=ssl_ca_certs, - ) - server = uvicorn.Server(uv_config) - - heartbeat_task = asyncio.create_task(_heartbeat_loop()) - try: - await asyncio.gather( - *(adapter.start() for adapter in adapters.values()), - server.serve(), - ) - finally: - heartbeat_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await heartbeat_task - - await asyncio.to_thread(storage.deregister_service, "channel", service_id) - log.info("channel.service_deregistered", service_id=service_id) - with contextlib.suppress(KeyboardInterrupt): - asyncio.run(_run_all()) + asyncio.run(_run_gateway(adapters, channel_app, storage, args)) if __name__ == "__main__": diff --git a/turnstone/channels/discord/bot.py b/turnstone/channels/discord/bot.py index c138c178..fb108f38 100644 --- a/turnstone/channels/discord/bot.py +++ b/turnstone/channels/discord/bot.py @@ -13,15 +13,17 @@ from __future__ import annotations import asyncio import contextlib -import json import time +from collections import OrderedDict from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any import httpx +from turnstone.channels._config import MAX_NOTIFY_TRACKING from turnstone.channels._formatter import chunk_message from turnstone.channels._routing import ChannelRouter +from turnstone.channels._sse import run_sse_stream from turnstone.core.log import get_logger from turnstone.sdk.events import ( ApprovalResolvedEvent, @@ -49,9 +51,26 @@ if TYPE_CHECKING: log = get_logger(__name__) -# SSE reconnection parameters -_SSE_RECONNECT_DELAY: float = 2.0 -_SSE_MAX_RECONNECT_DELAY: float = 30.0 + +_THREAD_INVOKER_CAP: int = 4096 + + +def _thread_owner_id(thread: discord.abc.Messageable) -> str: + """Return the Discord user ID who owns the thread / DM target. + + Used to gate approval / plan-review button clicks to the session + owner. For Discord threads this is the thread creator + (``thread.owner_id``). For DM channels we use ``recipient.id``. + Returns ``""`` when the owner cannot be determined — the views + then refuse the interaction. + """ + owner = getattr(thread, "owner_id", None) + if owner: + return str(owner) + recipient = getattr(thread, "recipient", None) + if recipient is not None and getattr(recipient, "id", None): + return str(recipient.id) + return "" # --------------------------------------------------------------------------- @@ -73,11 +92,39 @@ class StreamingMessage: edit_interval: float = 1.5 _message: discord.Message | None = field(default=None, init=False, repr=False) _buffer: list[str] = field(default_factory=list, init=False, repr=False) + # Rolling truncated in-progress display string; stops growing at + # max_length so per-flush cost is O(max_length) instead of + # O(total_streamed_chars). + _display: str = field(default="", init=False, repr=False) _last_edit: float = field(default=0.0, init=False, repr=False) + _finalized_text: str | None = field(default=None, init=False, repr=False) + + @property + def message(self) -> discord.Message | None: + """The underlying Discord message, once posted.""" + return self._message + + @message.setter + def message(self, value: discord.Message | None) -> None: + self._message = value + + @property + def accumulated_text(self) -> str: + """The joined text of everything appended so far. + + Cached after ``finalize()`` so the StreamEnd DM-forward path + (``discord/bot.py::_handle_stream_end``) doesn't re-join a + multi-MB buffer a second time. + """ + if self._finalized_text is not None: + return self._finalized_text + return "".join(self._buffer) async def append(self, text: str) -> None: """Add *text* to the buffer and edit the message if the interval has elapsed.""" self._buffer.append(text) + if len(self._display) < self.max_length: + self._display = (self._display + text)[: self.max_length] now = time.monotonic() if now - self._last_edit >= self.edit_interval: await self._flush() @@ -85,6 +132,7 @@ class StreamingMessage: async def finalize(self) -> None: """Flush any remaining buffered content, chunking if necessary.""" content = "".join(self._buffer) + self._finalized_text = content if not content: return @@ -104,14 +152,11 @@ class StreamingMessage: await self.channel.send(chunk) async def _flush(self) -> None: - """Edit or create the message with the current buffer contents.""" - content = "".join(self._buffer) - if not content: + """Edit or create the message with the current display slice.""" + display = self._display + if not display: return - # Truncate to max_length for the in-progress edit (finalize handles overflow). - display = content[: self.max_length] - try: if self._message is None: self._message = await self.channel.send(display) @@ -143,7 +188,7 @@ class TurnstoneBot: """ channel_type: str = "discord" - _MAX_NOTIFY_TRACKING: int = 100 + _MAX_NOTIFY_TRACKING: int = MAX_NOTIFY_TRACKING def __init__( self, @@ -204,6 +249,13 @@ class TurnstoneBot: # 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]] = {} + # Explicit thread-invoker map so the sec-3 owner check can admit + # `/ask` follow-ups (Discord sets `thread.owner_id = bot` when the + # thread is created via `channel.create_thread(...)` without a + # starter message, so `channel.owner_id` alone would reject every + # legitimate follow-up from the human who ran the slash command). + # Bounded LRU to prevent unbounded growth across long bot uptime. + self._thread_invokers: OrderedDict[int, int] = OrderedDict() # Shared HTTP client for SSE connections. # Read timeout detects half-open connections (server sends ping=5s @@ -348,13 +400,13 @@ class TurnstoneBot: self._subscribed_ws.add(ws_id) log.info("discord.subscribed", ws_id=ws_id) - async def unsubscribe_ws(self, ws_id: str) -> None: - """Cancel the SSE listener for *ws_id* and clean up streaming state.""" - task = self._sse_tasks.pop(ws_id, None) - if task is not None: - task.cancel() - with contextlib.suppress(asyncio.CancelledError, Exception): - await task + async def _clear_ws_state(self, ws_id: str) -> None: + """Drop all in-memory state keyed by *ws_id*. + + Does not cancel or await the SSE task — callers handle task + lifecycle differently (``unsubscribe_ws`` cancels and awaits; + ``_cleanup_stale_route`` is itself invoked from inside the task). + """ self._subscribed_ws.discard(ws_id) self._streaming.pop(ws_id, None) thinking_msg = self._thinking_msgs.pop(ws_id, None) @@ -368,129 +420,58 @@ class TurnstoneBot: 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] + + async def unsubscribe_ws(self, ws_id: str) -> None: + """Cancel the SSE listener for *ws_id* and clean up streaming state.""" + task = self._sse_tasks.pop(ws_id, None) + if task is not None: + task.cancel() + # Await the cancelled task so CancelledError propagates out of + # the SSE loop before we clear the per-ws state below. + # CancelledError is expected; other exceptions from the SSE + # loop are already logged there and must not block shutdown. + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + await self._clear_ws_state(ws_id) log.info("discord.unsubscribed", ws_id=ws_id) async def _cleanup_stale_route(self, ws_id: str) -> None: """Remove a channel route whose workstream no longer exists.""" route = await asyncio.to_thread(self.storage.get_channel_route_by_ws, ws_id) if route: - await asyncio.to_thread( - self.storage.delete_channel_route, route["channel_type"], route["channel_id"] - ) + # Route deletion must go through the router so the TTL cache + # of (channel_type, channel_id) → ws_id is also invalidated. + await self.router.delete_route(route["channel_type"], route["channel_id"]) log.info("discord.stale_route_removed", ws_id=ws_id) - # Clear same per-ws state that unsubscribe_ws() clears. - self._subscribed_ws.discard(ws_id) + # Called from inside the SSE task itself — don't await the task here. self._sse_tasks.pop(ws_id, None) - self._streaming.pop(ws_id, None) - thinking_msg = self._thinking_msgs.pop(ws_id, None) - if thinking_msg is not None: - with contextlib.suppress(Exception): - await thinking_msg.delete() - self._tool_info_msgs.pop(ws_id, None) - self._pending_approval_msgs.pop(ws_id, None) - self._notify_reply_channels.pop(ws_id, None) - 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] + await self._clear_ws_state(ws_id) # -- SSE listener -------------------------------------------------------- async def _sse_listener(self, ws_id: str, thread: discord.abc.Messageable) -> None: """SSE listener task for a workstream. - Connects to the server's per-workstream SSE endpoint, parses events - using :meth:`ServerEvent.from_dict`, and dispatches to - :meth:`_on_ws_event`. Reconnects with exponential backoff on - connection failures. + Delegates the reconnect/backoff loop to :func:`run_sse_stream`; + this wrapper just translates events to ``_on_ws_event`` calls and + handles the 404 stale-route case. """ - import httpx_sse - delay = _SSE_RECONNECT_DELAY - url = "" # set before loop so exception handlers can reference it + async def _on_event(event: ServerEvent) -> None: + await self._on_ws_event(ws_id, thread, event) - while True: - try: - # Re-resolve node URL on each attempt so reconnects pick up - # changes after bot restarts or router cache expiry. - node_base = await self.router.get_node_url(ws_id) - url = f"{node_base}/v1/api/events" - # Refresh auth header per-connection (token may have rotated) - sse_headers: dict[str, str] | None = None - if self._token_factory is not None: - sse_headers = {"Authorization": f"Bearer {self._token_factory()}"} - async with httpx_sse.aconnect_sse( - self._http_client, - "GET", - url, - params={"ws_id": ws_id}, - headers=sse_headers, - ) as event_source: - # Bail on non-retryable responses (404 = workstream gone). - status = event_source.response.status_code - if status == 404: - log.info("discord.sse_ws_gone", ws_id=ws_id) - await self._cleanup_stale_route(ws_id) - return - if status >= 400: - log.warning( - "discord.sse_upstream_error", - ws_id=ws_id, - status=status, - ) - # Don't try to parse a non-SSE error body — - # fall through to backoff/retry below. - raise httpx.HTTPStatusError( - f"SSE upstream {status}", - request=event_source.response.request, - response=event_source.response, - ) - delay = _SSE_RECONNECT_DELAY # reset on successful connect - async for sse in event_source.aiter_sse(): - if sse.event == "message" or not sse.event: - try: - data = json.loads(sse.data) - except json.JSONDecodeError: - log.debug( - "discord.sse_invalid_json", - ws_id=ws_id, - data=sse.data[:200], - ) - continue - event = ServerEvent.from_dict(data) - try: - await self._on_ws_event(ws_id, thread, event) - except Exception: - # Discord API failures (rate limits, outages) - # must not kill the SSE connection. - log.warning( - "discord.event_dispatch_failed", - ws_id=ws_id, - exc_info=True, - ) - except httpx.HTTPStatusError: - pass # already logged above; fall through to backoff - except httpx.RemoteProtocolError: - # Server closed connection (normal on stream_end or shutdown). - log.debug("discord.sse_remote_closed", ws_id=ws_id) - except asyncio.CancelledError: - return # unsubscribe or shutdown - except httpx.ReadTimeout: - # No data received within read timeout — likely a half-open - # connection. Reconnect to recover. - log.info("discord.sse_read_timeout", ws_id=ws_id) - except (httpx.ConnectError, httpx.ConnectTimeout) as exc: - log.warning( - "discord.sse_connect_failed", - ws_id=ws_id, - url=url, - error=str(exc), - ) - except Exception: - log.warning("discord.sse_error", ws_id=ws_id, exc_info=True) + async def _on_stale() -> None: + await self._cleanup_stale_route(ws_id) - # Exponential backoff before reconnecting. - await asyncio.sleep(delay) - delay = min(delay * 2, _SSE_MAX_RECONNECT_DELAY) + await run_sse_stream( + http_client=self._http_client, + log_prefix="discord", + ws_id=ws_id, + node_url_fn=self.router.get_node_url, + token_factory=self._token_factory, + on_event=_on_event, + on_stale=_on_stale, + ) # -- event dispatch ------------------------------------------------------ @@ -500,325 +481,357 @@ class TurnstoneBot: thread: discord.abc.Messageable, event: ServerEvent, ) -> None: - """Handle a typed server event for a subscribed workstream.""" - import discord - - from turnstone.channels._formatter import ( - format_approval_request, - format_plan_review, - format_verdict, - ) - from turnstone.channels.discord.views import ApprovalView, PlanReviewView - + """Dispatch a typed server event to its per-event handler.""" if isinstance(event, ThinkingStartEvent): - # Clean up any prior thinking message (consecutive starts without stop). - prev = self._thinking_msgs.pop(ws_id, None) - if prev is not None: - with contextlib.suppress(Exception): - await prev.delete() - try: - msg = await thread.send("*Thinking...*") - self._thinking_msgs[ws_id] = msg - except Exception: - log.debug("discord.thinking_start_send_failed", ws_id=ws_id) - + await self._handle_thinking_start(ws_id, thread) elif isinstance(event, ThinkingStopEvent): # Leave the thinking message in place — the next visible event # (ContentEvent, ToolInfoEvent, StreamEndEvent) will edit or # clean it up, avoiding a delete→gap→new-message flicker. pass - elif isinstance(event, ContentEvent): - # Reuse thinking message as the initial streaming message so the - # first flush edits it in-place (no delete→gap→send flicker). - thinking_msg = self._thinking_msgs.pop(ws_id, None) - sm = self._streaming.get(ws_id) - if sm is None: - sm = StreamingMessage( - channel=thread, - max_length=self.config.max_message_length, - edit_interval=self.config.streaming_edit_interval, - ) - if thinking_msg is not None: - sm._message = thinking_msg - self._streaming[ws_id] = sm - elif thinking_msg is not None: - with contextlib.suppress(Exception): - await thinking_msg.delete() - await sm.append(event.text) - + await self._handle_content(ws_id, thread, event) elif isinstance(event, ToolInfoEvent): - from turnstone.channels._formatter import truncate - - # Reuse the thinking message for the first tool embed. - thinking_msg = self._thinking_msgs.pop(ws_id, None) - - # Show a "running" embed for every tool. The approval dialog - # (ApproveRequestEvent) is a separate concern — it asks "do you - # authorize this?" while the running embed says "this tool is - # executing." Both can coexist in the thread. - for it in event.items: - raw_name = it.get("func_name") or it.get("approval_label") or "tool" - display_name = discord.utils.escape_markdown(raw_name) - raw_preview = it.get("preview", "") - # Escape backticks to prevent markdown breakout and - # strip @-mentions. - raw_preview = raw_preview.replace("`", "\\`") - raw_preview = discord.utils.escape_mentions(raw_preview) - preview = truncate(raw_preview, max_length=120) or None - embed = discord.Embed( - title=display_name, - description=preview, - color=discord.Color.light_grey(), - ) - # Edit thinking message into first tool embed to avoid flicker. - if thinking_msg is not None: - try: - await thinking_msg.edit(content=None, embed=embed) - msg = thinking_msg - except Exception: - msg = await thread.send(embed=embed) - thinking_msg = None - else: - msg = await thread.send(embed=embed) - call_id = it.get("call_id", "") - # Store raw (unescaped) name for matching against ToolResultEvent.name - self._tool_info_msgs.setdefault(ws_id, []).append( - (call_id, raw_name, preview or "", msg) - ) - - # If no items consumed the thinking message (empty event), clean up. - if thinking_msg is not None: - with contextlib.suppress(Exception): - await thinking_msg.delete() - + await self._handle_tool_info(ws_id, thread, event) elif isinstance(event, ToolResultEvent): - from turnstone.channels._formatter import format_tool_result + await self._handle_tool_result(ws_id, thread, event) + elif isinstance(event, ApproveRequestEvent): + await self._handle_approve_request(ws_id, thread, event) + elif isinstance(event, PlanReviewEvent): + await self._handle_plan_review(ws_id, thread, event) + elif isinstance(event, IntentVerdictEvent): + await self._handle_intent_verdict(ws_id, event) + elif isinstance(event, ApprovalResolvedEvent): + await self._handle_approval_resolved(ws_id, event) + elif isinstance(event, StreamEndEvent): + await self._handle_stream_end(ws_id) + elif isinstance(event, ErrorEvent): + safe_msg = event.message[:500] if event.message else "An error occurred" + await thread.send(f"**Error:** {safe_msg}") - # Mark the matching "running" embed as complete/errored. - # Prefer call_id match (deterministic); fall back to name (FIFO). - info_list = self._tool_info_msgs.get(ws_id, []) - matched_preview = "" - matched_msg: discord.Message | None = None - if event.call_id: - for i, (cid, _tname, _prev, _tmsg) in enumerate(info_list): - if cid == event.call_id: - entry = info_list.pop(i) - matched_preview, matched_msg = entry[2], entry[3] - break - if matched_msg is None: - for i, (_cid, tname, _prev, _tmsg) in enumerate(info_list): - if tname == event.name: - entry = info_list.pop(i) - matched_preview, matched_msg = entry[2], entry[3] - break - if matched_msg is not None: - status = "Error" if event.is_error else "Done" - status_color = discord.Color.red() if event.is_error else discord.Color.dark_grey() - status_embed = discord.Embed( - title=f"{discord.utils.escape_markdown(event.name)} \u2014 {status}", - description=matched_preview or None, - color=status_color, + async def _handle_thinking_start(self, ws_id: str, thread: discord.abc.Messageable) -> None: + # Clean up any prior thinking message (consecutive starts without stop). + prev = self._thinking_msgs.pop(ws_id, None) + if prev is not None: + with contextlib.suppress(Exception): + await prev.delete() + try: + msg = await thread.send("*Thinking...*") + self._thinking_msgs[ws_id] = msg + except Exception: + log.debug("discord.thinking_start_send_failed", ws_id=ws_id) + + async def _handle_content( + self, + ws_id: str, + thread: discord.abc.Messageable, + event: ContentEvent, + ) -> None: + # Reuse thinking message as the initial streaming message so the + # first flush edits it in-place (no delete→gap→send flicker). + thinking_msg = self._thinking_msgs.pop(ws_id, None) + sm = self._streaming.get(ws_id) + if sm is None: + sm = StreamingMessage( + channel=thread, + max_length=self.config.max_message_length, + edit_interval=self.config.streaming_edit_interval, + ) + if thinking_msg is not None: + sm.message = thinking_msg + self._streaming[ws_id] = sm + elif thinking_msg is not None: + with contextlib.suppress(Exception): + await thinking_msg.delete() + await sm.append(event.text) + + async def _handle_tool_info( + self, + ws_id: str, + thread: discord.abc.Messageable, + event: ToolInfoEvent, + ) -> None: + import discord + + from turnstone.channels._formatter import truncate + + # Reuse the thinking message for the first tool embed. + thinking_msg = self._thinking_msgs.pop(ws_id, None) + + # Show a "running" embed for every tool. The approval dialog + # (ApproveRequestEvent) is a separate concern — it asks "do you + # authorize this?" while the running embed says "this tool is + # executing." Both can coexist in the thread. + for it in event.items: + raw_name = it.get("func_name") or it.get("approval_label") or "tool" + display_name = discord.utils.escape_markdown(raw_name) + raw_preview = it.get("preview", "") + # Escape backticks to prevent markdown breakout and strip @-mentions. + raw_preview = raw_preview.replace("`", "\\`") + raw_preview = discord.utils.escape_mentions(raw_preview) + preview = truncate(raw_preview, max_length=120) or None + embed = discord.Embed( + title=display_name, + description=preview, + color=discord.Color.light_grey(), + ) + # Edit thinking message into first tool embed to avoid flicker. + if thinking_msg is not None: + try: + await thinking_msg.edit(content=None, embed=embed) + msg = thinking_msg + except Exception: + msg = await thread.send(embed=embed) + thinking_msg = None + else: + msg = await thread.send(embed=embed) + call_id = it.get("call_id", "") + # Store raw (unescaped) name for matching against ToolResultEvent.name + self._tool_info_msgs.setdefault(ws_id, []).append( + (call_id, raw_name, preview or "", msg) + ) + + # If no items consumed the thinking message (empty event), clean up. + if thinking_msg is not None: + with contextlib.suppress(Exception): + await thinking_msg.delete() + + async def _handle_tool_result( + self, + ws_id: str, + thread: discord.abc.Messageable, + event: ToolResultEvent, + ) -> None: + import discord + + from turnstone.channels._formatter import format_tool_result + + # Mark the matching "running" embed as complete/errored. + # Prefer call_id match (deterministic); fall back to name (FIFO). + info_list = self._tool_info_msgs.get(ws_id, []) + matched_preview = "" + matched_msg: discord.Message | None = None + if event.call_id: + for i, (cid, _tname, _prev, _tmsg) in enumerate(info_list): + if cid == event.call_id: + entry = info_list.pop(i) + matched_preview, matched_msg = entry[2], entry[3] + break + if matched_msg is None: + for i, (_cid, tname, _prev, _tmsg) in enumerate(info_list): + if tname == event.name: + entry = info_list.pop(i) + matched_preview, matched_msg = entry[2], entry[3] + break + if matched_msg is not None: + status = "Error" if event.is_error else "Done" + status_color = discord.Color.red() if event.is_error else discord.Color.dark_grey() + status_embed = discord.Embed( + title=f"{discord.utils.escape_markdown(event.name)} \u2014 {status}", + description=matched_preview or None, + color=status_color, + ) + try: + await matched_msg.edit(content=None, embed=status_embed) + except Exception: + log.debug("discord.tool_info_status_edit_failed", ws_id=ws_id) + + # Send the result as a separate message. + if not event.is_error: + from turnstone.channels._formatter import try_build_media_embed + + media_result = None + try: + media_result = await try_build_media_embed( + event.name, + event.output, + http=self._http_client, ) - try: - await matched_msg.edit(content=None, embed=status_embed) - except Exception: - log.debug("discord.tool_info_status_edit_failed", ws_id=ws_id) - - # Send the result as a separate message. - if not event.is_error: - from turnstone.channels._formatter import try_build_media_embed - - media_result = None - try: - media_result = await try_build_media_embed( - event.name, - event.output, - http=self._http_client, - ) - except Exception: - log.debug("discord.media_embed_failed", ws_id=ws_id, tool=event.name) - if media_result is not None: - embed, file = media_result - kwargs: dict[str, Any] = {"embed": embed} - if file is not None: - kwargs["file"] = file - await thread.send(**kwargs) - else: - desc = format_tool_result(event.output) - result_embed = discord.Embed( - title=event.name, - description=desc, - color=discord.Color.dark_grey(), - ) - await thread.send(embed=result_embed) + except Exception: + log.debug("discord.media_embed_failed", ws_id=ws_id, tool=event.name) + if media_result is not None: + embed, file = media_result + kwargs: dict[str, Any] = {"embed": embed} + if file is not None: + kwargs["file"] = file + await thread.send(**kwargs) else: desc = format_tool_result(event.output) result_embed = discord.Embed( title=event.name, description=desc, - color=discord.Color.red(), + color=discord.Color.dark_grey(), ) await thread.send(embed=result_embed) - - elif isinstance(event, ApproveRequestEvent): - # Evaluate admin tool policies before auto-approve. - _policy_handled = False - if self.storage is not None: - try: - from turnstone.core.policy import evaluate_tool_policies_batch - - _tool_names = [ - it.get("approval_label", "") or it.get("func_name", "") - for it in event.items - if it.get("needs_approval") and it.get("func_name") and not it.get("error") - ] - _tool_names = [n for n in _tool_names if n] - if _tool_names: - verdicts = await asyncio.to_thread( - evaluate_tool_policies_batch, - self.storage, - _tool_names, - ) - if any(v == "deny" for v in verdicts.values()): - denied = [n for n, v in verdicts.items() if v == "deny"] - await self.router.send_approval( - ws_id, - "", - approved=False, - feedback=f"Blocked by tool policy: {', '.join(denied)}", - ) - await thread.send( - f"*Tool blocked by admin policy: {', '.join(denied)}*" - ) - _policy_handled = True - elif all(verdicts.get(n) == "allow" for n in _tool_names): - await self.router.send_approval( - ws_id, - "", - approved=True, - ) - await thread.send("*Tool approved by policy.*") - _policy_handled = True - except Exception: - log.debug("Tool policy evaluation failed for ws %s", ws_id, exc_info=True) - if not _policy_handled and ( - self.config.auto_approve or self._should_auto_approve(event) - ): - # correlation_id is empty because the server's /api/approve - # endpoint resolves approvals by ws_id alone (one pending - # approval per workstream at a time). - await self.router.send_approval(ws_id, "", approved=True) - await thread.send("*Tool auto-approved.*") - elif not _policy_handled: - text = format_approval_request(event.items) - embed = discord.Embed( - title="Tool Approval Required", - description=text, - color=discord.Color.orange(), - ) - # Include heuristic verdicts from approval items. - for item in event.items: - verdict = item.get("verdict") - if verdict: - name = item.get("func_name") or item.get("approval_label") or "tool" - embed.add_field( - name=f"Verdict: {name}", - value=format_verdict(verdict), - inline=False, - ) - embed.set_footer(text=f"{ws_id}|") - msg = await thread.send(embed=embed, view=ApprovalView(self)._view) - self._pending_approval_msgs[ws_id] = msg - - elif isinstance(event, PlanReviewEvent): - text = format_plan_review(event.content) - embed = discord.Embed( - title="Plan Review", - description=text, - color=discord.Color.blue(), + else: + desc = format_tool_result(event.output) + result_embed = discord.Embed( + title=event.name, + description=desc, + color=discord.Color.red(), ) - embed.set_footer(text=f"{ws_id}|") - await thread.send(embed=embed, view=PlanReviewView(self)._view) + await thread.send(embed=result_embed) - elif isinstance(event, IntentVerdictEvent): - # LLM judge verdict arrived — update the pending approval embed. - approval_msg = self._pending_approval_msgs.get(ws_id) - if approval_msg and approval_msg.embeds: - embed = approval_msg.embeds[0] - verdict_data = { - "risk_level": event.risk_level, - "recommendation": event.recommendation, - "confidence": event.confidence, - "intent_summary": event.intent_summary, - "tier": event.tier, - } - name = event.func_name or "tool" - # Update embed color based on LLM judge risk level. - risk = (event.risk_level or "medium").upper() - color_map = { - "LOW": discord.Color.green(), - "MEDIUM": discord.Color.orange(), - "HIGH": discord.Color.red(), - "CRITICAL": discord.Color.dark_red(), - } - embed.color = color_map.get(risk, discord.Color.orange()) - embed.add_field( - name=f"Judge Verdict: {name}", - value=format_verdict(verdict_data), - inline=False, - ) - try: - await approval_msg.edit(embed=embed) - except Exception: - log.debug("discord.verdict_embed_edit_failed", ws_id=ws_id) + async def _handle_approve_request( + self, + ws_id: str, + thread: discord.abc.Messageable, + event: ApproveRequestEvent, + ) -> None: + import discord - elif isinstance(event, ApprovalResolvedEvent): - # Server resolved the approval (timeout, external approve/reject). - # Disable the buttons so they can't be clicked stale. - approval_msg = self._pending_approval_msgs.pop(ws_id, None) - if approval_msg is not None: - from turnstone.channels.discord.views import disable_message_buttons + from turnstone.channels._formatter import format_approval_request, format_verdict + from turnstone.channels.discord.views import ApprovalView - label = "Approved" if event.approved else "Denied" - try: - await disable_message_buttons(approval_msg, label) - except Exception: - log.debug("discord.approval_resolved_edit_failed", ws_id=ws_id) + # Evaluate admin tool policies before auto-approve. + policy_verdict = await self.router.evaluate_tool_policies(event.items) + policy_handled = False + if policy_verdict.kind == "deny": + denied = ", ".join(policy_verdict.denied_tools) + await self.router.send_approval( + ws_id, + "", + approved=False, + feedback=f"Blocked by tool policy: {denied}", + ) + await thread.send(f"*Tool blocked by admin policy: {denied}*") + policy_handled = True + elif policy_verdict.kind == "allow": + await self.router.send_approval(ws_id, "", approved=True) + await thread.send("*Tool approved by policy.*") + policy_handled = True - elif isinstance(event, StreamEndEvent): - # Edge-case cleanup: clear any lingering thinking indicator. - thinking_msg = self._thinking_msgs.pop(ws_id, None) - if thinking_msg is not None: - with contextlib.suppress(Exception): - await thinking_msg.delete() - self._tool_info_msgs.pop(ws_id, None) - sm = self._streaming.pop(ws_id, None) - if sm is not None: - await sm.finalize() - # Forward accumulated response to notification reply DM if active. - dm_entry = self._notify_reply_channels.pop(ws_id, None) - if dm_entry is not None and sm is not None: - content = "".join(sm._buffer) - if content: - dm_channel, target_user_id = dm_entry - last_msg: discord.Message | None = None - for dm_chunk in chunk_message(content, self.config.max_message_length): - try: - last_msg = await dm_channel.send(dm_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) + if not policy_handled and (self.config.auto_approve or self._should_auto_approve(event)): + # correlation_id is empty because the server's /api/approve + # endpoint resolves approvals by ws_id alone (one pending + # approval per workstream at a time). + await self.router.send_approval(ws_id, "", approved=True) + await thread.send("*Tool auto-approved.*") + elif not policy_handled: + text = format_approval_request(event.items) + embed = discord.Embed( + title="Tool Approval Required", + description=text, + color=discord.Color.orange(), + ) + # Include heuristic verdicts from approval items. + for item in event.items: + verdict = item.get("verdict") + if verdict: + name = item.get("func_name") or item.get("approval_label") or "tool" + embed.add_field( + name=f"Verdict: {name}", + value=format_verdict(verdict), + inline=False, + ) + embed.set_footer(text=f"{ws_id}||{_thread_owner_id(thread)}") + msg = await thread.send(embed=embed, view=ApprovalView(self)._view) + self._pending_approval_msgs[ws_id] = msg - elif isinstance(event, ErrorEvent): - safe_msg = event.message[:500] if event.message else "An error occurred" - await thread.send(f"**Error:** {safe_msg}") + async def _handle_plan_review( + self, + ws_id: str, + thread: discord.abc.Messageable, + event: PlanReviewEvent, + ) -> None: + import discord + + from turnstone.channels.discord.views import PlanReviewView + + embed = discord.Embed( + title="Plan Review", + description=f"**Plan review requested:**\n\n{event.content}", + color=discord.Color.blue(), + ) + embed.set_footer(text=f"{ws_id}||{_thread_owner_id(thread)}") + await thread.send(embed=embed, view=PlanReviewView(self)._view) + + async def _handle_intent_verdict( + self, + ws_id: str, + event: IntentVerdictEvent, + ) -> None: + import discord + + from turnstone.channels._formatter import format_verdict + + # LLM judge verdict arrived — update the pending approval embed. + approval_msg = self._pending_approval_msgs.get(ws_id) + if approval_msg and approval_msg.embeds: + embed = approval_msg.embeds[0] + verdict_data = { + "risk_level": event.risk_level, + "recommendation": event.recommendation, + "confidence": event.confidence, + "intent_summary": event.intent_summary, + "tier": event.tier, + } + name = event.func_name or "tool" + # Update embed color based on LLM judge risk level. + risk = (event.risk_level or "medium").upper() + color_map = { + "LOW": discord.Color.green(), + "MEDIUM": discord.Color.orange(), + "HIGH": discord.Color.red(), + "CRITICAL": discord.Color.dark_red(), + } + embed.color = color_map.get(risk, discord.Color.orange()) + embed.add_field( + name=f"Judge Verdict: {name}", + value=format_verdict(verdict_data), + inline=False, + ) + try: + await approval_msg.edit(embed=embed) + except Exception: + log.debug("discord.verdict_embed_edit_failed", ws_id=ws_id) + + async def _handle_approval_resolved( + self, + ws_id: str, + event: ApprovalResolvedEvent, + ) -> None: + # Server resolved the approval (timeout, external approve/reject). + # Disable the buttons so they can't be clicked stale. + approval_msg = self._pending_approval_msgs.pop(ws_id, None) + if approval_msg is not None: + from turnstone.channels.discord.views import disable_message_buttons + + label = "Approved" if event.approved else "Denied" + try: + await disable_message_buttons(approval_msg, label) + except Exception: + log.debug("discord.approval_resolved_edit_failed", ws_id=ws_id) + + async def _handle_stream_end(self, ws_id: str) -> None: + # Edge-case cleanup: clear any lingering thinking indicator. + thinking_msg = self._thinking_msgs.pop(ws_id, None) + if thinking_msg is not None: + with contextlib.suppress(Exception): + await thinking_msg.delete() + self._tool_info_msgs.pop(ws_id, None) + sm = self._streaming.pop(ws_id, None) + if sm is not None: + await sm.finalize() + # Forward accumulated response to notification reply DM if active. + dm_entry = self._notify_reply_channels.pop(ws_id, None) + if dm_entry is not None and sm is not None: + content = sm.accumulated_text + if content: + dm_channel, target_user_id = dm_entry + last_msg: discord.Message | None = None + for dm_chunk in chunk_message(content, self.config.max_message_length): + try: + last_msg = await dm_channel.send(dm_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) # -- helpers ------------------------------------------------------------- @@ -890,23 +903,67 @@ 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. + def register_thread_invoker(self, thread_id: int, discord_user_id: int) -> None: + """Record the Discord user who caused *thread_id* to be created. - 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. + Sec-3 gates inbound messages on thread ownership. When the bot + itself creates the thread via ``channel.create_thread(...)`` + Discord sets ``thread.owner_id`` to the bot, so ``channel.owner_id`` + alone would silently drop every legitimate follow-up from the + human who triggered the session. """ - 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) + self._thread_invokers[thread_id] = discord_user_id + self._thread_invokers.move_to_end(thread_id) + while len(self._thread_invokers) > _THREAD_INVOKER_CAP: + self._thread_invokers.popitem(last=False) + + def get_thread_invoker(self, thread_id: int) -> int | None: + """Return the recorded invoker for *thread_id*, or ``None``.""" + invoker = self._thread_invokers.get(thread_id) + if invoker is not None: + self._thread_invokers.move_to_end(thread_id) + return invoker + + async def send_notification(self, channel_id: str, content: str, ws_id: str) -> str: + """Send a notification and track the message for reply routing (DMs only). + + Like :meth:`send` but, when the target resolves to a DM channel, + records a mapping from the outgoing Discord message ID to + ``(ws_id, user_id)`` so that the user's reply can be routed back + to the originating workstream. Notifications delivered to guild + channels are NOT tracked — the reply-channel_id check would treat + the channel ID as a user ID and reject every legitimate reply. + """ + import discord + + int_id = int(channel_id) + target: discord.abc.Messageable | None = self._bot.get_channel(int_id) # type: ignore[assignment] + target_user_id: str = "" + if target is None: + try: + user = await self._bot.fetch_user(int_id) + target = await user.create_dm() + target_user_id = str(user.id) + except discord.NotFound as exc: + raise ValueError(f"Discord channel/user {channel_id} not found") from exc + + content = discord.utils.escape_mentions(content) + chunks = chunk_message(content, self.config.max_message_length) + msg: discord.Message | None = None + for chunk in chunks: + msg = await target.send(chunk) # type: ignore[union-attr] + + if msg is None: + return "" + + msg_id_str = str(msg.id) + if ws_id and target_user_id: + self._track_notification(int(msg_id_str), ws_id, target_user_id) log.debug( "discord.notification_tracked", message_id=msg_id_str, ws_id=ws_id, - target_user=channel_id, + target_user=target_user_id, ) return msg_id_str diff --git a/turnstone/channels/discord/cog.py b/turnstone/channels/discord/cog.py index 29e5e1f1..fdc83ecb 100644 --- a/turnstone/channels/discord/cog.py +++ b/turnstone/channels/discord/cog.py @@ -7,6 +7,8 @@ Handles ``on_message`` events and slash commands (``/link``, ``/unlink``, from __future__ import annotations import asyncio +import time +from collections import OrderedDict, deque from typing import TYPE_CHECKING from turnstone.core.log import get_logger @@ -23,6 +25,13 @@ log = get_logger(__name__) _THREAD_NAME_MAX = 100 _DM_REPLY_MAX_LENGTH = 4096 # Discord's own message limit +# /link is the only flow that reads Turnstone API tokens out of user +# input; throttle aggressively so an attacker with throw-away Discord +# accounts can't online-enumerate valid tokens. +_LINK_RATE_WINDOW_S: float = 3600.0 +_LINK_RATE_LIMIT: int = 5 +_LINK_RATE_CAP: int = 2048 + class MessageCog: """Cog that processes messages and registers slash commands. @@ -38,6 +47,10 @@ class MessageCog: self.bot = bot self.ts: TurnstoneBot = bot.turnstone # type: ignore[attr-defined] + # Per-Discord-user sliding-window rate limit on /link, to block + # online enumeration of Turnstone API tokens. + self._link_buckets: OrderedDict[str, deque[float]] = OrderedDict() + # -- Cog wiring (manual since we can't use decorators with guarded imports) -- # We build the cog dynamically so discord.py's import is fully deferred. @@ -132,14 +145,36 @@ class MessageCog: if not self.ts._is_allowed_channel(parent_id): return - # Check if this thread has an existing route. - route = await asyncio.to_thread( - self.ts.storage.get_channel_route, "discord", str(channel.id) - ) - if route is None: + # Check if this thread has an existing route (TTL-cached). + existing_ws_id = await self.ts.router.lookup_ws_id("discord", str(channel.id)) + if existing_ws_id is None: # Not our thread — ignore. return + # Owner check: only the thread creator (who initiated the + # workstream) can inject messages. Without this gate, any + # linked user in a public / multi-member thread could + # redirect someone else's assistant and bill their quota, + # because the gateway forwards with its service-scoped JWT + # and the server bypasses ownership on service scope. + # + # We prefer the explicitly-recorded invoker over + # `thread.owner_id`: `/ask` creates threads via + # `channel.create_thread(...)` which reports the bot as + # owner, so the Discord-reported value alone would reject + # every legitimate follow-up. + effective_owner_id = self.ts.get_thread_invoker(channel.id) + if effective_owner_id is None: + effective_owner_id = channel.owner_id + if effective_owner_id is None or message.author.id != effective_owner_id: + log.debug( + "discord.thread_message_rejected_non_owner", + thread_id=channel.id, + author_id=message.author.id, + owner_id=effective_owner_id, + ) + return + # Resolve user. user_id = await self.ts.router.resolve_user("discord", str(message.author.id)) if user_id is None: @@ -199,6 +234,9 @@ class MessageCog: name=thread_name, auto_archive_duration=self.ts.config.thread_auto_archive, # type: ignore[arg-type] ) + # Record invoker so the sec-3 gate admits follow-ups even if + # Discord's reported thread.owner_id diverges. + self.ts.register_thread_invoker(thread.id, message.author.id) # Create workstream WITHOUT initial_message — subscribe to events # first, then send the message. With SSE the event stream is @@ -288,10 +326,48 @@ class MessageCog: # -- slash commands ------------------------------------------------------ + def _allow_link_attempt(self, discord_user_id: str) -> bool: + """Return True when this Discord user is under the /link rate limit. + + Sliding window: up to ``_LINK_RATE_LIMIT`` attempts per + ``_LINK_RATE_WINDOW_S`` seconds. Each attempt — success or + failure — consumes a slot. The bucket map is LRU-bounded. + """ + now = time.monotonic() + window_start = now - _LINK_RATE_WINDOW_S + bucket = self._link_buckets.get(discord_user_id) + if bucket is None: + bucket = deque() + self._link_buckets[discord_user_id] = bucket + while len(self._link_buckets) > _LINK_RATE_CAP: + self._link_buckets.popitem(last=False) + else: + self._link_buckets.move_to_end(discord_user_id) + while bucket and bucket[0] < window_start: + bucket.popleft() + if len(bucket) >= _LINK_RATE_LIMIT: + return False + bucket.append(now) + return True + async def _cmd_link(self, interaction: discord.Interaction, token: str) -> None: """Link a Discord user to a turnstone account via API token.""" from turnstone.core.auth import hash_token + if not self._allow_link_attempt(str(interaction.user.id)): + log.warning( + "discord.link_rate_limited", + discord_user=str(interaction.user), + ) + await interaction.response.send_message( + ( + f"Too many /link attempts. Try again later — limit is " + f"{_LINK_RATE_LIMIT} per hour." + ), + ephemeral=True, + ) + return + # Check if already linked. existing = await asyncio.to_thread( self.ts.storage.get_channel_user, "discord", str(interaction.user.id) @@ -381,6 +457,10 @@ class MessageCog: auto_archive_duration=self.ts.config.thread_auto_archive, # type: ignore[arg-type] type=discord.ChannelType.public_thread, ) + # `channel.create_thread` without a starter message makes the + # bot the thread owner, so the sec-3 gate needs to see the + # real invoker here — otherwise `/ask` follow-ups get dropped. + self.ts.register_thread_invoker(thread.id, interaction.user.id) else: await interaction.followup.send( "Cannot create a thread in this channel type.", diff --git a/turnstone/channels/discord/views.py b/turnstone/channels/discord/views.py index bcd89558..71c422f9 100644 --- a/turnstone/channels/discord/views.py +++ b/turnstone/channels/discord/views.py @@ -19,15 +19,32 @@ if TYPE_CHECKING: log = get_logger(__name__) -def _parse_footer(interaction: discord.Interaction) -> tuple[str, str] | None: - """Extract ``(ws_id, correlation_id)`` from the first embed's footer.""" +def _parse_footer(interaction: discord.Interaction) -> tuple[str, str, str] | None: + """Extract ``(ws_id, correlation_id, owner_id)`` from the first embed's footer. + + Footer format is ``"{ws_id}|{correlation_id}|{owner_id}"``. Older + posts that pre-date the owner-check upgrade may have only two + fields; in that case ``owner_id`` is returned as an empty string + and the caller rejects the interaction (fail-closed). + """ if not interaction.message or not interaction.message.embeds: return None footer = interaction.message.embeds[0].footer.text if not footer or "|" not in footer: return None - parts = footer.split("|", 1) - return parts[0], parts[1] + parts = footer.split("|", 2) + ws_id = parts[0] + correlation_id = parts[1] if len(parts) > 1 else "" + owner_id = parts[2] if len(parts) > 2 else "" + return ws_id, correlation_id, owner_id + + +async def _deny_non_owner(interaction: discord.Interaction, verb: str) -> None: + """Reply with an ephemeral non-owner rejection.""" + await interaction.response.send_message( + f"Only the session owner can {verb} this.", + ephemeral=True, + ) async def disable_message_buttons(message: discord.Message, label: str) -> None: @@ -137,10 +154,19 @@ class ApprovalView: ) return - ws_id, correlation_id = parsed + ws_id, correlation_id, owner_id = parsed - # Verify user is linked. Scope enforcement (approve) happens - # server-side when the tool approval is executed. + # Owner check: the gateway forwards approvals using its own + # service-scoped JWT, and the server short-circuits scope checks + # for service tokens — so the adapter is the only place this + # can be enforced. Reject any other clicker, including linked + # users, with an ephemeral message. + if not owner_id or str(interaction.user.id) != owner_id: + verb = "always-approve" if always else ("approve" if approved else "reject") + await _deny_non_owner(interaction, verb) + return + + # Verify user is linked (session owner should already be linked). user_id = await self.bot.router.resolve_user("discord", str(interaction.user.id)) if user_id is None: await interaction.response.send_message( @@ -254,7 +280,11 @@ class PlanReviewView: ) return - ws_id, correlation_id = parsed + ws_id, correlation_id, owner_id = parsed + + if not owner_id or str(interaction.user.id) != owner_id: + await _deny_non_owner(interaction, "approve") + return user_id = await self.bot.router.resolve_user("discord", str(interaction.user.id)) if user_id is None: @@ -291,7 +321,12 @@ class PlanReviewView: ) return - ws_id, correlation_id = parsed + ws_id, correlation_id, owner_id = parsed + + if not owner_id or str(interaction.user.id) != owner_id: + await _deny_non_owner(interaction, "request changes on") + return + modal = self._modal_cls(ws_id, correlation_id) await interaction.response.send_modal(modal) diff --git a/turnstone/channels/slack/bot.py b/turnstone/channels/slack/bot.py index 734aabcf..a0d3406d 100644 --- a/turnstone/channels/slack/bot.py +++ b/turnstone/channels/slack/bot.py @@ -24,8 +24,8 @@ from __future__ import annotations import asyncio import contextlib -import json import time +from collections import OrderedDict, deque from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, cast @@ -33,8 +33,10 @@ import httpx from slack_bolt.async_app import AsyncApp from slack_sdk.web.async_client import AsyncWebClient +from turnstone.channels._config import MAX_NOTIFY_TRACKING from turnstone.channels._formatter import chunk_message from turnstone.channels._routing import ChannelRouter +from turnstone.channels._sse import run_sse_stream from turnstone.channels.slack.routes import SlackRoute from turnstone.core.log import get_logger from turnstone.sdk._types import TurnstoneAPIError @@ -63,15 +65,26 @@ log = get_logger(__name__) _GREETING = "Hey! Let me know what I can help with." -_SSE_RECONNECT_DELAY: float = 2.0 -_SSE_MAX_RECONNECT_DELAY: float = 30.0 - # Slack section.text caps at 3000 chars. Reserve headroom for the heading # and the "+N more" suffix; cap each preview individually so a single huge # tool can't crowd out the rest of a multi-tool batch. _APPROVAL_TEXT_BUDGET: int = 2700 _APPROVAL_PER_ITEM_PREVIEW: int = 600 +# Inbound-message limits — one user can't exhaust the LLM budget on +# their own. +_PER_USER_RATE_WINDOW_S: float = 60.0 +_PER_USER_RATE_LIMIT: int = 10 +_PER_USER_RATE_CAP: int = 4096 # LRU bound on the per-user deque map +_MAX_INBOUND_MESSAGE_LEN: int = 8192 # chars; roughly 8 KiB + +# /turnstone link has a much tighter throttle than regular messages so +# throw-away Slack accounts can't online-enumerate Turnstone API tokens +# (mirrors Discord's /link cap in discord/cog.py). +_LINK_RATE_WINDOW_S: float = 3600.0 +_LINK_RATE_LIMIT: int = 5 +_LINK_RATE_CAP: int = 2048 + def _sanitize_slack_preview(text: str, max_length: int = 1200) -> str: """Escape Slack mrkdwn-sensitive content for safe fenced display. @@ -92,9 +105,15 @@ def _sanitize_slack_preview(text: str, max_length: int = 1200) -> str: def _parse_ts(ts: str) -> tuple[int, int]: - parts = ts.split(".", 1) - seconds = int(parts[0]) - micros = int(parts[1]) if len(parts) > 1 else 0 + """Parse a Slack ``"seconds.microseconds"`` timestamp into ordered parts. + + Pads (or truncates) the fractional field to 6 digits before converting + so that e.g. ``"1.2"`` and ``"1.000002"`` compare correctly — ``int`` + strips leading zeros and otherwise both collapse to ``(1, 2)``. + """ + seconds_str, _sep, frac_str = ts.partition(".") + seconds = int(seconds_str) + micros = int(frac_str.ljust(6, "0")[:6]) if frac_str else 0 return (seconds, micros) @@ -103,6 +122,12 @@ class PendingApproval: channel: str message_ts: str owner_user_id: str | None = None + # Block Kit payload posted to Slack so IntentVerdictEvent can edit + # the message without an extra conversations_history fetch. A + # mutable list so repeat IntentVerdictEvents stack on the same + # approval message (matches the pre-refactor "fetch live blocks + # and append" behavior). + blocks: list[dict[str, Any]] = field(default_factory=list) @dataclass @@ -117,15 +142,39 @@ class StreamingMessage: _ts: str = field(default="", init=False, repr=False) _buffer: list[str] = field(default_factory=list, init=False, repr=False) + # Rolling truncated in-progress display string; stops growing at + # max_length so per-flush cost is O(max_length) instead of + # O(total_streamed_chars). + _display: str = field(default="", init=False, repr=False) _last_edit: float = field(default=0.0, init=False, repr=False) + _finalized_text: str | None = field(default=None, init=False, repr=False) + + @property + def message_ts(self) -> str: + """Slack ``ts`` of the initial posted message, or ``""`` if not yet sent.""" + return self._ts + + @property + def accumulated_text(self) -> str: + """The joined text of everything appended so far. + + Cached after ``finalize()`` so downstream callers that read the + full message content don't re-join the buffer twice. + """ + if self._finalized_text is not None: + return self._finalized_text + return "".join(self._buffer) async def append(self, text: str) -> None: self._buffer.append(text) + if len(self._display) < self.max_length: + self._display = (self._display + text)[: self.max_length] if time.monotonic() - self._last_edit >= self.edit_interval: await self._flush() async def finalize(self) -> None: content = "".join(self._buffer) + self._finalized_text = content if not content: return @@ -148,20 +197,23 @@ class StreamingMessage: ) else: for chunk in chunks: + # After the first chunk posts successfully, reuse its ts as + # the thread parent so the remaining chunks group into one + # thread instead of fragmenting as independent top-level + # messages (Slack DMs otherwise inherit thread_ts=""). resp = await self.client.chat_postMessage( channel=self.channel, - thread_ts=self.thread_ts or None, + thread_ts=self.thread_ts or self._ts or None, text=chunk, ) if not self._ts and resp.get("ok"): self._ts = resp["ts"] async def _flush(self) -> None: - content = "".join(self._buffer) - if not content: + display = self._display + if not display: return - display = content[: self.max_length] try: if not self._ts: resp = await self.client.chat_postMessage( @@ -187,7 +239,7 @@ class TurnstoneSlackBot: """Slack bot bridging Slack channels/threads to turnstone workstreams.""" channel_type: str = "slack" - _MAX_NOTIFY_TRACKING: int = 100 + _MAX_NOTIFY_TRACKING: int = MAX_NOTIFY_TRACKING def __init__( self, @@ -224,13 +276,22 @@ class TurnstoneSlackBot: self._streaming: dict[str, StreamingMessage] = {} self._pending_approval: dict[str, PendingApproval] = {} - self._pending_plan_review_ts: dict[str, tuple[str, str]] = {} + # (channel, message_ts, owner_user_id) per pending plan review, so + # _on_plan_* handlers can reject clicks from non-owners. + self._pending_plan_review_ts: dict[str, tuple[str, str, str]] = {} self._notify_ws_map: dict[str, tuple[str, SlackRoute]] = {} # Per-workstream override used to route the next streamed assistant # response back into a Slack notification reply thread instead of the # default session thread. self._notify_reply_routes: dict[str, SlackRoute] = {} self._channel_sessions: dict[tuple[str, str], tuple[str, str]] = {} + # Per-user sliding-window rate limit; maps slack_user_id → deque + # of recent allow timestamps within _PER_USER_RATE_WINDOW_S. + self._rate_buckets: OrderedDict[str, deque[float]] = OrderedDict() + # Separate sliding-window throttle for /turnstone link so an + # abusive user can't enumerate tokens while still being allowed + # to chat at the baseline message rate. + self._link_buckets: OrderedDict[str, deque[float]] = OrderedDict() headers: dict[str, str] = {} if api_token and not server_token_factory: @@ -248,8 +309,15 @@ class TurnstoneSlackBot: self._app.command(config.slash_command)(self._on_slash_command) self._app.event("message")(self._on_message) - self._app.action("ts_approve")(self._on_approve) - self._app.action("ts_deny")(self._on_deny) + + async def _approve_cb(ack: Any, body: dict[str, Any]) -> None: + await self._resolve_approval(ack, body, approved=True) + + async def _deny_cb(ack: Any, body: dict[str, Any]) -> None: + await self._resolve_approval(ack, body, approved=False) + + self._app.action("ts_approve")(_approve_cb) + self._app.action("ts_deny")(_deny_cb) self._app.action("ts_plan_approve")(self._on_plan_approve) self._app.action("ts_plan_request_changes")(self._on_plan_request_changes) self._app.view("ts_plan_feedback_modal")(self._on_plan_feedback_modal) @@ -308,11 +376,35 @@ class TurnstoneSlackBot: ) async def _on_slash_command(self, ack: Any, body: dict[str, Any]) -> None: - """Start or restart a per-user session in this channel.""" + """Start or restart a per-user session in this channel. + + Subcommands (parsed from the slash-command ``text`` field): + + - ``link ``: link the caller's Slack user ID to the + Turnstone account that owns *api_token*. + - ``unlink``: remove the link. + - (default): open or restart a Turnstone session in this channel. + """ await ack() slack_channel = body["channel_id"] user_id = body["user_id"] + raw_text = (body.get("text") or "").strip() + + # Subcommand routing — always available regardless of channel + # allowlist because linking is the prerequisite for everything + # else. Match on the first whitespace-delimited word so an + # innocuous prompt like `linking up the docs` doesn't hijack + # into the link handler with `"ing up the docs"` as the token. + parts = raw_text.split(None, 1) + subcmd = parts[0] if parts else "" + if subcmd == "link": + token = parts[1].strip() if len(parts) > 1 else "" + await self._handle_link(slack_channel, user_id, token) + return + if subcmd == "unlink": + await self._handle_unlink(slack_channel, user_id) + return if self.config.allowed_channels and slack_channel not in self.config.allowed_channels: await self._client.chat_postEphemeral( @@ -322,6 +414,12 @@ class TurnstoneSlackBot: ) return + # Only linked Slack users can start sessions — otherwise every + # Slack workspace member would be creating workstreams under + # the shared gateway identity with no attribution. + if not await self._require_linked(slack_channel, user_id): + return + existing = self._channel_sessions.get((slack_channel, user_id)) if existing is not None: old_ws_id, old_thread_ts = existing @@ -371,6 +469,195 @@ class TurnstoneSlackBot: thread_ts=opener_ts, ) + def _allow_user_send(self, slack_user_id: str) -> bool: + """Return True when *slack_user_id* is within the send rate limit. + + Sliding-window counter: up to ``_PER_USER_RATE_LIMIT`` allowed + sends per ``_PER_USER_RATE_WINDOW_S``. Stale entries are evicted + lazily, and the bucket map is LRU-bounded so a workspace with + churning user IDs can't grow the map unboundedly. + """ + now = time.monotonic() + window_start = now - _PER_USER_RATE_WINDOW_S + bucket = self._rate_buckets.get(slack_user_id) + if bucket is None: + bucket = deque() + self._rate_buckets[slack_user_id] = bucket + while len(self._rate_buckets) > _PER_USER_RATE_CAP: + self._rate_buckets.popitem(last=False) + else: + self._rate_buckets.move_to_end(slack_user_id) + while bucket and bucket[0] < window_start: + bucket.popleft() + if len(bucket) >= _PER_USER_RATE_LIMIT: + return False + bucket.append(now) + return True + + async def _check_inbound( + self, + slack_channel: str, + slack_user_id: str, + text: str, + ) -> bool: + """Validate an inbound message against size + rate limits. + + Returns True when the caller should proceed. Posts an ephemeral + message on rejection so the user understands why their message + was dropped. + """ + if len(text) > _MAX_INBOUND_MESSAGE_LEN: + await self._client.chat_postEphemeral( + channel=slack_channel, + user=slack_user_id, + text=( + f"Message too long ({len(text)} chars). Please keep " + f"messages under {_MAX_INBOUND_MESSAGE_LEN} characters." + ), + ) + log.info( + "slack.inbound_oversize", + user_id=slack_user_id, + length=len(text), + ) + return False + if not self._allow_user_send(slack_user_id): + await self._client.chat_postEphemeral( + channel=slack_channel, + user=slack_user_id, + text=( + f"You're sending messages too fast. Limit is " + f"{_PER_USER_RATE_LIMIT} per " + f"{int(_PER_USER_RATE_WINDOW_S)} seconds." + ), + ) + log.warning("slack.inbound_rate_limited", user_id=slack_user_id) + return False + return True + + async def _require_linked(self, slack_channel: str, slack_user_id: str) -> bool: + """Return True if *slack_user_id* is mapped to a Turnstone account. + + Posts an ephemeral "please /link first" message on False so the + caller can early-return. + """ + mapped = await self.router.resolve_user("slack", slack_user_id) + if mapped: + return True + await self._client.chat_postEphemeral( + channel=slack_channel, + user=slack_user_id, + text=( + "You need to link your Slack account to a Turnstone user before " + "using this bot. Run `/turnstone link ` — tokens " + "are managed in the Turnstone console." + ), + ) + return False + + def _allow_link_attempt(self, slack_user_id: str) -> bool: + """Return True when *slack_user_id* is under the /link rate limit. + + Mirrors the Discord helper at `discord/cog.py::_allow_link_attempt` + — 5 attempts per hour per Slack user, LRU-bounded. + """ + now = time.monotonic() + window_start = now - _LINK_RATE_WINDOW_S + bucket = self._link_buckets.get(slack_user_id) + if bucket is None: + bucket = deque() + self._link_buckets[slack_user_id] = bucket + while len(self._link_buckets) > _LINK_RATE_CAP: + self._link_buckets.popitem(last=False) + else: + self._link_buckets.move_to_end(slack_user_id) + while bucket and bucket[0] < window_start: + bucket.popleft() + if len(bucket) >= _LINK_RATE_LIMIT: + return False + bucket.append(now) + return True + + async def _handle_link(self, slack_channel: str, slack_user_id: str, token: str) -> None: + """Link the caller's Slack ID to the Turnstone user that owns *token*.""" + from turnstone.core.auth import hash_token + + if not self._allow_link_attempt(slack_user_id): + log.warning("slack.link_rate_limited", slack_user_id=slack_user_id) + await self._client.chat_postEphemeral( + channel=slack_channel, + user=slack_user_id, + text=( + f"Too many `/turnstone link` attempts. Try again " + f"later — limit is {_LINK_RATE_LIMIT} per hour." + ), + ) + return + + if not token: + await self._client.chat_postEphemeral( + channel=slack_channel, + user=slack_user_id, + text="Usage: `/turnstone link `", + ) + return + + existing = await asyncio.to_thread(self.storage.get_channel_user, "slack", slack_user_id) + if existing: + await self._client.chat_postEphemeral( + channel=slack_channel, + user=slack_user_id, + text="Your Slack account is already linked. Use `/turnstone unlink` first.", + ) + return + + record = await asyncio.to_thread(self.storage.get_api_token_by_hash, hash_token(token)) + if record is None or not record.get("user_id"): + # Generic failure message — don't leak whether the token + # existed but was e.g. revoked. + await self._client.chat_postEphemeral( + channel=slack_channel, + user=slack_user_id, + text="Invalid token. Please provide a valid Turnstone API token.", + ) + log.info("slack.link_invalid_token", slack_user_id=slack_user_id) + return + + await asyncio.to_thread( + self.storage.create_channel_user, + "slack", + slack_user_id, + record["user_id"], + ) + await self._client.chat_postEphemeral( + channel=slack_channel, + user=slack_user_id, + text="Linked. You can now use Turnstone from Slack.", + ) + log.info( + "slack.link_succeeded", + slack_user_id=slack_user_id, + turnstone_user_id=record["user_id"], + ) + + async def _handle_unlink(self, slack_channel: str, slack_user_id: str) -> None: + """Drop the Slack → Turnstone mapping for the caller.""" + existing = await asyncio.to_thread(self.storage.get_channel_user, "slack", slack_user_id) + if not existing: + await self._client.chat_postEphemeral( + channel=slack_channel, + user=slack_user_id, + text="Your Slack account was not linked.", + ) + return + await asyncio.to_thread(self.storage.delete_channel_user, "slack", slack_user_id) + await self._client.chat_postEphemeral( + channel=slack_channel, + user=slack_user_id, + text="Unlinked.", + ) + log.info("slack.unlink_succeeded", slack_user_id=slack_user_id) + async def _archive_session( self, slack_channel: str, @@ -409,6 +696,10 @@ class TurnstoneSlackBot: ) await self.unsubscribe_ws(ws_id) + # close_workstream is the only path that drops the ws_id from + # ChannelRouter._node_urls; skipping it leaks one cache entry per + # archived session over long bot uptime. + await self.router.close_workstream(ws_id) self._channel_sessions.pop((slack_channel, user_id), None) log.info( "slack.session_archived", @@ -434,6 +725,15 @@ class TurnstoneSlackBot: if thread_ts and thread_ts in self._notify_ws_map: origin_ws_id, notify_route = self._notify_ws_map[thread_ts] if channel_id == notify_route.channel and user_id == (notify_route.user_id or ""): + # Same identity + rate/size gates the other two routing + # paths use — the author-match above narrows the surface + # to the originally notified user, but that user may have + # unlinked or their token may have been revoked since the + # notification was recorded. + if not await self._require_linked(channel_id, user_id): + return + if not await self._check_inbound(channel_id, user_id, text): + return try: reply_route = SlackRoute( channel=channel_id, @@ -496,6 +796,15 @@ class TurnstoneSlackBot: if thread_ts != session_thread_ts: return + # Gate message routing on a linked Slack user so all LLM usage + # can be attributed to a Turnstone account. Unlinked users + # see an ephemeral prompt to run `/turnstone link`. + if not await self._require_linked(channel_id, user_id): + return + + if not await self._check_inbound(channel_id, user_id, text): + return + try: await self.router.send_message(ws_id, text) log.info("slack.message_dispatched", ws_id=ws_id, channel=channel_id) @@ -512,12 +821,23 @@ class TurnstoneSlackBot: thread_ts = event.get("thread_ts", "") user_id = event.get("user", "") text = event.get("text", "").strip() - msg_ts = event.get("ts", "") - thread_key = thread_ts or msg_ts + + # Same identity gate as `/turnstone` and channel messages: a + # Slack user must be linked to a Turnstone account before they + # can drive LLM calls via DM. + if not await self._require_linked(channel_id, user_id): + return + + if not await self._check_inbound(channel_id, user_id, text): + return + # Slack DMs don't auto-thread, so when no explicit thread_ts is + # present we route by (channel, user) alone. Falling back to the + # per-message ``ts`` would make every top-level DM spawn a new + # workstream because the key would be unique per message. route = SlackRoute( channel=channel_id, user_id=user_id or None, - thread_ts=thread_key or None, + thread_ts=thread_ts or None, ) try: @@ -536,7 +856,14 @@ class TurnstoneSlackBot: log.exception("slack.dm_dispatch_failed", user=user_id) await say(text="Sorry, something went wrong routing your message.") - async def _on_approve(self, ack: Any, body: dict[str, Any]) -> None: + async def _resolve_approval( + self, + ack: Any, + body: dict[str, Any], + *, + approved: bool, + ) -> None: + """Handle an approve or deny button click from an approval prompt.""" await ack() value = body["actions"][0].get("value", "") parts = value.split("|", 1) @@ -547,6 +874,7 @@ class TurnstoneSlackBot: entry = self._pending_approval.get(ws_id) actor_user_id = body.get("user", {}).get("id", "") channel = body["container"]["channel_id"] + verb = "approve" if approved else "deny" log.info( "slack.approval_actor_check", @@ -573,76 +901,60 @@ class TurnstoneSlackBot: await self._client.chat_postEphemeral( channel=channel, user=actor_user_id, - text="Only the session owner can approve this tool call.", + text=f"Only the session owner can {verb} this tool call.", ) return - await self.router.send_approval(ws_id, correlation_id, approved=True) + await self.router.send_approval(ws_id, correlation_id, approved=approved) + # Drop the pending entry now that we've handled it locally. Otherwise + # the subsequent ApprovalResolvedEvent will rewrite the message a + # second time ("Tool approved" → "Approved") — wasted chat_update and + # a visible edit flicker. + self._pending_approval.pop(ws_id, None) ts = body["container"]["message_ts"] + update_text = "Tool approved" if approved else "Tool denied" + event_key = "approve" if approved else "deny" try: await self._client.chat_update( channel=channel, ts=ts, - text="Tool approved", + text=update_text, blocks=[], ) except Exception: - log.debug("slack.approve_message_update_failed", exc_info=True) + log.debug(f"slack.{event_key}_message_update_failed", exc_info=True) - async def _on_deny(self, ack: Any, body: dict[str, Any]) -> None: - await ack() - value = body["actions"][0].get("value", "") - parts = value.split("|", 1) - if len(parts) != 2: - return + async def _ensure_plan_review_owner( + self, + entry: tuple[str, str, str] | None, + actor_user_id: str, + channel: str, + verb: str, + ) -> bool: + """Return True when *actor_user_id* owns the pending plan review. - ws_id, correlation_id = parts - entry = self._pending_approval.get(ws_id) - actor_user_id = body.get("user", {}).get("id", "") - channel = body["container"]["channel_id"] - - log.info( - "slack.approval_actor_check", - ws_id=ws_id, - actor_user_id=actor_user_id, - owner_user_id=entry.owner_user_id if entry else None, - has_entry=entry is not None, - ) - - if entry is None or not entry.owner_user_id: - log.warning( - "slack.approval_missing_owner", - ws_id=ws_id, - actor_user_id=actor_user_id, - ) + The gateway forwards ``/plan`` calls with its service-scoped JWT, + and the server bypasses ownership checks on service scope — so + every plan-review interaction needs an adapter-side owner gate. + """ + if entry is None: + log.warning("slack.plan_review_missing_entry", actor_user_id=actor_user_id) await self._client.chat_postEphemeral( channel=channel, user=actor_user_id, - text="This approval can no longer be verified. Please retry from the active session.", + text="This plan review can no longer be verified. Please retry from the active session.", ) - return - - if actor_user_id != entry.owner_user_id: + return False + _channel, _ts, owner_user_id = entry + if not owner_user_id or actor_user_id != owner_user_id: await self._client.chat_postEphemeral( channel=channel, user=actor_user_id, - text="Only the session owner can deny this tool call.", + text=f"Only the session owner can {verb} this plan.", ) - return - - await self.router.send_approval(ws_id, correlation_id, approved=False) - ts = body["container"]["message_ts"] - - try: - await self._client.chat_update( - channel=channel, - ts=ts, - text="Tool denied", - blocks=[], - ) - except Exception: - log.debug("slack.deny_message_update_failed", exc_info=True) + return False + return True async def _on_plan_approve(self, ack: Any, body: dict[str, Any]) -> None: await ack() @@ -651,11 +963,16 @@ class TurnstoneSlackBot: if not ws_id: return + actor_user_id = body.get("user", {}).get("id", "") + channel = body["container"]["channel_id"] + entry = self._pending_plan_review_ts.get(ws_id) + if not await self._ensure_plan_review_owner(entry, actor_user_id, channel, "approve"): + return + self._streaming.pop(ws_id, None) await self.router.send_plan_feedback(ws_id, "", "") log.info("slack.plan_feedback_sent", ws_id=ws_id, feedback="") - channel = body["container"]["channel_id"] ts = body["container"]["message_ts"] self._pending_plan_review_ts.pop(ws_id, None) @@ -676,6 +993,14 @@ class TurnstoneSlackBot: if not ws_id: return + actor_user_id = body.get("user", {}).get("id", "") + channel = body["container"]["channel_id"] + entry = self._pending_plan_review_ts.get(ws_id) + if not await self._ensure_plan_review_owner( + entry, actor_user_id, channel, "request changes on" + ): + return + trigger_id = body.get("trigger_id", "") if not trigger_id: return @@ -713,6 +1038,22 @@ class TurnstoneSlackBot: if not ws_id: return + actor_user_id = body.get("user", {}).get("id", "") + entry = self._pending_plan_review_ts.get(ws_id) + if entry is None: + # No pending review (stale modal) — silently drop. + return + _ignored_channel, _ignored_ts, owner_user_id = entry + if not owner_user_id or actor_user_id != owner_user_id: + # Modal submit can't emit ephemeral; just log and drop. + log.warning( + "slack.plan_feedback_non_owner", + ws_id=ws_id, + actor_user_id=actor_user_id, + owner_user_id=owner_user_id, + ) + return + feedback = ( view.get("state", {}) .get("values", {}) @@ -733,7 +1074,7 @@ class TurnstoneSlackBot: entry = self._pending_plan_review_ts.pop(ws_id, None) if entry is not None: - channel, ts = entry + channel, ts, _owner = entry try: await self._client.chat_update( channel=channel, @@ -745,6 +1086,16 @@ class TurnstoneSlackBot: log.debug("slack.plan_review_modal_update_failed", exc_info=True) async def subscribe_ws(self, ws_id: str, channel_id: str) -> None: + # Recover from a dead SSE task before the early-return check. If the + # prior listener died with an unhandled exception, its ws_id is still + # in _subscribed_ws — without this purge subscribe_ws would silently + # no-op forever. + prior = self._sse_tasks.get(ws_id) + if prior is not None and prior.done(): + log.info("slack.sse_task_recovered", ws_id=ws_id) + self._sse_tasks.pop(ws_id, None) + self._subscribed_ws.discard(ws_id) + if ws_id in self._subscribed_ws: return @@ -756,105 +1107,59 @@ class TurnstoneSlackBot: self._subscribed_ws.add(ws_id) log.info("slack.subscribed", ws_id=ws_id, channel_id=channel_id) - async def unsubscribe_ws(self, ws_id: str) -> None: - task = self._sse_tasks.pop(ws_id, None) - if task is not None: - task.cancel() - with contextlib.suppress(asyncio.CancelledError, Exception): - await task + def _clear_ws_state(self, ws_id: str) -> None: + """Drop all in-memory state keyed by *ws_id*. + Does not cancel or await the SSE task — callers handle task + lifecycle differently (``unsubscribe_ws`` cancels and awaits; + ``_cleanup_stale_route`` is itself invoked from inside the task). + """ self._subscribed_ws.discard(ws_id) self._streaming.pop(ws_id, None) self._pending_approval.pop(ws_id, None) self._pending_plan_review_ts.pop(ws_id, None) self._clear_notification_tracking_for_ws(ws_id) + async def unsubscribe_ws(self, ws_id: str) -> None: + task = self._sse_tasks.pop(ws_id, None) + if task is not None: + task.cancel() + # Await the cancelled task so CancelledError propagates out of + # the SSE loop before we clear the per-ws state below. + # CancelledError is expected; other exceptions from the SSE + # loop are already logged there and must not block shutdown. + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + + self._clear_ws_state(ws_id) + log.info("slack.unsubscribed", ws_id=ws_id) async def _sse_listener(self, ws_id: str, channel_id: str) -> None: - """Connect to the server SSE endpoint and dispatch events.""" - import httpx_sse + """Connect to the server SSE endpoint and dispatch events. - delay = _SSE_RECONNECT_DELAY - url = "" + Delegates the reconnect/backoff loop to :func:`run_sse_stream`. + """ - while True: - try: - node_base = await self.router.get_node_url(ws_id) - url = f"{node_base}/v1/api/events" + async def _on_event(event: ServerEvent) -> None: + effective_route = self._notify_reply_routes.get( + ws_id, + SlackRoute.parse(channel_id), + ) + await self._on_ws_event(ws_id, effective_route, event) - sse_headers: dict[str, str] | None = None - if self._token_factory is not None: - sse_headers = {"Authorization": f"Bearer {self._token_factory()}"} + async def _on_stale() -> None: + await self._cleanup_stale_route(ws_id, channel_id) - async with httpx_sse.aconnect_sse( - self._http_client, - "GET", - url, - params={"ws_id": ws_id}, - headers=sse_headers, - ) as event_source: - status = event_source.response.status_code - if status == 404: - log.info("slack.sse_ws_gone", ws_id=ws_id) - await self._cleanup_stale_route(ws_id, channel_id) - return - - if status >= 400: - log.warning( - "slack.sse_upstream_error", - ws_id=ws_id, - status=status, - ) - raise httpx.HTTPStatusError( - f"SSE upstream {status}", - request=event_source.response.request, - response=event_source.response, - ) - - delay = _SSE_RECONNECT_DELAY - async for sse in event_source.aiter_sse(): - if sse.event == "message" or not sse.event: - try: - data = json.loads(sse.data) - except json.JSONDecodeError: - log.debug("slack.sse_invalid_json", ws_id=ws_id, exc_info=True) - continue - - event = ServerEvent.from_dict(data) - try: - effective_route = self._notify_reply_routes.get( - ws_id, - SlackRoute.parse(channel_id), - ) - await self._on_ws_event(ws_id, effective_route, event) - except Exception: - log.warning( - "slack.event_dispatch_failed", - ws_id=ws_id, - exc_info=True, - ) - - except httpx.HTTPStatusError: - pass - except httpx.RemoteProtocolError: - log.debug("slack.sse_remote_closed", ws_id=ws_id) - except asyncio.CancelledError: - return - except httpx.ReadTimeout: - log.info("slack.sse_read_timeout", ws_id=ws_id) - except (httpx.ConnectError, httpx.ConnectTimeout) as exc: - log.warning( - "slack.sse_connect_failed", - ws_id=ws_id, - url=url, - error=str(exc), - ) - except Exception: - log.warning("slack.sse_error", ws_id=ws_id, exc_info=True) - - await asyncio.sleep(delay) - delay = min(delay * 2, _SSE_MAX_RECONNECT_DELAY) + await run_sse_stream( + http_client=self._http_client, + log_prefix="slack", + ws_id=ws_id, + node_url_fn=self.router.get_node_url, + token_factory=self._token_factory, + on_event=_on_event, + on_stale=_on_stale, + ) async def _cleanup_stale_route(self, ws_id: str, channel_id: str) -> None: """Remove a channel route whose workstream no longer exists.""" @@ -863,12 +1168,9 @@ class TurnstoneSlackBot: self._channel_sessions.pop((route.channel, route.user_id), None) await self.router.delete_route("slack", channel_id) - self._subscribed_ws.discard(ws_id) + # Called from inside the SSE task itself — don't await the task here. self._sse_tasks.pop(ws_id, None) - self._streaming.pop(ws_id, None) - self._pending_approval.pop(ws_id, None) - self._pending_plan_review_ts.pop(ws_id, None) - self._clear_notification_tracking_for_ws(ws_id) + self._clear_ws_state(ws_id) log.info("slack.stale_route_removed", ws_id=ws_id) @@ -878,219 +1180,225 @@ class TurnstoneSlackBot: route: SlackRoute, event: ServerEvent, ) -> None: - """Handle a typed server event for a subscribed workstream.""" + """Dispatch a typed server event to its per-event handler.""" + if isinstance(event, (ThinkingStartEvent, ThinkingStopEvent)): + return + + if isinstance(event, ContentEvent): + await self._handle_content(ws_id, route, event) + elif isinstance(event, ApproveRequestEvent): + await self._handle_approve_request(ws_id, route, event) + elif isinstance(event, IntentVerdictEvent): + await self._handle_intent_verdict(ws_id, event) + elif isinstance(event, PlanReviewEvent): + await self._handle_plan_review(ws_id, route, event) + elif isinstance(event, ApprovalResolvedEvent): + await self._handle_approval_resolved(ws_id, event) + elif isinstance(event, StreamEndEvent): + await self._handle_stream_end(ws_id) + elif isinstance(event, ErrorEvent): + await self._handle_error(route, event) + + async def _handle_content(self, ws_id: str, route: SlackRoute, event: ContentEvent) -> None: + slack_channel = route.channel + thread_ts = route.thread_ts or "" + sm = self._streaming.get(ws_id) + if sm is None or sm.channel != slack_channel or sm.thread_ts != thread_ts: + # Finalize the outgoing StreamingMessage before swapping so any + # already-buffered tokens still land on the old thread instead + # of being dropped (happens when the effective route flips + # mid-stream — e.g. a user replies to a notification and pins + # the conversation to a different thread). + if sm is not None: + await sm.finalize() + sm = StreamingMessage( + client=self._client, + channel=slack_channel, + thread_ts=thread_ts, + max_length=self.config.max_message_length, + edit_interval=self.config.streaming_edit_interval, + ) + self._streaming[ws_id] = sm + await sm.append(event.text) + + async def _handle_approve_request( + self, + ws_id: str, + route: SlackRoute, + event: ApproveRequestEvent, + ) -> None: slack_channel = route.channel thread_ts = route.thread_ts or "" owner_user_id = route.user_id - if isinstance(event, (ThinkingStartEvent, ThinkingStopEvent)): - pass - - elif isinstance(event, ContentEvent): - sm = self._streaming.get(ws_id) - if sm is None or sm.channel != slack_channel or sm.thread_ts != thread_ts: - sm = StreamingMessage( - client=self._client, - channel=slack_channel, - thread_ts=thread_ts, - max_length=self.config.max_message_length, - edit_interval=self.config.streaming_edit_interval, - ) - self._streaming[ws_id] = sm - - await sm.append(event.text) - - elif isinstance(event, ApproveRequestEvent): - _policy_handled = False - if self.storage is not None: - try: - from turnstone.core.policy import evaluate_tool_policies_batch - - _tool_names = [ - it.get("approval_label", "") or it.get("func_name", "") - for it in event.items - if it.get("needs_approval") and it.get("func_name") and not it.get("error") - ] - _tool_names = [n for n in _tool_names if n] - if _tool_names: - verdicts = await asyncio.to_thread( - evaluate_tool_policies_batch, - self.storage, - _tool_names, - ) - if any(v == "deny" for v in verdicts.values()): - denied = [n for n, v in verdicts.items() if v == "deny"] - await self.router.send_approval( - ws_id, - "", - approved=False, - feedback=f"Blocked by tool policy: {', '.join(denied)}", - ) - await self._client.chat_postMessage( - channel=slack_channel, - thread_ts=thread_ts or None, - text=f"_Tool blocked by admin policy: {', '.join(denied)}_", - ) - _policy_handled = True - elif all(verdicts.get(n) == "allow" for n in _tool_names): - await self.router.send_approval(ws_id, "", approved=True) - await self._client.chat_postMessage( - channel=slack_channel, - thread_ts=thread_ts or None, - text="_Tool approved by policy._", - ) - _policy_handled = True - except Exception: - log.debug( - "Tool policy evaluation failed for ws %s", - ws_id, - exc_info=True, - ) - - if not _policy_handled and ( - self.config.auto_approve or self._should_auto_approve(event) - ): - await self.router.send_approval(ws_id, "", approved=True) - await self._client.chat_postMessage( - channel=slack_channel, - thread_ts=thread_ts or None, - text="_Tool auto-approved._", - ) - elif not _policy_handled: - await self._send_approval_request( - ws_id, - "", - event.items, - slack_channel, - thread_ts, - owner_user_id, - ) - - elif isinstance(event, IntentVerdictEvent): - entry = self._pending_approval.get(ws_id) - if entry is not None: - pending_channel = entry.channel - pending_ts = entry.message_ts - risk = (event.risk_level or "medium").upper() - verdict_text = ( - f"*Judge Verdict: {event.func_name or 'tool'}*\n" - f"Risk: {risk} | Confidence: {event.confidence or 'N/A'}\n" - f"_{event.intent_summary or ''}_" - ) - try: - result = await self._client.conversations_history( - channel=pending_channel, - latest=pending_ts, - limit=1, - inclusive=True, - ) - existing_blocks = [] - if result.get("ok") and result.get("messages"): - existing_blocks = result["messages"][0].get("blocks", []) - existing_blocks.append( - { - "type": "section", - "text": { - "type": "mrkdwn", - "text": verdict_text, - }, - } - ) - await self._client.chat_update( - channel=pending_channel, - ts=pending_ts, - blocks=existing_blocks, - text="Tool approval required", - ) - except Exception: - log.debug("slack.verdict_message_update_failed", ws_id=ws_id, exc_info=True) - - elif isinstance(event, PlanReviewEvent): - log.info("slack.plan_review_received", ws_id=ws_id) - plan_preview = _sanitize_slack_preview(event.content, max_length=2000) - blocks = [ - { - "type": "section", - "text": { - "type": "mrkdwn", - "text": f"*Plan Review*\n```{plan_preview}```", - }, - }, - { - "type": "actions", - "elements": [ - { - "type": "button", - "text": {"type": "plain_text", "text": "Approve"}, - "style": "primary", - "action_id": "ts_plan_approve", - "value": ws_id, - }, - { - "type": "button", - "text": {"type": "plain_text", "text": "Request changes"}, - "style": "danger", - "action_id": "ts_plan_request_changes", - "value": ws_id, - }, - ], - }, - ] - - resp = await self._client.chat_postMessage( - channel=slack_channel, - thread_ts=thread_ts or None, - text="Plan review required", - blocks=cast("list[dict[str, Any]]", blocks), + verdict = await self.router.evaluate_tool_policies(event.items) + policy_handled = False + if verdict.kind == "deny": + denied = ", ".join(verdict.denied_tools) + await self.router.send_approval( + ws_id, + "", + approved=False, + feedback=f"Blocked by tool policy: {denied}", ) - if resp.get("ok"): - self._pending_plan_review_ts[ws_id] = (slack_channel, resp["ts"]) - - elif isinstance(event, ApprovalResolvedEvent): - entry = self._pending_approval.pop(ws_id, None) - if entry is not None: - pending_channel = entry.channel - pending_ts = entry.message_ts - label = "Approved" if event.approved else "Denied" - try: - await self._client.chat_update( - channel=pending_channel, - ts=pending_ts, - text=label, - blocks=cast("list[dict[str, Any]]", []), - ) - except Exception: - log.debug("slack.approval_resolved_edit_failed", ws_id=ws_id, exc_info=True) - - elif isinstance(event, StreamEndEvent): - sm = self._streaming.pop(ws_id, None) - - if sm is not None: - await sm.finalize() - - # Pop the notification-reply override so subsequent turns on - # this ws default back to the session route. Without the pop, - # one notification reply pins all future responses to the - # notification thread until the bot restarts. - reply_route = self._notify_reply_routes.pop(ws_id, None) - if ( - reply_route is not None - and sm is not None - and sm._ts - and reply_route.channel - and reply_route.user_id - ): - self._track_notification(sm._ts, ws_id, reply_route) - - self._pending_approval.pop(ws_id, None) - - elif isinstance(event, ErrorEvent): - safe_msg = event.message[:500] if event.message else "An error occurred" await self._client.chat_postMessage( channel=slack_channel, thread_ts=thread_ts or None, - text=f"*Error:* {safe_msg}", + text=f"_Tool blocked by admin policy: {denied}_", ) + policy_handled = True + elif verdict.kind == "allow": + await self.router.send_approval(ws_id, "", approved=True) + await self._client.chat_postMessage( + channel=slack_channel, + thread_ts=thread_ts or None, + text="_Tool approved by policy._", + ) + policy_handled = True + + if not policy_handled and (self.config.auto_approve or self._should_auto_approve(event)): + await self.router.send_approval(ws_id, "", approved=True) + await self._client.chat_postMessage( + channel=slack_channel, + thread_ts=thread_ts or None, + text="_Tool auto-approved._", + ) + elif not policy_handled: + await self._send_approval_request( + ws_id, + "", + event.items, + slack_channel, + thread_ts, + owner_user_id, + ) + + async def _handle_intent_verdict(self, ws_id: str, event: IntentVerdictEvent) -> None: + entry = self._pending_approval.get(ws_id) + if entry is None: + return + + pending_channel = entry.channel + pending_ts = entry.message_ts + risk = (event.risk_level or "medium").upper() + verdict_text = ( + f"*Judge Verdict: {event.func_name or 'tool'}*\n" + f"Risk: {risk} | Confidence: {event.confidence or 'N/A'}\n" + f"_{event.intent_summary or ''}_" + ) + # Append the verdict section in-place on the cached blocks so + # repeat IntentVerdictEvents stack on the same approval message + # (matches the pre-refactor "fetch live blocks and append" + # behavior, but without the extra conversations_history call). + entry.blocks.append( + { + "type": "section", + "text": {"type": "mrkdwn", "text": verdict_text}, + } + ) + try: + await self._client.chat_update( + channel=pending_channel, + ts=pending_ts, + blocks=entry.blocks, + text="Tool approval required", + ) + except Exception: + log.debug("slack.verdict_message_update_failed", ws_id=ws_id, exc_info=True) + + async def _handle_plan_review( + self, ws_id: str, route: SlackRoute, event: PlanReviewEvent + ) -> None: + slack_channel = route.channel + thread_ts = route.thread_ts or "" + owner_user_id = route.user_id or "" + + log.info("slack.plan_review_received", ws_id=ws_id) + plan_preview = _sanitize_slack_preview(event.content, max_length=2000) + blocks = [ + { + "type": "section", + "text": { + "type": "mrkdwn", + "text": f"*Plan Review*\n```{plan_preview}```", + }, + }, + { + "type": "actions", + "elements": [ + { + "type": "button", + "text": {"type": "plain_text", "text": "Approve"}, + "style": "primary", + "action_id": "ts_plan_approve", + "value": ws_id, + }, + { + "type": "button", + "text": {"type": "plain_text", "text": "Request changes"}, + "style": "danger", + "action_id": "ts_plan_request_changes", + "value": ws_id, + }, + ], + }, + ] + + resp = await self._client.chat_postMessage( + channel=slack_channel, + thread_ts=thread_ts or None, + text="Plan review required", + blocks=cast("list[dict[str, Any]]", blocks), + ) + if resp.get("ok"): + self._pending_plan_review_ts[ws_id] = (slack_channel, resp["ts"], owner_user_id) + + async def _handle_approval_resolved(self, ws_id: str, event: ApprovalResolvedEvent) -> None: + entry = self._pending_approval.pop(ws_id, None) + if entry is None: + return + + pending_channel = entry.channel + pending_ts = entry.message_ts + label = "Approved" if event.approved else "Denied" + try: + await self._client.chat_update( + channel=pending_channel, + ts=pending_ts, + text=label, + blocks=cast("list[dict[str, Any]]", []), + ) + except Exception: + log.debug("slack.approval_resolved_edit_failed", ws_id=ws_id, exc_info=True) + + async def _handle_stream_end(self, ws_id: str) -> None: + sm = self._streaming.pop(ws_id, None) + if sm is not None: + await sm.finalize() + + # Pop the notification-reply override so subsequent turns on this ws + # default back to the session route. Without the pop, one notification + # reply pins all future responses to the notification thread until the + # bot restarts. + reply_route = self._notify_reply_routes.pop(ws_id, None) + if ( + reply_route is not None + and sm is not None + and sm.message_ts + and reply_route.channel + and reply_route.user_id + ): + self._track_notification(sm.message_ts, ws_id, reply_route) + + self._pending_approval.pop(ws_id, None) + + async def _handle_error(self, route: SlackRoute, event: ErrorEvent) -> None: + safe_msg = event.message[:500] if event.message else "An error occurred" + await self._client.chat_postMessage( + channel=route.channel, + thread_ts=route.thread_ts or None, + text=f"*Error:* {safe_msg}", + ) async def _send_approval_request( self, @@ -1166,6 +1474,7 @@ class TurnstoneSlackBot: channel=channel, message_ts=resp["ts"], owner_user_id=owner_user_id, + blocks=list(cast("list[dict[str, Any]]", blocks)), ) log.info( "slack.pending_approval_stored", diff --git a/turnstone/channels/slack/routes.py b/turnstone/channels/slack/routes.py index ec68b750..0f429df4 100644 --- a/turnstone/channels/slack/routes.py +++ b/turnstone/channels/slack/routes.py @@ -1,3 +1,11 @@ +"""Slack channel-routing key dataclass. + +A :class:`SlackRoute` is the triple that uniquely identifies where a Slack +conversation lives: ``(channel, user_id?, thread_ts?)``. It round-trips +through a ``channel:user_id:thread_ts`` string so it can be used as the +opaque ``channel_id`` value stored in ``channel_routes``. +""" + from __future__ import annotations from dataclasses import dataclass @@ -5,6 +13,24 @@ from dataclasses import dataclass @dataclass(frozen=True) class SlackRoute: + """Routing key for a Slack conversation (channel / DM / thread). + + ``to_channel_id`` and ``parse`` are inverses only for values the + parser is willing to emit: + + - ``SlackRoute("C1")`` ↔ ``"C1"`` + - ``SlackRoute("C1", "U1")`` ↔ ``"C1:U1"`` + - ``SlackRoute("C1", "U1", "ts")`` ↔ ``"C1:U1:ts"`` + + Lax cases (not produced by ``to_channel_id`` but accepted by + ``parse``): + + - Trailing colons drop to ``None`` — ``"C1:"`` → ``SlackRoute("C1")``. + - Extra colons fold into ``thread_ts`` via ``split(":", 2)``, so + ``"C1:U1:ts:extra"`` → ``thread_ts="ts:extra"``. No Slack ts or + channel/user ID contains ``:`` so this is safe in practice. + """ + channel: str user_id: str | None = None thread_ts: str | None = None