Merge pull request #87 from superdesigndev/feat/mcp-idempotency-key

feat(mcp): call takes an idempotency_key, plus docs on every surface

.agents/skills/tools-registry-context/MAP.md

@@ -14,7 +14,9 @@ Regenerate via `scripts/build-map.py`.
 | `external:meetings/2026-06-30-jason-tools-registry.md` | foundation/charter.md, reference/glossary.md |
 | `render.yaml` | ops/deploy.md |
 | `src/treg/__main__.py` | ops/deploy.md |
-| `src/treg/api.py` | architecture/multi-tenancy.md, architecture/proxy-model.md, architecture/super-admin.md, interface/api.md, interface/dashboard.md, interface/landing-sandbox.md |
+| `src/treg/api.py` | architecture/multi-tenancy.md, architecture/proxy-model.md, architecture/super-admin.md, architecture/money.md, interface/api.md, interface/dashboard.md, interface/landing-sandbox.md |
+| `src/treg/mcp.py` | architecture/mcp-oauth.md |
+| `src/treg/mcp_oauth.py` | architecture/mcp-oauth.md |
 | `src/treg/audit.py` | architecture/data-model.md, ops/deploy.md |
 | `src/treg/cli.py` | interface/cli.md, interface/onboarding.md, interface/shell.md |
 | `src/treg/config.py` | architecture/super-admin.md, ops/deploy.md |
@@ -54,6 +56,7 @@ Regenerate via `scripts/build-map.py`.
 | `architecture/data-model.md` | `models.py`, `db.py`, `audit.py`, `ratestore.py` |
 | `architecture/local-run.md` | `localrun.py`, `egress.py`, `fsjail.py` |
 | `architecture/multi-tenancy.md` | `models.py`, `api.py`, `db.py` |
+| `architecture/mcp-oauth.md` | `mcp.py`, `mcp_oauth.py`, `connect-demo.html` |
 | `architecture/proxy-model.md` | `proxy.py`, `api.py` |
 | `architecture/super-admin.md` | `api.py`, `config.py` |
 | `foundation/charter.md` | `2026-06-30-jason-tools-registry.md`, `README.md` |

docs/context/architecture/mcp-oauth.md

@@ -47,6 +47,23 @@ drifting quietly.
 The catalog is the one exception: read straight from `catalog_store`, which is already parsed in
 memory, so a search answers in about a millisecond. That is a **speed** choice, not a permission one.
 
+## `call` takes an `idempotency_key`
+
+Optional, and it is the caller's. Pass the same key when repeating a call whose answer never arrived:
+treg replays the stored response, does not reach the provider, and charges nothing, with
+`replayed: true` on the result.
+
+It exists because the feature was built for agents and MCP is the agent path. Without it the whole
+thing was unreachable from the surface it was for.
+
+Deliberately NOT derived server-side from the endpoint and parameters, which was proposed and
+rejected: two identical searches an hour apart are new work, treg cannot tell that from a retry, and
+a server-invented key would quietly serve stale data. The tool description therefore spends its words
+on WHEN to use it, because that is where a mistake costs something. Reuse for a different request is
+refused rather than answered, so the failure is loud.
+
+Full reasoning, storage rules and the concurrency guard: `architecture/money.md`.
+
 ## Authentication: eager, every request
 
 `RequireAuthForProtectedTools` answers **any** uncredentialed MCP request with **401** and a

docs/context/architecture/money.md

@@ -5,6 +5,7 @@ sources:
   - src/treg/ledger.py
   - src/treg/billing.py
   - src/treg/reconcile.py
+  - src/treg/api.py
 related:
   - architecture/catalog.md
   - architecture/proxy-model.md
@@ -172,3 +173,73 @@ Closing the hold runs on its **own session** (the request's may be mid-rollback
 being released for) and **never raises** — the caller already has their answer, and a ledger hiccup
 must not turn a served call into a 500. A hold that fails to close is not lost money either: the
 reaper releases it, which errs in the org's favour.
