diff --git a/src/beaver_gateway/backends/claude_sdk.py b/src/beaver_gateway/backends/claude_sdk.py index facc3ea..2f3ee4d 100644 --- a/src/beaver_gateway/backends/claude_sdk.py +++ b/src/beaver_gateway/backends/claude_sdk.py @@ -111,6 +111,7 @@ warnings.filterwarnings("ignore", category=CanUseToolShadowedWarning) ASK_TOOL = "AskUserQuestion" __all__ = [ + "ORIGIN_HEADER", "AuditSink", "ClaudeSdkBackend", "RunnerConfig", @@ -120,6 +121,9 @@ __all__ = [ "fingerprint", ] +ORIGIN_HEADER = "X-Conversation-Origin" +"""Rides on every exposed MCP request; the gateway proxy passes it upstream.""" + ENV_KEEP: tuple[str, ...] = ( "PATH", "HOME", @@ -724,7 +728,14 @@ class ClaudeSdkBackend: system_prompt = ( prompt_assembly.assemble(sources) if sources else agent.system_prompt ) - servers: dict[str, Any] = dict(self._servers) if spec.tools else {} + servers: dict[str, Any] = ( + { + name: {**server, "headers": {ORIGIN_HEADER: key}} + for name, server in self._servers.items() + } + if spec.tools + else {} + ) gateway = ( self._tool_server(key, spec.kind) if spec.tools and self._tool_server is not None and agent.gateway_tools diff --git a/src/beaver_gateway/conversations/spawning.py b/src/beaver_gateway/conversations/spawning.py index 828cddd..71521c6 100644 --- a/src/beaver_gateway/conversations/spawning.py +++ b/src/beaver_gateway/conversations/spawning.py @@ -148,8 +148,7 @@ class Spawning(Messaging): ) if result.text.strip(): await self.inject(parent, result.text, urgency="normal", origin="merge") - await self.set_status(conv, "merged") - await self.mark_closed(conv) + await self.mark_closed(await self.set_status(conv, "merged")) self._bus.publish( "conversation.merged", conversation_id=conv.external_id, diff --git a/src/beaver_gateway/frontends/telegram/frontend.py b/src/beaver_gateway/frontends/telegram/frontend.py index a8f0c94..2893483 100644 --- a/src/beaver_gateway/frontends/telegram/frontend.py +++ b/src/beaver_gateway/frontends/telegram/frontend.py @@ -268,7 +268,10 @@ class TelegramFrontend(Frontend): ) async def mark_closed(self, conv: Conversation) -> bool: - return await self.mark_topic(conv) + merged = conv.status == "merged" + return await self.mark_topic( + conv, self.texts.closed_prefix if merged else self.texts.idle_prefix + ) async def mark_topic(self, conv: Conversation, prefix: str | None = None) -> bool: """Rename the topic; closeForumTopic does not exist in private chats.""" @@ -279,10 +282,7 @@ class TelegramFrontend(Frontend): name = self._topic_names.get(target[1]) or conv.title or self.texts.branch if name.startswith(prefix): return True - await self.bot.edit_forum_topic( - target[0], target[1], name=f"{prefix}{name}"[:128] - ) - self._topic_names[target[1]] = f"{prefix}{name}" + await self._rename(target[1], f"{prefix}{self._unmarked(name)}") return True def file_note(self) -> str: @@ -462,6 +462,36 @@ class TelegramFrontend(Frontend): attachments=attachments, ) + async def _branch_in_topic( + self, topic: int, *, text: str, attachments: list[dict[str, Any]] | None + ) -> None: + """A fresh topic is named by the first line, a closed one loses its mark.""" + current = self._topic_names.get(topic) + previous = None + if current is None: + previous = await self.conversations.find_bound( + frontend=FRONTEND, external_id=self._ext(topic) + ) + known = current or (previous.title if previous is not None else None) + title = (self._unmarked(known) if known else "") or self._title_from(text) + await self._branch(topic, title=title, text=text, attachments=attachments) + if current == title: + return + try: + await self._rename(topic, title) + except TelegramAPIError as exc: + _log.warning("topic %s not renamed to %r: %s", topic, title, exc) + + async def _rename(self, topic: int, name: str) -> None: + name = name[:128] + await self.bot.edit_forum_topic(self.chat_id, topic, name=name) + self._topic_names[topic] = name + + def _unmarked(self, name: str) -> str: + for prefix in (self.texts.closed_prefix, self.texts.idle_prefix): + name = name.removeprefix(prefix) + return name + async def _handle(self, update: Update) -> None: if update.message is not None: await self._on_message(update.message) @@ -592,9 +622,8 @@ class TelegramFrontend(Frontend): topic = cast("int", thread_id) conv = await self._live(topic) if conv is None: - title = self._topic_names.get(topic) or self._title_from(text) - await self._branch( - topic, title=title, text=text, attachments=attachments or None + await self._branch_in_topic( + topic, text=text, attachments=attachments or None ) return pending = self.conversations.pending_question(conv.external_id) diff --git a/src/beaver_gateway/frontends/telegram/texts.py b/src/beaver_gateway/frontends/telegram/texts.py index edf2116..3fbfbb4 100644 --- a/src/beaver_gateway/frontends/telegram/texts.py +++ b/src/beaver_gateway/frontends/telegram/texts.py @@ -42,6 +42,7 @@ class TelegramTexts: help_head: str = "General is the master, any other topic is a branch." branch: str = "branch" closed_prefix: str = "✅ " + idle_prefix: str = "💤 " waiting: str = "⏳ thinking" thinking: str = "🤔 thinking" writing: str = "✍️ writing" diff --git a/src/beaver_gateway/jobs/job.py b/src/beaver_gateway/jobs/job.py index 0b77654..04d162f 100644 --- a/src/beaver_gateway/jobs/job.py +++ b/src/beaver_gateway/jobs/job.py @@ -5,7 +5,7 @@ from __future__ import annotations import logging from dataclasses import dataclass from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast if TYPE_CHECKING: from collections.abc import Awaitable, Callable, Coroutine @@ -73,6 +73,23 @@ class JobRun: ) return True + async def inject_into( + self, + public_id: str | None, + text: str, + *, + urgency: Priority = "normal", + origin: str | None = None, + ) -> bool: + """Inject into that conversation, or the master once it is gone or over.""" + conv = await self.conversations.get(public_id) if public_id else None + if conv is None or conv.status != "open": + return await self.inject_master(text, urgency=urgency, origin=origin) + await self.conversations.inject( + conv, text, urgency=urgency, origin=origin or self.job.name + ) + return True + async def spawn_job( self, *, @@ -113,6 +130,21 @@ class JobRun: _log.exception("closing idle %s failed", conv.external_id) return out + async def close_quiet( + self, *, kind: str = "branch", days: int = 3, limit: int | None = None + ) -> list[Conversation]: + """Close chats quiet for ``days`` without a digest, their windows marked.""" + out: list[Conversation] = [] + for conv in await self.conversations.idle(kind=kind, days=days, limit=limit): + if await self.conversations.busy( + conv + ) or await self.conversations.queue.pending(cast("int", conv.id)): + continue + row = await self.conversations.close(conv) + await self.conversations.mark_closed(row) + out.append(row) + return out + async def retry_in(self, delay: timedelta) -> None: await self.scheduler.trigger( self.job, self.payload, delay=delay, trigger=self.trigger diff --git a/tests/test_claude_sdk_backend.py b/tests/test_claude_sdk_backend.py index 4920727..f65b7ab 100644 --- a/tests/test_claude_sdk_backend.py +++ b/tests/test_claude_sdk_backend.py @@ -250,7 +250,11 @@ async def test_stream_envelope_and_index_rebase(cwd: Path) -> None: assert opts.strict_mcp_config is True assert opts.permission_mode == "bypassPermissions" assert opts.mcp_servers == { - "firefly": {"type": "http", "url": "http://127.0.0.1:1/mcp/firefly/"} + "firefly": { + "type": "http", + "url": "http://127.0.0.1:1/mcp/firefly/", + "headers": {"X-Conversation-Origin": "conv-1"}, + } } assert opts.session_store is not None assert opts.resume is None diff --git a/tests/test_jobrun.py b/tests/test_jobrun.py new file mode 100644 index 0000000..7a443d2 --- /dev/null +++ b/tests/test_jobrun.py @@ -0,0 +1,85 @@ +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from typing import Any, cast + +from test_conversations import StubFrontend, World, world + +from beaver_gateway.jobs.job import Job, JobRun +from beaver_gateway.storage.models import Conversation + +__all__ = ["world"] + + +async def _noop(_: JobRun) -> None: + return None + + +def _run(world: World) -> JobRun: + return JobRun( + job=Job("test", _noop), + trigger="test", + payload={}, + scheduler=cast("Any", SimpleNamespace(conversations=world.conversations)), + ) + + +async def _quiet_since(world: World, conv: Conversation, ago: timedelta) -> None: + async def apply(row: Conversation) -> None: + row.last_activity_at = datetime.now(UTC) - ago + + await world.conversations._update(conv, apply) # noqa: SLF001 + + +async def test_close_quiet_closes_and_marks_only_idle_branches(world: World) -> None: + marked: list[tuple[str, str]] = [] + + class MarkingFrontend(StubFrontend): + async def mark_closed(self, conv: Conversation) -> bool: + marked.append((conv.external_id, conv.status)) + return True + + tg = MarkingFrontend("tg", ("master", "branch")) + tg.conversations = world.conversations + world.conversations._frontends = [tg] # noqa: SLF001 + convs = world.conversations + master = await convs.create(kind="master", agent="a", session_id="m") + idle = await convs.create(kind="branch", agent="a", parent=master, session_id="1") + fresh = await convs.create(kind="branch", agent="a", parent=master, session_id="2") + waiting = await convs.create( + kind="branch", agent="a", parent=master, session_id="3" + ) + await _quiet_since(world, idle, timedelta(days=4)) + await _quiet_since(world, fresh, timedelta(days=1)) + await _quiet_since(world, waiting, timedelta(days=4)) + await convs.queue.push( + conversation_id=cast("int", waiting.id), + priority="normal", + origin="cron", + text="later", + ) + + closed = await _run(world).close_quiet(kind="branch", days=3) + + assert [c.external_id for c in closed] == [idle.external_id] + assert marked == [(idle.external_id, "closed")] + assert (await convs.get(fresh.external_id)).status == "open" + assert (await convs.get(waiting.external_id)).status == "open" + + +async def test_inject_into_falls_back_to_the_master(world: World) -> None: + convs = world.conversations + master = await convs.create(kind="master", agent="a", session_id="m") + branch = await convs.create(kind="branch", agent="a", parent=master) + gone = await convs.create(kind="branch", agent="a", parent=master) + await convs.set_status(gone, "closed") + run = _run(world) + + assert await run.inject_into(branch.external_id, "to branch", urgency="normal") + assert await run.inject_into(gone.external_id, "to closed", urgency="normal") + assert await run.inject_into(None, "to nobody", urgency="normal") + + assert [i.text for i in await convs.queue.pending(branch.id)] == ["to branch"] + assert [i.text for i in await convs.queue.pending(master.id)] == [ + "to closed", + "to nobody", + ] diff --git a/tests/test_mcp_proxy_headers.py b/tests/test_mcp_proxy_headers.py new file mode 100644 index 0000000..22b6c7b --- /dev/null +++ b/tests/test_mcp_proxy_headers.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import asyncio +import contextlib +import socket +from typing import TYPE_CHECKING, Any + +import pytest +import uvicorn +from fastmcp import Client, FastMCP +from fastmcp.client.transports import StreamableHttpTransport +from fastmcp.server.dependencies import get_http_headers + +from beaver_gateway.backends.claude_sdk import ORIGIN_HEADER +from beaver_gateway.mcp.internal_app import build_internal_app +from beaver_gateway.mcp.types import McpServer + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +@contextlib.asynccontextmanager +async def _serving(app: Any, port: int) -> AsyncIterator[None]: + server = uvicorn.Server( + uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning") + ) + task = asyncio.create_task(server.serve()) + while not server.started: + await asyncio.sleep(0.02) + try: + yield + finally: + server.should_exit = True + await task + + +@pytest.mark.anyio +async def test_the_origin_header_reaches_the_upstream_mcp() -> None: + upstream = FastMCP("upstream") + + @upstream.tool + def caller() -> str: + return get_http_headers().get(ORIGIN_HEADER.lower(), "") + + up, gateway = _free_port(), _free_port() + app, urls, _ = build_internal_app( + [McpServer.http(name="t3", url=f"http://127.0.0.1:{up}/mcp")], + host="127.0.0.1", + port=gateway, + ) + seen = [] + async with ( + _serving(upstream.http_app(transport="http", path="/mcp"), up), + _serving(app, gateway), + ): + for conversation in ("conv-a", "conv-b"): + transport = StreamableHttpTransport( + urls["t3"], headers={ORIGIN_HEADER: conversation} + ) + async with Client(transport) as client: + result = await client.call_tool("caller", {}) + seen.append(result.data) + assert seen == ["conv-a", "conv-b"] diff --git a/tests/test_telegram.py b/tests/test_telegram.py index 441c83c..8ce8ac7 100644 --- a/tests/test_telegram.py +++ b/tests/test_telegram.py @@ -519,10 +519,10 @@ async def test_commands_status_merge_and_new(stack: Stack) -> None: frontend="telegram", external_id=f"{USER}/11" ) assert branch.status == "merged" - assert stack.bot.topics == ["🦫 General", "edit:11:✅ work"] + assert stack.bot.topics == ["🦫 General", "edit:11:work", "edit:11:✅ work"] stack.bot.message("/new отчёт") await stack.until(lambda: stack.sent_with("in a new topic"), what="new topic") - assert stack.bot.topics == ["🦫 General", "edit:11:✅ work", "отчёт"] + assert stack.bot.topics[-1] == "отчёт" child = await stack.world.conversations.find_bound( frontend="telegram", external_id=f"{USER}/902" ) @@ -1298,3 +1298,34 @@ def test_a_tool_call_between_progress_updates_starts_a_new_paragraph() -> None: ) feed(_thinking_delta("now writing it back")) assert draft.text == "reading the note\n\nnow writing it back" + + +async def test_a_closed_topic_is_marked_asleep_and_wakes_with_a_new_branch( + stack: Stack, +) -> None: + stack.bot.message("кухня на неделю\nи ещё строка", thread=12) + await stack.until(lambda: stack.sent_with("кухня на неделю"), what="reply") + assert stack.bot.topics[-1] == "edit:12:кухня на неделю" + branch = await stack.world.conversations.find_bound( + frontend="telegram", external_id=f"{USER}/12" + ) + closed = await stack.world.conversations.close(branch) + assert await stack.world.conversations.mark_closed(closed) + assert stack.bot.topics[-1] == "edit:12:💤 кухня на неделю" + stack.tg._topic_names.clear() # noqa: SLF001 + stack.bot.message("а в четверг?", thread=12) + await stack.until(lambda: stack.sent_with("а в четверг?"), what="new branch") + assert stack.bot.topics[-1] == "edit:12:кухня на неделю" + fresh = await stack.world.conversations.find_bound( + frontend="telegram", external_id=f"{USER}/12" + ) + assert fresh.id != branch.id + assert fresh.status == "open" + assert fresh.title == "кухня на неделю" + + +async def test_a_named_topic_keeps_its_name(stack: Stack) -> None: + stack.bot.topic_created("план", 13) + stack.bot.message("hello", thread=13) + await stack.until(lambda: stack.sent_with("hello"), what="reply") + assert not any(t.startswith("edit:13:") for t in stack.bot.topics)