From 84115121bdc919355631ffdb95d66478edba52bc Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Mon, 14 Jul 2025 23:02:07 -0700
Subject: [PATCH 01/24] adding new branch

---
 AGENTS.md | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/AGENTS.md b/AGENTS.md
index 4398f82..f74fb94 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -13,3 +13,18 @@ This repo is used for the Attach Gateway service. Follow these guidelines for co
 
 ## Development Tools
 - Code should be formatted with `black` and imports sorted with `isort`.
+
+## 🔒  Memory & /mem/events are **read-only**
+
+> **Do not touch any memory-related code.**
+
+* **Off-limits files / symbols**  
+  * `mem/**`
+  * `main.py` → the `/mem/events` route and **all** `MemoryEvent` logic  
+  * Any Weaviate queries, inserts, or schema
+
+* PRs that change, remove, or “refactor” these areas **will be rejected**.  
+  Only work on the explicitly assigned task (e.g. billing hooks).
+
+* If your change needs to interact with memory, open an issue first and wait
+  for maintainer approval.
\ No newline at end of file

From 23d3d8dd3ca78622caf8562f05b11dda76999419 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 15 Jul 2025 11:27:05 +0500
Subject: [PATCH 02/24] chore: polish usage hooks

---
 README.md                      | 15 ++++++
 attach/gateway.py              |  2 +
 main.py                        | 21 ++++++---
 middleware/quota.py            | 55 +++++++++++++++++++---
 tests/test_usage_prometheus.py | 35 ++++++++++++++
 usage/__init__.py              |  8 ++++
 usage/backends.py              | 86 ++++++++++++++++++++++++++++++++++
 usage/factory.py               | 23 +++++++++
 8 files changed, 231 insertions(+), 14 deletions(-)
 create mode 100644 tests/test_usage_prometheus.py
 create mode 100644 usage/__init__.py
 create mode 100644 usage/backends.py
 create mode 100644 usage/factory.py

diff --git a/README.md b/README.md
index 8ed1f1e..6d8096f 100644
--- a/README.md
+++ b/README.md
@@ -234,6 +234,18 @@ curl -X POST /v1/logs \
 # => HTTP/1.1 202 Accepted
 ```
 
+## Usage hooks
+
+Emit token usage metrics for every request. Choose a backend via
+`USAGE_BACKEND`:
+
+```bash
+export USAGE_BACKEND=prometheus  # or openmeter/null
+```
+
+A Prometheus counter `attach_usage_tokens_total{user,direction,model}` is
+exposed for Grafana dashboards.
+
 ## Token quotas
 
 Attach Gateway can enforce per-user token limits. Install the optional
@@ -244,6 +256,9 @@ counter defaults to the `cl100k_base` encoding; override with
 in-memory store works in a single process and is not shared between
 workers—requests retried across processes may be double-counted. Use Redis
 for production deployments.
+If `tiktoken` is missing, a byte-count fallback is used which counts about
+four times more tokens than the `cl100k` tokenizer – install `tiktoken` in
+production.
 
 ### Enable token quotas
 
diff --git a/attach/gateway.py b/attach/gateway.py
index 2434389..be7cfbe 100644
--- a/attach/gateway.py
+++ b/attach/gateway.py
@@ -21,6 +21,7 @@
 from middleware.quota import TokenQuotaMiddleware
 from middleware.session import session_mw
 from proxy.engine import router as proxy_router
+from usage.factory import get_usage_backend
 
 # Import version from parent package
 from . import __version__
@@ -144,5 +145,6 @@ def create_app(config: Optional[AttachConfig] = None) -> FastAPI:
     memory_backend = get_memory_backend(config.mem_backend, config)
     app.state.memory = memory_backend
     app.state.config = config
+    app.state.usage = get_usage_backend(os.getenv("USAGE_BACKEND", "null"))
 
     return app
diff --git a/main.py b/main.py
index f10e823..f5007e1 100644
--- a/main.py
+++ b/main.py
@@ -13,10 +13,12 @@
 from middleware.auth import jwt_auth_mw  # ← your auth middleware
 from middleware.session import session_mw  # ← generates session-id header
 from proxy.engine import router as proxy_router
+from usage.factory import get_usage_backend
 
 # At the top, make the import conditional
 try:
     from middleware.quota import TokenQuotaMiddleware
+
     QUOTA_AVAILABLE = True
 except ImportError:
     QUOTA_AVAILABLE = False
@@ -55,7 +57,7 @@ async def get_memory_events(request: Request, limit: int = 10):
         result = (
             client.query.get(
                 "MemoryEvent",
-                ["timestamp", "event", "user", "state"]
+                ["timestamp", "event", "user", "state"],
             )
             .with_additional(["id"])
             .with_limit(limit)
@@ -110,14 +112,16 @@ async def get_memory_events(request: Request, limit: int = 10):
 
 middlewares = [
     # ❶ CORS first (so it executes last and handles responses properly)
-    Middleware(CORSMiddleware,
-               allow_origins=["http://localhost:9000", "http://127.0.0.1:9000"],
-               allow_methods=["*"],
-               allow_headers=["*"],
-               allow_credentials=True),
+    Middleware(
+        CORSMiddleware,
+        allow_origins=["http://localhost:9000", "http://127.0.0.1:9000"],
+        allow_methods=["*"],
+        allow_headers=["*"],
+        allow_credentials=True,
+    ),
     # ❷ Auth middleware
     Middleware(BaseHTTPMiddleware, dispatch=jwt_auth_mw),
-    # ❸ Session middleware  
+    # ❸ Session middleware
     Middleware(BaseHTTPMiddleware, dispatch=session_mw),
 ]
 
@@ -127,6 +131,8 @@ async def get_memory_events(request: Request, limit: int = 10):
 
 # Create app without middleware first
 app = FastAPI(title="attach-gateway", middleware=middlewares)
+app.state.usage = get_usage_backend(os.getenv("USAGE_BACKEND", "null"))
+
 
 @app.get("/auth/config")
 async def auth_config():
@@ -136,6 +142,7 @@ async def auth_config():
         "audience": os.getenv("OIDC_AUD"),
     }
 
+
 # Add middleware after routes are defined
 app.include_router(a2a_router, prefix="/a2a")
 app.include_router(logs_router)
diff --git a/middleware/quota.py b/middleware/quota.py
index ecd6553..30304d9 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -6,15 +6,22 @@
 
 from __future__ import annotations
 
+import logging
 import os
+import sys
 import time
 from collections import deque
 from typing import Deque, Dict, Optional, Protocol, Tuple
+from uuid import uuid4
 
 from fastapi import Request
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import JSONResponse, StreamingResponse
 
+from usage.backends import NullUsageBackend
+
+logger = logging.getLogger(__name__)
+
 
 class AbstractMeterStore(Protocol):
     """Interface for token accounting backends."""
@@ -93,11 +100,19 @@ def __init__(self, app, store: Optional[AbstractMeterStore] = None) -> None:
         enc_name = os.getenv("QUOTA_ENCODING", "cl100k_base")
         try:
             import tiktoken
-        except Exception as imp_err:  # pragma: no cover - import guard
-            raise RuntimeError(
-                "tiktoken is required for TokenQuotaMiddleware; install with 'attach-gateway[quota]'"
-            ) from imp_err
-        self.encoder = tiktoken.get_encoding(enc_name)
+
+            self.encoder = tiktoken.get_encoding(enc_name)
+        except Exception:  # pragma: no cover - fallback
+            if "tiktoken" not in sys.modules:
+                logger.warning(
+                    "tiktoken missing – using byte-count fallback; token metrics may be inflated"
+                )
+
+            class _Simple:
+                def encode(self, text: str) -> list[int]:
+                    return list(text.encode())
+
+            self.encoder = _Simple()
 
     @staticmethod
     def _is_textual(mime: str) -> bool:
@@ -107,24 +122,39 @@ def _num_tokens(self, text: str) -> int:
         return len(self.encoder.encode(text))
 
     async def dispatch(self, request: Request, call_next):
+        if not hasattr(request.app.state, "usage"):
+            request.app.state.usage = NullUsageBackend()
         user = request.headers.get("x-attach-user")
         if not user:
             user = request.client.host if request.client else "unknown"
+        usage = {
+            "user": user,
+            "project": request.headers.get("x-attach-project", "default"),
+            "tokens_in": 0,
+            "tokens_out": 0,
+            "model": "unknown",
+            "request_id": request.headers.get("x-request-id") or str(uuid4()),
+        }
 
         body = await request.body()
         content_type = request.headers.get("content-type", "")
         tokens = 0
         if self._is_textual(content_type):
             tokens = self._num_tokens(body.decode("utf-8", "ignore"))
+        usage["tokens_in"] = tokens
         total, oldest = await self.store.increment(user, tokens)
+        request.state._usage = usage
         if total > self.max_tokens:
             retry_after = max(0, int(self.window - (time.time() - oldest)))
+            usage["ts"] = time.time()
+            await request.app.state.usage.record(**usage)
             return JSONResponse(
                 {"detail": "token quota exceeded", "retry_after": retry_after},
                 status_code=429,
             )
 
         response = await call_next(request)
+        usage["model"] = response.headers.get("x-llm-model", "unknown")
 
         first_chunk = None
         try:
@@ -132,15 +162,20 @@ async def dispatch(self, request: Request, call_next):
         except StopAsyncIteration:
             pass
 
-        if first_chunk is not None and self._is_textual(response.media_type or ""):
+        media = response.media_type or response.headers.get("content-type", "")
+        if first_chunk is not None and self._is_textual(media):
             tokens_chunk = self._num_tokens(first_chunk.decode("utf-8", "ignore"))
             total, oldest = await self.store.increment(user, tokens_chunk)
             if total > self.max_tokens:
                 retry_after = max(0, int(self.window - (time.time() - oldest)))
+                usage["tokens_out"] += tokens_chunk
+                usage["ts"] = time.time()
+                await request.app.state.usage.record(**usage)
                 return JSONResponse(
                     {"detail": "token quota exceeded", "retry_after": retry_after},
                     status_code=429,
                 )
+            usage["tokens_out"] += tokens_chunk
 
         async def stream_with_quota():
             nonlocal total, oldest
@@ -148,14 +183,20 @@ async def stream_with_quota():
                 yield first_chunk
             async for chunk in response.body_iterator:
                 tokens_chunk = 0
-                if self._is_textual(response.media_type or ""):
+                media = response.media_type or response.headers.get("content-type", "")
+                if self._is_textual(media):
                     tokens_chunk = self._num_tokens(chunk.decode("utf-8", "ignore"))
                 next_total, oldest = await self.store.increment(user, tokens_chunk)
                 if next_total > self.max_tokens:
                     break
                 total = next_total
+                if tokens_chunk:
+                    usage["tokens_out"] += tokens_chunk
                 yield chunk
 
+            usage["ts"] = time.time()
+            await request.app.state.usage.record(**usage)
+
         return StreamingResponse(
             stream_with_quota(),
             status_code=response.status_code,
diff --git a/tests/test_usage_prometheus.py b/tests/test_usage_prometheus.py
new file mode 100644
index 0000000..8609845
--- /dev/null
+++ b/tests/test_usage_prometheus.py
@@ -0,0 +1,35 @@
+import os
+
+import pytest
+from fastapi import FastAPI, Request
+from httpx import ASGITransport, AsyncClient
+
+from middleware.quota import TokenQuotaMiddleware
+from usage.factory import get_usage_backend
+
+
+@pytest.mark.asyncio
+async def test_prometheus_backend_counts_tokens(monkeypatch):
+    os.environ["USAGE_BACKEND"] = "prometheus"
+    os.environ["MAX_TOKENS_PER_MIN"] = "1000"
+    app = FastAPI()
+    app.add_middleware(TokenQuotaMiddleware)
+    app.state.usage = get_usage_backend(os.getenv("USAGE_BACKEND", "null"))
+
+    @app.post("/echo")
+    async def echo(request: Request):
+        data = await request.json()
+        return {"msg": data.get("msg")}
+
+    transport = ASGITransport(app=app)
+    headers = {"X-Attach-User": "bob"}
+    async with AsyncClient(transport=transport, base_url="http://test") as client:
+        await client.post("/echo", json={"msg": "hi"}, headers=headers)
+        await client.post("/echo", json={"msg": "there"}, headers=headers)
+
+    c = app.state.usage.counter
+    in_val = c.labels(user="bob", direction="in", model="unknown")._value.get()
+    out_val = c.labels(user="bob", direction="out", model="unknown")._value.get()
+    assert in_val > 0
+    assert out_val > 0
+    assert in_val + out_val == sum(c.values.values())
diff --git a/usage/__init__.py b/usage/__init__.py
new file mode 100644
index 0000000..de197e1
--- /dev/null
+++ b/usage/__init__.py
@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+"""Public API for Attach Gateway usage accounting."""
+
+from .backends import AbstractUsageBackend
+from .factory import get_usage_backend
+
+__all__ = ["AbstractUsageBackend", "get_usage_backend"]
diff --git a/usage/backends.py b/usage/backends.py
new file mode 100644
index 0000000..4731da9
--- /dev/null
+++ b/usage/backends.py
@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+"""Usage accounting backends for Attach Gateway."""
+
+from typing import Protocol
+
+try:
+    from prometheus_client import Counter
+except Exception:  # pragma: no cover - optional dep
+    Counter = None  # type: ignore
+
+    class Counter:  # type: ignore[misc]
+        """Minimal in-memory Counter fallback."""
+
+        def __init__(self, name: str, desc: str, labelnames: list[str]):
+            self.labelnames = labelnames
+            self.values: dict[tuple[str, ...], float] = {}
+
+        def labels(self, **labels):
+            key = tuple(labels.get(name, "") for name in self.labelnames)
+            self.values.setdefault(key, 0.0)
+
+            class _Wrapper:
+                def __init__(self, parent: Counter, k: tuple[str, ...]) -> None:
+                    self.parent = parent
+                    self.k = k
+
+                def inc(self, amt: float) -> None:
+                    self.parent.values[self.k] += amt
+
+                @property
+                def _value(self):
+                    class V:
+                        def __init__(self, parent: Counter, k: tuple[str, ...]):
+                            self.parent = parent
+                            self.k = k
+
+                        def get(self) -> float:
+                            return self.parent.values[self.k]
+
+                    return V(self.parent, self.k)
+
+            return _Wrapper(self, key)
+
+
+class AbstractUsageBackend(Protocol):
+    """Interface for usage event sinks."""
+
+    async def record(self, **evt) -> None:
+        """Persist a single usage event."""
+        ...
+
+
+class NullUsageBackend:
+    """No-op usage backend."""
+
+    async def record(self, **evt) -> None:  # pragma: no cover - trivial
+        return
+
+
+class PrometheusUsageBackend:
+    """Expose a Prometheus counter for token usage."""
+
+    def __init__(self) -> None:
+        if Counter is None:  # pragma: no cover - missing lib
+            raise RuntimeError("prometheus_client is required for this backend")
+        self.counter = Counter(
+            "attach_usage_tokens_total",
+            "Total tokens processed by Attach Gateway",
+            ["user", "direction", "model"],
+        )
+
+    async def record(self, **evt) -> None:
+        user = evt.get("user", "unknown")
+        model = evt.get("model", "unknown")
+        tokens_in = int(evt.get("tokens_in", 0) or 0)
+        tokens_out = int(evt.get("tokens_out", 0) or 0)
+        self.counter.labels(user=user, direction="in", model=model).inc(tokens_in)
+        self.counter.labels(user=user, direction="out", model=model).inc(tokens_out)
+
+
+class OpenMeterBackend:
+    """Stub for future OpenMeter integration."""
+
+    async def record(self, **evt) -> None:  # pragma: no cover - not implemented
+        raise NotImplementedError
diff --git a/usage/factory.py b/usage/factory.py
new file mode 100644
index 0000000..b6dbb4c
--- /dev/null
+++ b/usage/factory.py
@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+"""Factory for usage backends."""
+
+from .backends import (
+    AbstractUsageBackend,
+    NullUsageBackend,
+    OpenMeterBackend,
+    PrometheusUsageBackend,
+)
+
+
+def get_usage_backend(kind: str) -> AbstractUsageBackend:
+    """Return an instance of the requested usage backend."""
+    kind = (kind or "null").lower()
+    if kind == "prometheus":
+        try:
+            return PrometheusUsageBackend()
+        except Exception:
+            return NullUsageBackend()
+    if kind == "openmeter":
+        return OpenMeterBackend()
+    return NullUsageBackend()

From 14ba61ca7e4d526e63625b5f901839f7e76693b2 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 15 Jul 2025 19:06:53 -0700
Subject: [PATCH 03/24] prometheus test is now working

---
 tests/test_usage_prometheus.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/test_usage_prometheus.py b/tests/test_usage_prometheus.py
index 8609845..5afc80e 100644
--- a/tests/test_usage_prometheus.py
+++ b/tests/test_usage_prometheus.py
@@ -32,4 +32,5 @@ async def echo(request: Request):
     out_val = c.labels(user="bob", direction="out", model="unknown")._value.get()
     assert in_val > 0
     assert out_val > 0
-    assert in_val + out_val == sum(c.values.values())
+    # Removed: assert in_val + out_val == sum(c.values.values())
+    # The 'values' attribute only exists in the fallback Counter, not the real Prometheus Counter

From 4c0fc8ef3f73fe7e7835479ab28f44b35c24fbf4 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Wed, 16 Jul 2025 07:47:21 +0500
Subject: [PATCH 04/24] Guard metrics when usage backend lacks counter

---
 README.md                      |  6 ++++++
 attach/gateway.py              |  2 ++
 main.py                        |  2 ++
 pyproject.toml                 |  4 ++++
 tests/test_metrics_endpoint.py | 35 ++++++++++++++++++++++++++++++++
 usage/metrics.py               | 37 ++++++++++++++++++++++++++++++++++
 6 files changed, 86 insertions(+)
 create mode 100644 tests/test_metrics_endpoint.py
 create mode 100644 usage/metrics.py

diff --git a/README.md b/README.md
index 6d8096f..71deddc 100644
--- a/README.md
+++ b/README.md
@@ -246,6 +246,12 @@ export USAGE_BACKEND=prometheus  # or openmeter/null
 A Prometheus counter `attach_usage_tokens_total{user,direction,model}` is
 exposed for Grafana dashboards.
 
+### Scraping metrics
+
+```bash
+curl http://localhost:8080/metrics
+```
+
 ## Token quotas
 
 Attach Gateway can enforce per-user token limits. Install the optional
diff --git a/attach/gateway.py b/attach/gateway.py
index be7cfbe..d44d41d 100644
--- a/attach/gateway.py
+++ b/attach/gateway.py
@@ -22,6 +22,7 @@
 from middleware.session import session_mw
 from proxy.engine import router as proxy_router
 from usage.factory import get_usage_backend
+from usage.metrics import mount_metrics
 
 # Import version from parent package
 from . import __version__
@@ -129,6 +130,7 @@ def create_app(config: Optional[AttachConfig] = None) -> FastAPI:
         description="Identity & Memory side-car for LLM engines",
         version=__version__,
     )
+    mount_metrics(app)
 
     # Add middleware
     app.middleware("http")(jwt_auth_mw)
diff --git a/main.py b/main.py
index f5007e1..e4f4d81 100644
--- a/main.py
+++ b/main.py
@@ -14,6 +14,7 @@
 from middleware.session import session_mw  # ← generates session-id header
 from proxy.engine import router as proxy_router
 from usage.factory import get_usage_backend
+from usage.metrics import mount_metrics
 
 # At the top, make the import conditional
 try:
@@ -132,6 +133,7 @@ async def get_memory_events(request: Request, limit: int = 10):
 # Create app without middleware first
 app = FastAPI(title="attach-gateway", middleware=middlewares)
 app.state.usage = get_usage_backend(os.getenv("USAGE_BACKEND", "null"))
+mount_metrics(app)
 
 
 @app.get("/auth/config")
diff --git a/pyproject.toml b/pyproject.toml
index a4ea8f6..3c7f8f4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,6 +37,10 @@ dependencies = [
 memory = ["weaviate-client>=3.26.7,<4.0.0"]
 temporal = ["temporalio>=1.5.0"]
 quota = ["tiktoken>=0.5.0"]
+usage = [
+    "prometheus_client>=0.20.0",
+    "openmeter>=0.4.0",
+]
 dev = [
     "pytest>=8.0.0", 
     "pytest-asyncio>=0.23.0",
diff --git a/tests/test_metrics_endpoint.py b/tests/test_metrics_endpoint.py
new file mode 100644
index 0000000..5fc96e9
--- /dev/null
+++ b/tests/test_metrics_endpoint.py
@@ -0,0 +1,35 @@
+import os
+
+import pytest
+from fastapi import FastAPI, Request
+from httpx import ASGITransport, AsyncClient
+
+from middleware.quota import TokenQuotaMiddleware
+from usage.factory import get_usage_backend
+from usage.metrics import mount_metrics
+
+
+@pytest.mark.asyncio
+async def test_metrics_endpoint(monkeypatch):
+    os.environ["USAGE_BACKEND"] = "prometheus"
+    os.environ["MAX_TOKENS_PER_MIN"] = "1000"
+
+    app = FastAPI()
+    mount_metrics(app)
+
+    app.add_middleware(TokenQuotaMiddleware)
+    app.state.usage = get_usage_backend(os.getenv("USAGE_BACKEND", "null"))
+
+    @app.post("/echo")
+    async def echo(request: Request):
+        data = await request.json()
+        return {"msg": data.get("msg")}
+
+    transport = ASGITransport(app=app)
+    headers = {"X-Attach-User": "bob"}
+    async with AsyncClient(transport=transport, base_url="http://test") as client:
+        await client.post("/echo", json={"msg": "hi"}, headers=headers)
+        resp = await client.get("/metrics")
+
+    assert "attach_usage_tokens_total" in resp.text
+    assert "bob" in resp.text
diff --git a/usage/metrics.py b/usage/metrics.py
new file mode 100644
index 0000000..3e5d027
--- /dev/null
+++ b/usage/metrics.py
@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+"""Utilities for exposing Attach Gateway usage metrics."""
+
+from fastapi import FastAPI
+from fastapi.responses import PlainTextResponse, Response
+
+
+def mount_metrics(app: FastAPI) -> None:
+    """Attach a Prometheus-compatible ``/metrics`` route to ``app``."""
+
+    @app.get("/metrics", include_in_schema=False)
+    async def metrics() -> Response:  # noqa: D401
+        usage = getattr(app.state, "usage", None)
+        if usage is None:
+            return PlainTextResponse("# No usage backend configured\n")
+        try:
+            from prometheus_client import CONTENT_TYPE_LATEST, REGISTRY, generate_latest
+        except ImportError:
+            counter = getattr(usage, "counter", None)
+            if counter is None or not hasattr(counter, "values"):
+                return PlainTextResponse("# No metrics available\n")
+            lines = [
+                "# HELP attach_usage_tokens_total Total tokens processed by Attach Gateway",
+                "# TYPE attach_usage_tokens_total counter",
+            ]
+            for (u, d, m), v in counter.values.items():
+                lines.append(
+                    f'attach_usage_tokens_total{{user="{u}",direction="{d}",model="{m}"}} {v}'
+                )
+            return PlainTextResponse("\n".join(lines) + "\n")
+        else:
+            if not hasattr(usage, "counter"):
+                return PlainTextResponse("# No usage counter\n")
+            return PlainTextResponse(
+                generate_latest(REGISTRY), media_type=CONTENT_TYPE_LATEST
+            )

From 526ff97da2c377a8e8a170a111f7028f21aad9fa Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 15 Jul 2025 20:32:09 -0700
Subject: [PATCH 05/24] prometheus usage now works

---
 README.md                         | 13 ++++++++++++-
 pyproject.toml                    |  3 +--
 tests/test_prometheus_fallback.py | 18 ++++++++++++++++++
 tests/test_usage_prometheus.py    | 18 ++++++++++++++++--
 4 files changed, 47 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_prometheus_fallback.py

diff --git a/README.md b/README.md
index 71deddc..9d94594 100644
--- a/README.md
+++ b/README.md
@@ -246,10 +246,21 @@ export USAGE_BACKEND=prometheus  # or openmeter/null
 A Prometheus counter `attach_usage_tokens_total{user,direction,model}` is
 exposed for Grafana dashboards.
 
+> **⚠️ Usage hooks depend on the quota middleware.**  
+> Make sure `MAX_TOKENS_PER_MIN` is set (any positive number) so the  
+> `TokenQuotaMiddleware` is enabled; the middleware is what records usage  
+> events that feed Prometheus.
+
+```bash
+# Enable usage tracking (set any reasonable limit)
+export MAX_TOKENS_PER_MIN=60000
+export USAGE_BACKEND=prometheus
+```
+
 ### Scraping metrics
 
 ```bash