+
+## Retries: a call must not be paid for twice
+
+Prompted by a public question — *"how does result pricing handle retries, agents need idempotent
+billing before this works"* — and it matters more here than for a human-facing API, because agents
+retry far more than people do.
+
+**Most retries were already free**, which is what makes the real gap narrow. `_platform_billable`
+never bills a 5xx, a 3xx, a timeout or a network error, and bills a 4xx only under `per_call` where
+the provider charges for accepting the request at all. Result pricing settles on the provider's own
+reported number (`_observed_cost_micro`), so a `per_success` lookup that finds nothing costs nothing.
+
+The gap is one case: treg reached the provider, the provider succeeded **and charged us**, and the
+response was lost on the way back. The agent retries and we pay twice.
+
+### Why remembering the charge is not enough
+
+The cheap fix is to note that a key was already billed and skip the second charge. It does not work:
+treg would still make the second upstream call, so we would still pay the provider and would simply
+move the double cost onto ourselves. The second request has to not reach the provider at all, which
+means storing the first response and replaying it.
+
+### The surface
+
+`Idempotency-Key: <label>` on `/call/`, or the `idempotency_key` argument to the MCP `call` tool. A
+replay answers with `X-Treg-Idempotent-Replay: true` and `X-Treg-Cost-Micro` set to what the FIRST
+call cost, so a caller can report the charge honestly rather than implying a second one. Over MCP the
+result carries `replayed: true`.
+
+Nothing happens without it. A caller who sends no label sees byte-identical behaviour to before the
+feature existed, which is what made it safe to ship.
+
+### The key belongs to the caller
+
+`IdempotentCall` is keyed on `(membership_id, key)`. Every door — a personal token, an agent token,
+an OAuth grant — resolves to one `Membership`, so a single rule covers a human, an agent, and two
+agents in the same team: **the label belongs to whoever called**.
+
+Not `key` alone: clients choose their own labels, the same string will be picked twice, and that
+collision would serve one team's stored response to another. It is the only failure in this feature
+that leaks data rather than money. Not per-org either — two lazily written agents in one team both
+reach for `retry-1` and would collide for no reason.
+
+The key is never derived from the request. That was proposed and rejected: two identical searches an
+hour apart are new work, and treg cannot tell that from a retry. A server-invented key would turn a
+correctness feature into a 24-hour cache that quietly serves stale data.
+
+### What is stored, and for how long
+
+Metered successes only, for 24 hours. A team calling on its **own** key is billed by the provider, so
+there is nothing to protect and no reason to hold their response. A failure is never billed, and
+replaying one would freeze an error the caller should be free to retry out of — so a failed call
+frees its label immediately.
+
+That is also what bounds storage: bodies are kept only for calls that actually cost money, for a day.
+
+### Concurrency, and giving the label back
+
+A `pending` row is written **before** the upstream call, and that row is the lock: two retries
+arriving together race on the unique constraint, and the loser is told to wait (409) instead of
+duplicating the spend. Same reasoning as the conditional UPDATE in `ledger.reserve` — where two paths
+can read before either writes, the database has to arbitrate.
+
+A request that dies after claiming must give the label back, or a single bad parameter would hold it
+for the full window and answer every retry with 409 — worse than the problem this solves. The release
+happens in the `StarletteHTTPException` handler, the one place every refusal passes through; the call
+handler has a dozen raise points and releasing at each would be a dozen chances to miss one.
+
+Expired rows are swept lazily at claim time, scoped to the calling caller, exactly like the hold
+reaper: no scheduler, no leader election on a multi-instance deploy.

docs/context/interface/api.md

@@ -459,3 +459,21 @@ treg is an OAuth authorization server for its own MCP endpoint. Detail in
 `/call/` gained one thing for this: a metered response now carries `X-Treg-Cost-Micro`, so a caller
 can report what it spent instead of diffing the balance. Absent on an unmetered call — a team's own
 key is not ours to bill, and `0` would read as "free".
