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.
- langgraph_tenancy-0.1.0/.github/workflows/ci.yml +50 -0
- langgraph_tenancy-0.1.0/.github/workflows/release.yml +71 -0
- langgraph_tenancy-0.1.0/.gitignore +21 -0
- langgraph_tenancy-0.1.0/LICENSE +21 -0
- langgraph_tenancy-0.1.0/PKG-INFO +143 -0
- langgraph_tenancy-0.1.0/README.md +114 -0
- langgraph_tenancy-0.1.0/langgraph_tenancy/__init__.py +65 -0
- langgraph_tenancy-0.1.0/langgraph_tenancy/_checkpointer.py +254 -0
- langgraph_tenancy-0.1.0/langgraph_tenancy/_errors.py +39 -0
- langgraph_tenancy-0.1.0/langgraph_tenancy/_store.py +124 -0
- langgraph_tenancy-0.1.0/langgraph_tenancy/_usage.py +95 -0
- langgraph_tenancy-0.1.0/pyproject.toml +61 -0
- langgraph_tenancy-0.1.0/tests/conftest.py +53 -0
- langgraph_tenancy-0.1.0/tests/test_isolation.py +180 -0
|
@@ -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
|
+
[](https://github.com/ac12644/langgraph-tenancy/actions/workflows/ci.yml)
|
|
33
|
+
[](https://pypi.org/project/langgraph-tenancy/)
|
|
34
|
+
[](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
|
+
[](https://github.com/ac12644/langgraph-tenancy/actions/workflows/ci.yml)
|
|
4
|
+
[](https://pypi.org/project/langgraph-tenancy/)
|
|
5
|
+
[](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}
|