pytest, guide

README.md

@@ -1,10 +1,22 @@
 # rust-python-lab
 
-A Rust async execution controller for Python, exposed through PyO3. Python remains the authoring surface; Rust owns scheduling, limits, and lifecycle. Python's asyncio event loop still executes coroutine bodies.
+A Tokio-backed **execution controller** for Python AI/agent apps. Python stays the authoring surface (`await handle`, `async with`, `async for`). Rust owns admission, concurrency, timeouts, cancellation, retries, streaming, and shutdown. Python's asyncio event loop still executes coroutine bodies — LLM SDKs, HTTP clients, DB drivers, and tools.
+
+This is not a LangChain/LangGraph rewrite and not a Rust HTTP/LLM client. Rust waits on bridged asyncio futures; it does not step Python bytecode.
+
+## Execution split
+
+| Layer | Owns |
+| --- | --- |
+| **Rust** | queues, permits, when a task may start, timeout, cancel, retry/backoff, gather/stream aggregation, graceful shutdown, in-memory journal |
+| **Python asyncio** | coroutine bodies, sockets, SDK calls, object lifetimes, GIL |
+| **Bridge** | `pyo3-async-runtimes`: capture `TaskLocals`, convert awaitables both ways, never hold the GIL across `.await` |
+
+Cancel and timeout must cancel the **asyncio Task**, not only drop a Tokio waiter. Retries re-invoke the Python callable to get a new coroutine.
 
 ## Setup
 
-Requires Rust (stable) and Python 3.12 (local development uses 3.12.3). CI installs CPython 3.12.
+Requires Rust (stable) and Python 3.12. Local development uses 3.12.3.
 
 ```bash
 python -m venv .venv
@@ -13,10 +25,85 @@ pip install maturin
 maturin develop --extras dev
 ```
 
+## API
+
+```python
+from rust_python_lab import OnFull, RetryPolicy, Runtime, Task
+
+async with Runtime(
+    max_concurrency=32,
+    queue_capacity=128,
+    on_full=OnFull.REJECT,  # or OnFull.WAIT
+    default_timeout=30.0,
+    retry=RetryPolicy(max_attempts=3, backoff_ms=100),
+    idle_ttl=None,
+) as rt:
+    handle = rt.submit(call_llm, "hello", timeout=10.0)
+    result = await handle
+    handle.cancel()
+
+    results = await rt.gather(
+        [Task(call_llm, prompt) for prompt in prompts],
+        return_exceptions=True,
+    )
+    async for event in rt.as_completed(tasks):
+        ...
+    async for token in rt.stream(llm_tokens, "hello"):
+        ...
+
+    print(rt.stats())
+    await rt.shutdown(timeout=5.0)
+```
+
+After close, `submit` raises `RuntimeClosed`. A full queue with `OnFull.REJECT` raises `QueueFull`.
+
 ## Test
 
+CI and local default: **localhost HTTP only**, no API keys, no public internet.
+
 ```bash
 pytest
 ```
 
-CI runs `maturin develop` and `pytest` on pull requests. No API keys are required.
+The suite drives a local `aiohttp` server (`tests/http_harness.py`) through **httpx** so policy is proven against real sockets:
+
+- concurrency cap, timeout, and cancel against `/work` and `/slow`
+- retries against `/flaky`
+- `gather` / `as_completed` / SSE `stream` against mixed latencies and `/sse`
+- shutdown cancels in-flight `/slow` and rejects later submits
+
+Short soak (2k tiny GETs; `SOAK=1` runs 20k):
+
+```bash
+pytest tests/test_soak.py
+SOAK=1 pytest tests/test_soak.py
+```
+
+Agent-shaped example (tools + SSE “LLM”) against the same harness:
+
+```bash
+python examples/agent_tools.py
+```
+
+Control-plane baseline (not a CI gate): `Runtime.gather` vs `asyncio.gather` + `asyncio.Semaphore` on local HTTP.
+
+```bash
+python scripts/bench_fanout.py
+```
+
+### Optional live LLM
+
+Skipped unless a key is present. CI never sets keys. Cheap smoke: 2–3 short prompts, `as_completed`, one token stream, one cancelled handle.
+
+```bash
+python examples/live_llm.py
+```
+
+Uses `OPENAI_API_KEY` (preferred) or `ANTHROPIC_API_KEY`. Optional `OPENAI_MODEL` / `ANTHROPIC_MODEL`. To also run it under pytest: `LIVE_LLM=1 pytest tests/test_examples.py`. Do not log secrets.
+
+## What this does not prove
+
+- Rust is not faster at token generation. The model and Python SDK still dominate latency. Benchmarks measure **scheduler** cost on local HTTP, not provider SLA.
+- No crash-durable journal replay, no multi-process supervisor, no sandbox for tool code.
+- Live tests need a key, cost a few tokens, and will be rate-limited. They are a smoke, not a soak.
+- No built-in Rust HTTP or LLM adapters in this lab. Coroutine bodies stay on asyncio on purpose.