+
+## `Idempotency-Key` on `/call/`
+
+A caller-supplied label that makes a retry free. Sent on `/call/`, honoured only when present, so a
+caller who omits it sees byte-identical behaviour to before the feature existed.
+
+    Idempotency-Key: <caller's label>        → replay if we already answered this label
+    X-Treg-Idempotent-Replay: true           → on the response, when it came from store
+    X-Treg-Cost-Micro: <original charge>     → what the FIRST call cost, not a new charge
+
+Refusals: `422` when a key is reused for a different request (a caller bug, and answering it would
+hand them a response to a question they did not ask), `409` while the first call with that key is
+still in flight.
+
+Over MCP the same thing is the optional `idempotency_key` argument to the `call` tool, and a replayed
+result carries `replayed: true`.
+
+Reasoning, storage rules and the concurrency guard: `architecture/money.md`.

plugin/skills/treg/SKILL.md

@@ -134,6 +134,17 @@ Rules for spending someone's balance:
     you hold, and treg relays rather than rewrites your request.
 - An endpoint with no published price is refused rather than served free; connect your own key.
 
+## Retrying a call without paying twice
+
+If a call times out or you never see its answer, repeat it with the same `idempotency_key` (over MCP)
+or `Idempotency-Key` header (over HTTP). treg returns the stored answer, does not call the provider
+again, and charges nothing. The result says `replayed: true`.
+
+Only for a genuine retry. Asking the same question again to see what changed is NEW work: use a new
+key or none, or you will get the old answer back. Reusing one key for a different request is refused.
+
+Most retries need none of this — a failed call was never billed.
+
 ## Task — your own tools: call one the team registered
 **You already know the upstream API. Just build the real request and prefix it.** No treg
 vocabulary, no special params — use the API exactly as its own docs say:

src/treg/mcp.py

@@ -137,6 +137,7 @@ class CatalogGetOut(TypedDict, total=False):
 class CallOut(TypedDict, total=False):
     status: int | None              # the UPSTREAM status, relayed
     endpoint_id: str | None
+    replayed: bool | None           # answered from an earlier call with the same idempotency_key
     body: Any                       # the provider's response, verbatim
     cost_usd: float
     whose_error: str                # "treg" or "provider" — who to blame, and whether to retry
@@ -433,13 +434,21 @@ async def catalog_get(endpoint_id: str, ctx: Context) -> CatalogGetOut:
         "'render/v1/services'. treg injects the credential server-side and relays the provider's "
         "response unchanged, so you never hold an API key. Catalog calls on treg's key are metered "
         "from the team's prepaid balance; a team's own tool is never metered. Tell the human the "
-        "price (from catalog_get) before calling anything that costs more than a cent."
+        "price (from catalog_get) before calling anything that costs more than a cent.\n\n"
+        "`idempotency_key`: pass the SAME key when you are repeating a call whose answer you did "
+        "not receive — a timeout, a dropped connection, an error on your side after the request "
+        "went out. treg returns the stored answer, does not call the provider again, and charges "
+        "nothing the second time; the result carries `replayed: true`. Use a NEW key (or none) for "
+        "genuinely new work, even when the parameters are identical: repeating a search to see "
+        "what changed is a new call, not a retry, and reusing the key would hand you the old answer. "
+        "Reusing one key for a DIFFERENT request is refused rather than answered."
     ),
     annotations=_CALLS,
     structured_output=True
 )
 async def call(endpoint_id: str, params: dict | list | None = None,
-               method: str | None = None, ctx: Context = None) -> CallOut:  # type: ignore[assignment]
+               method: str | None = None, idempotency_key: str | None = None,
+               ctx: Context = None) -> CallOut:  # type: ignore[assignment]
     token = _bearer(ctx) if ctx else ""
     if not token:
         return _need_token()
