zettcode-im 0.4.2__py3-none-any.whl

zettcode/im/link.py

@@ -0,0 +1,229 @@
+"""The bridge's own session on the broker: one place to send a prompt, one place to hear the answer.
+
+The daemon is a normal ZettCode session, so the bridge talks to it the way a chat window does: it registers a
+session of its own, sends the text with ``direct`` set (the daemon runs it as its reader's prompt and answers
+with ``send_back``), and waits for that answer here.
+
+Only the daemon's answers are taken. Everything else on this session — another session writing to it, a group
+message — is refused in words, because a bridge that answers everything would turn a stray message into a turn
+of the agent the reader did not ask for.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+
+from zettcode.broker.models import Incoming, Receipt, SendError, SessionInfo
+from zettcode.broker.protocol import VIA_USER, ImageData
+
+
+class LinkError(RuntimeError):
+    """The message could not be handed to the daemon, and the message is what to tell the reader."""
+
+
+logger = logging.getLogger(__name__)
+
+
+#: What the answer queue carries when the connection, not the agent, ended the wait.
+_LOST = object()
+
+
+@dataclass
+class SessionLink:
+    """The bridge's line to the daemon, over the session broker.
+
+    Args:
+        session_id: The id this bridge registers under. It only has to be stable; the roster shows the handle.
+        handle: What other sessions call the bridge, in the roster and in the daemon's prompt.
+        workspace: The directory named in the roster beside this session.
+        timeout: Seconds to wait for the daemon's answer after it has taken a message.
+    """
+
+    session_id: str
+    handle: str = "chat-bridge"
+    workspace: str = ""
+    timeout: float = 600.0
+    _client: object | None = field(default=None, init=False, repr=False)
+    _answers: asyncio.Queue[object] = field(default_factory=asyncio.Queue, init=False, repr=False)
+    _waiting: bool = field(default=False, init=False, repr=False)
+    #: The session this bridge is waiting on, so another session's message cannot be taken for its answer.
+    _asking: str = field(default="", init=False, repr=False)
+    #: What the broker said when the connection went out of reach, while a turn was waiting.
+    _problem: str = field(default="", init=False, repr=False)
+
+    # -- the life of the line -------------------------------------------------------------------------
+    async def start(self) -> None:
+        """Join the broker. Without one, every send says so rather than hanging."""
+        from zettcode.broker.client import BrokerClient
+
+        client = BrokerClient(on_message=self._heard, on_problem=self._trouble, on_conflict=self._conflict)
+        self._client = client
+        await client.start()
+        client.set_session(
+            SessionInfo(session_id=self.session_id, title=self.handle, workspace=self.workspace, status="idle")
+        )
+
+    async def close(self) -> None:
+        """Leave the broker. Safe to call twice, and safe when it never joined."""
+        client = self._client
+        self._client = None
+        if client is not None:
+            await client.close()  # type: ignore[attr-defined]
+
+    async def connected(self, timeout: float = 10.0) -> bool:
+        """Return whether the broker is reachable, waiting a moment for a connection that is on its way."""
+        client = self._client
+        if client is None:
+            return False
+        return bool(await client.wait_connected(timeout))  # type: ignore[attr-defined]
+
+    def status(self, status: str) -> None:
+        """Say whether a turn is running, so the roster tells the truth about this session."""
+        client = self._client
+        if client is None:
+            return
+        client.set_session(  # type: ignore[attr-defined]
+            SessionInfo(session_id=self.session_id, title=self.handle, workspace=self.workspace, status=status)
+        )
+
+    # -- one prompt, one answer -----------------------------------------------------------------------
+    async def ask(self, session: str, text: str, *, images: Sequence[ImageData] = ()) -> str:
+        """Send ``text`` to ``session`` as this reader's own words and return the answer it sends back.
+
+        Pictures travel with it. A ``direct`` message is the only one the broker lets carry them, which is what
+        this is: the daemon puts them into the same prompt the text becomes.
+
+        Raises:
+            LinkError: The broker or the daemon refused the message, or the answer did not come in time.
+        """
+        client = self._client
+        if client is None:
+            raise LinkError("the chat bridge is not connected to the session broker")
+        while not self._answers.empty():  # an answer that arrived after its timeout belongs to nobody
+            self._answers.get_nowait()
+        self._waiting = True
+        self._asking = session
+        self._problem = ""
+        try:
+            await self._hand_over(client, session, text, images)
+        except SendError as error:
+            self._waiting = False
+            self._asking = ""
+            raise LinkError(_sentence(error)) from error
+        except Exception as error:  # noqa: BLE001 - a lost connection reads as a send that did not happen
+            self._waiting = False
+            self._asking = ""
+            raise LinkError(f"the message did not reach the daemon: {error}") from error
+        try:
+            answer = await asyncio.wait_for(self._answers.get(), timeout=self.timeout)
+        except asyncio.TimeoutError as error:
+            raise LinkError(
+                f"the daemon took the message but did not answer in {int(self.timeout)} seconds; it may still be "
+                "working on it, and an answer after this is not forwarded"
+            ) from error
+        finally:
+            self._waiting = False
+            self._asking = ""
+        if answer is _LOST:
+            raise LinkError(self._problem or "the connection to the session broker was lost while the agent answered")
+        return str(answer)
+
+    async def _hand_over(self, client: object, session: str, text: str, images: Sequence[ImageData]) -> None:
+        """Send the prompt, giving a daemon that has just restarted a moment to appear in the roster.
+
+        A daemon that was started or resumed is in its record before it is registered with the broker, so the first
+        send can be refused as ``unknown`` while the process is perfectly healthy. Three tries a second apart is
+        the difference between that and a sentence about a session that does not exist.
+        """
+        tries = 3
+        for attempt in range(tries):
+            try:
+                await client.send(  # type: ignore[attr-defined]
+                    session,
+                    self._shorten(text, client),
+                    direct=True,
+                    mode="followup",
+                    via=VIA_USER,
+                    images=list(images),
+                )
+                return
+            except SendError as error:
+                if error.reason in {"unknown", "gone"} and attempt + 1 < tries:
+                    await asyncio.sleep(1.0 * (attempt + 1))
+                    continue
+                raise
+
+    async def send(self, session: str, text: str) -> None:
+        """Tell a session something without waiting for an answer — a note, not a prompt."""
+        client = self._client
+        if client is None:
+            return
+        try:
+            await client.send(session, self._shorten(text, client), mode="followup", via=VIA_USER)  # type: ignore[attr-defined]
+        except Exception:  # noqa: BLE001 - a note that cannot be delivered is not worth failing a turn over
+            return
+
+    # -- what the broker calls ------------------------------------------------------------------------
+    def _heard(self, message: Incoming) -> Receipt:
+        """Take one message: the daemon's answer goes to the turn waiting for it, anything else is refused.
+
+        Only the session that was asked may answer. Another session — a stray ``send_message`` from an agent, a
+        group message — would otherwise be forwarded to the chat as if the agent had said it, and the real answer
+        would be refused as a second one.
+        """
+        if not self._waiting:
+            return Receipt(accepted=False, reason="this is a chat bridge; it is not reading messages right now")
+        if message.group or message.sender.session_id != self._asking:
+            return Receipt(
+                accepted=False,
+                reason="this bridge is waiting for the daemon's answer, not for another session's message",
+            )
+        self._answers.put_nowait(message.text)
+        return Receipt(accepted=True, disposition="delivered")
+
+    def _trouble(self, problem: str | None) -> None:
+        """A connection that is out of reach ends the wait instead of leaving the chat hanging.
+
+        The wait is ended with a marker, not with an empty answer: an empty string is what a silent agent looks
+        like, and telling the reader the agent finished when the broker merely dropped is a lie.
+        """
+        if problem is not None and self._waiting:
+            self._problem = problem
+            self._answers.put_nowait(_LOST)
+
+    def _conflict(self, message: str | None) -> None:
+        """Another process holds this bridge's session on the broker; say so where the worker's log is."""
+        if message is not None:
+            logger.warning("another process holds this bridge's session on the broker: %s", message)
+
+    @staticmethod
+    def _shorten(text: str, client: object) -> str:
+        """Return ``text`` fitted to what the broker takes in one message."""
+        limits = getattr(client, "limits", None)
+        room = int(getattr(limits, "max_text", 0) or 0)
+        if room <= 0 or len(text) <= room:
+            return text
+        marker = f"\n[\u2026 {len(text) - room} characters not sent: the broker takes {room} in one message]"
+        return text[: max(0, room - len(marker))] + marker
+
+
+def _sentence(error: SendError) -> str:
+    """Return what a refused send means for the reader, rather than the broker's code alone."""
+    reason = getattr(error, "reason", "") or "refused"
+    if reason == "gone":
+        return "the daemon is not connected to the session broker"
+    if reason == "unknown":
+        return "the daemon's session is not registered with the session broker"
+    if reason == "denied":
+        return "the daemon refused the message"
+    if reason == "too_long" and "picture" in str(error):
+        return f"the pictures were refused: {error}"
+    if reason == "invalid" and "picture" in str(error):
+        return f"the pictures were refused: {error}"
+    return f"the message was not delivered ({reason}): {error}"
+
+
+__all__ = ["LinkError", "SessionLink"]

