PR-B1 (ADR 0008 Phase B): gRPC server stub — Create / Close / GetSessionInfo

First PR of Phase B. Lands the gRPC RuntimeService surface
implementing three of the five ADR 0008 \u00a72.2 RPCs against the
SessionStore from PR-A2 / PR-A3b. AppendTokens and Generate are
explicitly NOT implemented yet (PR-B2 / PR-B3 territory); calling
them returns gRPC UNIMPLEMENTED, which is the framework default for
non-overridden servicer methods and the right placeholder per
\u00a72.10 'no graceful degradation'.

The server is asyncio (grpc.aio) per \u00a72.5: all RPCs run on a single
event loop, serializing SessionStore access at that layer.

New files:

  inference_engine/server/grpc_app.py (215 lines)
    - RuntimeServiceServicer with CreateSession / CloseSession /
      GetSessionInfo implementations.
    - Error mapping per \u00a72.6 / \u00a72.10:
        SessionNotFoundError -> NOT_FOUND
        PoolExhausted        -> RESOURCE_EXHAUSTED
      InvariantViolation / ValueError mapping land in PR-B2 when
      AppendTokens triggers them; not wired here to keep the
      diff minimal and 100%-tested.
    - GrpcServerConfig (frozen dataclass) with bind_address default
      127.0.0.1:50051 (\u00a78 OQ-5 default — loopback only).
    - create_grpc_server(session_store, config) factory; built but
      not started, so callers control the start / stop lifecycle.

  inference_engine/server/proto_gen/ (generated)
    - kakeya/v1/runtime_pb2.py, runtime_pb2.pyi, runtime_pb2_grpc.py
    - Empty __init__.py at every package level so
      'from inference_engine.server.proto_gen.kakeya.v1 import runtime_pb2'
      works under Python's package layout.
    - Generated by scripts/regenerate_proto_stubs.sh; CI's
      proto-stub-drift job re-runs the script and 'git diff
      --exit-code' to catch silent drift between proto/ and stubs.

  scripts/regenerate_proto_stubs.sh
    - Canonical regeneration command. Patches protoc's absolute
      'from kakeya.v1 import runtime_pb2' to relative
      'from . import runtime_pb2' (a known protoc/Python layout
      issue: protocolbuffers/protobuf#1491).

  tests/inference_engine/server/test_grpc_app.py (22 tests)
    - Real grpc.aio.server bound to 127.0.0.1:0 (random free port);
      real grpc.aio.insecure_channel client. End-to-end coverage
      of every reachable code path including PoolExhausted and
      SessionNotFoundError mappings.
    - Two regression tests for AppendTokens / Generate returning
      UNIMPLEMENTED (so a future PR-B2 / PR-B3 that forgets to
      implement them is caught at PR review time, not in production).

Modified files:

  requirements.txt
    + grpcio>=1.65,<2.0
    + grpcio-tools>=1.65,<2.0  (regen + drift-check)
  .coveragerc
    omit += inference_engine/server/proto_gen/*
    (generated stubs are not the surface we own; coverage on them
    is not meaningful or stable across protoc versions)
  .github/workflows/ci.yaml
    + proto-stub-drift job
    + grpc_app + proto_gen.kakeya.v1.runtime_pb2{,_grpc} imports
      added to package-import-smoke

Local verification (Linux VM, py3.12):
  Linux CI gate: 562 passed (was 540 + 22 new), coverage 100.00 %
                 on 1382 stmts (was 1336 + 46 new in grpc_app.py).
  Regen script idempotent: scripts/regenerate_proto_stubs.sh
                           produces byte-identical stubs to the
                           committed ones.
  Servicer methods exercised end-to-end via real gRPC channel
  (no mocks of the SUT; ServicerContext is the framework's, not a
  test double).

Per ADR 0008 \u00a79: this PR is Linux-only — no MLX paths touched, no
hardware-specific code. \u00a79 last-paragraph carve-out invoked:
'Linux-only path' justification, no Mac M4 integration test report
needed. The Mac-M4-only suite (tests/backends/mlx/test_verifier.py
etc.) is unaffected by this PR's diff.

Next PR after merge:
  PR-B2 (ADR 0008 \u00a76.2): wire AppendTokens through SessionStore
        + the \u00a72.3 byte-exact prefill-incremental contract. Adds
        InvariantViolation -> FAILED_PRECONDITION mapping (not
        reachable in PR-B1's RPC surface but reachable in
        AppendTokens). Linux-only path; \u00a79 carve-out continues to
        apply until PR-B3's Generate touches the verifier sampler
        on real MLX.

Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

Author: cursoragent

.coveragerc

@@ -30,6 +30,14 @@ branch = false
 # tested library code).
 omit =
     kv_cache_proposer/run_demo.py
+    # Generated gRPC / protobuf stubs (PR-B1 of ADR 0008). The
+    # canonical regeneration command is
+    # ``scripts/regenerate_proto_stubs.sh``; CI's
+    # ``proto-stub-drift`` step verifies that the committed stubs
+    # match what the script would produce. The stubs are not the
+    # surface we own — protoc owns them — so coverage on them is
+    # neither meaningful nor stable across protoc versions.
+    inference_engine/server/proto_gen/*
 
 [report]
 exclude_lines =

.github/workflows/ci.yaml

@@ -139,6 +139,9 @@ jobs:
             import inference_engine.pipeline.coordinator; \
             import inference_engine.session; \
             import inference_engine.session.store; \
+            import inference_engine.server.grpc_app; \
+            import inference_engine.server.proto_gen.kakeya.v1.runtime_pb2; \
+            import inference_engine.server.proto_gen.kakeya.v1.runtime_pb2_grpc; \
             import inference_engine.proposer; \
             import inference_engine.proposer.sparse_logits; \
             import inference_engine.backends.mlx.env; \
@@ -214,3 +217,42 @@ jobs:
             printf '%s\n' "$diff_output"
             exit 1
           fi
+
+  proto-stub-drift:
+    name: proto stub drift
+    runs-on: ubuntu-latest
+    # ADR 0008 PR-B1: the committed Python stubs under
+    # inference_engine/server/proto_gen/ MUST be byte-identical to what
+    # scripts/regenerate_proto_stubs.sh produces from the .proto. This
+    # job re-runs the script, then `git diff --exit-code` fails CI if
+    # anything changed. Catches drift between the .proto contract and
+    # the SDK-consumed Python API.
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+
+      - name: Set up Python 3.12
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+
+      - name: Install grpcio-tools
+        run: |
+          python -m pip install --upgrade pip
+          # Pin the same grpcio-tools range as requirements.txt so the
+          # stubs we regenerate match what production uses. If
+          # grpcio-tools updates and starts producing different stub
+          # bytes, this job catches it as a drift before merge.
+          pip install 'grpcio>=1.65,<2.0' 'grpcio-tools>=1.65,<2.0'
+
+      - name: Regenerate stubs
+        run: bash scripts/regenerate_proto_stubs.sh
+
+      - name: Fail if regenerated stubs differ from committed stubs
+        run: |
+          if ! git diff --exit-code -- inference_engine/server/proto_gen/; then
+            echo "::error::Committed stubs are out of date with proto/."
+            echo "Run scripts/regenerate_proto_stubs.sh locally and commit."
+            exit 1
+          fi

inference_engine/server/grpc_app.py

@@ -0,0 +1,215 @@
+"""gRPC runtime service for Kakeya — ADR 0008 PR-B1 (Phase B).
+
+Implements three RPCs from `proto/kakeya/v1/runtime.proto` against
+the :class:`SessionStore` from PR-A2 / PR-A3b:
+
+- ``CreateSession``
+- ``CloseSession``
+- ``GetSessionInfo``
+
+The remaining RPCs (``AppendTokens``, ``Generate``) are intentionally
+*not* implemented in this PR; calling them returns
+``UNIMPLEMENTED`` (the gRPC framework's default for un-overridden
+servicer methods, which is the correct stop-gap behavior per ADR
+0008 §2.10 "no graceful degradation"). They land in PR-B2 / PR-B3.
+
+The server is **asyncio-based** (``grpc.aio``) per ADR 0008 §2.5:
+all RPCs run on a single event loop and serialize SessionStore
+access at the asyncio layer. The store itself is single-threaded;
+multi-worker support is v0.4 scope (ADR 0008 §4.5).
+
+This module deliberately does *not* depend on FastAPI or the
+deprecated HTTP shim. The two surfaces share only the
+``SessionStore`` (or, in the deprecated shim's case, will share it
+once PR-D1 lands). They can be started together by the same CLI
+entry point or separately; the factory ``create_grpc_server`` is
+self-contained so any of those wirings is possible without code
+change here.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+import grpc
+
+from inference_engine.memory.pool import PoolExhausted
+from inference_engine.server.proto_gen.kakeya.v1 import (
+    runtime_pb2,
+    runtime_pb2_grpc,
+)
+from inference_engine.session import (
+    InvariantViolation,
+    SessionNotFoundError,
+    SessionStore,
+)
+
+_logger = logging.getLogger(__name__)
+
+DEFAULT_BIND_ADDRESS = "127.0.0.1:50051"
+"""Default bind address for the gRPC server.
+
+Per ADR 0008 §8 OQ-5 default while unresolved: bind to loopback
+only. Multi-tenant deployments that need to reach the runtime from
+another host configure a different bind address explicitly and add
+the appropriate auth (also OQ-5 scope; v0.4)."""
+
+
+@dataclass(frozen=True)
+class GrpcServerConfig:
+    """Configuration for the gRPC runtime server.
+
+    Kept as a frozen dataclass so the configuration is auditable at
+    construction time and accidentally re-binding mid-session is a
+    structural impossibility.
+    """
+
+    bind_address: str = DEFAULT_BIND_ADDRESS
+    """`host:port` to bind. Defaults to loopback per ADR 0008 §8 OQ-5."""
+
+    max_concurrent_rpcs: Optional[int] = None
+    """Per-server cap on in-flight RPCs.
+
+    ``None`` defers to grpc.aio's default. PR-B1 leaves this
+    unconstrained because the per-session concurrency cap (ADR 0008
+    §2.5 ``max_concurrent``) is enforced at the SessionStore /
+    Generate level, not at the RPC dispatch level. Set explicitly
+    when running on a constrained host where the gRPC handler
+    threads themselves are the bottleneck."""
+
+
+class RuntimeServiceServicer(runtime_pb2_grpc.RuntimeServiceServicer):
+    """RuntimeService implementation backed by a :class:`SessionStore`.
+
+    Error mapping (per ADR 0008 §2.6 / §2.10 — every typed
+    SessionStore error becomes a typed gRPC status, no silent
+    fallback):
+
+    +-----------------------------+----------------------------------+
+    | SessionStoreError subclass  | gRPC status                      |
+    +=============================+==================================+
+    | SessionNotFoundError        | NOT_FOUND                        |
+    +-----------------------------+----------------------------------+
+    | InvariantViolation          | FAILED_PRECONDITION              |
+    +-----------------------------+----------------------------------+
+    | (PoolExhausted from pool)   | RESOURCE_EXHAUSTED               |
+    +-----------------------------+----------------------------------+
+    | ValueError (token-id range) | INVALID_ARGUMENT                 |
+    +-----------------------------+----------------------------------+
+
+    Of these, PR-B1's three RPCs only ever raise
+    SessionNotFoundError or PoolExhausted (Create can raise
+    PoolExhausted; Close and GetSessionInfo can raise
+    SessionNotFoundError). InvariantViolation and ValueError become
+    reachable in PR-B2 (``AppendTokens``) and are wired here for
+    forward-compatibility — but un-tested in PR-B1 because the
+    RPC paths that trigger them do not exist yet.
+    """
+
+    def __init__(self, session_store: SessionStore) -> None:
+        self._store = session_store
+
+    async def CreateSession(  # noqa: N802 — gRPC-generated method casing
+        self,
+        request: runtime_pb2.CreateSessionRequest,
+        context: grpc.aio.ServicerContext,
+    ) -> runtime_pb2.CreateSessionResponse:
+        """Allocate a new session; return its server-issued id.
+
+        ADR 0008 §2.2 contract item 1: clients have no input on the
+        ``session_id`` value; this RPC is the only producer.
+        """
+        try:
+            session = self._store.create_session(
+                eos_token_ids=list(request.eos_token_ids),
+                client_label=request.client_label,
+            )
+        except PoolExhausted as exc:
+            await context.abort(
+                grpc.StatusCode.RESOURCE_EXHAUSTED,
+                f"slab pool exhausted: {exc}",
+            )
+        return runtime_pb2.CreateSessionResponse(
+            session_id=session.session_id,
+        )
+
+    async def CloseSession(  # noqa: N802
+        self,
+        request: runtime_pb2.CloseSessionRequest,
+        context: grpc.aio.ServicerContext,
+    ) -> runtime_pb2.CloseSessionResponse:
+        """Close a session and return its final history length.
+
+        Returns NOT_FOUND if the session is unknown (closed,
+        evicted, never existed — caller cannot distinguish, by
+        ADR 0008 §2.6 design).
+        """
+        try:
+            final_length = self._store.close_session(request.session_id)
+        except SessionNotFoundError as exc:
+            await context.abort(grpc.StatusCode.NOT_FOUND, str(exc))
+        return runtime_pb2.CloseSessionResponse(
+            final_history_length=final_length,
+        )
+
+    async def GetSessionInfo(  # noqa: N802
+        self,
+        request: runtime_pb2.GetSessionInfoRequest,
+        context: grpc.aio.ServicerContext,
+    ) -> runtime_pb2.GetSessionInfoResponse:
+        """Return diagnostic counters for a session.
+
+        Surfaces ADR 0008 §2.8's anomaly-invariant counters; healthy
+        operation reports zero for both INV-1 and INV-2. Non-zero
+        values are paging-grade — the session itself has by then
+        been removed from the store, and a follow-up
+        ``GetSessionInfo`` on the same id will return NOT_FOUND.
+        """
+        try:
+            session = self._store.get_session(request.session_id)
+        except SessionNotFoundError as exc:
+            await context.abort(grpc.StatusCode.NOT_FOUND, str(exc))
+        return runtime_pb2.GetSessionInfoResponse(
+            history_length=session.history_length,
+            kv_live_bytes=session.kv_live_bytes(),
+            cache_invariant_inv1_violations=session.inv1_violations,
+            cache_invariant_inv2_violations=session.inv2_violations,
+            idle_seconds=session.idle_seconds,
+        )
+
+
+def create_grpc_server(
+    *,
+    session_store: SessionStore,
+    config: Optional[GrpcServerConfig] = None,
+) -> grpc.aio.Server:
+    """Build, but do not start, a configured gRPC asyncio server.
+
+    The caller invokes ``await server.start()`` and ``await
+    server.wait_for_termination()`` (or ``await server.stop(grace)``
+    for shutdown). This split is intentional: tests construct
+    servers without starting them, and the eventual production
+    entry point may want to wire signal handlers between
+    construction and start.
+
+    The bound port is observable via the returned server's
+    ``add_insecure_port`` return value; callers that need the port
+    should use the lower-level ``grpc.aio.server()`` directly,
+    because PR-B1 returns the constructed server with the port
+    already bound (so the asyncio event loop sees the listen socket
+    immediately).
+    """
+    if config is None:
+        config = GrpcServerConfig()
+    server = grpc.aio.server(
+        maximum_concurrent_rpcs=config.max_concurrent_rpcs,
+    )
+    runtime_pb2_grpc.add_RuntimeServiceServicer_to_server(
+        RuntimeServiceServicer(session_store),
+        server,
+    )
+    server.add_insecure_port(config.bind_address)
+    _logger.info("gRPC RuntimeService bound to %s", config.bind_address)
+    return server