-curl http://localhost:8080/metrics
+curl -H "Authorization: Bearer $JWT" http://localhost:8080/metrics
 ```
 
 ## Token quotas
diff --git a/pyproject.toml b/pyproject.toml
index 3c7f8f4..52dc93a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,8 +38,7 @@ memory = ["weaviate-client>=3.26.7,<4.0.0"]
 temporal = ["temporalio>=1.5.0"]
 quota = ["tiktoken>=0.5.0"]
 usage = [
-    "prometheus_client>=0.20.0",
-    "openmeter>=0.4.0",
+    "prometheus_client>=0.20.0"
 ]
 dev = [
     "pytest>=8.0.0", 
diff --git a/tests/test_prometheus_fallback.py b/tests/test_prometheus_fallback.py
new file mode 100644
index 0000000..f1772aa
--- /dev/null
+++ b/tests/test_prometheus_fallback.py
@@ -0,0 +1,18 @@
+import os
+import pytest
+from usage.factory import get_usage_backend  
+
+pytest.importorskip("prometheus_client")
+
+def test_prometheus_backend_falls_back_to_null_when_unavailable():
+    os.environ["USAGE_BACKEND"] = "prometheus"
+    
+    # This should return NullUsageBackend when prometheus_client unavailable
+    backend = get_usage_backend("prometheus")
+    
+    if hasattr(backend, 'counter'):
+        # prometheus_client available - got PrometheusUsageBackend
+        assert backend.__class__.__name__ == 'PrometheusUsageBackend'
+    else:
+        # prometheus_client unavailable - got NullUsageBackend fallback
+        assert backend.__class__.__name__ == 'NullUsageBackend' 
\ No newline at end of file
diff --git a/tests/test_usage_prometheus.py b/tests/test_usage_prometheus.py
index 5afc80e..ba1573a 100644
--- a/tests/test_usage_prometheus.py
+++ b/tests/test_usage_prometheus.py
@@ -1,20 +1,34 @@
 import os
-
+import sys
 import pytest
 from fastapi import FastAPI, Request
 from httpx import ASGITransport, AsyncClient
 
+# Force reload to pick up prometheus_client if just installed
+if 'usage.backends' in sys.modules:
+    del sys.modules['usage.backends']
+if 'usage.factory' in sys.modules:
+    del sys.modules['usage.factory']
+
 from middleware.quota import TokenQuotaMiddleware
 from usage.factory import get_usage_backend
 
+# At the top, skip test if prometheus_client not available
+pytest.importorskip("prometheus_client")  # ← Skip entire test if not installed
+
 
 @pytest.mark.asyncio
 async def test_prometheus_backend_counts_tokens(monkeypatch):
+    # Verify we actually get PrometheusUsageBackend
+    backend = get_usage_backend("prometheus")
+    if not hasattr(backend, 'counter'):
+        pytest.skip("prometheus_client not available, got NullUsageBackend")
+    
     os.environ["USAGE_BACKEND"] = "prometheus"
     os.environ["MAX_TOKENS_PER_MIN"] = "1000"
     app = FastAPI()
     app.add_middleware(TokenQuotaMiddleware)
-    app.state.usage = get_usage_backend(os.getenv("USAGE_BACKEND", "null"))
+    app.state.usage = backend  # Use the backend we verified
 
     @app.post("/echo")
     async def echo(request: Request):

From 99567fea59cc48ce62e4ef7d95f4f1d54f0ef767 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Sat, 19 Jul 2025 12:03:31 +0500
Subject: [PATCH 06/24] Polish OpenMeter backend and docs

---
 README.md                         | 21 ++++++++--
 attach/gateway.py                 |  5 ++-
 main.py                           |  5 ++-
 pyproject.toml                    |  3 +-
 tests/test_prometheus_fallback.py | 19 ++++-----
 tests/test_usage_openmeter.py     | 66 +++++++++++++++++++++++++++++++
 tests/test_usage_prometheus.py    | 17 ++++----
 usage/backends.py                 | 55 ++++++++++++++++++++++++--
 usage/factory.py                  | 21 +++++++++-
 9 files changed, 183 insertions(+), 29 deletions(-)
 create mode 100644 tests/test_usage_openmeter.py

diff --git a/README.md b/README.md
index 9d94594..48e8155 100644
--- a/README.md
+++ b/README.md
@@ -237,14 +237,15 @@ curl -X POST /v1/logs \
 ## Usage hooks
 
 Emit token usage metrics for every request. Choose a backend via
-`USAGE_BACKEND`:
+`USAGE_METERING` (alias `USAGE_BACKEND`):
 
 ```bash
-export USAGE_BACKEND=prometheus  # or openmeter/null
+export USAGE_METERING=prometheus  # or null
 ```
 
 A Prometheus counter `attach_usage_tokens_total{user,direction,model}` is
 exposed for Grafana dashboards.
+Set `USAGE_METERING=null` (the default) to disable metering entirely.
 
 > **⚠️ Usage hooks depend on the quota middleware.**  
 > Make sure `MAX_TOKENS_PER_MIN` is set (any positive number) so the  
@@ -254,9 +255,23 @@ exposed for Grafana dashboards.
 ```bash
 # Enable usage tracking (set any reasonable limit)
 export MAX_TOKENS_PER_MIN=60000
-export USAGE_BACKEND=prometheus
+export USAGE_METERING=prometheus
 ```
 
+#### OpenMeter (Stripe / ClickHouse)
+
+```bash
+pip install "attach-gateway[usage]"
+export MAX_TOKENS_PER_MIN=60000
+export USAGE_METERING=openmeter
+export OPENMETER_API_KEY=...
+export OPENMETER_URL=http://localhost:8888   # optional self-host, defaults to https://openmeter.cloud
+```
+
+Events land in the tokens meter of OpenMeter and can sync to Stripe.
+
+The gateway runs fine without these vars; metering activates only when both USAGE_METERING=openmeter and OPENMETER_API_KEY are set.
+
 ### Scraping metrics
 
 ```bash