zettcode/im/state.py

@@ -0,0 +1,193 @@
+"""What the bridge has to remember between two starts.
+
+The IM side is stateless on purpose (``agim`` keeps nothing), so this is the only place the continuity lives:
+the login handshake and its credentials, the cursor a ``receive`` call answered with, the session the daemon
+runs, and the reply context of each chat. ``send`` needs that context back, and a restart mid-batch re-delivers
+the batch rather than dropping the messages behind the first one, which is why the cursor is written only after
+a message has been answered.
+
+Nothing here talks to the network. A file that cannot be read is treated as a first run — the worst it costs is
+a login — and a file that cannot be written is an error the caller has to see, because a bridge that silently
+loses its cursor re-answers messages the reader has already read.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+#: How many event ids are remembered, so a re-delivered batch is answered once. The platform hands a batch back
+#: whole after a restart, and a batch is small; a hundred is several restarts of slack.
+SEEN_LIMIT = 200
+
+#: How many chats keep their reply context. One entry per chat is tiny, and a chat that falls out of it simply
+#: loses the quoting, not the answer.
+CHATS_LIMIT = 100
+
+
+@dataclass
+class BridgeState:
+    """Everything the bridge persists, in one file."""
+
+    #: The login handshake as the IM client gave it, verbatim: its QR payload and the provider state behind it.
+    handshake: dict[str, Any] = field(default_factory=dict)
+    #: What a finished login handed back; the token the client needs to send and poll.
+    credentials: dict[str, Any] = field(default_factory=dict)
+    #: The cursor the last answered message came with.
+    cursor: str = ""
+    #: The session the daemon runs. Empty until the first message starts one.
+    session_id: str = ""
+    #: The reply context of the last message of each chat, by chat id.
+    chats: dict[str, str] = field(default_factory=dict)
+    #: Event ids already answered, oldest first.
+    seen: list[str] = field(default_factory=list)
+    #: What the bridge was started with, so the worker needs no flags: where the daemon works, the model it runs on,
+    #: its shell level, and who may drive it.
+    workspace: str = ""
+    model: str = ""
+    effort: str = ""
+    permission: str = "safe"
+    allow: list[str] = field(default_factory=list)
+    allow_all: bool = False
+    #: Where this state came from. Not written into the file; it is what :meth:`save` writes back to by default.
+    path: Path | None = field(default=None, compare=False, repr=False)
+
+    # -- what the bridge asks -------------------------------------------------------------------------
+    @property
+    def logged_in(self) -> bool:
+        """Return whether a login has finished, which is what a run needs: a login without a token is not one."""
+        return bool(self.token())
+
+    def token(self) -> str:
+        """Return the bot token a finished login produced, or an empty string.
+
+        The IM client hands it over inside ``credentials`` — ``{"config": {...}, "secrets": {...}}`` — and a token
+        read from the wrong level is an empty string that looks like a login until the first send fails, so both
+        the nested shape and a flat one are read.
+        """
+        secrets = self.credentials.get("secrets")
+        value = secrets.get("bot_token") if isinstance(secrets, dict) else None
+        if not isinstance(value, str) or not value:
+            value = self.credentials.get("bot_token")
+        return value if isinstance(value, str) else ""
+
+    def answered(self, event_id: str) -> bool:
+        """Return whether this message has already been answered."""
+        return event_id in self.seen
+
+    def remember(self, event_id: str, *, cursor: str, chat_id: str, reply_token: str | None) -> None:
+        """Record that a message has been answered, and the context the next answer of that chat needs."""
+        if event_id and event_id not in self.seen:
+            self.seen.append(event_id)
+            del self.seen[:-SEEN_LIMIT]
+        self.cursor = cursor
+        if reply_token:
+            self.chats[chat_id] = reply_token
+            while len(self.chats) > CHATS_LIMIT:
+                self.chats.pop(next(iter(self.chats)))
+        else:
+            self.chats.pop(chat_id, None)
+
+    def reply_token(self, chat_id: str) -> str | None:
+        """Return the context token of the last message of this chat, if it had one."""
+        return self.chats.get(chat_id)
+
+    # -- the file -------------------------------------------------------------------------------------
+    @classmethod
+    def load(cls, path: Path) -> BridgeState:
+        """Read the state at ``path``; a missing or unreadable file is a first run."""
+        try:
+            raw = json.loads(path.read_text(encoding="utf-8"))
+        except (OSError, ValueError):
+            return cls(path=path)
+        if not isinstance(raw, dict):
+            return cls(path=path)
+        state = cls(path=path)
+        handshake = raw.get("handshake")
+        state.handshake = handshake if isinstance(handshake, dict) else {}
+        credentials = raw.get("credentials")
+        state.credentials = credentials if isinstance(credentials, dict) else {}
+        cursor = raw.get("cursor")
+        state.cursor = cursor if isinstance(cursor, str) else ""
+        session_id = raw.get("session_id")
+        state.session_id = session_id if isinstance(session_id, str) else ""
+        chats = raw.get("chats")
+        if isinstance(chats, dict):
+            state.chats = {str(key): value for key, value in chats.items() if isinstance(value, str)}
+        seen = raw.get("seen")
+        if isinstance(seen, list):
+            state.seen = [item for item in seen if isinstance(item, str)][-SEEN_LIMIT:]
+        state.workspace = _text(raw.get("workspace"))
+        state.model = _text(raw.get("model"))
+        state.effort = _text(raw.get("effort"))
+        state.permission = _text(raw.get("permission")) or "safe"
+        allow = raw.get("allow")
+        state.allow = [item for item in allow if isinstance(item, str)] if isinstance(allow, list) else []
+        state.allow_all = raw.get("allow_all") is True
+        return state
+
+    def save_runtime(self, path: Path | None = None) -> None:
+        """Write what a running bridge produces, keeping the settings whoever is looking at the file set.
+
+        A worker holds the state it started with; a reader who changes the allow-list, the workspace or the shell
+        level while it runs must not have that change written back to the old value on the next message. So the
+        settings are taken from the file as it is now, and only the cursor, the answers and the login are written.
+        """
+        target = path or self.path
+        on_disk = BridgeState.load(target) if target is not None else None
+        if on_disk is not None:
+            self.allow = on_disk.allow
+            self.allow_all = on_disk.allow_all
+            self.workspace = on_disk.workspace or self.workspace
+            self.model = on_disk.model
+            self.effort = on_disk.effort
+            self.permission = on_disk.permission or self.permission
+        self.save(target)
+
+    def save(self, path: Path | None = None) -> None:
+        """Write the state to ``path`` (this state's own file by default), all at once, readable by this user only.
+
+        The write goes to a temporary file in the same directory and is moved into place, so a process killed
+        mid-write leaves the previous state rather than half of this one — a torn cursor is a re-answered batch.
+        """
+        target = path or self.path
+        if target is None:
+            raise ValueError("this state has no file; pass the path to save it to")
+        target.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
+        payload = {
+            "handshake": self.handshake,
+            "credentials": self.credentials,
+            "cursor": self.cursor,
+            "session_id": self.session_id,
+            "chats": self.chats,
+            "seen": self.seen,
+            "workspace": self.workspace,
+            "model": self.model,
+            "effort": self.effort,
+            "permission": self.permission,
+            "allow": self.allow,
+            "allow_all": self.allow_all,
+        }
+        handle, temporary = tempfile.mkstemp(dir=target.parent, prefix=".state-", suffix=".json")
+        try:
+            with os.fdopen(handle, "w", encoding="utf-8") as stream:
+                json.dump(payload, stream, ensure_ascii=False, indent=2)
+                stream.flush()
+                os.fsync(stream.fileno())
+            os.chmod(temporary, 0o600)
+            os.replace(temporary, target)
+        except BaseException:
+            Path(temporary).unlink(missing_ok=True)
+            raise
+
+
+def _text(value: object) -> str:
+    """Return ``value`` when it is a string, and an empty string for anything a damaged file may hold instead."""
+    return value if isinstance(value, str) else ""
+
+
+__all__ = ["BridgeState", "CHATS_LIMIT", "SEEN_LIMIT"]

zettcode/im/weixin/__init__.py

@@ -0,0 +1,5 @@
+"""Personal WeChat as the chat side: login by scanning, then a background process that runs the bridge."""
+
+from __future__ import annotations
+
+__all__: list[str] = []