langgraph-tenancy 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,50 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+
8
+ jobs:
9
+ test:
10
+ runs-on: ubuntu-latest
11
+ strategy:
12
+ fail-fast: false
13
+ matrix:
14
+ python-version: ["3.10", "3.11", "3.12", "3.13"]
15
+ services:
16
+ postgres:
17
+ image: postgres:16
18
+ env:
19
+ POSTGRES_USER: postgres
20
+ POSTGRES_PASSWORD: postgres
21
+ POSTGRES_DB: langgraph_tenancy_test
22
+ ports:
23
+ - 5432:5432
24
+ options: >-
25
+ --health-cmd "pg_isready -U postgres"
26
+ --health-interval 5s
27
+ --health-timeout 5s
28
+ --health-retries 10
29
+ env:
30
+ LG_TENANCY_PG_URI: postgresql://postgres:postgres@localhost:5432/langgraph_tenancy_test
31
+ LG_TENANCY_PG_REQUIRED: "1"
32
+ steps:
33
+ - uses: actions/checkout@v4
34
+ - uses: astral-sh/setup-uv@v5
35
+ with:
36
+ python-version: ${{ matrix.python-version }}
37
+ - name: Install
38
+ run: uv pip install -e ".[test]"
39
+ - name: Test
40
+ run: uv run pytest -v
41
+
42
+ lint:
43
+ runs-on: ubuntu-latest
44
+ steps:
45
+ - uses: actions/checkout@v4
46
+ - uses: astral-sh/setup-uv@v5
47
+ - name: Ruff check
48
+ run: uvx ruff check .
49
+ - name: Ruff format check
50
+ run: uvx ruff format --check .
@@ -0,0 +1,71 @@
1
+ name: Release
2
+
3
+ on:
4
+ release:
5
+ types: [published]
6
+
7
+ jobs:
8
+ test:
9
+ runs-on: ubuntu-latest
10
+ services:
11
+ postgres:
12
+ image: postgres:16
13
+ env:
14
+ POSTGRES_USER: postgres
15
+ POSTGRES_PASSWORD: postgres
16
+ POSTGRES_DB: langgraph_tenancy_test
17
+ ports:
18
+ - 5432:5432
19
+ options: >-
20
+ --health-cmd "pg_isready -U postgres"
21
+ --health-interval 5s
22
+ --health-timeout 5s
23
+ --health-retries 10
24
+ env:
25
+ LG_TENANCY_PG_URI: postgresql://postgres:postgres@localhost:5432/langgraph_tenancy_test
26
+ LG_TENANCY_PG_REQUIRED: "1"
27
+ steps:
28
+ - uses: actions/checkout@v4
29
+ - uses: astral-sh/setup-uv@v5
30
+ with:
31
+ python-version: "3.12"
32
+ - name: Install
33
+ run: uv pip install -e ".[test]"
34
+ - name: Test
35
+ run: uv run pytest -v
36
+
37
+ build:
38
+ needs: test
39
+ runs-on: ubuntu-latest
40
+ steps:
41
+ - uses: actions/checkout@v4
42
+ - uses: astral-sh/setup-uv@v5
43
+ - name: Build distributions
44
+ run: uv build
45
+ - name: Check release tag matches package version
46
+ run: |
47
+ version=$(uvx hatchling version 2>/dev/null || grep -m1 '^version' pyproject.toml | cut -d'"' -f2)
48
+ if [ "v${version}" != "${{ github.event.release.tag_name }}" ]; then
49
+ echo "Release tag ${{ github.event.release.tag_name }} does not match package version v${version}"
50
+ exit 1
51
+ fi
52
+ - uses: actions/upload-artifact@v4
53
+ with:
54
+ name: dist
55
+ path: dist/
56
+
57
+ publish:
58
+ needs: build
59
+ runs-on: ubuntu-latest
60
+ environment:
61
+ name: pypi
62
+ url: https://pypi.org/p/langgraph-tenancy
63
+ permissions:
64
+ id-token: write
65
+ steps:
66
+ - uses: actions/download-artifact@v4
67
+ with:
68
+ name: dist
69
+ path: dist/
70
+ - name: Publish to PyPI (trusted publishing)
71
+ uses: pypa/gh-action-pypi-publish@release/v1
@@ -0,0 +1,21 @@
1
+ # python
2
+ __pycache__/
3
+ *.py[cod]
4
+ *.egg-info/
5
+ dist/
6
+ build/
7
+
8
+ # environments
9
+ .venv/
10
+ venv/
11
+ .env
12
+
13
+ # tooling caches
14
+ .pytest_cache/
15
+ .ruff_cache/
16
+ .mypy_cache/
17
+
18
+ # editors / os
19
+ .idea/
20
+ .vscode/
21
+ .DS_Store
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Abhishek Chauhan
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,143 @@
1
+ Metadata-Version: 2.4
2
+ Name: langgraph-tenancy
3
+ Version: 0.1.0
4
+ Summary: Tenant isolation and per-tenant usage metering for LangGraph checkpointers and stores
5
+ Project-URL: Source, https://github.com/ac12644/langgraph-tenancy
6
+ Project-URL: Issues, https://github.com/ac12644/langgraph-tenancy/issues
7
+ Author-email: Abhishek Chauhan <ac12644@gmail.com>
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Keywords: agents,checkpointer,langgraph,llm,multi-tenant,tenant-isolation
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Programming Language :: Python :: 3.10
15
+ Classifier: Programming Language :: Python :: 3.11
16
+ Classifier: Programming Language :: Python :: 3.12
17
+ Classifier: Programming Language :: Python :: 3.13
18
+ Classifier: Topic :: Security
19
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
20
+ Requires-Python: >=3.10
21
+ Requires-Dist: langchain-core>=0.2.38
22
+ Requires-Dist: langgraph-checkpoint>=2.0.0
23
+ Provides-Extra: test
24
+ Requires-Dist: langgraph-checkpoint-postgres>=2.0.0; extra == 'test'
25
+ Requires-Dist: langgraph>=0.4; extra == 'test'
26
+ Requires-Dist: psycopg[binary]>=3.2; extra == 'test'
27
+ Requires-Dist: pytest>=8; extra == 'test'
28
+ Description-Content-Type: text/markdown
29
+
30
+ # langgraph-tenancy
31
+
32
+ [![CI](https://github.com/ac12644/langgraph-tenancy/actions/workflows/ci.yml/badge.svg)](https://github.com/ac12644/langgraph-tenancy/actions/workflows/ci.yml)
33
+ [![PyPI](https://img.shields.io/pypi/v/langgraph-tenancy)](https://pypi.org/project/langgraph-tenancy/)
34
+ [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
35
+
36
+ **Tenant isolation for LangGraph persistence — as a drop-in wrapper.**
37
+
38
+ LangGraph's own [threat model](https://github.com/langchain-ai/langgraph/blob/main/.github/THREAT_MODEL.md) says it plainly:
39
+
40
+ > Checkpoint savers index by `thread_id`. Without application-level auth, any
41
+ > caller with a valid thread_id can access that thread's state. [...] Users
42
+ > embedding LangGraph directly must implement their own access controls.
43
+
44
+ If you run a multi-tenant product on open-source LangGraph, the only thing
45
+ between Customer A's agent state and Customer B's is a query filter in your
46
+ application code. This package replaces that convention with enforcement.
47
+
48
+ ## Install
49
+
50
+ ```bash
51
+ pip install langgraph-tenancy
52
+ ```
53
+
54
+ ## Usage
55
+
56
+ Wrap your existing checkpointer and store. Nothing else changes.
57
+
58
+ ```python
59
+ from langgraph_tenancy import (
60
+ TenantScopedCheckpointer,
61
+ TenantScopedStore,
62
+ InMemoryUsageLedger,
63
+ )
64
+
65
+ ledger = InMemoryUsageLedger()
66
+ checkpointer = TenantScopedCheckpointer(PostgresSaver(...), usage_ledger=ledger)
67
+ store = TenantScopedStore(InMemoryStore())
68
+
69
+ graph = builder.compile(checkpointer=checkpointer, store=store)
70
+
71
+ # tenant_id is now REQUIRED on every invocation
72
+ graph.invoke(
73
+ {"messages": ["hello"]},
74
+ config={"configurable": {"thread_id": "t1", "tenant_id": "acme"}},
75
+ )
76
+
77
+ # free per-tenant token metering, extracted from checkpointed messages
78
+ ledger.totals("acme") # TenantUsage(input_tokens=..., output_tokens=..., by_model={...})
79
+ ```
80
+
81
+ ## What it enforces
82
+
83
+ | Raw LangGraph behavior | With `langgraph-tenancy` |
84
+ |---|---|
85
+ | Any caller with a `thread_id` reads that thread | Threads are physically keyed `tenant::thread`; wrong-thread_id bugs cannot cross tenants |
86
+ | Missing filter → silent unscoped query | Missing `tenant_id` → `TenantRequiredError`, nothing read or written |
87
+ | `checkpointer.list(None)` enumerates **every** tenant's threads | Refused with `UnscopedAccessError` |
88
+ | Store namespaces are convention; any node can read any namespace | Every op is rooted at the tenant segment, resolved from the run config automatically |
89
+ | `delete_thread("t1")` deletes whoever owns `t1` | Requires an explicit `for_tenant("acme").delete_thread("t1")` handle |
90
+ | `usage_metadata` buried in checkpoint blobs, unqueryable | Aggregated per tenant (and per model), deduped by message id |
91
+
92
+ ## No magic
93
+
94
+ The entire mechanism is key prefixing plus mandatory-context checks, in two
95
+ small files you can audit in ten minutes:
96
+
97
+ - thread ids become `"{tenant_id}::{thread_id}"` before reaching your
98
+ database; the prefix is stripped from everything returned.
99
+ - store namespaces `("memories",)` become `("{tenant_id}", "memories")`.
100
+ - tenant ids containing the separator are rejected, so `acme` can never craft
101
+ a key that collides with another tenant's space.
102
+
103
+ It composes with any `BaseCheckpointSaver` / `BaseStore` implementation —
104
+ Postgres, SQLite, Redis, MongoDB, in-memory — because it never touches
105
+ storage itself.
106
+
107
+ ## What it is not
108
+
109
+ - Not authentication. You decide which tenant a request belongs to; this
110
+ package guarantees that decision is enforced everywhere downstream.
111
+ - Not encryption. Combine with `EncryptedSerializer` for at-rest encryption.
112
+ - Not a replacement for database-level controls in high-assurance setups
113
+ (RLS, schema-per-tenant) — it's the layer that makes your *application*
114
+ unable to leak, whatever the database allows.
115
+
116
+ ## Tested
117
+
118
+ The adversarial test suite — every test attempts a cross-tenant access the
119
+ raw LangGraph API allows — runs against `InMemorySaver` **and** a real
120
+ `PostgresSaver` in CI. The isolation guarantees are proven on actual SQL
121
+ storage, not just the in-memory reference.
122
+
123
+ ## Development
124
+
125
+ ```bash
126
+ uv venv && uv pip install -e ".[test]"
127
+ uv run pytest # postgres tests skip if no server is reachable
128
+
129
+ # to run the postgres leg locally:
130
+ export LG_TENANCY_PG_URI=postgresql://user@localhost:5432/langgraph_tenancy_test
131
+ uv run pytest
132
+ ```
133
+
134
+ ## Status
135
+
136
+ Early (0.1.x). Covered today: sync + async checkpointer paths, sync store
137
+ paths, in-memory and Postgres backends. Not yet covered: subgraph
138
+ `checkpoint_ns` edge cases, `AsyncPostgresSaver`, `PostgresStore`, store TTL
139
+ ops. Issues and PRs welcome.
140
+
141
+ ## License
142
+
143
+ [MIT](LICENSE)
@@ -0,0 +1,114 @@
1
+ # langgraph-tenancy
2
+
3
+ [![CI](https://github.com/ac12644/langgraph-tenancy/actions/workflows/ci.yml/badge.svg)](https://github.com/ac12644/langgraph-tenancy/actions/workflows/ci.yml)
4
+ [![PyPI](https://img.shields.io/pypi/v/langgraph-tenancy)](https://pypi.org/project/langgraph-tenancy/)
5
+ [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
6
+
7
+ **Tenant isolation for LangGraph persistence — as a drop-in wrapper.**
8
+
9
+ LangGraph's own [threat model](https://github.com/langchain-ai/langgraph/blob/main/.github/THREAT_MODEL.md) says it plainly:
10
+
11
+ > Checkpoint savers index by `thread_id`. Without application-level auth, any
12
+ > caller with a valid thread_id can access that thread's state. [...] Users
13
+ > embedding LangGraph directly must implement their own access controls.
14
+
15
+ If you run a multi-tenant product on open-source LangGraph, the only thing
16
+ between Customer A's agent state and Customer B's is a query filter in your
17
+ application code. This package replaces that convention with enforcement.
18
+
19
+ ## Install
20
+
21
+ ```bash
22
+ pip install langgraph-tenancy
23
+ ```
24
+
25
+ ## Usage
26
+
27
+ Wrap your existing checkpointer and store. Nothing else changes.
28
+
29
+ ```python
30
+ from langgraph_tenancy import (
31
+ TenantScopedCheckpointer,
32
+ TenantScopedStore,
33
+ InMemoryUsageLedger,
34
+ )
35
+
36
+ ledger = InMemoryUsageLedger()
37
+ checkpointer = TenantScopedCheckpointer(PostgresSaver(...), usage_ledger=ledger)
38
+ store = TenantScopedStore(InMemoryStore())
39
+
40
+ graph = builder.compile(checkpointer=checkpointer, store=store)
41
+
42
+ # tenant_id is now REQUIRED on every invocation
43
+ graph.invoke(
44
+ {"messages": ["hello"]},
45
+ config={"configurable": {"thread_id": "t1", "tenant_id": "acme"}},
46
+ )
47
+
48
+ # free per-tenant token metering, extracted from checkpointed messages
49
+ ledger.totals("acme") # TenantUsage(input_tokens=..., output_tokens=..., by_model={...})
50
+ ```
51
+
52
+ ## What it enforces
53
+
54
+ | Raw LangGraph behavior | With `langgraph-tenancy` |
55
+ |---|---|
56
+ | Any caller with a `thread_id` reads that thread | Threads are physically keyed `tenant::thread`; wrong-thread_id bugs cannot cross tenants |
57
+ | Missing filter → silent unscoped query | Missing `tenant_id` → `TenantRequiredError`, nothing read or written |
58
+ | `checkpointer.list(None)` enumerates **every** tenant's threads | Refused with `UnscopedAccessError` |
59
+ | Store namespaces are convention; any node can read any namespace | Every op is rooted at the tenant segment, resolved from the run config automatically |
60
+ | `delete_thread("t1")` deletes whoever owns `t1` | Requires an explicit `for_tenant("acme").delete_thread("t1")` handle |
61
+ | `usage_metadata` buried in checkpoint blobs, unqueryable | Aggregated per tenant (and per model), deduped by message id |
62
+
63
+ ## No magic
64
+
65
+ The entire mechanism is key prefixing plus mandatory-context checks, in two
66
+ small files you can audit in ten minutes:
67
+
68
+ - thread ids become `"{tenant_id}::{thread_id}"` before reaching your
69
+ database; the prefix is stripped from everything returned.
70
+ - store namespaces `("memories",)` become `("{tenant_id}", "memories")`.
71
+ - tenant ids containing the separator are rejected, so `acme` can never craft
72
+ a key that collides with another tenant's space.
73
+
74
+ It composes with any `BaseCheckpointSaver` / `BaseStore` implementation —
75
+ Postgres, SQLite, Redis, MongoDB, in-memory — because it never touches
76
+ storage itself.
77
+
78
+ ## What it is not
79
+
80
+ - Not authentication. You decide which tenant a request belongs to; this
81
+ package guarantees that decision is enforced everywhere downstream.
82
+ - Not encryption. Combine with `EncryptedSerializer` for at-rest encryption.
83
+ - Not a replacement for database-level controls in high-assurance setups
84
+ (RLS, schema-per-tenant) — it's the layer that makes your *application*
85
+ unable to leak, whatever the database allows.
86
+
87
+ ## Tested
88
+
89
+ The adversarial test suite — every test attempts a cross-tenant access the
90
+ raw LangGraph API allows — runs against `InMemorySaver` **and** a real
91
+ `PostgresSaver` in CI. The isolation guarantees are proven on actual SQL
92
+ storage, not just the in-memory reference.
93
+
94
+ ## Development
95
+
96
+ ```bash
97
+ uv venv && uv pip install -e ".[test]"
98
+ uv run pytest # postgres tests skip if no server is reachable
99
+
100
+ # to run the postgres leg locally:
101
+ export LG_TENANCY_PG_URI=postgresql://user@localhost:5432/langgraph_tenancy_test
102
+ uv run pytest
103
+ ```
104
+
105
+ ## Status
106
+
107
+ Early (0.1.x). Covered today: sync + async checkpointer paths, sync store
108
+ paths, in-memory and Postgres backends. Not yet covered: subgraph
109
+ `checkpoint_ns` edge cases, `AsyncPostgresSaver`, `PostgresStore`, store TTL
110
+ ops. Issues and PRs welcome.
111
+
112
+ ## License
113
+
114
+ [MIT](LICENSE)
@@ -0,0 +1,65 @@
1
+ """Tenant isolation for LangGraph persistence.
2
+
3
+ LangGraph's own threat model: "Without application-level auth, any caller
4
+ with a valid thread_id can access that thread's state." This package is that
5
+ application-level wall, as a drop-in wrapper:
6
+
7
+ ```python
8
+ from langgraph_tenancy import (
9
+ TenantScopedCheckpointer, TenantScopedStore, InMemoryUsageLedger,
10
+ )
11
+
12
+ ledger = InMemoryUsageLedger()
13
+ checkpointer = TenantScopedCheckpointer(PostgresSaver(...), usage_ledger=ledger)
14
+ store = TenantScopedStore(InMemoryStore())
15
+
16
+ graph = builder.compile(checkpointer=checkpointer, store=store)
17
+
18
+ graph.invoke(
19
+ {"messages": [...]},
20
+ config={"configurable": {"thread_id": "t1", "tenant_id": "acme"}},
21
+ )
22
+
23
+ ledger.totals("acme") # per-tenant token usage, for free
24
+ ```
25
+
26
+ Guarantees:
27
+ - No tenant_id in config -> the call raises; nothing is read or written.
28
+ - Thread ids and store namespaces are physically tenant-prefixed in storage.
29
+ - `list(None)` (enumerate all tenants' threads) is refused.
30
+ - Maintenance ops (delete/copy/prune) require an explicit `for_tenant()` handle.
31
+ """
32
+
33
+ from langgraph_tenancy._checkpointer import (
34
+ TenantCheckpointerHandle,
35
+ TenantScopedCheckpointer,
36
+ )
37
+ from langgraph_tenancy._errors import (
38
+ InvalidTenantError,
39
+ TenancyError,
40
+ TenantRequiredError,
41
+ UnscopedAccessError,
42
+ )
43
+ from langgraph_tenancy._store import TenantScopedStore
44
+ from langgraph_tenancy._usage import (
45
+ InMemoryUsageLedger,
46
+ TenantUsage,
47
+ UsageLedger,
48
+ UsageRecord,
49
+ extract_usage,
50
+ )
51
+
52
+ __all__ = [
53
+ "InMemoryUsageLedger",
54
+ "InvalidTenantError",
55
+ "TenancyError",
56
+ "TenantCheckpointerHandle",
57
+ "TenantRequiredError",
58
+ "TenantScopedCheckpointer",
59
+ "TenantScopedStore",
60
+ "TenantUsage",
61
+ "UnscopedAccessError",
62
+ "UsageLedger",
63
+ "UsageRecord",
64
+ "extract_usage",
65
+ ]
@@ -0,0 +1,254 @@
1
+ """Tenant-scoped wrapper around any `BaseCheckpointSaver`.
2
+
3
+ Design:
4
+ - Reads `tenant_id` from `config["configurable"]` on every call. Missing
5
+ tenant -> `TenantRequiredError`. There is no unscoped fallback.
6
+ - Physically prefixes `thread_id` with the tenant (``acme::thread-1``) before
7
+ it reaches the inner saver, and strips the prefix from everything returned.
8
+ A wrong-thread_id bug in app code therefore cannot cross a tenant boundary:
9
+ the key the database sees is always tenant-qualified.
10
+ - Blocks the dangerous raw-API escape hatches: `list(None)` and
11
+ `delete_thread(...)` without a tenant.
12
+ - Optionally records per-tenant token usage from checkpointed messages into a
13
+ `UsageLedger` (see `_usage.py`) — same integration point, free metering.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ from collections.abc import AsyncIterator, Iterator, Sequence
19
+ from typing import Any
20
+
21
+ from langchain_core.runnables import RunnableConfig
22
+ from langgraph.checkpoint.base import (
23
+ BaseCheckpointSaver,
24
+ ChannelVersions,
25
+ Checkpoint,
26
+ CheckpointMetadata,
27
+ CheckpointTuple,
28
+ )
29
+
30
+ from langgraph_tenancy._errors import (
31
+ InvalidTenantError,
32
+ TenantRequiredError,
33
+ UnscopedAccessError,
34
+ )
35
+ from langgraph_tenancy._usage import UsageLedger, extract_usage
36
+
37
+ SEP = "::"
38
+
39
+
40
+ def _validate_tenant(tenant: Any, where: str) -> str:
41
+ if not isinstance(tenant, str) or not tenant:
42
+ raise TenantRequiredError(where)
43
+ if SEP in tenant:
44
+ raise InvalidTenantError(f"tenant_id may not contain '{SEP}': {tenant!r}")
45
+ return tenant
46
+
47
+
48
+ class TenantScopedCheckpointer(BaseCheckpointSaver):
49
+ """Wrap a checkpointer so every operation is scoped to one tenant."""
50
+
51
+ def __init__(
52
+ self,
53
+ inner: BaseCheckpointSaver,
54
+ *,
55
+ usage_ledger: UsageLedger | None = None,
56
+ ) -> None:
57
+ super().__init__(serde=inner.serde)
58
+ self.inner = inner
59
+ self.usage_ledger = usage_ledger
60
+
61
+ # -- scoping helpers ----------------------------------------------------
62
+
63
+ def _scope(self, config: RunnableConfig, where: str) -> tuple[str, RunnableConfig]:
64
+ conf = config.get("configurable") or {}
65
+ tenant = _validate_tenant(conf.get("tenant_id"), where)
66
+ thread_id = conf.get("thread_id")
67
+ scoped_thread = f"{tenant}{SEP}{thread_id}"
68
+ return tenant, {
69
+ **config,
70
+ "configurable": {**conf, "thread_id": scoped_thread},
71
+ }
72
+
73
+ def _unscope_config(
74
+ self, tenant: str, config: RunnableConfig | None
75
+ ) -> RunnableConfig | None:
76
+ if config is None:
77
+ return None
78
+ conf = config.get("configurable") or {}
79
+ thread_id = conf.get("thread_id")
80
+ prefix = f"{tenant}{SEP}"
81
+ if isinstance(thread_id, str) and thread_id.startswith(prefix):
82
+ conf = {**conf, "thread_id": thread_id[len(prefix) :], "tenant_id": tenant}
83
+ return {**config, "configurable": conf}
84
+
85
+ def _unscope_tuple(
86
+ self, tenant: str, tup: CheckpointTuple | None
87
+ ) -> CheckpointTuple | None:
88
+ if tup is None:
89
+ return None
90
+ return CheckpointTuple(
91
+ config=self._unscope_config(tenant, tup.config),
92
+ checkpoint=tup.checkpoint,
93
+ metadata=tup.metadata,
94
+ parent_config=self._unscope_config(tenant, tup.parent_config),
95
+ pending_writes=tup.pending_writes,
96
+ )
97
+
98
+ # -- core protocol (sync) -----------------------------------------------
99
+
100
+ @property
101
+ def config_specs(self) -> list:
102
+ return self.inner.config_specs
103
+
104
+ def get_next_version(self, current: Any, channel: None = None) -> Any:
105
+ return self.inner.get_next_version(current, channel)
106
+
107
+ def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
108
+ tenant, scoped = self._scope(config, "get_tuple()")
109
+ return self._unscope_tuple(tenant, self.inner.get_tuple(scoped))
110
+
111
+ def list(
112
+ self,
113
+ config: RunnableConfig | None,
114
+ *,
115
+ filter: dict[str, Any] | None = None,
116
+ before: RunnableConfig | None = None,
117
+ limit: int | None = None,
118
+ ) -> Iterator[CheckpointTuple]:
119
+ if config is None:
120
+ raise UnscopedAccessError(
121
+ "list(None) would enumerate every tenant's threads; "
122
+ "pass a config with tenant_id (and optionally thread_id)."
123
+ )
124
+ tenant, scoped = self._scope(config, "list()")
125
+ scoped_before = self._scope(before, "list(before=...)")[1] if before else None
126
+ for tup in self.inner.list(
127
+ scoped, filter=filter, before=scoped_before, limit=limit
128
+ ):
129
+ yield self._unscope_tuple(tenant, tup)
130
+
131
+ def put(
132
+ self,
133
+ config: RunnableConfig,
134
+ checkpoint: Checkpoint,
135
+ metadata: CheckpointMetadata,
136
+ new_versions: ChannelVersions,
137
+ ) -> RunnableConfig:
138
+ tenant, scoped = self._scope(config, "put()")
139
+ if self.usage_ledger is not None:
140
+ for record in extract_usage(checkpoint):
141
+ self.usage_ledger.record(tenant, record)
142
+ result = self.inner.put(scoped, checkpoint, metadata, new_versions)
143
+ return self._unscope_config(tenant, result)
144
+
145
+ def put_writes(
146
+ self,
147
+ config: RunnableConfig,
148
+ writes: Sequence[tuple[str, Any]],
149
+ task_id: str,
150
+ task_path: str = "",
151
+ ) -> None:
152
+ _, scoped = self._scope(config, "put_writes()")
153
+ self.inner.put_writes(scoped, writes, task_id, task_path)
154
+
155
+ # -- blocked / redirected escape hatches ----------------------------------
156
+
157
+ def delete_thread(self, thread_id: str) -> None:
158
+ raise UnscopedAccessError(
159
+ "delete_thread() has no tenant context; "
160
+ "use for_tenant(tenant_id).delete_thread(thread_id)."
161
+ )
162
+
163
+ def delete_for_runs(self, run_ids: Sequence[str]) -> None:
164
+ raise UnscopedAccessError(
165
+ "delete_for_runs() cannot be tenant-scoped (run ids are global); "
166
+ "call it on the inner saver explicitly if you accept that."
167
+ )
168
+
169
+ def for_tenant(self, tenant_id: str) -> TenantCheckpointerHandle:
170
+ """Admin/maintenance handle pinned to one tenant."""
171
+ return TenantCheckpointerHandle(
172
+ self, _validate_tenant(tenant_id, "for_tenant()")
173
+ )
174
+
175
+ # -- core protocol (async) ------------------------------------------------
176
+
177
+ async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
178
+ tenant, scoped = self._scope(config, "aget_tuple()")
179
+ return self._unscope_tuple(tenant, await self.inner.aget_tuple(scoped))
180
+
181
+ async def alist(
182
+ self,
183
+ config: RunnableConfig | None,
184
+ *,
185
+ filter: dict[str, Any] | None = None,
186
+ before: RunnableConfig | None = None,
187
+ limit: int | None = None,
188
+ ) -> AsyncIterator[CheckpointTuple]:
189
+ if config is None:
190
+ raise UnscopedAccessError(
191
+ "alist(None) would enumerate every tenant's threads; "
192
+ "pass a config with tenant_id (and optionally thread_id)."
193
+ )
194
+ tenant, scoped = self._scope(config, "alist()")
195
+ scoped_before = self._scope(before, "alist(before=...)")[1] if before else None
196
+ async for tup in self.inner.alist(
197
+ scoped, filter=filter, before=scoped_before, limit=limit
198
+ ):
199
+ yield self._unscope_tuple(tenant, tup)
200
+
201
+ async def aput(
202
+ self,
203
+ config: RunnableConfig,
204
+ checkpoint: Checkpoint,
205
+ metadata: CheckpointMetadata,
206
+ new_versions: ChannelVersions,
207
+ ) -> RunnableConfig:
208
+ tenant, scoped = self._scope(config, "aput()")
209
+ if self.usage_ledger is not None:
210
+ for record in extract_usage(checkpoint):
211
+ self.usage_ledger.record(tenant, record)
212
+ result = await self.inner.aput(scoped, checkpoint, metadata, new_versions)
213
+ return self._unscope_config(tenant, result)
214
+
215
+ async def aput_writes(
216
+ self,
217
+ config: RunnableConfig,
218
+ writes: Sequence[tuple[str, Any]],
219
+ task_id: str,
220
+ task_path: str = "",
221
+ ) -> None:
222
+ _, scoped = self._scope(config, "aput_writes()")
223
+ await self.inner.aput_writes(scoped, writes, task_id, task_path)
224
+
225
+ async def adelete_thread(self, thread_id: str) -> None:
226
+ self.delete_thread(thread_id)
227
+
228
+
229
+ class TenantCheckpointerHandle:
230
+ """Maintenance operations pre-bound to a single tenant.
231
+
232
+ Exists because `BaseCheckpointSaver.delete_thread/copy_thread/prune` take
233
+ bare thread ids with no config, so there is no per-call tenant to read.
234
+ """
235
+
236
+ def __init__(self, parent: TenantScopedCheckpointer, tenant: str) -> None:
237
+ self._inner = parent.inner
238
+ self._tenant = tenant
239
+
240
+ def _scoped(self, thread_id: str) -> str:
241
+ return f"{self._tenant}{SEP}{thread_id}"
242
+
243
+ def delete_thread(self, thread_id: str) -> None:
244
+ self._inner.delete_thread(self._scoped(thread_id))
245
+
246
+ def copy_thread(self, source_thread_id: str, target_thread_id: str) -> None:
247
+ self._inner.copy_thread(
248
+ self._scoped(source_thread_id), self._scoped(target_thread_id)
249
+ )
250
+
251
+ def prune(
252
+ self, thread_ids: Sequence[str], *, strategy: str = "keep_latest"
253
+ ) -> None:
254
+ self._inner.prune([self._scoped(t) for t in thread_ids], strategy=strategy)
@@ -0,0 +1,39 @@
1
+ """Errors raised by langgraph-tenancy.
2
+
3
+ Every error here exists to turn a silent cross-tenant leak into a loud failure.
4
+ """
5
+
6
+
7
+ class TenancyError(Exception):
8
+ """Base class for all tenancy errors."""
9
+
10
+
11
+ class TenantRequiredError(TenancyError):
12
+ """Raised when an operation runs without a tenant_id in scope.
13
+
14
+ The wrapper never falls back to an unscoped read or write: no tenant, no data.
15
+ """
16
+
17
+ def __init__(self, where: str) -> None:
18
+ super().__init__(
19
+ f"{where} requires a tenant. Pass it in the run config: "
20
+ "config={'configurable': {'thread_id': ..., 'tenant_id': ...}} "
21
+ "or use .for_tenant(tenant_id) for out-of-band access."
22
+ )
23
+
24
+
25
+ class UnscopedAccessError(TenancyError):
26
+ """Raised for operations that would touch data across tenant boundaries.
27
+
28
+ Example: `checkpointer.list(None)` on a raw saver enumerates every
29
+ customer's threads. This wrapper refuses that call instead.
30
+ """
31
+
32
+
33
+ class InvalidTenantError(TenancyError):
34
+ """Raised when a tenant_id could be used to escape its scope.
35
+
36
+ The tenant id becomes part of storage keys, so it must not contain the
37
+ separator (or be empty) — otherwise tenant "a" could craft ids that
38
+ collide with tenant "a::b".
39
+ """
@@ -0,0 +1,124 @@
1
+ """Tenant-scoped wrapper around any `BaseStore`.
2
+
3
+ Raw `BaseStore` namespaces are pure convention — any caller can `get()`,
4
+ `search()`, or `list_namespaces()` across all of them. This wrapper prepends
5
+ the tenant id as the root namespace segment on every operation and strips it
6
+ from every result, so two tenants using the identical namespace tuple
7
+ (e.g. ``("memories",)``) land in physically distinct locations.
8
+
9
+ Tenant resolution, in order:
10
+ 1. A pinned tenant from `.for_tenant(tenant_id)` (out-of-band/admin access).
11
+ 2. The ambient run config via `langgraph.config.get_config()` — inside a node,
12
+ the `tenant_id` you passed to `graph.invoke(...)` is picked up automatically.
13
+ 3. Otherwise: `TenantRequiredError`. Never an unscoped operation.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ from collections.abc import Iterable
19
+ from typing import Any
20
+
21
+ from langgraph.store.base import (
22
+ BaseStore,
23
+ GetOp,
24
+ ListNamespacesOp,
25
+ MatchCondition,
26
+ Op,
27
+ PutOp,
28
+ Result,
29
+ SearchOp,
30
+ )
31
+
32
+ from langgraph_tenancy._errors import InvalidTenantError, TenantRequiredError
33
+
34
+ SEP = "::"
35
+
36
+
37
+ class TenantScopedStore(BaseStore):
38
+ def __init__(self, inner: BaseStore, *, _tenant: str | None = None) -> None:
39
+ self.inner = inner
40
+ self._tenant = _tenant
41
+
42
+ def for_tenant(self, tenant_id: str) -> TenantScopedStore:
43
+ """A view of the store pinned to one tenant, for use outside a run."""
44
+ self._validate(tenant_id)
45
+ return TenantScopedStore(self.inner, _tenant=tenant_id)
46
+
47
+ # -- tenant resolution ----------------------------------------------------
48
+
49
+ @staticmethod
50
+ def _validate(tenant: Any) -> str:
51
+ if not isinstance(tenant, str) or not tenant:
52
+ raise TenantRequiredError("store access")
53
+ if SEP in tenant:
54
+ raise InvalidTenantError(f"tenant_id may not contain '{SEP}': {tenant!r}")
55
+ return tenant
56
+
57
+ def _current_tenant(self) -> str:
58
+ if self._tenant is not None:
59
+ return self._tenant
60
+ try:
61
+ from langgraph.config import get_config
62
+
63
+ config = get_config()
64
+ except Exception:
65
+ config = None
66
+ tenant = ((config or {}).get("configurable") or {}).get("tenant_id")
67
+ return self._validate(tenant)
68
+
69
+ # -- op rewriting -----------------------------------------------------------
70
+
71
+ def _scope_op(self, tenant: str, op: Op) -> Op:
72
+ if isinstance(op, (GetOp, PutOp)):
73
+ return op._replace(namespace=(tenant, *op.namespace))
74
+ if isinstance(op, SearchOp):
75
+ return op._replace(namespace_prefix=(tenant, *op.namespace_prefix))
76
+ if isinstance(op, ListNamespacesOp):
77
+ conditions = list(op.match_conditions or ())
78
+ for i, cond in enumerate(conditions):
79
+ if cond.match_type == "prefix":
80
+ conditions[i] = MatchCondition(
81
+ match_type="prefix", path=(tenant, *cond.path)
82
+ )
83
+ break
84
+ else:
85
+ conditions.insert(
86
+ 0, MatchCondition(match_type="prefix", path=(tenant,))
87
+ )
88
+ return op._replace(
89
+ match_conditions=tuple(conditions),
90
+ max_depth=None if op.max_depth is None else op.max_depth + 1,
91
+ )
92
+ raise TypeError(f"unsupported op: {op!r}")
93
+
94
+ def _unscope_result(self, tenant: str, op: Op, result: Result) -> Result:
95
+ if isinstance(op, GetOp) and result is not None:
96
+ result.namespace = tuple(result.namespace[1:])
97
+ return result
98
+ if isinstance(op, SearchOp):
99
+ for item in result:
100
+ item.namespace = tuple(item.namespace[1:])
101
+ return result
102
+ if isinstance(op, ListNamespacesOp):
103
+ return [tuple(ns[1:]) for ns in result if ns and ns[0] == tenant]
104
+ return result
105
+
106
+ # -- BaseStore protocol -------------------------------------------------------
107
+
108
+ def batch(self, ops: Iterable[Op]) -> list[Result]:
109
+ tenant = self._current_tenant()
110
+ ops = list(ops)
111
+ results = self.inner.batch([self._scope_op(tenant, op) for op in ops])
112
+ return [
113
+ self._unscope_result(tenant, op, result)
114
+ for op, result in zip(ops, results, strict=True)
115
+ ]
116
+
117
+ async def abatch(self, ops: Iterable[Op]) -> list[Result]:
118
+ tenant = self._current_tenant()
119
+ ops = list(ops)
120
+ results = await self.inner.abatch([self._scope_op(tenant, op) for op in ops])
121
+ return [
122
+ self._unscope_result(tenant, op, result)
123
+ for op, result in zip(ops, results, strict=True)
124
+ ]
@@ -0,0 +1,95 @@
1
+ """Per-tenant token usage, extracted at the checkpoint boundary.
2
+
3
+ LangGraph already persists `usage_metadata` on every AI message inside
4
+ `checkpoint["channel_values"]` — it is just never indexed or aggregated.
5
+ Since the tenant-scoped checkpointer sees every checkpoint anyway, it can
6
+ pull usage out and attribute it to the tenant for free.
7
+
8
+ Messages are deduplicated by message id, so re-checkpointing the same
9
+ conversation does not double-count.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ from collections import defaultdict
15
+ from dataclasses import dataclass, field
16
+ from threading import Lock
17
+ from typing import Any, Protocol
18
+
19
+
20
+ @dataclass(frozen=True)
21
+ class UsageRecord:
22
+ message_id: str
23
+ model: str | None
24
+ input_tokens: int
25
+ output_tokens: int
26
+ total_tokens: int
27
+
28
+
29
+ class UsageLedger(Protocol):
30
+ """Anything that can receive per-tenant usage records.
31
+
32
+ Swap the in-memory default for one backed by your own Postgres table,
33
+ StatsD, OpenMeter, etc.
34
+ """
35
+
36
+ def record(self, tenant_id: str, record: UsageRecord) -> None: ...
37
+
38
+
39
+ @dataclass
40
+ class TenantUsage:
41
+ input_tokens: int = 0
42
+ output_tokens: int = 0
43
+ total_tokens: int = 0
44
+ messages: int = 0
45
+ by_model: dict[str, int] = field(default_factory=lambda: defaultdict(int))
46
+
47
+
48
+ class InMemoryUsageLedger:
49
+ """Reference ledger: per-tenant totals, deduped by message id."""
50
+
51
+ def __init__(self) -> None:
52
+ self._seen: set[tuple[str, str]] = set()
53
+ self._totals: dict[str, TenantUsage] = defaultdict(TenantUsage)
54
+ self._lock = Lock()
55
+
56
+ def record(self, tenant_id: str, record: UsageRecord) -> None:
57
+ with self._lock:
58
+ key = (tenant_id, record.message_id)
59
+ if key in self._seen:
60
+ return
61
+ self._seen.add(key)
62
+ usage = self._totals[tenant_id]
63
+ usage.input_tokens += record.input_tokens
64
+ usage.output_tokens += record.output_tokens
65
+ usage.total_tokens += record.total_tokens
66
+ usage.messages += 1
67
+ if record.model:
68
+ usage.by_model[record.model] += record.total_tokens
69
+
70
+ def totals(self, tenant_id: str) -> TenantUsage:
71
+ with self._lock:
72
+ return self._totals[tenant_id]
73
+
74
+
75
+ def extract_usage(checkpoint: dict[str, Any]) -> list[UsageRecord]:
76
+ """Pull usage records out of message objects in a checkpoint's channels."""
77
+ records: list[UsageRecord] = []
78
+ for value in (checkpoint.get("channel_values") or {}).values():
79
+ items = value if isinstance(value, (list, tuple)) else [value]
80
+ for item in items:
81
+ usage = getattr(item, "usage_metadata", None)
82
+ message_id = getattr(item, "id", None)
83
+ if not usage or not message_id:
84
+ continue
85
+ model = (getattr(item, "response_metadata", None) or {}).get("model_name")
86
+ records.append(
87
+ UsageRecord(
88
+ message_id=message_id,
89
+ model=model,
90
+ input_tokens=usage.get("input_tokens", 0),
91
+ output_tokens=usage.get("output_tokens", 0),
92
+ total_tokens=usage.get("total_tokens", 0),
93
+ )
94
+ )
95
+ return records
@@ -0,0 +1,61 @@
1
+ [project]
2
+ name = "langgraph-tenancy"
3
+ version = "0.1.0"
4
+ description = "Tenant isolation and per-tenant usage metering for LangGraph checkpointers and stores"
5
+ readme = "README.md"
6
+ requires-python = ">=3.10"
7
+ license = "MIT"
8
+ license-files = ["LICENSE"]
9
+ authors = [{ name = "Abhishek Chauhan", email = "ac12644@gmail.com" }]
10
+ keywords = [
11
+ "langgraph",
12
+ "multi-tenant",
13
+ "tenant-isolation",
14
+ "checkpointer",
15
+ "agents",
16
+ "llm",
17
+ ]
18
+ classifiers = [
19
+ "Development Status :: 4 - Beta",
20
+ "Intended Audience :: Developers",
21
+ "Programming Language :: Python :: 3",
22
+ "Programming Language :: Python :: 3.10",
23
+ "Programming Language :: Python :: 3.11",
24
+ "Programming Language :: Python :: 3.12",
25
+ "Programming Language :: Python :: 3.13",
26
+ "Topic :: Software Development :: Libraries :: Python Modules",
27
+ "Topic :: Security",
28
+ ]
29
+ dependencies = [
30
+ "langgraph-checkpoint>=2.0.0",
31
+ "langchain-core>=0.2.38",
32
+ ]
33
+
34
+ [project.urls]
35
+ Source = "https://github.com/ac12644/langgraph-tenancy"
36
+ Issues = "https://github.com/ac12644/langgraph-tenancy/issues"
37
+
38
+ [project.optional-dependencies]
39
+ test = [
40
+ "pytest>=8",
41
+ "langgraph>=0.4",
42
+ "langgraph-checkpoint-postgres>=2.0.0",
43
+ "psycopg[binary]>=3.2",
44
+ ]
45
+
46
+ [build-system]
47
+ requires = ["hatchling"]
48
+ build-backend = "hatchling.build"
49
+
50
+ [tool.hatch.build.targets.wheel]
51
+ packages = ["langgraph_tenancy"]
52
+
53
+ [tool.pytest.ini_options]
54
+ testpaths = ["tests"]
55
+
56
+ [tool.ruff]
57
+ line-length = 88
58
+ target-version = "py310"
59
+
60
+ [tool.ruff.lint]
61
+ select = ["E", "F", "I", "UP", "B", "SIM"]
@@ -0,0 +1,53 @@
1
+ """Checkpointer backends for the isolation suite.
2
+
3
+ Every checkpointer test runs against each backend, so the isolation
4
+ guarantees are proven on real storage, not just the in-memory reference.
5
+ Postgres tests skip automatically if no server is reachable.
6
+ """
7
+
8
+ import os
9
+
10
+ import pytest
11
+
12
+ PG_URI = os.environ.get(
13
+ "LG_TENANCY_PG_URI",
14
+ "postgresql://postgres:postgres@localhost:5432/langgraph_tenancy_test",
15
+ )
16
+ # In CI we want a missing database to FAIL the postgres leg, not skip it —
17
+ # otherwise a service misconfiguration silently halves the suite.
18
+ PG_REQUIRED = os.environ.get("LG_TENANCY_PG_REQUIRED") == "1"
19
+
20
+
21
+ @pytest.fixture(params=["memory", "postgres"])
22
+ def make_inner(request):
23
+ """Factory producing fresh inner savers on a clean backend."""
24
+ if request.param == "memory":
25
+ from langgraph.checkpoint.memory import InMemorySaver
26
+
27
+ yield InMemorySaver
28
+ return
29
+
30
+ psycopg = pytest.importorskip("psycopg")
31
+ from langgraph.checkpoint.postgres import PostgresSaver
32
+ from psycopg.rows import dict_row
33
+
34
+ try:
35
+ conn = psycopg.Connection.connect(
36
+ PG_URI, autocommit=True, prepare_threshold=0, row_factory=dict_row
37
+ )
38
+ except Exception as exc: # pragma: no cover - environment dependent
39
+ if PG_REQUIRED:
40
+ raise
41
+ pytest.skip(f"postgres unavailable: {exc}")
42
+
43
+ PostgresSaver(conn).setup()
44
+ # clean slate per test (keep the migrations table)
45
+ tables = conn.execute(
46
+ "select tablename from pg_tables where schemaname='public' "
47
+ "and tablename like 'checkpoint%' and tablename not like '%migrations'"
48
+ ).fetchall()
49
+ for row in tables:
50
+ conn.execute(f'TRUNCATE "{row["tablename"]}"')
51
+
52
+ yield lambda: PostgresSaver(conn)
53
+ conn.close()
@@ -0,0 +1,180 @@
1
+ """Adversarial isolation tests.
2
+
3
+ Every test here attempts a cross-tenant access the raw LangGraph API allows,
4
+ and asserts the wrapper makes it either impossible or an explicit error.
5
+ Run against the real InMemorySaver/InMemoryStore from langgraph.
6
+ """
7
+
8
+ import operator
9
+ from typing import Annotated, TypedDict
10
+ from uuid import uuid4
11
+
12
+ import pytest
13
+ from langchain_core.messages import AIMessage
14
+ from langgraph.checkpoint.memory import InMemorySaver
15
+ from langgraph.config import get_store
16
+ from langgraph.graph import StateGraph
17
+ from langgraph.store.memory import InMemoryStore
18
+
19
+ from langgraph_tenancy import (
20
+ InMemoryUsageLedger,
21
+ InvalidTenantError,
22
+ TenantRequiredError,
23
+ TenantScopedCheckpointer,
24
+ TenantScopedStore,
25
+ UnscopedAccessError,
26
+ )
27
+
28
+
29
+ class State(TypedDict):
30
+ messages: Annotated[list, operator.add]
31
+
32
+
33
+ def make_graph(checkpointer, store=None, model_name="claude-sonnet-4-6"):
34
+ """One-node graph that fakes an LLM reply (with usage) and writes a memory."""
35
+
36
+ def agent(state: State) -> State:
37
+ # real chat models always set an id; the ledger dedupes on it
38
+ reply = AIMessage(
39
+ id=str(uuid4()),
40
+ content=f"echo: {state['messages'][-1]}",
41
+ usage_metadata={
42
+ "input_tokens": 10,
43
+ "output_tokens": 5,
44
+ "total_tokens": 15,
45
+ },
46
+ response_metadata={"model_name": model_name},
47
+ )
48
+ if store is not None:
49
+ get_store().put(
50
+ ("memories",),
51
+ f"note-{len(state['messages'])}",
52
+ {"last": str(state["messages"][-1])},
53
+ )
54
+ return {"messages": [reply]}
55
+
56
+ builder = StateGraph(State)
57
+ builder.add_node("agent", agent)
58
+ builder.set_entry_point("agent")
59
+ builder.set_finish_point("agent")
60
+ return builder.compile(checkpointer=checkpointer, store=store)
61
+
62
+
63
+ def cfg(tenant, thread="t1"):
64
+ return {"configurable": {"thread_id": thread, "tenant_id": tenant}}
65
+
66
+
67
+ # --- checkpointer isolation ---------------------------------------------------
68
+
69
+
70
+ def test_same_thread_id_different_tenants_do_not_collide(make_inner):
71
+ graph = make_graph(TenantScopedCheckpointer(make_inner()))
72
+ graph.invoke({"messages": ["from acme"]}, cfg("acme"))
73
+ graph.invoke({"messages": ["from globex"]}, cfg("globex"))
74
+
75
+ acme = graph.get_state(cfg("acme")).values["messages"]
76
+ globex = graph.get_state(cfg("globex")).values["messages"]
77
+ assert acme[0] == "from acme" and len(acme) == 2
78
+ assert globex[0] == "from globex" and len(globex) == 2
79
+ assert not any("globex" in str(m) for m in acme)
80
+
81
+
82
+ def test_missing_tenant_id_raises_instead_of_leaking(make_inner):
83
+ graph = make_graph(TenantScopedCheckpointer(make_inner()))
84
+ with pytest.raises(TenantRequiredError):
85
+ graph.invoke({"messages": ["hi"]}, {"configurable": {"thread_id": "t1"}})
86
+
87
+
88
+ def test_tenant_id_cannot_contain_separator(make_inner):
89
+ graph = make_graph(TenantScopedCheckpointer(make_inner()))
90
+ with pytest.raises(InvalidTenantError):
91
+ graph.invoke({"messages": ["hi"]}, cfg("acme::evil"))
92
+
93
+
94
+ def test_list_none_is_refused(make_inner):
95
+ saver = TenantScopedCheckpointer(make_inner())
96
+ make_graph(saver).invoke({"messages": ["hi"]}, cfg("acme"))
97
+ with pytest.raises(UnscopedAccessError):
98
+ list(saver.list(None))
99
+
100
+
101
+ def test_list_only_sees_own_tenant(make_inner):
102
+ saver = TenantScopedCheckpointer(make_inner())
103
+ graph = make_graph(saver)
104
+ graph.invoke({"messages": ["a"]}, cfg("acme"))
105
+ graph.invoke({"messages": ["b"]}, cfg("globex"))
106
+
107
+ seen = [t.config["configurable"]["thread_id"] for t in saver.list(cfg("acme"))]
108
+ assert seen and all(t == "t1" for t in seen)
109
+ values = [t for t in saver.list(cfg("acme"))]
110
+ assert not any("globex" in str(t.checkpoint.get("channel_values")) for t in values)
111
+
112
+
113
+ def test_delete_thread_requires_tenant_handle(make_inner):
114
+ saver = TenantScopedCheckpointer(make_inner())
115
+ graph = make_graph(saver)
116
+ graph.invoke({"messages": ["a"]}, cfg("acme"))
117
+ graph.invoke({"messages": ["b"]}, cfg("globex"))
118
+
119
+ with pytest.raises(UnscopedAccessError):
120
+ saver.delete_thread("t1")
121
+
122
+ saver.for_tenant("acme").delete_thread("t1")
123
+ assert saver.get_tuple(cfg("acme")) is None
124
+ assert saver.get_tuple(cfg("globex")) is not None # other tenant untouched
125
+
126
+
127
+ def test_storage_keys_are_physically_tenant_prefixed(make_inner):
128
+ inner = make_inner()
129
+ make_graph(TenantScopedCheckpointer(inner)).invoke({"messages": ["a"]}, cfg("acme"))
130
+ raw_threads = {t.config["configurable"]["thread_id"] for t in inner.list(None)}
131
+ assert raw_threads == {"acme::t1"}
132
+
133
+
134
+ # --- store isolation ----------------------------------------------------------
135
+
136
+
137
+ def test_store_namespaces_are_isolated_per_tenant():
138
+ store = TenantScopedStore(InMemoryStore())
139
+ graph = make_graph(TenantScopedCheckpointer(InMemorySaver()), store=store)
140
+ graph.invoke({"messages": ["secret-acme"]}, cfg("acme"))
141
+ graph.invoke({"messages": ["secret-globex"]}, cfg("globex"))
142
+
143
+ acme_items = store.for_tenant("acme").search(("memories",))
144
+ globex_items = store.for_tenant("globex").search(("memories",))
145
+ assert [i.value["last"] for i in acme_items] == ["secret-acme"]
146
+ assert [i.value["last"] for i in globex_items] == ["secret-globex"]
147
+ # returned namespaces are unprefixed — the tenant segment never leaks out
148
+ assert acme_items[0].namespace == ("memories",)
149
+
150
+
151
+ def test_store_outside_run_without_tenant_raises():
152
+ store = TenantScopedStore(InMemoryStore())
153
+ with pytest.raises(TenantRequiredError):
154
+ store.search(("memories",))
155
+
156
+
157
+ def test_list_namespaces_only_sees_own_tenant():
158
+ store = TenantScopedStore(InMemoryStore())
159
+ store.for_tenant("acme").put(("memories", "work"), "k", {"v": 1})
160
+ store.for_tenant("globex").put(("memories", "home"), "k", {"v": 2})
161
+ assert store.for_tenant("acme").list_namespaces() == [("memories", "work")]
162
+
163
+
164
+ # --- usage metering -------------------------------------------------------------
165
+
166
+
167
+ def test_usage_attributed_per_tenant_and_deduped(make_inner):
168
+ ledger = InMemoryUsageLedger()
169
+ graph = make_graph(TenantScopedCheckpointer(make_inner(), usage_ledger=ledger))
170
+
171
+ graph.invoke({"messages": ["q1"]}, cfg("acme"))
172
+ graph.invoke({"messages": ["q2"]}, cfg("acme")) # 2nd turn, same thread
173
+ graph.invoke({"messages": ["q1"]}, cfg("globex"))
174
+
175
+ acme, globex = ledger.totals("acme"), ledger.totals("globex")
176
+ # two LLM calls for acme, one for globex; old messages re-checkpointed
177
+ # on turn 2 must not double-count
178
+ assert acme.messages == 2 and acme.total_tokens == 30
179
+ assert globex.messages == 1 and globex.total_tokens == 15
180
+ assert acme.by_model == {"claude-sonnet-4-6": 30}