diff --git a/attach/gateway.py b/attach/gateway.py
index d44d41d..2683e41 100644
--- a/attach/gateway.py
+++ b/attach/gateway.py
@@ -21,7 +21,7 @@
 from middleware.quota import TokenQuotaMiddleware
 from middleware.session import session_mw
 from proxy.engine import router as proxy_router
-from usage.factory import get_usage_backend
+from usage.factory import _select_backend, get_usage_backend
 from usage.metrics import mount_metrics
 
 # Import version from parent package
@@ -147,6 +147,7 @@ def create_app(config: Optional[AttachConfig] = None) -> FastAPI:
     memory_backend = get_memory_backend(config.mem_backend, config)
     app.state.memory = memory_backend
     app.state.config = config
-    app.state.usage = get_usage_backend(os.getenv("USAGE_BACKEND", "null"))
+    backend_selector = _select_backend()
+    app.state.usage = get_usage_backend(backend_selector)
 
     return app
diff --git a/main.py b/main.py
index e4f4d81..04236c1 100644
--- a/main.py
+++ b/main.py
@@ -13,7 +13,7 @@
 from middleware.auth import jwt_auth_mw  # ← your auth middleware
 from middleware.session import session_mw  # ← generates session-id header
 from proxy.engine import router as proxy_router
-from usage.factory import get_usage_backend
+from usage.factory import _select_backend, get_usage_backend
 from usage.metrics import mount_metrics
 
 # At the top, make the import conditional
@@ -132,7 +132,8 @@ async def get_memory_events(request: Request, limit: int = 10):
 
 # Create app without middleware first
 app = FastAPI(title="attach-gateway", middleware=middlewares)
-app.state.usage = get_usage_backend(os.getenv("USAGE_BACKEND", "null"))
+backend_selector = _select_backend()
+app.state.usage = get_usage_backend(backend_selector)
 mount_metrics(app)
 
 
