pinecall-runtime 0.1.6__py3-none-any.whl

pinecall/gateway/_sockets.py

@@ -0,0 +1,444 @@
+"""Who holds each agent: the app sockets registered in this process, per scope, and the line."""
+
+import time
+from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import dataclass, replace
+from uuid import uuid4
+
+from pydantic import TypeAdapter
+
+from pinecall.domain.agent import AgentConfig
+from pinecall.domain.errors import DeclarationRefused
+from pinecall.domain.names import PRODUCTION, Env, JsonObject
+from pinecall.domain.scope import THE_ORGS_OWN, Scope
+from pinecall.log.logs import Logs
+from pinecall.process.shared import Assigned, Shared, newest
+from pinecall.wire.commands import AgentRegister
+from pinecall.wire.events import (
+    AgentConfigured,
+    AgentDetached,
+    AgentDraining,
+    AgentRegistered,
+)
+from pinecall.wire.frames import Entry
+
+# Minted, never id(websocket): CPython reuses addresses, and the id travels back as `?app=`.
+type SocketId = str
+
+
+# A slug is one org's: the scope names the org, the world and whose copy it is.
+type Held = tuple[Scope, str]
+
+
+AN_APP = "app_"
+
+
+NOT_HOLDING = "agent {slug} is not held in {env} by an app of yours that answers an unclaimed call"
+
+
+ANOTHER_ORGS = "agent {slug} belongs to another org: a slug is one org's"
+
+
+NOT_REGISTERED = "agent {slug} is not registered on this socket: register it first"
+
+
+NO_AGENT = "no app is holding agent {slug}"
+
+
+NOT_THAT_APP = "app {app} is not holding agent {slug}: it disconnected, or it never held it"
+
+
+NO_UNCLAIMED = (
+    "agent {slug} is held only by apps that take no call they did not open: run `pinecall start`"
+)
+
+
+NO_DEV_ANSWERER = (
+    "agent {slug} is held, but nothing answers the console for it: "
+    "run `pinecall start` in its directory"
+)
+
+
+type Stop = Callable[[str], Awaitable[None]]
+
+# Every gateway's sockets, lines and phones, said on this channel and merged by each.
+HELD_CHANNEL = "held"
+
+
+@dataclass(frozen=True)
+class Registration:
+    """One agent as one app socket holds it."""
+
+    slug: str
+    scope: Scope
+    owner: SocketId
+    config: AgentConfig
+    sdk: str | None = None
+    # False for a console: it takes only the calls that name it.
+    takes_unclaimed: bool = True
+    # True for the CLI's companion: the console's dev verbs go to it, never to the class's socket.
+    answers_dev: bool = False
+    # When it was claimed, on any gateway, so the newest holder is known across all of them.
+    claimed: float = 0.0
+    # Set by agent.drain: its tools still answer, it takes no new call.
+    draining: bool = False
+
+    @property
+    def held_as(self) -> Held:
+        """The registry's key for this agent in its scope."""
+        return (self.scope, self.slug)
+
+
+@dataclass(frozen=True)
+class Share:
+    """What one gateway holds: its sockets' agents, and the lines and phones it set."""
+
+    registrations: tuple[Registration, ...] = ()
+    lines: tuple[Assigned, ...] = ()
+    phones: tuple[Assigned, ...] = ()
+
+
+# Durable facts are in the log; this is which sockets are open now, on every gateway of the box:
+# each says its own share on the signal and reads the merged view. The line is where a ring
+# lands: the org's own scope in production, one developer's in the sandbox, which an org shares.
+class Sockets:
+    """The agents the app sockets hold, per scope, the line of each, and the developers' phones."""
+
+    def __init__(self, logs: Logs) -> None:
+        """Nobody holds anything yet."""
+        self.logs = logs
+        # This gateway's own sockets, and the lines and phones it set.
+        self.mine: dict[Held, list[Registration]] = {}
+        self.mine_lines: dict[tuple[str, str], Assigned] = {}
+        self.mine_phones: dict[tuple[str, str], Assigned] = {}
+        # The merged view every read answers from: every gateway's.
+        self.holders: dict[Held, list[Registration]] = {}
+        self.lines: dict[tuple[str, str], str] = {}
+        self.phones: dict[tuple[str, str], str] = {}
+        self.owned: dict[SocketId, set[Held]] = {}
+        self.shared = Shared(logs.relay.signal, HELD_CHANNEL, _SHARE, Share(), self._merged)
+        self._last_claim = 0.0
+
+    async def start(self) -> None:
+        """Hear the other gateways' sockets, and say this one's."""
+        await self.shared.start()
+
+    async def close(self) -> None:
+        """Stop hearing and saying."""
+        await self.shared.close()
+
+    # ── reading ──
+
+    def of(self, scope: Scope, slug: str) -> Registration | None:
+        """The newest holder in this scope, else in the org's own."""
+        return self._newest((scope, slug)) or self._newest((orgs_own(scope), slug))
+
+    def serving(self, scope: Scope, slug: str, app: SocketId | None) -> Registration | None:
+        """The socket a call opened in this scope goes to: the one named, else the newest taker."""
+        if app is not None:
+            found = self.on(app, scope.env, slug)
+            return found if found is not None and found.scope.org == scope.org else None
+        return self._taker((scope, slug)) or self._taker((orgs_own(scope), slug))
+
+    def answering_dev(self, scope: Scope, slug: str) -> Registration | None:
+        """The socket the console's dev verbs go to: the newest that answers them, scope first."""
+        return self._answerer((scope, slug)) or self._answerer((orgs_own(scope), slug))
+
+    # A caller's registered phone reaches their own scope first; then the line.
+    def taking(self, scope: Scope, slug: str, caller: str | None) -> Registration | None:
+        """The socket a ring at the agent lands in."""
+        if caller is not None and (holder := self.phones.get((scope.env, caller))) is not None:
+            theirs = self._taker((replace(scope, holder=holder), slug))
+            if theirs is not None:
+                return theirs
+        holder = self.lines.get((scope.env, slug))
+        if holder is None:
+            return None
+        return self._taker((replace(scope, holder=holder), slug))
+
+    def on(self, app: SocketId, env: Env, slug: str) -> Registration | None:
+        """The agent as this socket holds it in this world."""
+        for found in self.owned.get(app, ()):
+            if found[0].env == env and found[1] == slug:
+                return next(
+                    (item for item in self.holders.get(found, ()) if item.owner == app), None
+                )
+        return None
+
+    def owned_by(self, app: SocketId) -> list[Registration]:
+        """Every agent one socket holds, in claim order."""
+        found = (item for name in self.owned.get(app, ()) for item in self.holders.get(name, ()))
+        return sorted(
+            (holder for holder in found if holder.owner == app),
+            key=lambda item_found: item_found.claimed,
+        )
+
+    # One row per slug, the reader's scope over the org's; every scope for a team reader.
+    def holding(self, scope: Scope, *, every_corner: bool) -> list[Registration]:
+        """The org's held agents in the world, as this reader sees them."""
+        seen: dict[str | Held, Registration] = {}
+        for (owner, slug), holding in self.holders.items():
+            if owner.org != scope.org or owner.env != scope.env or not holding:
+                continue
+            mine = owner.holder in (THE_ORGS_OWN, scope.holder)
+            if not every_corner and not mine:
+                continue
+            under: str | Held = (owner, slug) if every_corner else slug
+            if under not in seen or owner.holder == scope.holder:
+                seen[under] = holding[-1]
+        return list(seen.values())
+
+    def slugs(self, org: str) -> frozenset[str]:
+        """Every slug the org holds, in any world and scope: the agents quota counts these."""
+        return frozenset(
+            slug
+            for (owner_scope, slug), kept in self.holders.items()
+            if owner_scope.org == org and kept
+        )
+
+    # A log reader knows the slug alone: production's declaration, else any world's.
+    def declared(self, slug: str) -> AgentConfig | None:
+        """The agent's declared config, production's first."""
+        found = sorted(
+            (kept[-1] for (_, name), kept in self.holders.items() if name == slug and kept),
+            key=lambda item_found: (item_found.scope.env != PRODUCTION, -item_found.claimed),
+        )
+        return found[0].config if found else None
+
+    # ── the line ──
+
+    def line_of(self, env: Env, slug: str) -> str | None:
+        """The scope holding the agent's line; None when nobody does."""
+        return self.lines.get((env, slug))
+
+    # Taking the line is said out loud: starting later never takes a colleague's calls.
+    def take_the_line(self, scope: Scope, slug: str) -> None:
+        """Give the agent's line to this scope, which must hold an app that answers a ring."""
+        if self._taker((scope, slug)) is None:
+            raise DeclarationRefused(NOT_HOLDING.format(slug=slug, env=scope.env))
+        self._assigned(self.mine_lines, scope.env, slug, scope.holder)
+
+    def drop_the_line(self, scope: Scope, slug: str) -> bool:
+        """Let the line go, to the newest other scope that could take it."""
+        if self.lines.get((scope.env, slug)) != scope.holder:
+            return False
+        self._assigned(self.mine_lines, scope.env, slug, None)
+        self._next_takes_the_line(scope, slug, leaving=scope.holder)
+        return True
+
+    def waiting_for_the_line(self, scope: Scope, slug: str) -> list[Registration]:
+        """Every scope of the org that could take the line, newest first."""
+        taking = [
+            item
+            for (owner_scope, name) in self.holders
+            if name == slug and owner_scope.org == scope.org and owner_scope.env == scope.env
+            if (item := self._taker((owner_scope, name))) is not None
+        ]
+        return sorted(taking, key=lambda item: item.claimed, reverse=True)
+
+    # ── the developers' phones ──
+
+    def calls_from(self, env: Env, number: str, holder: str) -> None:
+        """Send rings from this phone to this person's scope; the last one said wins."""
+        self._assigned(self.mine_phones, env, number, holder)
+
+    def forget_calls_from(self, env: Env, holder: str) -> list[str]:
+        """Forget every phone of this person, and say which."""
+        gone = sorted(
+            n
+            for (world, n), owner_scope in self.phones.items()
+            if world == env and owner_scope == holder
+        )
+        for number in gone:
+            self._assigned(self.mine_phones, env, number, None)
+        return gone
+
+    def phone_of(self, env: Env, number: str) -> str | None:
+        """Whose own phone the number is in the world, or None."""
+        return self.phones.get((env, number.strip()))
+
+    def calling(self, env: Env, holder: str) -> list[str]:
+        """The phones this person rings from."""
+        return sorted(
+            n
+            for (world, n), owner_scope in self.phones.items()
+            if world == env and owner_scope == holder
+        )
+
+    # ── claiming ──
+
+    async def register(
+        self, owner: SocketId, scope: Scope, slug: str, wanted: AgentRegister
+    ) -> Entry:
+        """This socket holds the agent from now on; agent.registered is written on its log."""
+        store = self.logs.store
+        owner_org = await store.owner(slug)
+        if owner_org is not None and owner_org != scope.org:
+            raise DeclarationRefused(ANOTHER_ORGS.format(slug=slug))
+        await store.claim(None, slug, scope.org)
+        # A call that arrives before agent.configure still runs on what the agent declared.
+        registration = self.on(owner, scope.env, slug) or self.of(scope, slug)
+        config = registration.config if registration is not None else AgentConfig(slug=slug)
+        self._replace(
+            Registration(
+                slug=slug,
+                scope=scope,
+                owner=owner,
+                config=config,
+                sdk=wanted.sdk,
+                takes_unclaimed=wanted.takes_unclaimed,
+                answers_dev=wanted.answers_dev,
+            )
+        )
+        data = AgentRegistered(routes=[], app=owner, sdk=wanted.sdk, env=scope.env)
+        return await self._written(scope.env, slug, "agent.registered", data.written())
+
+    async def configure(
+        self, owner: SocketId, env: Env, slug: str, config: AgentConfig, changed: Sequence[str]
+    ) -> Entry:
+        """The declaration this socket sent; agent.configured names what changed."""
+        found = self.on(owner, env, slug)
+        if found is None:
+            raise DeclarationRefused(NOT_REGISTERED.format(slug=slug))
+        self._replace(replace(found, config=config))
+        data = AgentConfigured(changed=list(changed))
+        return await self._written(env, slug, "agent.configured", data.written())
+
+    # A deploy: the socket keeps the agent so its tools still answer, and takes no new call.
+    def drain(self, owner: SocketId, env: Env, slug: str) -> None:
+        """Mark the socket as draining the agent."""
+        found = self.on(owner, env, slug)
+        if found is None:
+            raise DeclarationRefused(NOT_REGISTERED.format(slug=slug))
+        self._replace(replace(found, draining=True))
+
+    async def drained(
+        self, owner: SocketId, env: Env, slug: str, *, handed: int, parked: int
+    ) -> Entry:
+        """agent.draining, once the socket's calls moved."""
+        data = AgentDraining(app=owner, env=env, handed=handed, parked=parked)
+        return await self._written(env, slug, "agent.draining", data.written())
+
+    async def release(self, owner: SocketId) -> None:
+        """A socket closed: its agents let go, the line handed on, agent.detached for each."""
+        theirs = [key for key, kept in self.mine.items() if any(i.owner == owner for i in kept)]
+        for found in theirs:
+            kept = [item for item in self.mine[found] if item.owner != owner]
+            if kept:
+                self.mine[found] = kept
+            else:
+                del self.mine[found]
+            self._merged()
+            left = self.holders.get(found, [])
+            scope, slug = found
+            if self.lines.get((scope.env, slug)) == scope.holder and not left:
+                self._assigned(self.mine_lines, scope.env, slug, None)
+                self._next_takes_the_line(scope, slug)
+            data = AgentDetached(app=owner, env=scope.env, left=not left)
+            await self._written(scope.env, slug, "agent.detached", data.written())
+        self._say()
+
+    # Claimed at the gateway's clock, never earlier than its last claim: the newest is the last.
+    def _replace(self, registration: Registration) -> None:
+        self._last_claim = max(time.time(), self._last_claim + 1e-6)
+        claim = replace(registration, claimed=self._last_claim)
+        holding = self.mine.setdefault(claim.held_as, [])
+        at = next(
+            (n for n, item_found in enumerate(holding) if item_found.owner == claim.owner), None
+        )
+        if at is None:
+            holding.append(claim)
+        else:
+            holding[at] = claim
+        self._merged()
+        # The first scope able to take a ring gets the line; later ones claim it.
+        taking = claim.takes_unclaimed and not claim.draining
+        if taking and (claim.scope.env, claim.slug) not in self.lines:
+            self._assigned(self.mine_lines, claim.scope.env, claim.slug, claim.scope.holder)
+        self._say()
+
+    # A line or a phone set here: said to every gateway, the newest setting of a key standing.
+    def _assigned(
+        self, table: dict[tuple[str, str], Assigned], env: Env, key: str, holder: str | None
+    ) -> None:
+        table[(env, key)] = Assigned(env, key, holder, time.time())
+        self._merged()
+        self._say()
+
+    def _say(self) -> None:
+        self.shared.put(
+            Share(
+                registrations=tuple(item for kept in self.mine.values() for item in kept),
+                lines=tuple(self.mine_lines.values()),
+                phones=tuple(self.mine_phones.values()),
+            )
+        )
+
+    # Every gateway's sockets, oldest claim first; each line and phone as its newest setting says.
+    def _merged(self) -> None:
+        theirs = [heard.share for heard in self.shared.theirs.values()]
+        holders: dict[Held, list[Registration]] = {k: list(v) for k, v in self.mine.items()}
+        for share in theirs:
+            for item in share.registrations:
+                holders.setdefault(item.held_as, []).append(item)
+        owned: dict[SocketId, set[Held]] = {}
+        for key, kept in holders.items():
+            kept.sort(key=lambda item: item.claimed)
+            for item in kept:
+                owned.setdefault(item.owner, set()).add(key)
+        self.holders, self.owned = holders, owned
+        self.lines = newest(
+            self.mine_lines, (setting for share in theirs for setting in share.lines)
+        )
+        self.phones = newest(
+            self.mine_phones, (setting for share in theirs for setting in share.phones)
+        )
+
+    def _newest(self, holding_found: Held) -> Registration | None:
+        holding = self.holders.get(holding_found)
+        return holding[-1] if holding else None
+
+    def _taker(self, holding_found: Held) -> Registration | None:
+        holding = reversed(self.holders.get(holding_found, ()))
+        return next((item for item in holding if item.takes_unclaimed and not item.draining), None)
+
+    def _answerer(self, holding_found: Held) -> Registration | None:
+        holding = reversed(self.holders.get(holding_found, ()))
+        return next((item for item in holding if item.answers_dev and not item.draining), None)
+
+    def _next_takes_the_line(self, scope: Scope, slug: str, *, leaving: str | None = None) -> None:
+        waiting = [
+            item for item in self.waiting_for_the_line(scope, slug) if item.scope.holder != leaving
+        ]
+        if waiting:
+            self._assigned(self.mine_lines, scope.env, slug, waiting[0].scope.holder)
+
+    # A laptop restarts often: in the sandbox these entries are forgettable, or they flood the log.
+    async def _written(self, env: Env, slug: str, kind: str, data: JsonObject) -> Entry:
+        forgettable = None if env == PRODUCTION else True
+        return await self.logs.agent(slug).append(kind, data, ephemeral=forgettable)
+
+
+@dataclass(frozen=True)
+class Process:
+    """One connected app socket, as the org's list of processes shows it."""
+
+    app: SocketId
+    scope: Scope
+    address: str | None
+    connected_at: float
+    # Sends `error stopped` and closes: the SDK exits instead of reconnecting.
+    stop: Stop
+    host: str | None = None
+
+
+_SHARE: TypeAdapter[Share] = TypeAdapter(Share)
+
+
+def orgs_own(scope: Scope) -> Scope:
+    return replace(scope, holder=THE_ORGS_OWN)
+
+
+def new_socket_id() -> SocketId:
+    """A new id for an app socket."""
+    return f"{AN_APP}{uuid4().hex[:12]}"