@@ -461,6 +470,12 @@ async def call(endpoint_id: str, params: dict | list | None = None,
     # type. Query strings still need key/value pairs, so a list is only meaningful as a body.
     args = params if params is not None else {}
     async with _api(token) as client:
+        if idempotency_key:
+            # Straight through to the header the API already honours. Deliberately the CALLER's key
+            # and never derived from the request: two identical searches an hour apart are new work,
+            # not a retry, and a server-invented key would hand back the stale answer — a 24-hour
+            # cache wearing an idempotency badge.
+            client.headers["Idempotency-Key"] = idempotency_key[:200]
         # The SAME route the CLI and the proxy use, so the tool ACL, deny rules, both daily caps,
         # the balance reserve and the settle all happen exactly once, in one place.
         if method in ("GET", "HEAD", "DELETE"):
@@ -472,6 +487,10 @@ async def call(endpoint_id: str, params: dict | list | None = None,
             r = await client.request(method, f"/call/{endpoint_id}", json=args)
 
     out: dict[str, Any] = {"status": r.status_code, "endpoint_id": endpoint_id, "body": _body(r)}
+    if r.headers.get("X-Treg-Idempotent-Replay") == "true":
+        out["replayed"] = True
+        out["hint"] = ("this is the stored answer from the earlier call with the same "
+                       "idempotency_key — nothing was charged for it")
     # Set by /call/ on a METERED call only — a team's own key is never charged, and its absence
     # therefore means "not applicable" rather than "free". This header did not exist when the tool
     # first read it: I wrote against a convention I had invented, so `cost_usd` was always null and

src/treg/web/llms.txt

@@ -58,6 +58,29 @@ resolves the tool by host and injects the credential:
   `base64+gzip`); treg decodes it server-side, so the upstream still gets the real bytes. The
   `treg` CLI does this automatically on a WAF 403 - you only need it over raw HTTP.
 
+### Retrying a call: `Idempotency-Key`
+
+Agents retry, and a retry must not be paid for twice. Most already cost nothing: a 5xx, a timeout or
+a network error is never billed, and result pricing settles on what the provider actually charged, so
+a lookup that finds nothing is free.
+
+The case that needs a key is narrower — treg reached the provider, the provider charged us, and the
+answer was lost on the way back to you. Send the SAME key when you repeat that call:
+
+    curl "{BASE}/call/tikhub.tiktok.user.profile?uniqueId=tiktok" \
+      -H "X-Treg-Token: $TREG_TOKEN" -H "Idempotency-Key: 4f9c1b2e"
+
+treg returns the stored answer, does not call the provider again, and charges nothing. The reply
+carries `X-Treg-Idempotent-Replay: true` and `X-Treg-Cost-Micro` set to what the original cost.
+
+Use a NEW key, or none, for genuinely new work — even with identical parameters. Repeating a search
+to see what changed is a new call, not a retry, and reusing the key would hand you yesterday's
+answer. Reusing one key for a DIFFERENT request is refused (422) rather than answered.
+
+Keys are yours alone: scoped to the caller, kept 24 hours. Nothing happens without the header, and a
+call that fails frees its key at once. Over MCP the same thing is the `idempotency_key` argument to
+`call`, and a replayed result carries `replayed: true`.
+
 ## The catalog — tools you do not have a key for
 
 Curated endpoints across ~40 providers, grouped by what they DO: keyword and rank tracking,

src/treg/web/skill.md

@@ -110,6 +110,17 @@ Rules for spending someone's balance:
     you hold, and treg relays rather than rewrites your request.
 - An endpoint with no published price is refused rather than served free; connect your own key.
 
+## Retrying a call without paying twice
+
+If a call times out or you never see its answer, repeat it with the same `idempotency_key` (over MCP)
+or `Idempotency-Key` header (over HTTP). treg returns the stored answer, does not call the provider
+again, and charges nothing. The result says `replayed: true`.
+
+Only for a genuine retry. Asking the same question again to see what changed is NEW work: use a new
+key or none, or you will get the old answer back. Reusing one key for a different request is refused.
+
+Most retries need none of this — a failed call was never billed.
+
 ## Task — your own tools: call one the team registered
 **You already know the upstream API. Just build the real request and prefix it.** No treg
 vocabulary, no special params — use the API exactly as its own docs say:

tests/test_mcp.py

@@ -667,3 +667,83 @@ async def test_a_per_org_token_still_reaches_the_tool(clients):
             "params": {"name": "balance", "arguments": {}}},
             headers={**MCP_HEADERS, "Authorization": "Bearer not-an-oauth-token-shape"})
     assert r.status_code == 200, r.text  # the tool answers (with its own error prose) — not a 401