diff --git a/pyproject.toml b/pyproject.toml
index 52dc93a..f74f709 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,7 +38,8 @@ memory = ["weaviate-client>=3.26.7,<4.0.0"]
 temporal = ["temporalio>=1.5.0"]
 quota = ["tiktoken>=0.5.0"]
 usage = [
-    "prometheus_client>=0.20.0"
+    "prometheus_client>=0.20.0",
+    "openmeter>=1.0.0b188,<2",  # openmeter>=1.0.0b188 requires Python >= 3.9
 ]
 dev = [
     "pytest>=8.0.0", 
diff --git a/tests/test_prometheus_fallback.py b/tests/test_prometheus_fallback.py
index f1772aa..0536d0c 100644
--- a/tests/test_prometheus_fallback.py
+++ b/tests/test_prometheus_fallback.py
@@ -1,18 +1,19 @@
-import os
 import pytest
-from usage.factory import get_usage_backend  
+
+from usage.factory import get_usage_backend
 
 pytest.importorskip("prometheus_client")
 
-def test_prometheus_backend_falls_back_to_null_when_unavailable():
-    os.environ["USAGE_BACKEND"] = "prometheus"
-    
+
+def test_prometheus_backend_falls_back_to_null_when_unavailable(monkeypatch):
+    monkeypatch.setenv("USAGE_METERING", "prometheus")
+
     # This should return NullUsageBackend when prometheus_client unavailable
     backend = get_usage_backend("prometheus")
-    
-    if hasattr(backend, 'counter'):
+
+    if hasattr(backend, "counter"):
         # prometheus_client available - got PrometheusUsageBackend
-        assert backend.__class__.__name__ == 'PrometheusUsageBackend'
+        assert backend.__class__.__name__ == "PrometheusUsageBackend"
     else:
         # prometheus_client unavailable - got NullUsageBackend fallback
-        assert backend.__class__.__name__ == 'NullUsageBackend' 
\ No newline at end of file
+        assert backend.__class__.__name__ == "NullUsageBackend"
diff --git a/tests/test_usage_openmeter.py b/tests/test_usage_openmeter.py
new file mode 100644
index 0000000..171dd54
--- /dev/null
+++ b/tests/test_usage_openmeter.py
@@ -0,0 +1,66 @@
+import os
+import sys
+import types
+
+import pytest
+
+
+class DummyEvents:
+    def __init__(self):
+        self.called = None
+
+    async def create(self, **event):
+        self.called = event
+
+
+class DummyClient:
+    def __init__(self, api_key: str, base_url: str = "https://openmeter.cloud") -> None:
+        self.api_key = api_key
+        self.base_url = base_url
+        self.events = DummyEvents()
+
+    async def aclose(self) -> None:
+        pass
+
+
+dummy_module = types.SimpleNamespace(Client=DummyClient)
+
+
+@pytest.mark.asyncio
+async def test_openmeter_backend_create(monkeypatch):
+    monkeypatch.setitem(sys.modules, "openmeter", dummy_module)
+    if "usage.backends" in sys.modules:
+        del sys.modules["usage.backends"]
+    if "usage.factory" in sys.modules:
+        del sys.modules["usage.factory"]
+    from usage.factory import get_usage_backend
+
+    monkeypatch.setenv("OPENMETER_API_KEY", "k")
+    monkeypatch.setenv("OPENMETER_URL", "https://example.com")
+
+    backend = get_usage_backend("openmeter")
+    assert backend.__class__.__name__ == "OpenMeterBackend"
+    await backend.record(user="bob", tokens_in=1, tokens_out=2, model="m")
+    await backend.aclose()
+
+    called = backend.client.events.called
+    assert called["type"] == "tokens"
+    assert called["subject"] == "bob"
+    assert called["project"] is None
+    assert called["data"] == {"tokens_in": 1, "tokens_out": 2, "model": "m"}
+    assert "time" in called
+
+
+def test_openmeter_backend_missing_key(monkeypatch):
+    monkeypatch.setitem(sys.modules, "openmeter", dummy_module)
+    if "usage.backends" in sys.modules:
+        del sys.modules["usage.backends"]
+    if "usage.factory" in sys.modules:
+        del sys.modules["usage.factory"]
+    from usage.factory import get_usage_backend
+
+    monkeypatch.setenv("USAGE_METERING", "openmeter")
+    monkeypatch.delenv("OPENMETER_API_KEY", raising=False)
+
+    backend = get_usage_backend("openmeter")
+    assert backend.__class__.__name__ == "NullUsageBackend"
diff --git a/tests/test_usage_prometheus.py b/tests/test_usage_prometheus.py
index ba1573a..d009a0f 100644
--- a/tests/test_usage_prometheus.py
+++ b/tests/test_usage_prometheus.py
@@ -1,14 +1,15 @@
 import os
 import sys
+
 import pytest
 from fastapi import FastAPI, Request
 from httpx import ASGITransport, AsyncClient
 
 # Force reload to pick up prometheus_client if just installed
-if 'usage.backends' in sys.modules:
-    del sys.modules['usage.backends']
-if 'usage.factory' in sys.modules:
-    del sys.modules['usage.factory']
+if "usage.backends" in sys.modules:
+    del sys.modules["usage.backends"]
+if "usage.factory" in sys.modules:
+    del sys.modules["usage.factory"]
 
 from middleware.quota import TokenQuotaMiddleware
 from usage.factory import get_usage_backend
@@ -21,11 +22,11 @@
 async def test_prometheus_backend_counts_tokens(monkeypatch):
     # Verify we actually get PrometheusUsageBackend
     backend = get_usage_backend("prometheus")
-    if not hasattr(backend, 'counter'):
+    if not hasattr(backend, "counter"):
         pytest.skip("prometheus_client not available, got NullUsageBackend")
-    
-    os.environ["USAGE_BACKEND"] = "prometheus"
-    os.environ["MAX_TOKENS_PER_MIN"] = "1000"
+
+    monkeypatch.setenv("USAGE_METERING", "prometheus")
+    monkeypatch.setenv("MAX_TOKENS_PER_MIN", "1000")
     app = FastAPI()
     app.add_middleware(TokenQuotaMiddleware)
     app.state.usage = backend  # Use the backend we verified
diff --git a/usage/backends.py b/usage/backends.py
index 4731da9..734c503 100644
--- a/usage/backends.py
+++ b/usage/backends.py
@@ -2,8 +2,14 @@
 
 """Usage accounting backends for Attach Gateway."""
 
+import inspect
+import logging
+import os
+from datetime import datetime, timezone
 from typing import Protocol
 
+logger = logging.getLogger(__name__)
+
 try:
     from prometheus_client import Counter
 except Exception:  # pragma: no cover - optional dep
@@ -80,7 +86,50 @@ async def record(self, **evt) -> None:
 
 
 class OpenMeterBackend:
-    """Stub for future OpenMeter integration."""
+    """Send token usage events to OpenMeter."""
+
+    def __init__(self) -> None:
+        try:
+            from openmeter import Client  # type: ignore
+        except Exception as exc:  # pragma: no cover - optional dep
+            raise ImportError("openmeter package is required") from exc
+
+        api_key = os.getenv("OPENMETER_API_KEY")
+        if not api_key:
+            raise ImportError("OPENMETER_API_KEY is required for OpenMeter")
 
-    async def record(self, **evt) -> None:  # pragma: no cover - not implemented
-        raise NotImplementedError
+        url = os.getenv("OPENMETER_URL", "https://openmeter.cloud")
+        self.client = Client(api_key=api_key, base_url=url)
+
+    async def aclose(self) -> None:
+        """Close the underlying OpenMeter client."""
+        try:
+            await self.client.aclose()  # type: ignore[call-arg]
+        except Exception:  # pragma: no cover - optional
+            pass
+
+    async def record(self, **evt) -> None:
+        event = {
+            "type": "tokens",
+            "subject": evt.get("user"),
+            "project": evt.get("project"),
+            "time": datetime.now(timezone.utc)
+            .isoformat(timespec="milliseconds")
+            .replace("+00:00", "Z"),
+            "data": {
+                "tokens_in": int(evt.get("tokens_in", 0) or 0),
+                "tokens_out": int(evt.get("tokens_out", 0) or 0),
+                "model": evt.get("model"),
+            },
+        }
+
+        create_fn = self.client.events.create
+        import anyio
+
+        try:
+            if inspect.iscoroutinefunction(create_fn):
+                await create_fn(**event)  # type: ignore[arg-type]
+            else:
+                await anyio.to_thread.run_sync(create_fn, **event)
+        except Exception as exc:  # pragma: no cover - network errors
+            logger.warning("OpenMeter create failed: %s", exc)
diff --git a/usage/factory.py b/usage/factory.py
index b6dbb4c..3663db5 100644
--- a/usage/factory.py
+++ b/usage/factory.py
@@ -2,6 +2,9 @@
 
 """Factory for usage backends."""
 
+import os
+import warnings
+
 from .backends import (
     AbstractUsageBackend,
     NullUsageBackend,
@@ -10,6 +13,19 @@
 )
 
 
+def _select_backend() -> str:
+    """Return backend name from env vars with deprecation warning."""
+    if "USAGE_METERING" in os.environ:
+        return os.getenv("USAGE_METERING", "null")
+    if "USAGE_BACKEND" in os.environ:
+        warnings.warn(
+            "USAGE_BACKEND is deprecated; rename to USAGE_METERING",
+            UserWarning,
+            stacklevel=2,
+        )
+    return os.getenv("USAGE_BACKEND", "null")
+
+
 def get_usage_backend(kind: str) -> AbstractUsageBackend:
     """Return an instance of the requested usage backend."""
     kind = (kind or "null").lower()
@@ -19,5 +35,8 @@ def get_usage_backend(kind: str) -> AbstractUsageBackend:
         except Exception:
             return NullUsageBackend()
     if kind == "openmeter":
-        return OpenMeterBackend()
+        try:
+            return OpenMeterBackend()
+        except Exception:
+            return NullUsageBackend()
     return NullUsageBackend()

From 5079191918151701bd144ec80c3546e814045b0c Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Sun, 20 Jul 2025 21:04:55 -0700
Subject: [PATCH 07/24] excluded quota pathways

---
 middleware/quota.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/middleware/quota.py b/middleware/quota.py
index 30304d9..d5026d3 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -121,7 +121,21 @@ def _is_textual(mime: str) -> bool:
     def _num_tokens(self, text: str) -> int:
         return len(self.encoder.encode(text))
 
+    @staticmethod
+    def _is_monitoring_endpoint(path: str) -> bool:
+        """Monitoring endpoints that should skip token-based rate limiting."""
+        monitoring_paths = [
+            "/metrics",      # Prometheus metrics
+            "/mem/events",   # Memory events
+            "/auth/config",  # Auth configuration
+        ]
+        return any(path.startswith(p) for p in monitoring_paths)
+
     async def dispatch(self, request: Request, call_next):
+        # Skip token-based rate limiting for monitoring endpoints
+        if self._is_monitoring_endpoint(request.url.path):
+            return await call_next(request)
+            
         if not hasattr(request.app.state, "usage"):
             request.app.state.usage = NullUsageBackend()
         user = request.headers.get("x-attach-user")

From 929bc0719460f7eb91f64b602e830c33967ac4b2 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Mon, 21 Jul 2025 11:25:24 +0500
Subject: [PATCH 08/24] Add mid-stream quota log and expand tail buffer

---
 middleware/quota.py | 411 ++++++++++++++++++++++++++++++++------------
 1 file changed, 300 insertions(+), 111 deletions(-)

diff --git a/middleware/quota.py b/middleware/quota.py
index d5026d3..ba18179 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -1,41 +1,53 @@
-"""Token quota middleware for Attach Gateway.
-
-Enforces per-user token limits using a sliding 1-minute window. The
-quota is applied to both the request body and the response body.
-"""
-
 from __future__ import annotations
 
+"""Token quota middleware enforcing a sliding window budget."""
+
+import json
 import logging
 import os
-import sys
 import time
 from collections import deque
-from typing import Deque, Dict, Optional, Protocol, Tuple
+from typing import (
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Deque,
+    Dict,
+    Iterable,
+    Protocol,
+    Tuple,
+)
 from uuid import uuid4
 
 from fastapi import Request
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import JSONResponse, StreamingResponse
 
+try:
+    import tiktoken  # type: ignore
+except Exception:  # pragma: no cover
+    tiktoken = None
+
 from usage.backends import NullUsageBackend
 
 logger = logging.getLogger(__name__)
 
 
 class AbstractMeterStore(Protocol):
-    """Interface for token accounting backends."""
+    """Token accounting backend."""
 
     async def increment(self, user: str, tokens: int) -> Tuple[int, float]:
-        """Return the running total and timestamp of the oldest entry."""
+        """Increment ``user``'s counter and return ``(total, oldest_ts)``."""
 
+    async def adjust(self, user: str, delta: int) -> Tuple[int, float]:
+        """Atomically adjust ``user``'s total by ``delta``."""
 
-class InMemoryMeterStore:
-    """Simple in-memory sliding window counter.
+    async def peek_total(self, user: str) -> int:
+        """Return ``user``'s current total without mutating state."""
 
-    Not safe for multi-process deployments; each process keeps its own
-    counters. Use :class:`RedisMeterStore` in production.
-    """
+
+class InMemoryMeterStore:
+    """Simple in-memory meter for tests and development."""
 
     def __init__(self, window: int = 60) -> None:
         self.window = window
@@ -52,11 +64,21 @@ async def increment(self, user: str, tokens: int) -> Tuple[int, float]:
         oldest = dq[0][0] if dq else now
         return total, oldest
 
+    async def adjust(self, user: str, delta: int) -> Tuple[int, float]:
+        return await self.increment(user, delta)
+
+    async def peek_total(self, user: str) -> int:
+        total, _ = await self.increment(user, 0)
+        dq = self._data.get(user)
+        if total and dq:
+            dq.pop()
+        return total
+
 
 class RedisMeterStore:
-    """Redis backed sliding window counter."""
+    """Redis backed sliding window meter."""
 
-    def __init__(self, url: str = "redis://localhost:6379", window: int = 60) -> None:
+    def __init__(self, url: str, window: int = 60) -> None:
         import redis.asyncio as redis  # type: ignore
 
         self.window = window
@@ -74,73 +96,200 @@ async def increment(self, user: str, tokens: int) -> Tuple[int, float]:
         entries = results[-1]
         total = 0
         oldest = now
-        for member, ts in entries:
+        for m, ts in entries:
             try:
-                _, tok = member.split(":", 1)
+                _, tok = m.split(":", 1)
                 total += int(tok)
             except Exception:
                 pass
             oldest = min(oldest, ts)
         return total, oldest
 
+    async def adjust(self, user: str, delta: int) -> Tuple[int, float]:
+        now = time.time()
+        key = f"attach:quota:{user}"
+        member = f"{now}:{delta}"
+        async with self.redis.pipeline(transaction=True) as pipe:
+            await pipe.zadd(key, {member: now})
+            await pipe.zremrangebyscore(key, 0, now - self.window)
+            await pipe.zrange(key, 0, -1, withscores=True)
+            results = await pipe.execute()
+        entries = results[-1]
+        total = 0
+        oldest = now
+        for m, ts in entries:
+            try:
+                _, tok = m.split(":", 1)
+                total += int(tok)
+            except Exception:
+                pass
+            oldest = min(oldest, ts)
+        return total, oldest
 
-class TokenQuotaMiddleware(BaseHTTPMiddleware):
-    """FastAPI middleware enforcing per-user token quotas."""
+    async def peek_total(self, user: str) -> int:
+        now = time.time()
+        key = f"attach:quota:{user}"
+        member = f"{now}:0:{uuid4().hex}"
+        async with self.redis.pipeline(transaction=True) as pipe:
+            await pipe.zadd(key, {member: now})
+            await pipe.zremrangebyscore(key, 0, now - self.window)
+            await pipe.zrange(key, 0, -1, withscores=True)
+            pipe.zrem(key, member)
+            results = await pipe.execute()
+        entries = results[2]
+        total = 0
+        for m, _ in entries:
+            try:
+                _, tok = m.split(":", 1)
+                total += int(tok)
+            except Exception:
+                pass
+        return total
 
-    def __init__(self, app, store: Optional[AbstractMeterStore] = None) -> None:
-        """Create middleware.
 
-        Must be added **after** :func:`session_mw` so the ``X-Attach-User``
-        header or client IP is available for quota tracking.
-        """
-        super().__init__(app)
-        self.store = store or InMemoryMeterStore()
-        self.window = 60
-        self.max_tokens = int(os.getenv("MAX_TOKENS_PER_MIN", "60000"))
-        enc_name = os.getenv("QUOTA_ENCODING", "cl100k_base")
-        try:
-            import tiktoken
+def _is_textual(mime: str) -> bool:
+    mime = (mime or "").lower()
+    if not mime or mime == "*/*":
+        return False
+    return mime.startswith("text/") or "json" in mime
+
+
+def _encoder_for_model(model: str):
+    """Return a tiktoken encoder, falling back to byte count."""
+    if tiktoken is None:  # pragma: no cover - fallback
+        logger.warning(
+            "tiktoken missing – using byte-count fallback; token metrics may be inflated"
+        )
 
-            self.encoder = tiktoken.get_encoding(enc_name)
-        except Exception:  # pragma: no cover - fallback
-            if "tiktoken" not in sys.modules:
-                logger.warning(
-                    "tiktoken missing – using byte-count fallback; token metrics may be inflated"
-                )
+        class _Simple:
+            def encode(self, text: str) -> list[int]:
+                return list(text.encode())
+
+        return _Simple()
+
+    try:
+        return tiktoken.encoding_for_model(model)
+    except Exception:
+        try:
+            return tiktoken.get_encoding("cl100k_base")
+        except Exception:
 
             class _Simple:
                 def encode(self, text: str) -> list[int]:
                     return list(text.encode())
 
-            self.encoder = _Simple()
+            return _Simple()
+
+
+def _num_tokens(text: str, model: str = "cl100k_base") -> int:
+    return len(_encoder_for_model(model).encode(text))
 
-    @staticmethod
-    def _is_textual(mime: str) -> bool:
-        return mime.startswith("text/") or "json" in mime
 
-    def _num_tokens(self, text: str) -> int:
-        return len(self.encoder.encode(text))
+def num_tokens_from_messages(messages: Iterable[dict], model: str) -> int:
+    enc = _encoder_for_model(model)
+    total = 3
+    for msg in messages:
+        total += 4
+        for k, v in msg.items():
+            total += len(enc.encode(str(v)))
+            if k == "name":
+                total -= 1
+    return total
+
+
+async def async_iter(data: Iterable[bytes]) -> AsyncIterator[bytes]:
+    for chunk in data:
+        yield chunk
+
+
+_SKIP_PATHS = {
+    "/metrics",
+    "/mem/events",
+    "/auth/config",
+    "/health",
+    "/docs",
+    "/redoc",
+    "/openapi.json",
+}
+
+
+class TokenQuotaMiddleware(BaseHTTPMiddleware):
+    """Apply per-user LLM token quotas."""
+
+    def __init__(self, app, store: AbstractMeterStore | None = None) -> None:
+        super().__init__(app)
+        self.window = int(os.getenv("WINDOW", "60"))
+        self.max_tokens = int(os.getenv("MAX_TOKENS_PER_MIN", "60000"))
+        if store is not None:
+            self.store = store
+        else:
+            redis_url = os.getenv("REDIS_URL")
+            self.store = (
+                RedisMeterStore(redis_url, self.window)
+                if redis_url
+                else InMemoryMeterStore(self.window)
+            )
 
-    @staticmethod
-    def _is_monitoring_endpoint(path: str) -> bool:
-        """Monitoring endpoints that should skip token-based rate limiting."""
-        monitoring_paths = [
-            "/metrics",      # Prometheus metrics
-            "/mem/events",   # Memory events
-            "/auth/config",  # Auth configuration
-        ]
-        return any(path.startswith(p) for p in monitoring_paths)
+    class _Streamer:
+        def __init__(
+            self,
+            iterator: AsyncIterator[bytes],
+            *,
+            user: str,
+            store: AbstractMeterStore,
+            max_tokens: int,
+            is_textual: bool,
+        ) -> None:
+            self.iterator = iterator
+            self.tail = bytearray()
+            self.on_complete: Callable[[], Awaitable[None]] | None = None
+            self.user = user
+            self.store = store
+            self.limit = max_tokens
+            self.is_textual = is_textual
+            self.quota_exceeded = False
+
+        def __aiter__(self) -> AsyncIterator[bytes]:
+            return self._gen()
+
+        async def _gen(self) -> AsyncIterator[bytes]:
+            try:
+                async for chunk in self.iterator:
+                    self.tail.extend(chunk)
+                    if len(self.tail) > 8192:
+                        del self.tail[:-8192]
+                    if self.limit:
+                        chunk_tokens = (
+                            _num_tokens(chunk.decode("utf-8", "ignore"))
+                            if self.is_textual
+                            else 0
+                        )
+                        if chunk_tokens:
+                            total, _ = await self.store.adjust(self.user, chunk_tokens)
+                            if total > self.limit:
+                                self.quota_exceeded = True
+                                logger.warning(
+                                    "User %s quota breached mid-stream", self.user
+                                )
+                                return
+                    yield chunk
+            finally:
+                if self.on_complete:
+                    await self.on_complete()
+
+        def get_tail(self) -> bytes:
+            return bytes(self.tail)
 
     async def dispatch(self, request: Request, call_next):
-        # Skip token-based rate limiting for monitoring endpoints
-        if self._is_monitoring_endpoint(request.url.path):
+        if any(request.url.path.startswith(p) for p in _SKIP_PATHS):
             return await call_next(request)
-            
+
         if not hasattr(request.app.state, "usage"):
             request.app.state.usage = NullUsageBackend()
-        user = request.headers.get("x-attach-user")
-        if not user:
-            user = request.client.host if request.client else "unknown"
+
+        user = request.headers.get("x-attach-user") or (
+            request.client.host if request.client else "unknown"
+        )
         usage = {
             "user": user,
             "project": request.headers.get("x-attach-project", "default"),
@@ -149,16 +298,27 @@ async def dispatch(self, request: Request, call_next):
             "model": "unknown",
             "request_id": request.headers.get("x-request-id") or str(uuid4()),
         }
+        req_is_text = _is_textual(request.headers.get("content-type", ""))
+
+        tokens_in = 0
+        if req_is_text:
+            raw = await request.body()
+            try:
+                payload = json.loads(raw.decode())
+            except Exception:
+                payload = None
+            if isinstance(payload, dict) and "messages" in payload:
+                model = payload.get("model", "cl100k_base")
+                usage["model"] = model
+                tokens_in = num_tokens_from_messages(payload.get("messages", []), model)
+            else:
+                tokens_in = _num_tokens(raw.decode("utf-8", "ignore"))
+
+        usage["tokens_in"] = tokens_in
 
-        body = await request.body()
-        content_type = request.headers.get("content-type", "")
-        tokens = 0
-        if self._is_textual(content_type):
-            tokens = self._num_tokens(body.decode("utf-8", "ignore"))
-        usage["tokens_in"] = tokens
-        total, oldest = await self.store.increment(user, tokens)
-        request.state._usage = usage
+        total, oldest = await self.store.adjust(user, tokens_in)
         if total > self.max_tokens:
+            await self.store.adjust(user, -tokens_in)
             retry_after = max(0, int(self.window - (time.time() - oldest)))
             usage["ts"] = time.time()
             await request.app.state.usage.record(**usage)
@@ -167,53 +327,82 @@ async def dispatch(self, request: Request, call_next):
                 status_code=429,
             )
 
-        response = await call_next(request)
-        usage["model"] = response.headers.get("x-llm-model", "unknown")
+        resp = await call_next(request)
 
-        first_chunk = None
-        try:
-            first_chunk = await response.body_iterator.__anext__()
-        except StopAsyncIteration:
-            pass
-
-        media = response.media_type or response.headers.get("content-type", "")
-        if first_chunk is not None and self._is_textual(media):
-            tokens_chunk = self._num_tokens(first_chunk.decode("utf-8", "ignore"))
-            total, oldest = await self.store.increment(user, tokens_chunk)
-            if total > self.max_tokens:
-                retry_after = max(0, int(self.window - (time.time() - oldest)))
-                usage["tokens_out"] += tokens_chunk
+        media = resp.media_type or resp.headers.get("content-type", "")
+        resp_is_text = _is_textual(media)
+        streamer = self._Streamer(
+            resp.body_iterator,
+            user=user,
+            store=self.store,
+            max_tokens=self.max_tokens,
+            is_textual=resp_is_text,
+        )
+
+        headers = dict(resp.headers)
+        headers.pop("content-length", None)
+        response = StreamingResponse(
+            content=streamer,
+            status_code=resp.status_code,
+            headers=headers,
+            media_type=resp.media_type,
+        )
+
+        async def finalize() -> None:
+            nonlocal tokens_in
+            if getattr(streamer, "quota_exceeded", False):
+                usage["detail"] = "token quota exceeded mid-stream"
                 usage["ts"] = time.time()
                 await request.app.state.usage.record(**usage)
-                return JSONResponse(
-                    {"detail": "token quota exceeded", "retry_after": retry_after},
-                    status_code=429,
-                )
-            usage["tokens_out"] += tokens_chunk
-
-        async def stream_with_quota():
-            nonlocal total, oldest
-            if first_chunk is not None:
-                yield first_chunk
-            async for chunk in response.body_iterator:
-                tokens_chunk = 0
-                media = response.media_type or response.headers.get("content-type", "")
-                if self._is_textual(media):
-                    tokens_chunk = self._num_tokens(chunk.decode("utf-8", "ignore"))
-                next_total, oldest = await self.store.increment(user, tokens_chunk)
-                if next_total > self.max_tokens:
-                    break
-                total = next_total
-                if tokens_chunk:
-                    usage["tokens_out"] += tokens_chunk
-                yield chunk
+                return
+            tokens_out = 0
+            model = usage.get("model", "unknown")
+            parsed: dict | None = None
+            if resp_is_text:
+                text_tail = streamer.get_tail().decode("utf-8", "ignore")
+                idx = max(text_tail.rfind("{"), text_tail.rfind('{"'))
+                if idx != -1:
+                    try:
+                        parsed = json.loads(text_tail[idx:])
+                    except Exception:
+                        parsed = None
+                if isinstance(parsed, dict):
+                    model = parsed.get("model", model)
+                    if "usage" in parsed:
+                        u = parsed.get("usage") or {}
+                        tokens_out = int(u.get("completion_tokens", 0))
+                        prompt_tokens = int(u.get("prompt_tokens", tokens_in))
+                        delta_prompt = prompt_tokens - tokens_in
+                        tokens_in = prompt_tokens
+                        usage["tokens_in"] = tokens_in
+                        await self.store.adjust(user, delta_prompt)
+                    elif "choices" in parsed:
+                        msgs = [
+                            (c.get("message") or {}).get("content", "")
+                            for c in parsed.get("choices", [])
+                        ]
+                        tokens_out = num_tokens_from_messages(
+                            [{"content": m} for m in msgs], model
+                        )
+                if tokens_out == 0:
+                    tokens_out = _num_tokens(text_tail)
+
+            usage["tokens_out"] = tokens_out
+            usage["model"] = model
+
+            await self.store.adjust(user, tokens_out)
+
+            total = await self.store.peek_total(user)
+            if total > self.max_tokens:
+                usage["detail"] = "token quota exceeded post-stream"
+
+            response.headers["x-llm-model"] = model
+            response.headers["x-tokens-in"] = str(tokens_in)
+            response.headers["x-tokens-out"] = str(tokens_out)
 
             usage["ts"] = time.time()
             await request.app.state.usage.record(**usage)
+            logger.info(json.dumps(usage))
 
-        return StreamingResponse(
-            stream_with_quota(),
-            status_code=response.status_code,
-            headers=dict(response.headers),
-            media_type=response.media_type,
-        )
+        streamer.on_complete = finalize
+        return response

From 6a835524c794d41eea0779864f4763e1e50245b1 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Mon, 21 Jul 2025 11:56:01 +0500
Subject: [PATCH 09/24] fix token quota finalization

---
 middleware/quota.py | 32 ++++++++++++++++++++++++--------
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/middleware/quota.py b/middleware/quota.py
index ba18179..0037368 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -359,11 +359,21 @@ async def finalize() -> None:
             model = usage.get("model", "unknown")
             parsed: dict | None = None
             if resp_is_text:
+                # -- Robustly extract the last JSON object -----------------
                 text_tail = streamer.get_tail().decode("utf-8", "ignore")
-                idx = max(text_tail.rfind("{"), text_tail.rfind('{"'))
-                if idx != -1:
+                # 1. Split SSE frames if present: keep only the part after the final "data: "
+                if "data:" in text_tail:
+                    *_, last_frame = text_tail.strip().split("data:")
+                    text_tail = last_frame.strip()
+                # 2. Strip the trailing '[DONE]' token if it exists
+                if text_tail.endswith("[DONE]"):
+                    text_tail = text_tail[: text_tail.rfind("[DONE]")].rstrip()
+                # 3. Find the first '{' from the *left* (because the frame has been cleaned)
+                brace = text_tail.find("{")
+                parsed = None
+                if brace != -1:
                     try:
-                        parsed = json.loads(text_tail[idx:])
+                        parsed = json.loads(text_tail[brace:])
                     except Exception:
                         parsed = None
                 if isinstance(parsed, dict):
@@ -372,9 +382,10 @@ async def finalize() -> None:
                         u = parsed.get("usage") or {}
                         tokens_out = int(u.get("completion_tokens", 0))
                         prompt_tokens = int(u.get("prompt_tokens", tokens_in))
-                        delta_prompt = prompt_tokens - tokens_in
-                        tokens_in = prompt_tokens
+                        delta_prompt = prompt_tokens - tokens_in  # adjust quota window
+                        tokens_in = prompt_tokens  # ← canonical value
                         usage["tokens_in"] = tokens_in
+                        # Replace provisional count with canonical one in the meter
                         await self.store.adjust(user, delta_prompt)
                     elif "choices" in parsed:
                         msgs = [
@@ -390,15 +401,20 @@ async def finalize() -> None:
             usage["tokens_out"] = tokens_out
             usage["model"] = model
 
+            # All *out* tokens are new – add them once.
             await self.store.adjust(user, tokens_out)
 
             total = await self.store.peek_total(user)
             if total > self.max_tokens:
                 usage["detail"] = "token quota exceeded post-stream"
 
-            response.headers["x-llm-model"] = model
-            response.headers["x-tokens-in"] = str(tokens_in)
-            response.headers["x-tokens-out"] = str(tokens_out)
+            response.headers.update(
+                {
+                    "x-llm-model": model,
+                    "x-tokens-in": str(tokens_in),
+                    "x-tokens-out": str(tokens_out),
+                }
+            )
 
             usage["ts"] = time.time()
             await request.app.state.usage.record(**usage)

From f162f20ddd09bfe3c04b59233b3901a6f2c5a7ac Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Mon, 21 Jul 2025 00:01:02 -0700
Subject: [PATCH 10/24] token calculation matches between response output and
 metrics endpoint

---
 middleware/quota.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/middleware/quota.py b/middleware/quota.py
index 0037368..5477d05 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -406,6 +406,9 @@ async def finalize() -> None:
 
             total = await self.store.peek_total(user)
             if total > self.max_tokens:
+                retry_after = self.window              # simple worst-case
+                response.status_code = 429
+                response.headers["Retry-After"] = str(retry_after)
                 usage["detail"] = "token quota exceeded post-stream"
 
             response.headers.update(

From 38bf0c2d05f1517c186840ba0cc90a99e8b7c9e2 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 22 Jul 2025 05:43:53 +0500
Subject: [PATCH 11/24] Make token quota optional with shared int_env helper

---
 attach/gateway.py   |  5 ++++-
 main.py             |  6 ++++--
 middleware/quota.py | 11 ++++++-----
 utils/env.py        | 25 +++++++++++++++++++++++++
 4 files changed, 39 insertions(+), 8 deletions(-)
 create mode 100644 utils/env.py

diff --git a/attach/gateway.py b/attach/gateway.py
index 2683e41..ce27650 100644
--- a/attach/gateway.py
+++ b/attach/gateway.py
@@ -23,6 +23,7 @@
 from proxy.engine import router as proxy_router
 from usage.factory import _select_backend, get_usage_backend
 from usage.metrics import mount_metrics
+from utils.env import int_env
 
 # Import version from parent package
 from . import __version__
@@ -135,7 +136,9 @@ def create_app(config: Optional[AttachConfig] = None) -> FastAPI:
     # Add middleware
     app.middleware("http")(jwt_auth_mw)
     app.middleware("http")(session_mw)
-    app.add_middleware(TokenQuotaMiddleware)
+    limit = int_env("MAX_TOKENS_PER_MIN", 60000)
+    if limit is not None:
+        app.add_middleware(TokenQuotaMiddleware)
 
     # Add routes
     app.include_router(a2a_router)
diff --git a/main.py b/main.py
index 04236c1..0e3bf72 100644
--- a/main.py
+++ b/main.py
@@ -15,6 +15,7 @@
 from proxy.engine import router as proxy_router
 from usage.factory import _select_backend, get_usage_backend
 from usage.metrics import mount_metrics
+from utils.env import int_env
 
 # At the top, make the import conditional
 try:
@@ -126,8 +127,9 @@ async def get_memory_events(request: Request, limit: int = 10):
     Middleware(BaseHTTPMiddleware, dispatch=session_mw),
 ]
 
-# Only add quota middleware if tiktoken is available AND user configured it
-if QUOTA_AVAILABLE and os.getenv("MAX_TOKENS_PER_MIN"):
+# Only add quota middleware if tiktoken is available and a positive limit is set
+limit = int_env("MAX_TOKENS_PER_MIN", 60000)
+if QUOTA_AVAILABLE and limit is not None:
     middlewares.append(Middleware(TokenQuotaMiddleware))
 
 # Create app without middleware first
diff --git a/middleware/quota.py b/middleware/quota.py
index 5477d05..a6a81b9 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -29,6 +29,7 @@
     tiktoken = None
 
 from usage.backends import NullUsageBackend
+from utils.env import int_env
 
 logger = logging.getLogger(__name__)
 
@@ -219,7 +220,7 @@ class TokenQuotaMiddleware(BaseHTTPMiddleware):
     def __init__(self, app, store: AbstractMeterStore | None = None) -> None:
         super().__init__(app)
         self.window = int(os.getenv("WINDOW", "60"))
-        self.max_tokens = int(os.getenv("MAX_TOKENS_PER_MIN", "60000"))
+        self.max_tokens: int | None = int_env("MAX_TOKENS_PER_MIN", 60000)
         if store is not None:
             self.store = store
         else:
@@ -237,7 +238,7 @@ def __init__(
             *,
             user: str,
             store: AbstractMeterStore,
-            max_tokens: int,
+            max_tokens: int | None,
             is_textual: bool,
         ) -> None:
             self.iterator = iterator
@@ -317,7 +318,7 @@ async def dispatch(self, request: Request, call_next):
         usage["tokens_in"] = tokens_in
 
         total, oldest = await self.store.adjust(user, tokens_in)
-        if total > self.max_tokens:
+        if self.max_tokens is not None and total > self.max_tokens:
             await self.store.adjust(user, -tokens_in)
             retry_after = max(0, int(self.window - (time.time() - oldest)))
             usage["ts"] = time.time()
@@ -405,8 +406,8 @@ async def finalize() -> None:
             await self.store.adjust(user, tokens_out)
 
             total = await self.store.peek_total(user)
-            if total > self.max_tokens:
-                retry_after = self.window              # simple worst-case
+            if self.max_tokens is not None and total > self.max_tokens:
+                retry_after = self.window  # simple worst-case
                 response.status_code = 429
                 response.headers["Retry-After"] = str(retry_after)
                 usage["detail"] = "token quota exceeded post-stream"
diff --git a/utils/env.py b/utils/env.py
new file mode 100644
index 0000000..909534c
--- /dev/null
+++ b/utils/env.py
@@ -0,0 +1,25 @@
+"""Env helpers used across entry-points."""
+
+import logging
+import os
+
+log = logging.getLogger(__name__)
+
+
+def int_env(var: str, default: int | None = None) -> int | None:
+    """Read $VAR as positive int.
+    • '', 'null', 'none', 'false', 'infinite'  -> None
+    • invalid / non-positive -> default
+    """
+    val = os.getenv(var)
+    if val is None:
+        return default
+    val = val.strip().lower()
+    if val in {"", "null", "none", "false", "infinite"}:
+        return None
+    try:
+        num = int(val)
+        return num if num > 0 else default
+    except ValueError:
+        log.warning("⚠️  %s=%s is not a valid int; using default=%s", var, val, default)
+        return default

From 19b10d4201e7f1bbf549b1702efc5fc6167acb57 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 22 Jul 2025 06:20:01 +0500
Subject: [PATCH 12/24] Add content-type header for quota 429

---
 middleware/quota.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/middleware/quota.py b/middleware/quota.py
index a6a81b9..dfd7408 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -350,9 +350,14 @@ async def dispatch(self, request: Request, call_next):
         )
 
         async def finalize() -> None:
-            nonlocal tokens_in
+            nonlocal tokens_in, oldest
             if getattr(streamer, "quota_exceeded", False):
                 usage["detail"] = "token quota exceeded mid-stream"
+                retry_after = max(0, int(self.window - (time.time() - oldest)))
+                response.status_code = 429
+                response.body_iterator = async_iter([])
+                response.headers["Retry-After"] = str(retry_after)
+                response.headers["content-type"] = "application/json"
                 usage["ts"] = time.time()
                 await request.app.state.usage.record(**usage)
                 return
@@ -411,6 +416,7 @@ async def finalize() -> None:
                 response.status_code = 429
                 response.headers["Retry-After"] = str(retry_after)
                 usage["detail"] = "token quota exceeded post-stream"
+                response.headers["content-type"] = "application/json"
 
             response.headers.update(
                 {

From 9e31551f8732531c1ac694bb87a9f587ae5b76ab Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 22 Jul 2025 06:48:29 +0500
Subject: [PATCH 13/24] fix quota streaming rollback

---
 middleware/quota.py                  |  5 +----
 tests/test_token_quota_middleware.py | 31 +++++++++++++++++++++++++++-
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/middleware/quota.py b/middleware/quota.py
index dfd7408..1360891 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -268,6 +268,7 @@ async def _gen(self) -> AsyncIterator[bytes]:
                         if chunk_tokens:
                             total, _ = await self.store.adjust(self.user, chunk_tokens)
                             if total > self.limit:
+                                await self.store.adjust(self.user, -chunk_tokens)
                                 self.quota_exceeded = True
                                 logger.warning(
                                     "User %s quota breached mid-stream", self.user
@@ -354,10 +355,6 @@ async def finalize() -> None:
             if getattr(streamer, "quota_exceeded", False):
                 usage["detail"] = "token quota exceeded mid-stream"
                 retry_after = max(0, int(self.window - (time.time() - oldest)))
-                response.status_code = 429
-                response.body_iterator = async_iter([])
-                response.headers["Retry-After"] = str(retry_after)
-                response.headers["content-type"] = "application/json"
                 usage["ts"] = time.time()
                 await request.app.state.usage.record(**usage)
                 return
diff --git a/tests/test_token_quota_middleware.py b/tests/test_token_quota_middleware.py
index 00a638f..5d3aec3 100644
--- a/tests/test_token_quota_middleware.py
+++ b/tests/test_token_quota_middleware.py
@@ -6,7 +6,9 @@
 
 pytest.importorskip("tiktoken")
 
-from middleware.quota import TokenQuotaMiddleware
+from starlette.responses import StreamingResponse
+
+from middleware.quota import InMemoryMeterStore, TokenQuotaMiddleware
 
 
 @pytest.mark.asyncio
@@ -44,3 +46,30 @@ async def echo(request: Request):
         resp = await client.post("/echo", json={"msg": "hello"}, headers=headers)
     assert resp.status_code == 429
     assert "retry_after" in resp.json()
+
+
+@pytest.mark.asyncio
+async def test_midstream_over_limit_rolls_back(monkeypatch):
+    monkeypatch.setenv("MAX_TOKENS_PER_MIN", "5")
+    store = InMemoryMeterStore()
+    app = FastAPI()
+    app.add_middleware(TokenQuotaMiddleware, store=store)
+
+    @app.get("/stream")
+    async def stream():
+        async def gen():
+            yield b"hi"
+            yield b"aaaaaaa"
+
+        return StreamingResponse(gen(), media_type="text/plain")
+
+    headers = {"X-Attach-User": "carol"}
+    transport = ASGITransport(app=app)
+    async with AsyncClient(transport=transport, base_url="http://test") as client:
+        resp = await client.get("/stream", headers=headers)
+        text = resp.text
+
+    assert resp.status_code == 200
+    assert text == "hi"
+    total = await store.peek_total("carol")
+    assert total == 2

From 9b35225a2ad60ac8dc3419cc686128c978034e46 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 22 Jul 2025 07:30:55 +0500
Subject: [PATCH 14/24] Return proper 429 JSON when quota exceeded mid-stream

---
 middleware/quota.py                  | 57 +++++++++++++++++++++++++++-
 tests/test_token_quota_middleware.py |  5 +--
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/middleware/quota.py b/middleware/quota.py
index 1360891..21f8666 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -282,6 +282,57 @@ async def _gen(self) -> AsyncIterator[bytes]:
         def get_tail(self) -> bytes:
             return bytes(self.tail)
 
+    class _BufferedResponse(StreamingResponse):
+        def __init__(
+            self, streamer: "TokenQuotaMiddleware._Streamer", **kwargs
+        ) -> None:
+            super().__init__(streamer, **kwargs)
+            self._streamer = streamer
+
+        async def stream_response(self, send):
+            chunks: list[bytes] = []
+            async for chunk in self._streamer:
+                chunks.append(chunk)
+            if getattr(self._streamer, "quota_exceeded", False):
+                retry_after = max(
+                    0,
+                    int(
+                        self._streamer.window
+                        - (time.time() - getattr(self._streamer, "oldest", 0))
+                    ),
+                )
+                self.status_code = 429
+                self.raw_headers = [
+                    (b"content-type", b"application/json"),
+                    (b"retry-after", str(retry_after).encode()),
+                ]
+                iterator = async_iter(
+                    [
+                        json.dumps(
+                            {
+                                "detail": "token quota exceeded",
+                                "retry_after": retry_after,
+                            }
+                        ).encode()
+                    ]
+                )
+            else:
+                iterator = async_iter(chunks)
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": self.status_code,
+                    "headers": self.raw_headers,
+                }
+            )
+            async for chunk in iterator:
+                if not isinstance(chunk, (bytes, memoryview)):
+                    chunk = chunk.encode(self.charset)
+                await send(
+                    {"type": "http.response.body", "body": chunk, "more_body": True}
+                )
+            await send({"type": "http.response.body", "body": b"", "more_body": False})
+
     async def dispatch(self, request: Request, call_next):
         if any(request.url.path.startswith(p) for p in _SKIP_PATHS):
             return await call_next(request)
@@ -340,11 +391,13 @@ async def dispatch(self, request: Request, call_next):
             max_tokens=self.max_tokens,
             is_textual=resp_is_text,
         )
+        streamer.window = self.window
+        streamer.oldest = oldest
 
         headers = dict(resp.headers)
         headers.pop("content-length", None)
-        response = StreamingResponse(
-            content=streamer,
+        response = self._BufferedResponse(
+            streamer,
             status_code=resp.status_code,
             headers=headers,
             media_type=resp.media_type,
diff --git a/tests/test_token_quota_middleware.py b/tests/test_token_quota_middleware.py
index 5d3aec3..487b41b 100644
--- a/tests/test_token_quota_middleware.py
+++ b/tests/test_token_quota_middleware.py
@@ -67,9 +67,8 @@ async def gen():
     transport = ASGITransport(app=app)
     async with AsyncClient(transport=transport, base_url="http://test") as client:
         resp = await client.get("/stream", headers=headers)
-        text = resp.text
 
-    assert resp.status_code == 200
-    assert text == "hi"
+    assert resp.status_code == 429
+    assert resp.json()["detail"] == "token quota exceeded"
     total = await store.peek_total("carol")
     assert total == 2

From db6cd2811f65d50aba45c41a3d2595f89ea2469b Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 22 Jul 2025 19:17:07 -0700
Subject: [PATCH 15/24] open meter now works

---
 .env.example                   |  2 +
 main.py                        | 22 ++++++++---
 tests/test_metrics_endpoint.py | 17 ++++++++-
 usage/backends.py              | 69 ++++++++++++++++++++++------------
 4 files changed, 80 insertions(+), 30 deletions(-)

diff --git a/.env.example b/.env.example
index c6ef8ca..6019812 100644
--- a/.env.example
+++ b/.env.example
@@ -19,6 +19,8 @@ WEAVIATE_URL=http://localhost:8081
 MAX_TOKENS_PER_MIN=60000
 QUOTA_ENCODING=cl100k_base
 
+USAGE_METERING=null
+
 # Development: Auth0 credentials for dev_login script
 # AUTH0_DOMAIN=your-domain.auth0.com
 # AUTH0_CLIENT=your-client-id
diff --git a/main.py b/main.py
index 0e3bf72..76e1efc 100644
--- a/main.py
+++ b/main.py
@@ -5,6 +5,7 @@
 from fastapi.middleware import Middleware
 from fastapi.middleware.cors import CORSMiddleware
 from starlette.middleware.base import BaseHTTPMiddleware
+from contextlib import asynccontextmanager
 
 from a2a.routes import router as a2a_router
 from auth.oidc import verify_jwt  # Fixed: was auth.jwt, now auth.oidc
@@ -132,11 +133,22 @@ async def get_memory_events(request: Request, limit: int = 10):
 if QUOTA_AVAILABLE and limit is not None:
     middlewares.append(Middleware(TokenQuotaMiddleware))
 
-# Create app without middleware first
-app = FastAPI(title="attach-gateway", middleware=middlewares)
-backend_selector = _select_backend()
-app.state.usage = get_usage_backend(backend_selector)
-mount_metrics(app)
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Manage application lifespan - startup and shutdown."""
+    # Startup
+    backend_selector = _select_backend()
+    app.state.usage = get_usage_backend(backend_selector)
+    mount_metrics(app)
+    
+    yield
+    
+    # Shutdown
+    if hasattr(app.state.usage, 'aclose'):
+        await app.state.usage.aclose()
+
+# Create app with lifespan
+app = FastAPI(title="attach-gateway", middleware=middlewares, lifespan=lifespan)
 
 
 @app.get("/auth/config")
diff --git a/tests/test_metrics_endpoint.py b/tests/test_metrics_endpoint.py
index 5fc96e9..45af1a5 100644
--- a/tests/test_metrics_endpoint.py
+++ b/tests/test_metrics_endpoint.py
@@ -9,16 +9,29 @@
 from usage.metrics import mount_metrics
 
 
+def setup_module():
+    # ✅ Fixed: Use new variable name
+    os.environ["USAGE_METERING"] = "prometheus"
+
+
+@pytest.fixture
+def app():
+    app = FastAPI()
+    mount_metrics(app)
+    # ✅ Fixed: Use new variable name
+    app.state.usage = get_usage_backend(os.getenv("USAGE_METERING", "null"))
+    return app
+
+
 @pytest.mark.asyncio
 async def test_metrics_endpoint(monkeypatch):
-    os.environ["USAGE_BACKEND"] = "prometheus"
     os.environ["MAX_TOKENS_PER_MIN"] = "1000"
 
     app = FastAPI()
     mount_metrics(app)
 
     app.add_middleware(TokenQuotaMiddleware)
-    app.state.usage = get_usage_backend(os.getenv("USAGE_BACKEND", "null"))
+    app.state.usage = get_usage_backend(os.getenv("USAGE_METERING", "null"))
 
     @app.post("/echo")
     async def echo(request: Request):
diff --git a/usage/backends.py b/usage/backends.py
index 734c503..9e9fdf2 100644
--- a/usage/backends.py
+++ b/usage/backends.py
@@ -89,47 +89,70 @@ class OpenMeterBackend:
     """Send token usage events to OpenMeter."""
 
     def __init__(self) -> None:
-        try:
-            from openmeter import Client  # type: ignore
-        except Exception as exc:  # pragma: no cover - optional dep
-            raise ImportError("openmeter package is required") from exc
-
         api_key = os.getenv("OPENMETER_API_KEY")
         if not api_key:
             raise ImportError("OPENMETER_API_KEY is required for OpenMeter")
 
-        url = os.getenv("OPENMETER_URL", "https://openmeter.cloud")
-        self.client = Client(api_key=api_key, base_url=url)
+        self.api_key = api_key
+        self.base_url = os.getenv("OPENMETER_URL", "https://openmeter.cloud")
+        
+        # Use httpx instead of buggy OpenMeter SDK
+        try:
+            import httpx
+            self.client = httpx.AsyncClient(
+                timeout=30.0
+            )
+        except ImportError as exc:
+            raise ImportError("httpx is required for OpenMeter") from exc
 
     async def aclose(self) -> None:
-        """Close the underlying OpenMeter client."""
-        try:
-            await self.client.aclose()  # type: ignore[call-arg]
-        except Exception:  # pragma: no cover - optional
-            pass
+        """Close the underlying HTTP client."""
+        if hasattr(self.client, 'aclose'):
+            await self.client.aclose()
 
     async def record(self, **evt) -> None:
+        try:
+            from uuid import uuid4
+        except ImportError as exc:
+            print(f"❌ UUID import failed: {exc}")
+            return
+
+        # Create event in the same format as your working curl
         event = {
+            "specversion": "1.0",
             "type": "tokens",
-            "subject": evt.get("user"),
-            "project": evt.get("project"),
+            "id": str(uuid4()),
             "time": datetime.now(timezone.utc)
             .isoformat(timespec="milliseconds")
             .replace("+00:00", "Z"),
+            "source": "attach-gateway",
+            "subject": evt.get("user"),
             "data": {
                 "tokens_in": int(evt.get("tokens_in", 0) or 0),
                 "tokens_out": int(evt.get("tokens_out", 0) or 0),
                 "model": evt.get("model"),
-            },
+            }
         }
 