examples/agent_tools.py

@@ -0,0 +1,67 @@
+"""Agent-shaped fan-out against the local HTTP harness.
+
+Several tools are ordinary HTTP calls; the "LLM" is an SSE token stream.
+Run from the repo root after `maturin develop --extras dev`:
+
+    python examples/agent_tools.py
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from pathlib import Path
+
+from aiohttp.test_utils import TestServer
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from rust_python_lab import RetryPolicy, Runtime, Task  # noqa: E402
+from tests.http_harness import Counters, create_app, http_get, sse_tokens  # noqa: E402
+
+
+async def search_docs(base: str, query: str) -> str:
+    _ = query
+    return await http_get(f"{base}/work?delay=0.05")
+
+
+async def fetch_user(base: str, user_id: str) -> str:
+    _ = user_id
+    return await http_get(f"{base}/work?delay=0.05")
+
+
+async def flaky_write(base: str, note: str) -> str:
+    _ = note
+    return await http_get(f"{base}/flaky?fail_times=2")
+
+
+async def main() -> None:
+    counters = Counters()
+    async with TestServer(create_app(counters)) as server:
+        base = str(server.make_url("")).rstrip("/")
+        retry = RetryPolicy(max_attempts=3, backoff_ms=20, jitter=False)
+        async with Runtime(max_concurrency=8, retry=retry) as rt:
+            tool_results = await rt.gather(
+                [
+                    Task(search_docs, base, "runtime limits"),
+                    Task(fetch_user, base, "u-1"),
+                    Task(flaky_write, base, "remember the last tool result"),
+                ]
+            )
+            tokens = [token async for token in rt.stream(sse_tokens, f"{base}/sse")]
+            stats = rt.stats()
+
+    print("tools:", tool_results)
+    print("llm tokens:", " ".join(tokens))
+    print("stats:", stats)
+    print("server peak in-flight:", counters.peak_in_flight)
+    assert tool_results == ["ok", "ok", "ok"]
+    assert tokens == [f"token-{i}" for i in range(8)]
+    assert stats["in_flight"] == 0
+    assert counters.flaky_hits == 3
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/live_llm.py

@@ -0,0 +1,181 @@
+"""Optional OpenAI / Anthropic smoke. Skips unless a key is set.
+
+Never required for CI. Uses httpx (no official SDKs, no LangChain).
+
+    python examples/live_llm.py
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import sys
+
+import httpx
+
+from rust_python_lab import Runtime, Task
+
+OPENAI_URL = "https://api.openai.com/v1/chat/completions"
+ANTHROPIC_URL = "https://api.anthropic.com/v1/messages"
+PROMPTS = (
+    "Reply with the single word ping.",
+    "Reply with the single word pong.",
+    "Reply with the single word ok.",
+)
+
+
+def _provider() -> str | None:
+    if os.environ.get("OPENAI_API_KEY"):
+        return "openai"
+    if os.environ.get("ANTHROPIC_API_KEY"):
+        return "anthropic"
+    return None
+
+
+async def openai_complete(prompt: str) -> str:
+    key = os.environ["OPENAI_API_KEY"]
+    model = os.environ.get("OPENAI_MODEL", "gpt-4o-mini")
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": 32,
+    }
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        response = await client.post(
+            OPENAI_URL,
+            headers={"Authorization": f"Bearer {key}"},
+            json=payload,
+        )
+        response.raise_for_status()
+        data = response.json()
+    return data["choices"][0]["message"]["content"]
+
+
+async def openai_tokens(prompt: str):
+    key = os.environ["OPENAI_API_KEY"]
+    model = os.environ.get("OPENAI_MODEL", "gpt-4o-mini")
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": 32,
+        "stream": True,
+    }
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        async with client.stream(
+            "POST",
+            OPENAI_URL,
+            headers={"Authorization": f"Bearer {key}"},
+            json=payload,
+        ) as response:
+            response.raise_for_status()
+            async for line in response.aiter_lines():
+                if not line.startswith("data: "):
+                    continue
+                data = line[6:]
+                if data == "[DONE]":
+                    return
+                delta = json.loads(data)["choices"][0].get("delta", {})
+                text = delta.get("content")
+                if text:
+                    yield text
+
+
+async def anthropic_complete(prompt: str) -> str:
+    key = os.environ["ANTHROPIC_API_KEY"]
+    model = os.environ.get("ANTHROPIC_MODEL", "claude-3-haiku-20240307")
+    payload = {
+        "model": model,
+        "max_tokens": 32,
+        "messages": [{"role": "user", "content": prompt}],
+    }
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        response = await client.post(
+            ANTHROPIC_URL,
+            headers={
+                "x-api-key": key,
+                "anthropic-version": "2023-06-01",
+            },
+            json=payload,
+        )
+        response.raise_for_status()
+        data = response.json()
+    return "".join(
+        block.get("text", "")
+        for block in data.get("content", [])
+        if block.get("type") == "text"
+    )
+
+
+async def anthropic_tokens(prompt: str):
+    key = os.environ["ANTHROPIC_API_KEY"]
+    model = os.environ.get("ANTHROPIC_MODEL", "claude-3-haiku-20240307")
+    payload = {
+        "model": model,
+        "max_tokens": 32,
+        "stream": True,
+        "messages": [{"role": "user", "content": prompt}],
+    }
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        async with client.stream(
+            "POST",
+            ANTHROPIC_URL,
+            headers={
+                "x-api-key": key,
+                "anthropic-version": "2023-06-01",
+            },
+            json=payload,
+        ) as response:
+            response.raise_for_status()
+            async for line in response.aiter_lines():
+                if not line.startswith("data: "):
+                    continue
+                event = json.loads(line[6:])
+                if event.get("type") == "content_block_delta":
+                    text = event.get("delta", {}).get("text")
+                    if text:
+                        yield text
+
+
+async def main() -> int:
+    provider = _provider()
+    if provider is None:
+        print("skip: set OPENAI_API_KEY or ANTHROPIC_API_KEY to run the live LLM smoke")
+        return 0
+
+    complete = openai_complete if provider == "openai" else anthropic_complete
+    tokens = openai_tokens if provider == "openai" else anthropic_tokens
+    print(f"provider: {provider}")
+
+    async with Runtime(max_concurrency=3, default_timeout=30.0) as rt:
+        order: list[int] = []
+        async for event in rt.as_completed(
+            [Task(complete, prompt) for prompt in PROMPTS[:2]]
+        ):
+            if not event.ok:
+                raise RuntimeError(f"completion {event.index} failed")
+            order.append(event.index)
+            print(f"completed[{event.index}]: {event.value!r}")
+
+        streamed = []
+        async for piece in rt.stream(tokens, PROMPTS[2]):
+            streamed.append(piece)
+        print("streamed chars:", len("".join(streamed)))
+
+        handle = rt.submit(complete, PROMPTS[0])
+        handle.cancel()
+        try:
+            await handle
+        except asyncio.CancelledError:
+            print("cancelled: ok")
+
+    print("as_completed order:", order)
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        raise SystemExit(asyncio.run(main()))
+    except httpx.HTTPError as exc:
+        print(f"live llm request failed: {exc}", file=sys.stderr)
+        raise SystemExit(1)

pyproject.toml

@@ -22,6 +22,7 @@ dev = [
 
 [tool.maturin]
 features = ["pyo3/extension-module"]
+python-source = "python"
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]

python/rust_python_lab/__init__.py

@@ -0,0 +1,23 @@
+"""Tokio-backed async execution controller for Python."""
+
+from .rust_python_lab import (
+    Completion,
+    Handle,
+    OnFull,
+    QueueFull,
+    RetryPolicy,
+    Runtime,
+    RuntimeClosed,
+    Task,
+)
+
+__all__ = [
+    "Completion",
+    "Handle",
+    "OnFull",
+    "QueueFull",
+    "RetryPolicy",
+    "Runtime",
+    "RuntimeClosed",
+    "Task",
+]