+
+
+async def test_call_passes_an_idempotency_key_through(clients):
+    """The feature was built for agents and MCP is the agent path, so leaving `call` unable to send a
+    key made it unreachable from the surface it was for.
+
+    The key is the CALLER's, never derived from the request: two identical searches an hour apart are
+    new work, not a retry, and a server-invented key would hand back the stale answer — a 24-hour
+    cache wearing an idempotency badge."""
+    token = (await clients.post("/users", json={"email": "mcpidem@superdesign.dev"})).json()["token"]
+    prev = clients.headers.get("X-Treg-Token")
+    clients.headers["X-Treg-Token"] = token
+    made = await clients.post("/tools", json={"name": "echo", "base_url": "http://upstream"})
+    if prev:
+        clients.headers["X-Treg-Token"] = prev
+    assert made.status_code == 200, made.text
+
+    async with mcp_session(clients) as c:
+        out = await _call_tool(c, "call", {
+            "endpoint_id": "echo/anything", "method": "POST",
+            "params": {"x": 1}, "idempotency_key": "agent-retry-1"}, token=token)
+    assert out.get("status") == 200, out
+
+
+async def test_the_key_is_optional_and_described_for_the_model(clients):
+    """A model can only use it if the description says WHEN. The distinction that matters is retry
+    versus new work, because getting it wrong returns stale data rather than failing loudly."""
+    from treg.mcp import mcp as server
+
+    tool = [t for t in await server.list_tools() if t.name == "call"][0]
+    assert "idempotency_key" in tool.input_schema["properties"]
+    assert "idempotency_key" not in (tool.input_schema.get("required") or [])
+    desc = tool.description or ""
+    assert "repeating a call whose answer you did not receive" in desc
+    assert "new call, not a retry" in desc, "the model must be told when NOT to reuse a key"
+
+
+async def test_the_same_key_through_MCP_bills_once(clients, monkeypatch):
+    """End to end on the agent path: an agent retries with the same key, the provider is reached
+    once, and the balance moves once."""
+    from sqlmodel import select
+
+    from treg.config import get_settings
+    from treg.db import session_maker
+    from treg.models import Org
+
+    monkeypatch.setenv("TREG_PLATFORM_KEY_TIKHUB", "PLATKEY")
+    monkeypatch.setenv("TREG_PLATFORM_PROVIDERS", "tikhub")
+    get_settings.cache_clear()
+
+    token = (await clients.post("/users", json={"email": "mcponce@superdesign.dev"})).json()["token"]
+    prev = clients.headers.get("X-Treg-Token")
+    clients.headers["X-Treg-Token"] = token
+    org_id = (await clients.get("/orgs")).json()[0]["org_id"]
+    if prev:
+        clients.headers["X-Treg-Token"] = prev
+
+    async def balance() -> int:
+        async with session_maker() as db:
+            org = (await db.execute(select(Org).where(Org.id == org_id))).scalar_one()
+            return org.balance_micro or 0
+
+    args = {"endpoint_id": "tikhub.tiktok.video.comments", "params": {"aweme_id": "7"},
+            "idempotency_key": "one-piece-of-work"}
+    before = await balance()
+    async with mcp_session(clients) as c:
+        first = await _call_tool(c, "call", args, token=token)
+    assert first.get("status") == 200, first
+    charged = before - await balance()
+    assert charged > 0, "the first call must bill"
+
+    after_first = await balance()
+    async with mcp_session(clients) as c:
+        second = await _call_tool(c, "call", args, token=token)
+    get_settings.cache_clear()
+
+    assert second.get("status") == 200, second
+    assert second.get("replayed") is True, "the retry must be marked as a replay"
+    assert second.get("body") == first.get("body"), "and return the same answer"
+    assert await balance() == after_first, "and bill nothing"