-        create_fn = self.client.events.create
-        import anyio
-
         try:
-            if inspect.iscoroutinefunction(create_fn):
-                await create_fn(**event)  # type: ignore[arg-type]
+            # Use the same format as your working curl command
+            response = await self.client.post(
+                f"{self.base_url}/api/v1/events",
+                json=event,
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/cloudevents+json"
+                }
+            )
+            
+            if response.status_code in [200, 201, 202, 204]:  # ← Add 204
+                try:
+                    result = response.json()
+                    logger.info("OpenMeter event sent successfully")
+                except:
+                    logger.info(f"✅ OpenMeter success (no content): HTTP {response.status_code}")
             else:
-                await anyio.to_thread.run_sync(create_fn, **event)
-        except Exception as exc:  # pragma: no cover - network errors
-            logger.warning("OpenMeter create failed: %s", exc)
+                logger.warning(f"OpenMeter error: {response.status_code}")
+                
+        except Exception as exc:
+            logger.warning("OpenMeter request failed: %s", exc)

From fc5cfeddd732dadbab5f9e47bd4818bdba5fb00d Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 22 Jul 2025 19:50:12 -0700
Subject: [PATCH 16/24] updated readme and openmeter

---
 README.md         | 23 +++++++-----
 pyproject.toml    |  1 -
 usage/backends.py | 91 ++++++++++++++++++++++++++++-------------------
 3 files changed, 68 insertions(+), 47 deletions(-)