pinecall/gateway/_streams.py

@@ -0,0 +1,83 @@
+"""Server-sent events: the frames, the pings that keep a stream alive, its end on shutdown."""
+
+import asyncio
+import json
+from collections.abc import AsyncIterator
+
+from fastapi.responses import StreamingResponse
+
+from pinecall.domain.names import JsonObject
+
+SSE = "text/event-stream"
+
+# Short: a reconnect resumes losslessly from Last-Event-ID.
+RETRY_MS = 1000
+
+# Under the 30 s a proxy lets a quiet connection idle.
+PING_S = 25.0
+
+PING = ": ping\n\n"
+
+# x-accel-buffering: nginx would hold the stream in its buffer.
+SSE_HEADERS = {"cache-control": "no-store", "connection": "keep-alive", "x-accel-buffering": "no"}
+
+
+def wants_sse(accept: str | None) -> bool:
+    """Whether the Accept header asks for the stream rather than a page."""
+    return SSE in (accept or "")
+
+
+def frame(event: str, data: JsonObject, *, seq: int | None = None) -> str:
+    """One SSE frame: its id (the seq), its event, its data as one line of JSON."""
+    head = "" if seq is None else f"id: {seq}\n"
+    return f"{head}event: {event}\ndata: {json.dumps(data, separators=(',', ':'))}\n\n"
+
+
+def streamed(frames: AsyncIterator[str], closing: asyncio.Event) -> StreamingResponse:
+    """Frames as an SSE response, ended when they run out or the process stops."""
+    return StreamingResponse(until(closing, frames), media_type=SSE, headers=SSE_HEADERS)
+
+
+# One pending __anext__ across pings: a new one after a ping would lose the item the first waits
+# for.
+async def paced[T](coming: AsyncIterator[T], every_s: float = PING_S) -> AsyncIterator[T | None]:
+    """The items as they come, and None after every quiet stretch."""
+    pending: asyncio.Task[T] | None = None
+    try:
+        while True:
+            if pending is None:
+                pending = asyncio.ensure_future(anext(coming))
+            done, _ = await asyncio.wait({pending}, timeout=every_s)
+            if not done:
+                yield None
+                continue
+            finished, pending = pending, None
+            try:
+                yield finished.result()
+            except StopAsyncIteration:
+                return
+    finally:
+        if pending is not None:
+            pending.cancel()
+
+
+async def until[T](closing: asyncio.Event, coming: AsyncIterator[T]) -> AsyncIterator[T]:
+    """The items until they end or the process starts closing."""
+    stopping = asyncio.ensure_future(closing.wait())
+    pending: asyncio.Task[T] | None = None
+    try:
+        while True:
+            if pending is None:
+                pending = asyncio.ensure_future(anext(coming))
+            await asyncio.wait({pending, stopping}, return_when=asyncio.FIRST_COMPLETED)
+            if not pending.done():
+                return
+            finished, pending = pending, None
+            try:
+                yield finished.result()
+            except StopAsyncIteration:
+                return
+    finally:
+        stopping.cancel()
+        if pending is not None:
+            pending.cancel()