@theowlops/boombot 0.0.0-stage → 1.1.0

package/dist/python/agent-core-python/src/boombot_agent/rpc.py

@@ -0,0 +1,265 @@
+"""gRPC transport for the transport-free cognitive turn loop."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+from concurrent import futures
+from queue import Empty, Queue
+from threading import Event, Lock, Thread
+from typing import Any, TypeVar
+import argparse
+import json
+import os
+import sys
+
+import grpc
+
+from .generated import agent_core_pb2 as pb
+from .generated import agent_core_pb2_grpc as rpc
+from .turn import CancelledError, Host, Json, run_turn
+
+T = TypeVar("T")
+
+
+def _encode(value: Any) -> bytes:
+    return json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+
+
+def _decode(value: bytes) -> Json:
+    decoded = json.loads(value.decode("utf-8")) if value else {}
+    if not isinstance(decoded, dict):
+        raise ValueError("RPC JSON payload must be an object")
+    return decoded
+
+
+class RpcHost(Host):
+    def __init__(
+        self,
+        turn_id: str,
+        incoming: Iterator[pb.TurnFrame],
+        context: grpc.ServicerContext,
+    ):
+        self.turn_id = turn_id
+        self._incoming = incoming
+        self._context = context
+        self._outgoing: Queue[pb.TurnFrame | None] = Queue()
+        self._responses: dict[str, Queue[pb.CallbackResponse]] = {}
+        self._lock = Lock()
+        self._cancelled = Event()
+        self._next_id = 0
+
+    def start_reader(self) -> Thread:
+        thread = Thread(
+            target=self._read, name=f"boombot-rpc-read-{self.turn_id}", daemon=True
+        )
+        thread.start()
+        return thread
+
+    def _read(self) -> None:
+        try:
+            for frame in self._incoming:
+                if frame.HasField("cancel"):
+                    self._cancelled.set()
+                    break
+                if not frame.HasField("callback_response"):
+                    continue
+                with self._lock:
+                    target = self._responses.get(frame.request_id)
+                if target is not None:
+                    target.put(frame.callback_response)
+        except Exception:
+            self._cancelled.set()
+        finally:
+            self._cancelled.set()
+            with self._lock:
+                queues = list(self._responses.values())
+            for target in queues:
+                target.put(pb.CallbackResponse(error="turn transport closed", end=True))
+
+    def outgoing(self) -> Queue[pb.TurnFrame | None]:
+        return self._outgoing
+
+    def close(self) -> None:
+        self._outgoing.put(None)
+
+    def _request(
+        self, method: str, payload: Json, streaming: bool
+    ) -> tuple[str, Queue[pb.CallbackResponse]]:
+        with self._lock:
+            self._next_id += 1
+            request_id = f"cb-{self._next_id}"
+            responses: Queue[pb.CallbackResponse] = Queue()
+            self._responses[request_id] = responses
+        self._outgoing.put(
+            pb.TurnFrame(
+                turn_id=self.turn_id,
+                request_id=request_id,
+                callback_request=pb.CallbackRequest(
+                    method=method,
+                    payload_json=_encode(payload),
+                    streaming=streaming,
+                ),
+            )
+        )
+        return request_id, responses
+
+    def _next_response(
+        self, responses: Queue[pb.CallbackResponse]
+    ) -> pb.CallbackResponse:
+        while True:
+            if self.cancelled():
+                raise CancelledError("turn cancelled")
+            try:
+                response = responses.get(timeout=0.1)
+                if response.error:
+                    raise RuntimeError(response.error)
+                return response
+            except Empty:
+                continue
+
+    def _release(self, request_id: str) -> None:
+        with self._lock:
+            self._responses.pop(request_id, None)
+
+    def call(self, method: str, payload: Json) -> Json:
+        request_id, responses = self._request(method, payload, False)
+        try:
+            return _decode(self._next_response(responses).payload_json)
+        finally:
+            self._release(request_id)
+
+    def stream(self, method: str, payload: Json) -> Iterator[Json]:
+        request_id, responses = self._request(method, payload, True)
+        try:
+            while True:
+                response = self._next_response(responses)
+                if response.end:
+                    return
+                yield _decode(response.payload_json)
+        finally:
+            self._release(request_id)
+
+    def emit(self, event: Json) -> None:
+        self._outgoing.put(
+            pb.TurnFrame(
+                turn_id=self.turn_id,
+                event=pb.AgentEvent(event_json=_encode(event)),
+            )
+        )
+
+    def cancelled(self) -> bool:
+        return self._cancelled.is_set() or not self._context.is_active()
+
+    def phase(self, name: str, operation: Callable[[], T]) -> T:
+        begin = self.call("phase.begin", {"phase": name})
+        token = str(begin["token"])
+        try:
+            result = operation()
+        except Exception as error:
+            self.call("phase.fail", {"token": token, "error": str(error)})
+            raise
+        completed = self.call("phase.complete", {"token": token, "result": result})
+        return completed.get("result", result)  # type: ignore[return-value]
+
+
+class AgentCoreService(rpc.AgentCoreServicer):
+    def __init__(self, token: str) -> None:
+        self._token = token
+
+    def _authenticate(self, context: grpc.ServicerContext) -> None:
+        metadata = dict(context.invocation_metadata())
+        if metadata.get("x-boombot-token") != self._token:
+            context.abort(grpc.StatusCode.UNAUTHENTICATED, "invalid internal token")
+
+    def Health(
+        self, request: pb.HealthRequest, context: grpc.ServicerContext
+    ) -> pb.HealthResponse:
+        del request
+        self._authenticate(context)
+        return pb.HealthResponse(status="ok", runtime="python")
+
+    def RunTurn(
+        self, request_iterator: Iterator[pb.TurnFrame], context: grpc.ServicerContext
+    ):
+        self._authenticate(context)
+        try:
+            first = next(request_iterator)
+        except StopIteration:
+            context.abort(grpc.StatusCode.INVALID_ARGUMENT, "missing start frame")
+        if not first.HasField("start") or not first.turn_id:
+            context.abort(
+                grpc.StatusCode.INVALID_ARGUMENT, "first frame must start a turn"
+            )
+        host = RpcHost(first.turn_id, request_iterator, context)
+        host.start_reader()
+
+        def work() -> None:
+            try:
+                result = run_turn(host, _decode(first.start.context_json))
+                host.outgoing().put(
+                    pb.TurnFrame(
+                        turn_id=first.turn_id,
+                        result=pb.TurnResult(result_json=_encode(result)),
+                    )
+                )
+            except Exception as error:
+                host.outgoing().put(
+                    pb.TurnFrame(
+                        turn_id=first.turn_id,
+                        result=pb.TurnResult(
+                            result_json=_encode(
+                                {
+                                    "status": "error",
+                                    "messages": [],
+                                    "rounds": 0,
+                                    "toolCallsExecuted": 0,
+                                    "toolCallsDenied": 0,
+                                    "error": str(error),
+                                }
+                            )
+                        ),
+                    )
+                )
+            finally:
+                host.close()
+
+        Thread(target=work, name=f"boombot-turn-{first.turn_id}", daemon=True).start()
+        while True:
+            frame = host.outgoing().get()
+            if frame is None:
+                return
+            yield frame
+
+
+def serve(address: str, token: str) -> None:
+    # shortcut: 16 concurrent turn streams; make configurable when gateway load exceeds this bound.
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=16))
+    rpc.add_AgentCoreServicer_to_server(AgentCoreService(token), server)
+    port = server.add_insecure_port(address)
+    if port == 0:
+        raise RuntimeError(f"could not bind agent core to {address}")
+    server.start()
+    print(json.dumps({"status": "ready", "address": f"127.0.0.1:{port}"}), flush=True)
+    server.wait_for_termination()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--address", default="127.0.0.1:0")
+    args = parser.parse_args()
+    token = os.environ.get("BOOMBOT_PYTHON_CORE_TOKEN", "")
+    if len(token) < 32:
+        raise RuntimeError(
+            "BOOMBOT_PYTHON_CORE_TOKEN must contain at least 32 characters"
+        )
+    try:
+        serve(args.address, token)
+    except KeyboardInterrupt:
+        return
+    except Exception as error:
+        print(str(error), file=sys.stderr)
+        raise
+
+
+if __name__ == "__main__":
+    main()