From 7aad4090b3db7458d69976458b361eea03b8eecf Mon Sep 17 00:00:00 2001 From: kgridou <32600911+kgridou@users.noreply.github.com> Date: Wed, 10 Jun 2026 21:21:59 +0200 Subject: [PATCH] Multi-tenancy step 1: per-tenant thread + approval isolation Install the tenant-context seam and close the two places where per-request state leaked across tenants: graph checkpoint threads and the approval queue. - New agentforge/api/tenancy.py: resolve_tenant (X-Tenant-ID header, validated; falls back to DEFAULT_TENANT unless REQUIRE_TENANT is set), validate_thread_id, and scoped_thread(tenant, thread) -> "tenant:thread". The id charset forbids ':' so neither a tenant nor a thread id can forge another's scope. - Endpoints (chat, chat/stream, approve, approvals) take the tenant via a FastAPI dependency and use the composite key for the checkpointer config and the registry. The client still sees the bare thread_id. - Approval registry re-keyed by composite, gains a tenant_id field; list_pending(tenant) returns only the caller's queue. - Config: default_tenant ("default"), require_tenant (false) -> single-tenant deploys and the console are unchanged. Scope: this isolates conversation/approval state only. The knowledge base is still shared across tenants (next step), and the header is trusted rather than authenticated (a later step). Tests: test_tenancy.py covers resolution, the safe default, require_tenant, id validation, and per-tenant isolation of the approval queue; test_approvals.py updated for the new signatures. Co-Authored-By: Claude Opus 4.8 --- .env.example | 7 ++++ agentforge/api/approvals.py | 28 ++++++++----- agentforge/api/main.py | 37 ++++++++++------- agentforge/api/schemas.py | 1 + agentforge/api/tenancy.py | 52 ++++++++++++++++++++++++ agentforge/config.py | 10 +++++ tests/test_approvals.py | 22 ++++++---- tests/test_tenancy.py | 81 +++++++++++++++++++++++++++++++++++++ 8 files changed, 206 insertions(+), 32 deletions(-) create mode 100644 agentforge/api/tenancy.py create mode 100644 tests/test_tenancy.py diff --git a/.env.example b/.env.example index 983393f..d5864f5 100644 --- a/.env.example +++ b/.env.example @@ -44,6 +44,13 @@ CHUNK_OVERLAP=150 # Guardrails. REDACT_PII=true +# Multi-tenancy — tenant id comes from the X-Tenant-ID header and namespaces +# graph threads + the approval queue. Missing header falls back to +# DEFAULT_TENANT (single-tenant deploys need nothing here). Set REQUIRE_TENANT +# to reject requests without a valid tenant id. (Knowledge base is still shared.) +DEFAULT_TENANT=default +REQUIRE_TENANT=false + # Observability — "langsmith" | "langfuse" | "none". OBSERVABILITY_BACKEND=none LANGCHAIN_API_KEY= diff --git a/agentforge/api/approvals.py b/agentforge/api/approvals.py index 88bb128..fa8d969 100644 --- a/agentforge/api/approvals.py +++ b/agentforge/api/approvals.py @@ -5,6 +5,10 @@ same process that holds its checkpoint. Both become shared/durable together when you switch to ``PostgresSaver`` (see ``agents/graph.py``); until then run a single API replica for HITL. + +Entries are keyed by ``scoped_thread(tenant, thread)`` so two tenants reusing +the same bare ``thread_id`` stay isolated, and ``list_pending(tenant)`` only +ever returns the caller's own queue. """ from __future__ import annotations @@ -14,25 +18,28 @@ from dataclasses import dataclass from agentforge.api.schemas import PendingAction +from agentforge.api.tenancy import scoped_thread from agentforge.observability import metrics @dataclass class PendingApproval: - thread_id: str + tenant_id: str + thread_id: str # bare, client-facing id (the composite is the dict key) action: PendingAction question: str created_at: float -_pending: dict[str, PendingApproval] = {} +_pending: dict[str, PendingApproval] = {} # key: scoped_thread(tenant, thread) _lock = threading.Lock() -def register(thread_id: str, action: PendingAction, question: str) -> None: +def register(tenant_id: str, thread_id: str, action: PendingAction, question: str) -> None: """Record (or refresh) a thread that paused for approval.""" with _lock: - _pending[thread_id] = PendingApproval( + _pending[scoped_thread(tenant_id, thread_id)] = PendingApproval( + tenant_id=tenant_id, thread_id=thread_id, action=action, question=question, @@ -41,14 +48,17 @@ def register(thread_id: str, action: PendingAction, question: str) -> None: metrics.pending_approvals.set(len(_pending)) -def resolve(thread_id: str) -> None: +def resolve(tenant_id: str, thread_id: str) -> None: """Drop a thread once its approval has been decided.""" with _lock: - _pending.pop(thread_id, None) + _pending.pop(scoped_thread(tenant_id, thread_id), None) metrics.pending_approvals.set(len(_pending)) -def list_pending() -> list[PendingApproval]: - """Pending approvals, oldest first.""" +def list_pending(tenant_id: str | None = None) -> list[PendingApproval]: + """Pending approvals, oldest first; scoped to ``tenant_id`` when given.""" with _lock: - return sorted(_pending.values(), key=lambda p: p.created_at) + items = list(_pending.values()) + if tenant_id is not None: + items = [p for p in items if p.tenant_id == tenant_id] + return sorted(items, key=lambda p: p.created_at) diff --git a/agentforge/api/main.py b/agentforge/api/main.py index 2e1e509..aa621ac 100644 --- a/agentforge/api/main.py +++ b/agentforge/api/main.py @@ -23,7 +23,7 @@ from contextlib import asynccontextmanager from typing import Any -from fastapi import FastAPI, Request, Response +from fastapi import Depends, FastAPI, Request, Response from fastapi.middleware.cors import CORSMiddleware from langgraph.types import Command from sse_starlette.sse import EventSourceResponse @@ -39,6 +39,7 @@ PendingAction, PendingApprovalItem, ) +from agentforge.api.tenancy import resolve_tenant, scoped_thread, validate_thread_id from agentforge.config import get_settings from agentforge.observability import get_callbacks, metrics, setup_observability @@ -161,49 +162,55 @@ def documents() -> list[DocumentSummaryItem]: @app.get("/approvals", response_model=list[PendingApprovalItem]) -def approval_queue() -> list[PendingApprovalItem]: +def approval_queue(tenant: str = Depends(resolve_tenant)) -> list[PendingApprovalItem]: return [ PendingApprovalItem( + tenant_id=p.tenant_id, thread_id=p.thread_id, question=p.question, created_at=p.created_at, action=p.action, ) - for p in approvals.list_pending() + for p in approvals.list_pending(tenant) ] @app.post("/chat", response_model=ChatResponse) -def chat(req: ChatRequest) -> ChatResponse: +def chat(req: ChatRequest, tenant: str = Depends(resolve_tenant)) -> ChatResponse: metrics.chat_requests_total.inc() - thread_id = req.thread_id or str(uuid.uuid4()) + thread_id = validate_thread_id(req.thread_id) if req.thread_id else str(uuid.uuid4()) graph = get_compiled_graph() - result = graph.invoke({"question": req.message}, config=_run_config(thread_id)) + result = graph.invoke( + {"question": req.message}, config=_run_config(scoped_thread(tenant, thread_id)) + ) resp = _to_response(thread_id, result) _record_domain_metrics(resp) if resp.approval_required and resp.pending_action: - approvals.register(thread_id, resp.pending_action, req.message) + approvals.register(tenant, thread_id, resp.pending_action, req.message) return resp @app.post("/approve", response_model=ChatResponse) -def approve(req: ApprovalRequest) -> ChatResponse: +def approve(req: ApprovalRequest, tenant: str = Depends(resolve_tenant)) -> ChatResponse: metrics.approvals_total.labels(req.decision).inc() + thread_id = validate_thread_id(req.thread_id) graph = get_compiled_graph() result = graph.invoke( - Command(resume=req.decision), config=_run_config(req.thread_id) + Command(resume=req.decision), config=_run_config(scoped_thread(tenant, thread_id)) ) - approvals.resolve(req.thread_id) - resp = _to_response(req.thread_id, result) + approvals.resolve(tenant, thread_id) + resp = _to_response(thread_id, result) _record_domain_metrics(resp) return resp @app.post("/chat/stream") -async def chat_stream(req: ChatRequest) -> EventSourceResponse: - thread_id = req.thread_id or str(uuid.uuid4()) +async def chat_stream( + req: ChatRequest, tenant: str = Depends(resolve_tenant) +) -> EventSourceResponse: + thread_id = validate_thread_id(req.thread_id) if req.thread_id else str(uuid.uuid4()) graph = get_compiled_graph() - config = _run_config(thread_id) + config = _run_config(scoped_thread(tenant, thread_id)) # Only stream tokens from the answer-producing specialists — never the # supervisor, whose model emits a one-word routing label, not user output. @@ -227,7 +234,7 @@ async def event_generator(): if snapshot.next: # paused at an interrupt action = values.get("proposed_action") or {} if action: - approvals.register(thread_id, PendingAction(**action), req.message) + approvals.register(tenant, thread_id, PendingAction(**action), req.message) yield { "event": "approval_required", "data": json.dumps(action), diff --git a/agentforge/api/schemas.py b/agentforge/api/schemas.py index 0e67878..c4876d9 100644 --- a/agentforge/api/schemas.py +++ b/agentforge/api/schemas.py @@ -35,6 +35,7 @@ class ApprovalRequest(BaseModel): class PendingApprovalItem(BaseModel): + tenant_id: str thread_id: str question: str created_at: float diff --git a/agentforge/api/tenancy.py b/agentforge/api/tenancy.py new file mode 100644 index 0000000..0f14ba9 --- /dev/null +++ b/agentforge/api/tenancy.py @@ -0,0 +1,52 @@ +"""Tenant identity + per-tenant scoping for graph threads and approvals. + +Step 1 of multi-tenancy: install the seam. The tenant is taken from the +``X-Tenant-ID`` header (trusted for now — authentication arrives in a later +step) and used to namespace checkpoint threads and the approval registry, so +one tenant can never see or resume another's runs. + +What this does *not* yet do: isolate the knowledge base. All tenants still +share one corpus at this stage — that is a separate, later step. +""" + +from __future__ import annotations + +import re + +from fastapi import HTTPException, Request + +from agentforge.config import get_settings + +TENANT_HEADER = "X-Tenant-ID" +# Deliberately excludes ':' (the composite-key delimiter) so neither a tenant +# id nor a thread id can forge another tenant's scope. +_VALID_ID = re.compile(r"^[A-Za-z0-9_-]{1,64}$") + + +def resolve_tenant(request: Request) -> str: + """FastAPI dependency returning the validated tenant id for the request. + + Missing header -> ``default_tenant`` (unless ``require_tenant`` is set, in + which case it's a 400). An ill-formed id is always a 400. + """ + settings = get_settings() + raw = request.headers.get(TENANT_HEADER) + if not raw: + if settings.require_tenant: + raise HTTPException(status_code=400, detail=f"Missing {TENANT_HEADER} header") + return settings.default_tenant + if not _VALID_ID.match(raw): + raise HTTPException(status_code=400, detail="Invalid tenant id") + return raw + + +def validate_thread_id(thread_id: str) -> str: + """Reject client thread ids that could break out of their tenant scope.""" + if not _VALID_ID.match(thread_id): + raise HTTPException(status_code=400, detail="Invalid thread_id") + return thread_id + + +def scoped_thread(tenant_id: str, thread_id: str) -> str: + """The internal checkpoint/registry key. Never exposed to the client.""" + return f"{tenant_id}:{thread_id}" diff --git a/agentforge/config.py b/agentforge/config.py index 45d03da..154d9f8 100644 --- a/agentforge/config.py +++ b/agentforge/config.py @@ -70,6 +70,16 @@ class Settings(BaseSettings): # (see examples/). None keeps the built-in banking-compliance default. system_prompt: str | None = Field(default=None) + # --- Multi-tenancy --------------------------------------------------- + # Tenant identity comes from the X-Tenant-ID header and namespaces graph + # threads + the approval queue so tenants can't see/resume each other's + # runs. With require_tenant False (default), a request without the header + # falls back to default_tenant, so single-tenant deploys need no config. + # Set require_tenant True to reject requests that omit a valid tenant id. + # NOTE: the knowledge base is still shared across tenants at this stage. + default_tenant: str = Field(default="default") + require_tenant: bool = Field(default=False) + # --- Guardrails ------------------------------------------------------ redact_pii: bool = Field(default=True) # Tools that always require human approval before execution. diff --git a/tests/test_approvals.py b/tests/test_approvals.py index d813614..0db1241 100644 --- a/tests/test_approvals.py +++ b/tests/test_approvals.py @@ -10,24 +10,30 @@ def test_registry_register_list_resolve(): - approvals.register("t-1", PendingAction(name="file_sar", args={"customer": "C-9"}), "file a SAR") - pending = approvals.list_pending() + approvals.register( + "acme", "t-1", PendingAction(name="file_sar", args={"customer": "C-9"}), "file a SAR" + ) + pending = approvals.list_pending("acme") hit = next(p for p in pending if p.thread_id == "t-1") + assert hit.tenant_id == "acme" assert hit.action.name == "file_sar" assert hit.action.args == {"customer": "C-9"} - approvals.resolve("t-1") - assert all(p.thread_id != "t-1" for p in approvals.list_pending()) + approvals.resolve("acme", "t-1") + assert all(p.thread_id != "t-1" for p in approvals.list_pending("acme")) def test_approvals_endpoint_serializes(): - approvals.register("t-2", PendingAction(name="escalate_case", args={}), "escalate this") + approvals.register("acme", "t-2", PendingAction(name="escalate_case", args={}), "escalate this") client = TestClient(app) - item = next(d for d in client.get("/approvals").json() if d["thread_id"] == "t-2") + items = client.get("/approvals", headers={"X-Tenant-ID": "acme"}).json() + item = next(d for d in items if d["thread_id"] == "t-2") + assert item["tenant_id"] == "acme" assert item["action"]["name"] == "escalate_case" assert item["question"] == "escalate this" assert isinstance(item["created_at"], (int, float)) - approvals.resolve("t-2") - assert all(d["thread_id"] != "t-2" for d in client.get("/approvals").json()) + approvals.resolve("acme", "t-2") + items = client.get("/approvals", headers={"X-Tenant-ID": "acme"}).json() + assert all(d["thread_id"] != "t-2" for d in items) diff --git a/tests/test_tenancy.py b/tests/test_tenancy.py new file mode 100644 index 0000000..2f6075b --- /dev/null +++ b/tests/test_tenancy.py @@ -0,0 +1,81 @@ +"""Tenant resolution + per-tenant isolation of the approval queue. + +These cover the step-1 seam: identity from the X-Tenant-ID header, the safe +default-tenant fallback, id validation that forbids the composite delimiter, +and the guarantee that one tenant never sees another's pending approvals. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException +from fastapi.testclient import TestClient + +import agentforge.api.tenancy as tenancy +from agentforge.api import approvals +from agentforge.api.main import app +from agentforge.api.schemas import PendingAction + + +class _Req: + """Minimal stand-in for a Starlette Request (only .headers.get is used).""" + + def __init__(self, headers: dict | None = None): + self.headers = headers or {} + + +def test_missing_header_falls_back_to_default_tenant(): + from agentforge.config import get_settings + + assert tenancy.resolve_tenant(_Req()) == get_settings().default_tenant + + +def test_explicit_tenant_from_header(): + assert tenancy.resolve_tenant(_Req({"X-Tenant-ID": "acme"})) == "acme" + + +def test_invalid_tenant_id_rejected(): + with pytest.raises(HTTPException) as exc: + tenancy.resolve_tenant(_Req({"X-Tenant-ID": "bad:id"})) + assert exc.value.status_code == 400 + + +def test_require_tenant_rejects_missing_header(monkeypatch): + monkeypatch.setattr( + tenancy, "get_settings", lambda: SimpleNamespace(require_tenant=True, default_tenant="default") + ) + with pytest.raises(HTTPException) as exc: + tenancy.resolve_tenant(_Req()) + assert exc.value.status_code == 400 + + +def test_validate_thread_id_forbids_colon(): + with pytest.raises(HTTPException): + tenancy.validate_thread_id("acme:t-1") # can't smuggle a scope boundary + + +def test_scoped_thread_separates_same_bare_id(): + assert tenancy.scoped_thread("acme", "t-1") != tenancy.scoped_thread("globex", "t-1") + + +def test_approval_queue_is_isolated_per_tenant(): + # Two tenants reuse the same bare thread_id; their actions must not bleed. + approvals.register("acme", "shared", PendingAction(name="file_sar"), "q-acme") + approvals.register("globex", "shared", PendingAction(name="escalate_case"), "q-globex") + client = TestClient(app) + + acme = client.get("/approvals", headers={"X-Tenant-ID": "acme"}).json() + globex = client.get("/approvals", headers={"X-Tenant-ID": "globex"}).json() + + acme_shared = [d for d in acme if d["thread_id"] == "shared"] + globex_shared = [d for d in globex if d["thread_id"] == "shared"] + assert {d["action"]["name"] for d in acme_shared} == {"file_sar"} + assert {d["action"]["name"] for d in globex_shared} == {"escalate_case"} + + # Resolving under the wrong tenant must not drop the other's entry. + approvals.resolve("acme", "shared") + still = client.get("/approvals", headers={"X-Tenant-ID": "globex"}).json() + assert any(d["thread_id"] == "shared" for d in still) + approvals.resolve("globex", "shared")