mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
refactor(channels): backfill review of Slack/Discord adapters (#382)
* refactor(channels): backfill review of Slack/Discord adapters Retrospective multi-stage review of the Slack (PR #355) and Discord channel adapters — they shipped before the review pipeline existed, so this pass goes back and fixes everything the pipeline would have caught plus a follow-up round of ultrareview findings. ## Security (8 fixes) - Adapter-side owner checks on all interactive flows: Discord ApprovalView / PlanReviewView encode the owner Discord user ID in the embed footer (`{ws_id}|{corr_id}|{owner_id}`) and reject non-owner clicks; Slack plan-approve / request-changes / feedback-modal gain owner tracking in `_pending_plan_review_ts` and a shared `_ensure_plan_review_owner` gate. These closed the two critical authz gaps where the gateway's service-scoped JWT bypassed server-side ownership checks. - Discord thread-message gate: only the registered invoker can drive the workstream (prevents a linked user posting in another user's public thread from injecting into their assistant). Invoker recorded explicitly so `/ask` follow-ups survive the `channel.create_thread` bot-as-owner quirk. - Slack /link flow + per-user identity gate: unlinked Slack users see an ephemeral `/turnstone link <token>` prompt on every message instead of silently creating workstreams under the shared gateway identity. Rate-limited (5/hour) to block online token enumeration. - Gateway `/v1/api/notify` requires `write` scope on the validated JWT; low-scope tokens get 403 + audit. - Thumbnail URL validator DNS-resolves the hostname before fetch and rejects any resolved IP that's loopback / link-local / multicast / reserved, plus an explicit deny-list for IPv6 cloud metadata (`fd00:ec2::/32` — AWS Nitro IMDS + ECS task metadata) that would otherwise slip past the `is_private` allowance. - Per-user rate limit (10 msgs / 60s) + 8 KiB inbound size cap on Slack DMs / channels / notification-reply threads so one user can't exhaust the shared LLM budget. - Discord /link rate limit (5/hour) for token-enumeration defense. ## Bug fixes (9 correctness issues) - Slack DM routing: each top-level DM no longer spawns a fresh workstream (was using per-message `ts` as the route key). - Multi-chunk Slack responses thread correctly under the first chunk's ts instead of fragmenting as independent top-level messages. - Finalize the outgoing StreamingMessage before swapping channel / thread_ts mid-stream, so buffered tokens still land on the old thread. - Redundant `chat_update` on approve/deny eliminated by popping `_pending_approval[ws_id]` after local resolution. - Notification reply tracking on Discord only registers for DMs (guild-channel targets were storing channel IDs where user IDs were expected, so legitimate replies were always rejected). - `get_channel_default_alias` rolls `_channel_default_ts` back on `list_models()` failure so the next caller retries instead of serving an empty alias for the full TTL. - Slack `subscribe_ws` purges dead SSE tasks before the membership short-circuit (previously an unhandled exception left the ws_id in `_subscribed_ws` forever, silently no-opping subsequent subscribes). - ChannelRouter `_create_locks` is now an LRU-bounded OrderedDict that evicts only unheld locks (original dict grew unbounded; naive LRU could evict a held lock and let a second caller race through the critical section, creating duplicate workstreams). - Slack `_parse_ts` pads the fractional field to 6 digits so `"1.2"` and `"1.000002"` stop colliding as `(1, 2)` in the latest-session tiebreaker. ## Performance (6 fixes) - StreamingMessage keeps a rolling truncated display string capped at `max_length` so per-flush cost is O(max_length) instead of O(total_streamed_chars) — long streaming responses no longer do quadratic work every edit interval. - `StreamingMessage.finalize()` caches the joined content so the Discord stream-end DM-forward path doesn't re-join a multi-MB buffer twice. - `PendingApproval` stores the Block Kit payload posted to Slack; `IntentVerdictEvent` appends the verdict in-place and `chat_update`s, skipping an extra `conversations_history` round-trip. - ChannelRouter `lookup_ws_id()` TTL-caches the channel → ws_id resolution (30s TTL, 4096-entry LRU); hot inbound paths skip storage on every message. - Service-discovery startup retry uses exponential backoff (1s → 8s cap) with a 30s deadline instead of 30 × 1s fixed sleep. - `_archive_session` now calls `router.close_workstream` so the `_node_urls` cache entry is dropped (was leaking one entry per archived session). ## Quality / refactors (19 improvements) - `cli.main()` extracted from a 365-line function into focused helpers; imports carefully kept lazy where test patches target source-module paths. - `_run_gateway` finally block now awaits `adapter.stop()` on every adapter so SSE tasks, httpx clients, and the Slack socket handler close cleanly on shutdown. - Shared SSE reconnect loop extracted to `turnstone/channels/_sse.py` (`run_sse_stream` with `on_event` + `on_stale` callbacks); both adapters' `_sse_listener` methods just wire up callbacks. The "404 stops reconnect" invariant is enforced inside the helper so a broken `on_stale` can't livelock. - `_on_ws_event` god-dispatchers split into per-event `_handle_*` methods with a thin isinstance dispatcher at the top. - Slack `_on_approve` / `_on_deny` collapsed into a single `_resolve_approval(*, approved: bool)`. - `ApproveRequestEvent` policy evaluation hoisted into `ChannelRouter.evaluate_tool_policies` returning a `PolicyVerdict`; adapters switch on the verdict kind. - `ChannelAdapter` protocol trimmed to the four methods adapters actually implement; unused `ChannelEvent` dataclass removed. - Shared constants lifted to `turnstone/channels/_config.py`. - `_cleanup_stale_route` and `unsubscribe_ws` share a `_clear_ws_state` helper. - `StreamingMessage` private attrs promoted to `message` / `message_ts` / `accumulated_text` properties so callers don't reach past the `_`-prefix. - Various cleanups: dead var, noqa'd lambdas, renamed `_policy_handled` → `policy_handled`, inlined single-use helpers, added module docstrings, documented `SlackRoute.parse` edge cases. - `chunk_message` plain-text fast path (no backticks → skip fence bookkeeping). ## Test coverage Added 45 tests (178 → 223): - `tests/test_channel_sse.py` (new) — SSE reconnect / backoff / 404-stale-route / on-stale-exception / invalid-JSON-skip / on-event-exception-doesn't-kill-stream / per-connection token refresh / ConnectError retry. - ApprovalView + PlanReviewView owner-check regression tests (owner allowed, non-owner rejected, legacy 2-pipe footer fails closed, modal path rejected for non-owner, `/ask` bot-as-thread-owner follow-up allowed). - Slack `_recover_routes` latest-ts-wins, `_archive_session` drops route + closes workstream. - SSRF tests: DNS rebinding rejected, IPv4 link-local metadata rejected, IPv6 ULA metadata (fd00:ec2::254 / fd00:ec2::23) rejected. - Slack link prefix match (natural-language prompts don't hijack), link rate-limit ceiling. - SlackRoute round-trip across all three shapes + lax-parse behaviour. Lint (ruff) + mypy clean; 210 channel-focused tests pass. * chore(channels): address PR #382 review-bot feedback Three line-level findings from github-code-quality on the backfill review PR. Copilot had no line-level comments. - _sse.py:132 — the `except httpx.HTTPStatusError: pass` branch was flagged as an empty except. The original status was already logged at WARNING inside the try block (we re-raise ourselves after logging), so the handler has real intent. Added a debug log of the exception text + a comment explaining the control flow, so the empty-except lint stops firing and the next reader sees why we fall through to backoff. - discord/bot.py:430, cli.py:354, slack/bot.py:1127 — `await task` inside `contextlib.suppress` was flagged as "statement has no effect". It's a false positive (await is an effect) and the alternative try/except/pass triggers ruff SIM105. Kept the contextlib.suppress pattern and added an explanatory comment above each call so the intent (await CancelledError propagation before state cleanup) is obvious; will reply on the PR thread noting the false positive. No behavior change. Lint + mypy clean; 210 channel tests pass.
This commit is contained in:
Patrick Buckley
@@ -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)
|
||||
|
||||
+6
-6
@@ -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/<platform>/` package
|
||||
|
||||
+380
-46
@@ -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()
|
||||
|
||||
@@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
+184
-9
@@ -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 []
|
||||
|
||||
@@ -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]
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -25,8 +25,27 @@ def chunk_message(text: str, max_length: int = 2000) -> list[str]:
|
||||
if len(text) <= max_length:
|
||||
return [text]
|
||||
|
||||
# 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
|
||||
|
||||
while remaining:
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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."""
|
||||
...
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
+195
-139
@@ -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
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
# -- Auth config ---------------------------------------------------------
|
||||
jwt_secret = os.environ.get("TURNSTONE_JWT_SECRET", "").strip()
|
||||
|
||||
# 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
|
||||
|
||||
_scopes = frozenset({"read", "write", "approve", "service"})
|
||||
_console_mgr = ServiceTokenManager(
|
||||
scopes = frozenset({"read", "write", "approve", "service"})
|
||||
console_mgr = ServiceTokenManager(
|
||||
user_id="channel-gateway",
|
||||
scopes=_scopes,
|
||||
scopes=scopes,
|
||||
source="channel",
|
||||
secret=jwt_secret,
|
||||
audience=JWT_AUD_CONSOLE,
|
||||
expiry_hours=1,
|
||||
)
|
||||
_server_mgr = ServiceTokenManager(
|
||||
server_mgr = ServiceTokenManager(
|
||||
user_id="channel-gateway",
|
||||
scopes=_scopes,
|
||||
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
|
||||
|
||||
server_url: str = args.server_url
|
||||
console_url: str = args.console_url
|
||||
def console_factory() -> str:
|
||||
return console_mgr.token
|
||||
|
||||
# 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 server_factory() -> str:
|
||||
return server_mgr.token
|
||||
|
||||
return console_factory, server_factory
|
||||
|
||||
|
||||
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.
|
||||
|
||||
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
|
||||
|
||||
try:
|
||||
from turnstone.core.storage._registry import get_storage as _get_st
|
||||
|
||||
_disc_storage = _get_st()
|
||||
log.info("channel.discovering_services")
|
||||
for _attempt in range(30): # up to 30s
|
||||
deadline = time.monotonic() + _DISCOVERY_BUDGET_S
|
||||
delay = _DISCOVERY_INITIAL_DELAY_S
|
||||
while True:
|
||||
if not console_url:
|
||||
consoles = _disc_storage.list_services("console", max_age_seconds=3600)
|
||||
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 = _disc_storage.list_services("server", max_age_seconds=120)
|
||||
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
|
||||
_time.sleep(1)
|
||||
else:
|
||||
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
log.warning(
|
||||
"channel.discovery_timeout",
|
||||
console_url=console_url,
|
||||
server_url=server_url,
|
||||
)
|
||||
break
|
||||
|
||||
time.sleep(min(delay, remaining))
|
||||
delay = min(delay * 2, _DISCOVERY_MAX_DELAY_S)
|
||||
except Exception:
|
||||
log.warning("channel.discovery_failed", exc_info=True)
|
||||
|
||||
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)
|
||||
return console_url, server_url
|
||||
|
||||
# -- 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)
|
||||
|
||||
# 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,52 +274,27 @@ 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
|
||||
|
||||
log.info(
|
||||
"channel.starting",
|
||||
adapters=list(adapters.keys()),
|
||||
http_port=args.http_port,
|
||||
server_url=server_url,
|
||||
)
|
||||
|
||||
async def _run_all() -> None:
|
||||
"""Run all adapters + HTTP server + service heartbeat concurrently."""
|
||||
import uvicorn
|
||||
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
|
||||
|
||||
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
|
||||
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"
|
||||
advertise_url = f"{scheme}://{advertise_host}:{args.http_port}"
|
||||
service_url = advertise_url
|
||||
return f"{scheme}://{advertise_host}:{args.http_port}"
|
||||
|
||||
# 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."""
|
||||
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:
|
||||
@@ -341,11 +306,25 @@ def main() -> None:
|
||||
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):
|
||||
|
||||
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,
|
||||
@@ -357,13 +336,13 @@ def main() -> None:
|
||||
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,
|
||||
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())
|
||||
heartbeat_task = asyncio.create_task(_heartbeat_loop(storage, service_id))
|
||||
try:
|
||||
await asyncio.gather(
|
||||
*(adapter.start() for adapter in adapters.values()),
|
||||
@@ -371,14 +350,91 @@ def main() -> None:
|
||||
)
|
||||
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",
|
||||
adapters=list(adapters.keys()),
|
||||
http_port=args.http_port,
|
||||
server_url=server_url,
|
||||
)
|
||||
|
||||
with contextlib.suppress(KeyboardInterrupt):
|
||||
asyncio.run(_run_all())
|
||||
asyncio.run(_run_gateway(adapters, channel_app, storage, args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+264
-207
@@ -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
|
||||
|
||||
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:
|
||||
async def _on_event(event: ServerEvent) -> None:
|
||||
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)
|
||||
|
||||
# Exponential backoff before reconnecting.
|
||||
await asyncio.sleep(delay)
|
||||
delay = min(delay * 2, _SSE_MAX_RECONNECT_DELAY)
|
||||
async def _on_stale() -> None:
|
||||
await self._cleanup_stale_route(ws_id)
|
||||
|
||||
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,17 +481,35 @@ 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):
|
||||
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):
|
||||
await self._handle_content(ws_id, thread, event)
|
||||
elif isinstance(event, ToolInfoEvent):
|
||||
await self._handle_tool_info(ws_id, thread, event)
|
||||
elif isinstance(event, ToolResultEvent):
|
||||
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}")
|
||||
|
||||
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:
|
||||
@@ -522,13 +521,12 @@ class TurnstoneBot:
|
||||
except Exception:
|
||||
log.debug("discord.thinking_start_send_failed", ws_id=ws_id)
|
||||
|
||||
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):
|
||||
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)
|
||||
@@ -540,14 +538,21 @@ class TurnstoneBot:
|
||||
edit_interval=self.config.streaming_edit_interval,
|
||||
)
|
||||
if thinking_msg is not None:
|
||||
sm._message = thinking_msg
|
||||
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)
|
||||
|
||||
elif isinstance(event, ToolInfoEvent):
|
||||
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.
|
||||
@@ -561,8 +566,7 @@ class TurnstoneBot:
|
||||
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.
|
||||
# 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
|
||||
@@ -592,7 +596,14 @@ class TurnstoneBot:
|
||||
with contextlib.suppress(Exception):
|
||||
await thinking_msg.delete()
|
||||
|
||||
elif isinstance(event, ToolResultEvent):
|
||||
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.
|
||||
@@ -661,56 +672,42 @@ class TurnstoneBot:
|
||||
)
|
||||
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
|
||||
async def _handle_approve_request(
|
||||
self,
|
||||
ws_id: str,
|
||||
thread: discord.abc.Messageable,
|
||||
event: ApproveRequestEvent,
|
||||
) -> None:
|
||||
import discord
|
||||
|
||||
_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"]
|
||||
from turnstone.channels._formatter import format_approval_request, format_verdict
|
||||
from turnstone.channels.discord.views import ApprovalView
|
||||
|
||||
# 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: {', '.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,
|
||||
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
|
||||
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)
|
||||
):
|
||||
policy_handled = 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:
|
||||
elif not policy_handled:
|
||||
text = format_approval_request(event.items)
|
||||
embed = discord.Embed(
|
||||
title="Tool Approval Required",
|
||||
@@ -727,21 +724,37 @@ class TurnstoneBot:
|
||||
value=format_verdict(verdict),
|
||||
inline=False,
|
||||
)
|
||||
embed.set_footer(text=f"{ws_id}|")
|
||||
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, PlanReviewEvent):
|
||||
text = format_plan_review(event.content)
|
||||
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=text,
|
||||
description=f"**Plan review requested:**\n\n{event.content}",
|
||||
color=discord.Color.blue(),
|
||||
)
|
||||
embed.set_footer(text=f"{ws_id}|")
|
||||
embed.set_footer(text=f"{ws_id}||{_thread_owner_id(thread)}")
|
||||
await thread.send(embed=embed, view=PlanReviewView(self)._view)
|
||||
|
||||
elif isinstance(event, IntentVerdictEvent):
|
||||
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:
|
||||
@@ -773,7 +786,11 @@ class TurnstoneBot:
|
||||
except Exception:
|
||||
log.debug("discord.verdict_embed_edit_failed", ws_id=ws_id)
|
||||
|
||||
elif isinstance(event, ApprovalResolvedEvent):
|
||||
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)
|
||||
@@ -786,7 +803,7 @@ class TurnstoneBot:
|
||||
except Exception:
|
||||
log.debug("discord.approval_resolved_edit_failed", ws_id=ws_id)
|
||||
|
||||
elif isinstance(event, StreamEndEvent):
|
||||
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:
|
||||
@@ -799,7 +816,7 @@ class TurnstoneBot:
|
||||
# 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)
|
||||
content = sm.accumulated_text
|
||||
if content:
|
||||
dm_channel, target_user_id = dm_entry
|
||||
last_msg: discord.Message | None = None
|
||||
@@ -816,10 +833,6 @@ class TurnstoneBot:
|
||||
# Clean up pending approval message tracking.
|
||||
self._pending_approval_msgs.pop(ws_id, None)
|
||||
|
||||
elif isinstance(event, ErrorEvent):
|
||||
safe_msg = event.message[:500] if event.message else "An error occurred"
|
||||
await thread.send(f"**Error:** {safe_msg}")
|
||||
|
||||
# -- helpers -------------------------------------------------------------
|
||||
|
||||
def _should_auto_approve(self, event: ApproveRequestEvent) -> bool:
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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.",
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
+539
-230
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user