diff --git a/README.md b/README.md
index 48e8155..6a3a004 100644
--- a/README.md
+++ b/README.md
@@ -261,16 +261,21 @@ export USAGE_METERING=prometheus
 #### OpenMeter (Stripe / ClickHouse)
 
 ```bash
-pip install "attach-gateway[usage]"
-export MAX_TOKENS_PER_MIN=60000
-export USAGE_METERING=openmeter
-export OPENMETER_API_KEY=...
-export OPENMETER_URL=http://localhost:8888   # optional self-host, defaults to https://openmeter.cloud
+# No additional dependencies needed - uses direct HTTP API
+export MAX_TOKENS_PER_MIN=60000              # Required: enables quota middleware
+export USAGE_METERING=openmeter              # Required: activates OpenMeter backend  
+export OPENMETER_API_KEY=your-api-key-here   # Required: API authentication
+export OPENMETER_URL=https://openmeter.cloud # Optional: defaults to https://openmeter.cloud
 ```
 
-Events land in the tokens meter of OpenMeter and can sync to Stripe.
+Events are sent directly to OpenMeter's HTTP API and are processed by the LLM tokens meter for billing integration with Stripe.
+
+> **⚠️ All three variables are required for OpenMeter to work:**  
+> - `MAX_TOKENS_PER_MIN` enables the quota middleware that records usage events  
+> - `USAGE_METERING=openmeter` activates the OpenMeter backend  
+> - `OPENMETER_API_KEY` provides authentication to OpenMeter's API  
 
-The gateway runs fine without these vars; metering activates only when both USAGE_METERING=openmeter and OPENMETER_API_KEY are set.
+The gateway gracefully falls back to `NullUsageBackend` if any required variable is missing.
 
 ### Scraping metrics
 
@@ -281,7 +286,7 @@ curl -H "Authorization: Bearer $JWT" http://localhost:8080/metrics
 ## Token quotas
 
 Attach Gateway can enforce per-user token limits. Install the optional
-dependency with `pip install attach-gateway[quota]` and set
+dependency with `pip install attach-dev[quota]` and set
 `MAX_TOKENS_PER_MIN` in your environment to enable the middleware. The
 counter defaults to the `cl100k_base` encoding; override with
 `QUOTA_ENCODING` if your model uses a different tokenizer. The default
@@ -297,7 +302,7 @@ production.
 ```bash
 # Optional: Enable token quotas
 export MAX_TOKENS_PER_MIN=60000
-pip install tiktoken  # or pip install attach-gateway[quota]
+pip install tiktoken  # or pip install attach-dev[quota]
 ```
 
 To customize the tokenizer:
diff --git a/pyproject.toml b/pyproject.toml
index f74f709..760f14b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,6 @@ temporal = ["temporalio>=1.5.0"]
 quota = ["tiktoken>=0.5.0"]
 usage = [
     "prometheus_client>=0.20.0",
-    "openmeter>=1.0.0b188,<2",  # openmeter>=1.0.0b188 requires Python >= 3.9
 ]
 dev = [
     "pytest>=8.0.0", 
diff --git a/usage/backends.py b/usage/backends.py
index 9e9fdf2..f589ce3 100644
--- a/usage/backends.py
+++ b/usage/backends.py
@@ -114,45 +114,62 @@ async def record(self, **evt) -> None:
         try:
             from uuid import uuid4
         except ImportError as exc:
-            print(f"❌ UUID import failed: {exc}")
             return
 
-        # Create event in the same format as your working curl
-        event = {
-            "specversion": "1.0",
-            "type": "tokens",
-            "id": str(uuid4()),
-            "time": datetime.now(timezone.utc)
-            .isoformat(timespec="milliseconds")
-            .replace("+00:00", "Z"),
-            "source": "attach-gateway",
-            "subject": evt.get("user"),
-            "data": {
-                "tokens_in": int(evt.get("tokens_in", 0) or 0),
-                "tokens_out": int(evt.get("tokens_out", 0) or 0),
-                "model": evt.get("model"),
-            }
-        }
+        base_time = datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")
+        user = evt.get("user")
+        model = evt.get("model")
+        
+        tokens_in = int(evt.get("tokens_in", 0) or 0)
+        tokens_out = int(evt.get("tokens_out", 0) or 0)
 
-        try:
-            # Use the same format as your working curl command
-            response = await self.client.post(
-                f"{self.base_url}/api/v1/events",
-                json=event,
-                headers={
-                    "Authorization": f"Bearer {self.api_key}",
-                    "Content-Type": "application/cloudevents+json"
+        # Send separate events for input and output tokens
+        events_to_send = []
+        
+        if tokens_in > 0:
+            events_to_send.append({
+                "specversion": "1.0",
+                "type": "prompt",        # ← Changed from "tokens" to "prompt"
+                "id": str(uuid4()),
+                "time": base_time,
+                "source": "attach-gateway",
+                "subject": user,
+                "data": {
+                    "tokens": tokens_in,
+                    "model": model,
+                    "type": "input"      # ← This stays the same
                 }
-            )
-            
-            if response.status_code in [200, 201, 202, 204]:  # ← Add 204
-                try:
-                    result = response.json()
-                    logger.info("OpenMeter event sent successfully")
-                except:
-                    logger.info(f"✅ OpenMeter success (no content): HTTP {response.status_code}")
-            else:
-                logger.warning(f"OpenMeter error: {response.status_code}")
+            })
+        
+        if tokens_out > 0:
+            events_to_send.append({
+                "specversion": "1.0", 
+                "type": "prompt",
+                "id": str(uuid4()),
+                "time": base_time,
+                "source": "attach-gateway",
+                "subject": user,
+                "data": {
+                    "tokens": tokens_out,   # ← Single tokens field
+                    "model": model,
+                    "type": "output"        # ← Add type field
+                }
+            })
+
+        # Send each event
+        for event in events_to_send:
+            try:
+                response = await self.client.post(
+                    f"{self.base_url}/api/v1/events",
+                    json=event,
+                    headers={
+                        "Authorization": f"Bearer {self.api_key}",
+                        "Content-Type": "application/cloudevents+json"
+                    }
+                )
                 
-        except Exception as exc:
-            logger.warning("OpenMeter request failed: %s", exc)
+                if response.status_code not in [200, 201, 202, 204]:
+                    logger.warning(f"OpenMeter error: {response.status_code}")
+                    
+            except Exception as exc:
+                logger.warning("OpenMeter request failed: %s", exc)

From 29ba261a549ae647fb60f6b4024e8099b7843a1b Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 22 Jul 2025 23:30:39 -0700
Subject: [PATCH 17/24] fixed len(text) if tiktoken is not available

---
 middleware/quota.py | 45 ++++++++++++++++++++++++++++++++++-----------
 1 file changed, 34 insertions(+), 11 deletions(-)

diff --git a/middleware/quota.py b/middleware/quota.py
index 21f8666..2e5fc53 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -23,11 +23,17 @@
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import JSONResponse, StreamingResponse
 
+# ─────────────────────────────────────────────────────────
+# Tokenizer setup
+# ─────────────────────────────────────────────────────────
 try:
     import tiktoken  # type: ignore
 except Exception:  # pragma: no cover
     tiktoken = None
 
+# ≈ cl100k_base: ~4 bytes / token for typical English
+_APPROX_BYTES_PER_TOKEN = 4
+
 from usage.backends import NullUsageBackend
 from utils.env import int_env
 
@@ -155,18 +161,20 @@ def _is_textual(mime: str) -> bool:
     return mime.startswith("text/") or "json" in mime
 
 
+# ---------------------------------------------------------------------------
+# Token-count helpers
+# ---------------------------------------------------------------------------
+
 def _encoder_for_model(model: str):
     """Return a tiktoken encoder, falling back to byte count."""
-    if tiktoken is None:  # pragma: no cover - fallback
-        logger.warning(
-            "tiktoken missing – using byte-count fallback; token metrics may be inflated"
-        )
+    if tiktoken is None:  # fallback: 1 token ≈ 4 bytes
 
-        class _Simple:
+        class _Approx:
             def encode(self, text: str) -> list[int]:
-                return list(text.encode())
+                # Never return 0 → always count at least 1 token
+                return [0] * max(1, len(text) // _APPROX_BYTES_PER_TOKEN)
 
-        return _Simple()
+        return _Approx()
 
     try:
         return tiktoken.encoding_for_model(model)
@@ -175,11 +183,11 @@ def encode(self, text: str) -> list[int]:
             return tiktoken.get_encoding("cl100k_base")
         except Exception:
 
-            class _Simple:
+            class _Approx:
                 def encode(self, text: str) -> list[int]:
-                    return list(text.encode())
+                    return [0] * max(1, len(text) // _APPROX_BYTES_PER_TOKEN)
 
-            return _Simple()
+            return _Approx()
 
 
 def _num_tokens(text: str, model: str = "cl100k_base") -> int:
@@ -340,6 +348,21 @@ async def dispatch(self, request: Request, call_next):
         if not hasattr(request.app.state, "usage"):
             request.app.state.usage = NullUsageBackend()
 
+        # ── OPTIONAL request-size guard (default 1 MB) ───────────────
+        max_bytes = int(os.getenv("MAX_REQUEST_BYTES", "1000000"))
+        raw = await request.body()
+        if len(raw) > max_bytes:
+            return JSONResponse(
+                {
+                    "detail": "request too large",
+                    "limit_bytes": max_bytes,
+                },
+                status_code=413,
+            )
+
+        # Re-use the already-read body from here on
+        request._body = raw
+
         user = request.headers.get("x-attach-user") or (
             request.client.host if request.client else "unknown"
         )
@@ -355,7 +378,7 @@ async def dispatch(self, request: Request, call_next):
 
         tokens_in = 0
         if req_is_text:
-            raw = await request.body()
+            # raw already read above
             try:
                 payload = json.loads(raw.decode())
             except Exception:

From 3f479e8519071ebc1a9f274e4674fef57f12d5de Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Tue, 22 Jul 2025 23:34:06 -0700
Subject: [PATCH 18/24] updated readme roadmap

---
 README.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 6a3a004..f87441d 100644
--- a/README.md
+++ b/README.md
@@ -312,8 +312,8 @@ export QUOTA_ENCODING=cl100k_base  # default
 
 ## Roadmap
 
-* **v0.2** — Protected‑resource metadata endpoint (OAuth 2.1), enhanced DID resolvers.  
-* **v0.3** — Token‑exchange (RFC 8693) for on‑behalf‑of delegation.  
+* **v0.3.1** — Protected‑resource metadata endpoint (OAuth 2.1), enhanced DID resolvers.  
+* **v0.3.1** — Token‑exchange (RFC 8693) for on‑behalf‑of delegation.  
 * **v0.4** — Attach Store v1 (Git‑style, policy guards).
 
 ---

From 34427e2dc5c961a834dc1e75fa22906dfa7354f9 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Wed, 23 Jul 2025 15:49:59 -0700
Subject: [PATCH 19/24] gateway now mirrors main

---
 attach/gateway.py | 70 +++++++++++++++++++++++++++++++++++------------
 1 file changed, 53 insertions(+), 17 deletions(-)

diff --git a/attach/gateway.py b/attach/gateway.py
index ce27650..d0c1dff 100644
--- a/attach/gateway.py
+++ b/attach/gateway.py
@@ -3,28 +3,33 @@
 """
 
 import os
+from contextlib import asynccontextmanager
 from typing import Optional
 
 import weaviate
 from fastapi import APIRouter, FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from starlette.middleware.base import BaseHTTPMiddleware
 from pydantic import BaseModel
 
 from a2a.routes import router as a2a_router
-
-# Clean relative imports
-from auth import verify_jwt
 from auth.oidc import _require_env
-
-# from logs import router as logs_router
+from logs import router as logs_router
 from mem import get_memory_backend
 from middleware.auth import jwt_auth_mw
-from middleware.quota import TokenQuotaMiddleware
 from middleware.session import session_mw
 from proxy.engine import router as proxy_router
 from usage.factory import _select_backend, get_usage_backend
 from usage.metrics import mount_metrics
 from utils.env import int_env
 
+# Guard TokenQuotaMiddleware import (matches main.py pattern)
+try:
+    from middleware.quota import TokenQuotaMiddleware
+    QUOTA_AVAILABLE = True
+except ImportError:  # optional extra not installed
+    QUOTA_AVAILABLE = False
+
 # Import version from parent package
 from . import __version__
 
@@ -52,7 +57,7 @@ async def get_memory_events(request: Request, limit: int = 10):
             return {"data": {"Get": {"MemoryEvent": []}}}
 
         result = (
-            client.query.get("MemoryEvent", ["timestamp", "role", "content"])
+            client.query.get("MemoryEvent", ["timestamp", "event", "user", "state"])
             .with_additional(["id"])
             .with_limit(limit)
             .with_sort([{"path": ["timestamp"], "order": "desc"}])
@@ -100,6 +105,21 @@ class AttachConfig(BaseModel):
     auth0_client: Optional[str] = None
 
 
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Manage application lifespan - startup and shutdown."""
+    # Startup
+    backend_selector = _select_backend()
+    app.state.usage = get_usage_backend(backend_selector)
+    mount_metrics(app)
+    
+    yield
+    
+    # Shutdown
+    if hasattr(app.state.usage, 'aclose'):
+        await app.state.usage.aclose()
+
+
 def create_app(config: Optional[AttachConfig] = None) -> FastAPI:
     """
     Create a FastAPI app with Attach Gateway functionality
@@ -130,27 +150,43 @@ def create_app(config: Optional[AttachConfig] = None) -> FastAPI:
         title="Attach Gateway",
         description="Identity & Memory side-car for LLM engines",
         version=__version__,
+        lifespan=lifespan,
     )
-    mount_metrics(app)
 
-    # Add middleware
-    app.middleware("http")(jwt_auth_mw)
-    app.middleware("http")(session_mw)
-    limit = int_env("MAX_TOKENS_PER_MIN", 60000)
-    if limit is not None:
+    @app.get("/auth/config")
+    async def auth_config():
+        return {
+            "domain": config.auth0_domain,
+            "client_id": config.auth0_client,
+            "audience": config.oidc_audience,
+        }
+
+    # Add middleware in correct order (CORS outer-most)
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["http://localhost:9000", "http://127.0.0.1:9000"],
+        allow_methods=["*"],
+        allow_headers=["*"],
+        allow_credentials=True,
+    )
+    
+    app.add_middleware(BaseHTTPMiddleware, dispatch=jwt_auth_mw)
+    app.add_middleware(BaseHTTPMiddleware, dispatch=session_mw)
+    
+    # Only add quota middleware if available and explicitly configured
+    limit = int_env("MAX_TOKENS_PER_MIN", None)
+    if QUOTA_AVAILABLE and limit is not None:
         app.add_middleware(TokenQuotaMiddleware)
 
     # Add routes
-    app.include_router(a2a_router)
+    app.include_router(a2a_router, prefix="/a2a")
     app.include_router(proxy_router)
-    # app.include_router(logs_router)
+    app.include_router(logs_router)
     app.include_router(mem_router)
 
     # Setup memory backend
     memory_backend = get_memory_backend(config.mem_backend, config)
     app.state.memory = memory_backend
     app.state.config = config
-    backend_selector = _select_backend()
-    app.state.usage = get_usage_backend(backend_selector)
 
     return app

From 79e4e471405b096ea38057de365b2fc3197e70ba Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Thu, 24 Jul 2025 17:15:14 -0700
Subject: [PATCH 20/24] header bugs squashed

---
 .env.example        |  3 +-
 attach/gateway.py   |  2 +-
 main.py             | 71 +++++++++++++++------------------------------
 middleware/quota.py |  4 ++-
 4 files changed, 29 insertions(+), 51 deletions(-)

diff --git a/.env.example b/.env.example
index 6019812..b821c1c 100644
--- a/.env.example
+++ b/.env.example
@@ -19,7 +19,8 @@ WEAVIATE_URL=http://localhost:8081
 MAX_TOKENS_PER_MIN=60000
 QUOTA_ENCODING=cl100k_base
 
-USAGE_METERING=null
+# Metering Option (null, prometheus, openmeter)
+USAGE_METERING=null 
 
 # Development: Auth0 credentials for dev_login script
 # AUTH0_DOMAIN=your-domain.auth0.com
diff --git a/attach/gateway.py b/attach/gateway.py
index d0c1dff..ba286d9 100644
--- a/attach/gateway.py
+++ b/attach/gateway.py
@@ -174,7 +174,7 @@ async def auth_config():
     app.add_middleware(BaseHTTPMiddleware, dispatch=session_mw)
     
     # Only add quota middleware if available and explicitly configured
-    limit = int_env("MAX_TOKENS_PER_MIN", None)
+    limit = int_env("MAX_TOKENS_PER_MIN", 60000)
     if QUOTA_AVAILABLE and limit is not None:
         app.add_middleware(TokenQuotaMiddleware)
 
diff --git a/main.py b/main.py
index 76e1efc..5ac3826 100644
--- a/main.py
+++ b/main.py
@@ -8,45 +8,36 @@
 from contextlib import asynccontextmanager
 
 from a2a.routes import router as a2a_router
-from auth.oidc import verify_jwt  # Fixed: was auth.jwt, now auth.oidc
 from logs import router as logs_router
-from mem import write as mem_write  # Import memory write function
-from middleware.auth import jwt_auth_mw  # ← your auth middleware
-from middleware.session import session_mw  # ← generates session-id header
+from middleware.auth import jwt_auth_mw
+from middleware.session import session_mw
 from proxy.engine import router as proxy_router
 from usage.factory import _select_backend, get_usage_backend
 from usage.metrics import mount_metrics
 from utils.env import int_env
 
-# At the top, make the import conditional
 try:
     from middleware.quota import TokenQuotaMiddleware
-
     QUOTA_AVAILABLE = True
 except ImportError:
     QUOTA_AVAILABLE = False
 
-# Memory router
 mem_router = APIRouter(prefix="/mem", tags=["memory"])
 
 
 @mem_router.get("/events")
 async def get_memory_events(request: Request, limit: int = 10):
-    """Fetch recent MemoryEvent objects from Weaviate"""
+    """Fetch recent MemoryEvent objects from Weaviate."""
     try:
-        # Get user info from request state (set by jwt_auth_mw middleware)
         user_sub = getattr(request.state, "sub", None)
         if not user_sub:
             raise HTTPException(status_code=401, detail="User not authenticated")
 
-        # Use the exact same client setup as demo_view_memory.py
-        client = weaviate.Client("http://localhost:6666")
+        client = weaviate.Client(os.getenv("WEAVIATE_URL", "http://localhost:6666"))
 
-        # Test connection first
         if not client.is_ready():
             raise HTTPException(status_code=503, detail="Weaviate is not ready")
 
-        # Check schema the same way as demo_view_memory.py
         try:
             schema = client.schema.get()
             classes = {c["class"] for c in schema.get("classes", [])}
@@ -56,7 +47,6 @@ async def get_memory_events(request: Request, limit: int = 10):
         except Exception:
             return {"data": {"Get": {"MemoryEvent": []}}}
 
-        # Query with descending order by timestamp (newest first)
         result = (
             client.query.get(
                 "MemoryEvent",
@@ -68,7 +58,6 @@ async def get_memory_events(request: Request, limit: int = 10):
             .do()
         )
 
-        # Check for GraphQL errors like demo_view_memory.py does
         if "errors" in result:
             raise HTTPException(
                 status_code=500, detail=f"GraphQL error: {result['errors']}"
@@ -79,31 +68,26 @@ async def get_memory_events(request: Request, limit: int = 10):
 
         events = result["data"]["Get"]["MemoryEvent"]
 
-        # Add the result field from the raw objects for richer display
         try:
             raw_objects = client.data_object.get(class_name="MemoryEvent", limit=limit)
 
-            # Create a mapping of IDs to full objects
             id_to_full_object = {}
             for obj in raw_objects.get("objects", []):
                 obj_id = obj.get("id")
                 if obj_id:
                     id_to_full_object[obj_id] = obj.get("properties", {})
 
-            # Enrich the GraphQL results with data from raw objects
             for event in events:
                 event_id = event.get("_additional", {}).get("id")
                 if event_id and event_id in id_to_full_object:
                     full_props = id_to_full_object[event_id]
-                    # Add the result field if it exists
                     if "result" in full_props:
                         event["result"] = full_props["result"]
-                    # Add other useful fields
                     for field in ["event", "session_id", "task_id", "user"]:
                         if field in full_props:
                             event[field] = full_props[field]
         except Exception:
-            pass  # Silently fail if we can't enrich with raw object data
+            pass
 
         return result
 
@@ -113,43 +97,36 @@ async def get_memory_events(request: Request, limit: int = 10):
         )
 
 
-middlewares = [
-    # ❶ CORS first (so it executes last and handles responses properly)
-    Middleware(
-        CORSMiddleware,
-        allow_origins=["http://localhost:9000", "http://127.0.0.1:9000"],
-        allow_methods=["*"],
-        allow_headers=["*"],
-        allow_credentials=True,
-    ),
-    # ❷ Auth middleware
-    Middleware(BaseHTTPMiddleware, dispatch=jwt_auth_mw),
-    # ❸ Session middleware
-    Middleware(BaseHTTPMiddleware, dispatch=session_mw),
-]
-
-# Only add quota middleware if tiktoken is available and a positive limit is set
-limit = int_env("MAX_TOKENS_PER_MIN", 60000)
-if QUOTA_AVAILABLE and limit is not None:
-    middlewares.append(Middleware(TokenQuotaMiddleware))
-
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Manage application lifespan - startup and shutdown."""
-    # Startup
     backend_selector = _select_backend()
     app.state.usage = get_usage_backend(backend_selector)
     mount_metrics(app)
     
     yield
     
-    # Shutdown
     if hasattr(app.state.usage, 'aclose'):
         await app.state.usage.aclose()
 
-# Create app with lifespan
-app = FastAPI(title="attach-gateway", middleware=middlewares, lifespan=lifespan)
+app = FastAPI(title="attach-gateway", lifespan=lifespan)
 
+# Add middleware in correct order (CORS outer-most)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["http://localhost:9000", "http://127.0.0.1:9000"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+    allow_credentials=True,
+)
+
+# Only add quota middleware if available and explicitly configured
+limit = int_env("MAX_TOKENS_PER_MIN", 60000)
+if QUOTA_AVAILABLE and limit is not None:
+    app.add_middleware(TokenQuotaMiddleware)
+
+app.add_middleware(BaseHTTPMiddleware, dispatch=jwt_auth_mw)
+app.add_middleware(BaseHTTPMiddleware, dispatch=session_mw)
 
 @app.get("/auth/config")
 async def auth_config():
@@ -159,9 +136,7 @@ async def auth_config():
         "audience": os.getenv("OIDC_AUD"),
     }
 
-
-# Add middleware after routes are defined
 app.include_router(a2a_router, prefix="/a2a")
 app.include_router(logs_router)
 app.include_router(mem_router)
-app.include_router(proxy_router)  # ← ADD THIS BACK
+app.include_router(proxy_router)
diff --git a/middleware/quota.py b/middleware/quota.py
index 2e5fc53..15af934 100644
--- a/middleware/quota.py
+++ b/middleware/quota.py
@@ -363,9 +363,11 @@ async def dispatch(self, request: Request, call_next):
         # Re-use the already-read body from here on
         request._body = raw
 
-        user = request.headers.get("x-attach-user") or (
+        # Fix: Get user from request.state.sub (set by auth middleware)
+        user = getattr(request.state, "sub", None) or (
             request.client.host if request.client else "unknown"
         )
+
         usage = {
             "user": user,
             "project": request.headers.get("x-attach-project", "default"),

From 454556f77958d48b028237d7d64f0491c709343c Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Thu, 24 Jul 2025 17:44:05 -0700
Subject: [PATCH 21/24] fixed the auth.py issue with CORS options

---
 attach/gateway.py  | 6 +++---
 middleware/auth.py | 8 +++++---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/attach/gateway.py b/attach/gateway.py
index ba286d9..25064a1 100644
--- a/attach/gateway.py
+++ b/attach/gateway.py
@@ -170,14 +170,14 @@ async def auth_config():
         allow_credentials=True,
     )
     
-    app.add_middleware(BaseHTTPMiddleware, dispatch=jwt_auth_mw)
-    app.add_middleware(BaseHTTPMiddleware, dispatch=session_mw)
-    
     # Only add quota middleware if available and explicitly configured
     limit = int_env("MAX_TOKENS_PER_MIN", 60000)
     if QUOTA_AVAILABLE and limit is not None:
         app.add_middleware(TokenQuotaMiddleware)
 
+    app.add_middleware(BaseHTTPMiddleware, dispatch=jwt_auth_mw)
+    app.add_middleware(BaseHTTPMiddleware, dispatch=session_mw)
+
     # Add routes
     app.include_router(a2a_router, prefix="/a2a")
     app.include_router(proxy_router)
diff --git a/middleware/auth.py b/middleware/auth.py
index e9fd448..01d7db3 100644
--- a/middleware/auth.py
+++ b/middleware/auth.py
@@ -29,8 +29,12 @@ async def jwt_auth_mw(request: Request, call_next):
     • Verifies it with `auth.oidc.verify_jwt`.
     • Stores the `sub` claim in `request.state.sub` for downstream middleware.
     • Rejects the request with 401 on any failure.
-    • Skips authentication for excluded paths.
+    • Skips authentication for excluded paths and OPTIONS requests.
     """
+    # Skip authentication for OPTIONS requests (CORS preflight)
+    if request.method == "OPTIONS":
+        return await call_next(request)
+    
     # Skip authentication for excluded paths
     if request.url.path in EXCLUDED_PATHS:
         return await call_next(request)
@@ -48,6 +52,4 @@ async def jwt_auth_mw(request: Request, call_next):
 
     # attach the user id (sub) for the session-middleware
     request.state.sub = claims["sub"]
-
-    # continue down the middleware stack / route handler
     return await call_next(request)

From c6b81ee9f5182b731313900d7e1a04aa1058ba57 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Thu, 24 Jul 2025 17:48:44 -0700
Subject: [PATCH 22/24] bumped the version to 0.3.0

---
 attach/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/attach/__init__.py b/attach/__init__.py
index 994d780..5457a6d 100644
--- a/attach/__init__.py
+++ b/attach/__init__.py
@@ -4,7 +4,7 @@
 Add OIDC SSO, agent-to-agent handoff, and pluggable memory to any Python project.
 """
 
-__version__ = "0.2.2"
+__version__ = "0.3.0"
 __author__ = "Hammad Tariq"
 __email__ = "hammad@attach.dev"
 

From c5bf8262157497b5c86407b59b793ae68e83f93c Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Thu, 24 Jul 2025 19:02:53 -0700
Subject: [PATCH 23/24] trying to debug log error in the test package

---
 README.md          | 6 ------
 attach/__init__.py | 2 +-
 attach/gateway.py  | 3 ++-
 main.py            | 3 ++-
 pyproject.toml     | 2 +-
 5 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index f87441d..d28ba46 100644
--- a/README.md
+++ b/README.md
@@ -310,12 +310,6 @@ To customize the tokenizer:
 export QUOTA_ENCODING=cl100k_base  # default
 ```
 
-## Roadmap
-
-* **v0.3.1** — Protected‑resource metadata endpoint (OAuth 2.1), enhanced DID resolvers.  
-* **v0.3.1** — Token‑exchange (RFC 8693) for on‑behalf‑of delegation.  
-* **v0.4** — Attach Store v1 (Git‑style, policy guards).
-
 ---
 
 ## License
diff --git a/attach/__init__.py b/attach/__init__.py
index 5457a6d..04730ca 100644
--- a/attach/__init__.py
+++ b/attach/__init__.py
@@ -4,7 +4,7 @@
 Add OIDC SSO, agent-to-agent handoff, and pluggable memory to any Python project.
 """
 
-__version__ = "0.3.0"
+__version__ = "0.3.2"
 __author__ = "Hammad Tariq"
 __email__ = "hammad@attach.dev"
 
diff --git a/attach/gateway.py b/attach/gateway.py
index 25064a1..855ca9f 100644
--- a/attach/gateway.py
+++ b/attach/gateway.py
@@ -14,7 +14,8 @@
 
 from a2a.routes import router as a2a_router
 from auth.oidc import _require_env
-from logs import router as logs_router
+import logs
+logs_router = logs.router
 from mem import get_memory_backend
 from middleware.auth import jwt_auth_mw
 from middleware.session import session_mw
diff --git a/main.py b/main.py
index 5ac3826..1a99d96 100644
--- a/main.py
+++ b/main.py
@@ -8,7 +8,8 @@
 from contextlib import asynccontextmanager
 
 from a2a.routes import router as a2a_router
-from logs import router as logs_router
+import logs
+logs_router = logs.router
 from middleware.auth import jwt_auth_mw
 from middleware.session import session_mw
 from proxy.engine import router as proxy_router
diff --git a/pyproject.toml b/pyproject.toml
index 760f14b..c50bb8d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -56,7 +56,7 @@ attach-gateway = "attach.__main__:main"
 
 # Include your existing modules with cleaner targeting
 [tool.hatch.build.targets.wheel]
-packages = ["attach", "auth", "middleware", "mem", "proxy", "a2a", "attach_pydid"]
+packages = ["attach", "auth", "middleware", "mem", "proxy", "a2a", "attach_pydid", "usage", "utils"]
 
 
 # Dynamic version from attach/__init__.py

From 70992db651bd47412a6acd5b89158da923c70da1 Mon Sep 17 00:00:00 2001
From: Hammad Tariq <hammadtariq@yahoo.co.uk>
Date: Thu, 24 Jul 2025 20:29:00 -0700
Subject: [PATCH 24/24] 0.3.7 house-keeping for packages

---
 attach/__init__.py          | 15 ++++++++++++---
 attach/__main__.py          | 35 ++++++++++++++++++++++++++++++-----
 logs.py => logs/__init__.py |  0
 pyproject.toml              | 11 ++++++-----
 usage/factory.py            | 29 ++++++++++++++++++++++-------
 5 files changed, 70 insertions(+), 20 deletions(-)
 rename logs.py => logs/__init__.py (100%)

diff --git a/attach/__init__.py b/attach/__init__.py
index 04730ca..349db7c 100644
--- a/attach/__init__.py
+++ b/attach/__init__.py
@@ -4,11 +4,20 @@
 Add OIDC SSO, agent-to-agent handoff, and pluggable memory to any Python project.
 """
 
-__version__ = "0.3.2"
+__version__ = "0.3.7"
 __author__ = "Hammad Tariq"
 __email__ = "hammad@attach.dev"
 
-# Clean imports - no sys.path hacks needed since everything will be in the wheel
-from .gateway import create_app, AttachConfig
+# Remove this line that causes early failure:
+# from .gateway import create_app, AttachConfig
+
+# Optional: Add lazy import for convenience
+def create_app(*args, **kwargs):
+    from .gateway import create_app as _real
+    return _real(*args, **kwargs)
+
+def AttachConfig(*args, **kwargs):
+    from .gateway import AttachConfig as _real
+    return _real(*args, **kwargs)
 
 __all__ = ["create_app", "AttachConfig", "__version__"] 
\ No newline at end of file
diff --git a/attach/__main__.py b/attach/__main__.py
index 122d419..99490f2 100644
--- a/attach/__main__.py
+++ b/attach/__main__.py
@@ -2,7 +2,7 @@
 CLI entry point - replaces the need for main.py in wheel
 """
 import uvicorn
-from .gateway import create_app
+import click
 
 def main():
     """Run Attach Gateway server"""
@@ -13,17 +13,42 @@ def main():
     except ImportError:
         pass  # python-dotenv not installed, that's OK for production
     
-    import click
-    
     @click.command()
     @click.option("--host", default="0.0.0.0", help="Host to bind to")
     @click.option("--port", default=8080, help="Port to bind to") 
     @click.option("--reload", is_flag=True, help="Enable auto-reload")
     def cli(host: str, port: int, reload: bool):
-        app = create_app()
-        uvicorn.run(app, host=host, port=port, reload=reload)
+        try:
+            # Import here AFTER .env is loaded and CLI is parsed
+            from .gateway import create_app
+            app = create_app()
+            uvicorn.run(app, host=host, port=port, reload=reload)
+        except RuntimeError as e:
+            _friendly_exit(e)
+        except Exception as e:  # unexpected crash
+            click.echo(f"❌ Startup failed: {e}", err=True)
+            raise click.Abort()
     
     cli()
 
+def _friendly_exit(err):
+    """Convert RuntimeError to clean user message."""
+    err_str = str(err)
+    
+    if "OPENMETER_API_KEY" in err_str:
+        msg = (f"❌ {err}\n\n"
+               "💡 Fix:\n"
+               "   export OPENMETER_API_KEY=\"sk_live_...\"\n"
+               "   (or) export USAGE_METERING=null    # to disable metering\n\n"
+               "📖 See README.md for complete setup")
+    else:
+        msg = (f"❌ {err}\n\n"
+               "💡 Required environment variables:\n"
+               "   export OIDC_ISSUER=\"https://your-domain.auth0.com/\"\n"
+               "   export OIDC_AUD=\"your-api-identifier\"\n\n"
+               "📖 See README.md for complete setup instructions")
+    
+    raise click.ClickException(msg)
+
 if __name__ == "__main__":
     main() 
\ No newline at end of file
diff --git a/logs.py b/logs/__init__.py
similarity index 100%
rename from logs.py
rename to logs/__init__.py
diff --git a/pyproject.toml b/pyproject.toml
index c50bb8d..35daaa7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,14 +31,15 @@ dependencies = [
     "python-jose[cryptography]>=3.3.0",
     "click>=8.0.0",
     "python-dotenv>=1.0.0",
+    "weaviate-client>=3.26.7,<4.0.0",
 ]
 
 [project.optional-dependencies]
-memory = ["weaviate-client>=3.26.7,<4.0.0"]
-temporal = ["temporalio>=1.5.0"]
 quota = ["tiktoken>=0.5.0"]
-usage = [
-    "prometheus_client>=0.20.0",
+usage = ["prometheus_client>=0.20.0"]
+full = [
+    "tiktoken>=0.5.0", 
+    "prometheus_client>=0.20.0"
 ]
 dev = [
     "pytest>=8.0.0", 
@@ -56,7 +57,7 @@ attach-gateway = "attach.__main__:main"
 
 # Include your existing modules with cleaner targeting
 [tool.hatch.build.targets.wheel]
-packages = ["attach", "auth", "middleware", "mem", "proxy", "a2a", "attach_pydid", "usage", "utils"]
+packages = ["attach", "auth", "middleware", "mem", "proxy", "a2a", "attach_pydid", "usage", "utils", "logs"]
 
 
 # Dynamic version from attach/__init__.py
diff --git a/usage/factory.py b/usage/factory.py
index 3663db5..61ed842 100644
--- a/usage/factory.py
+++ b/usage/factory.py
@@ -4,6 +4,7 @@
 
 import os
 import warnings
+import logging
 
 from .backends import (
     AbstractUsageBackend,
@@ -12,14 +13,16 @@
     PrometheusUsageBackend,
 )
 
+log = logging.getLogger(__name__)
+
 
 def _select_backend() -> str:
     """Return backend name from env vars with deprecation warning."""
     if "USAGE_METERING" in os.environ:
         return os.getenv("USAGE_METERING", "null")
-    if "USAGE_BACKEND" in os.environ:
+    if "USAGE_BACKEND" in os.environ:  # old name, keep BC
         warnings.warn(
-            "USAGE_BACKEND is deprecated; rename to USAGE_METERING",
+            "USAGE_BACKEND is deprecated; use USAGE_METERING",
             UserWarning,
             stacklevel=2,
         )
@@ -29,14 +32,26 @@ def _select_backend() -> str:
 def get_usage_backend(kind: str) -> AbstractUsageBackend:
     """Return an instance of the requested usage backend."""
     kind = (kind or "null").lower()
+
     if kind == "prometheus":
         try:
             return PrometheusUsageBackend()
-        except Exception:
+        except ImportError as exc:
+            log.warning(
+                "Prometheus metering unavailable: %s – "
+                "falling back to NullUsageBackend. "
+                "Install with: pip install 'attach-dev[usage]'",
+                exc
+            )
             return NullUsageBackend()
+
     if kind == "openmeter":
-        try:
-            return OpenMeterBackend()
-        except Exception:
-            return NullUsageBackend()
+        # fail-fast on bad config
+        if not os.getenv("OPENMETER_API_KEY"):
+            raise RuntimeError(
+                "USAGE_METERING=openmeter requires OPENMETER_API_KEY. "
+                "Set the variable or change USAGE_METERING=null to disable."
+            )
+        return OpenMeterBackend()  # exceptions inside bubble up
+
     return NullUsageBackend()