keystone-request-context 0.3.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,13 @@
1
+ Metadata-Version: 2.4
2
+ Name: keystone-request-context
3
+ Version: 0.3.0
4
+ Summary: Shared request-scoped ContextVar slots + Starlette middleware for Keystone AI services.
5
+ Requires-Python: >=3.12
6
+ Requires-Dist: starlette<2,>=1.3.1
7
+ Requires-Dist: uuid-utils<1,>=0.14
8
+ Provides-Extra: dev
9
+ Requires-Dist: pytest<9,>=8; extra == "dev"
10
+ Requires-Dist: pytest-asyncio<1,>=0.24; extra == "dev"
11
+ Requires-Dist: pytest-cov<7,>=6; extra == "dev"
12
+ Requires-Dist: httpx<1,>=0.28; extra == "dev"
13
+ Requires-Dist: ruff==0.15.15; extra == "dev"
@@ -0,0 +1,85 @@
1
+ """Keystone AI — shared request-context library.
2
+
3
+ Provides the ``ContextVar`` slots and Starlette middleware that propagate
4
+ per-request identifiers (request_id, correlation_id, trace_id, user_id,
5
+ workspace_id, entity_id, stage_id, agent_id, traceparent) plus the full
6
+ header snapshot, available to any downstream code via plain
7
+ ``ContextVar.get()`` calls.
8
+
9
+ ``request_id`` (``X-Request-ID``) is per-API-call — each service mints a
10
+ fresh one on every inbound hop, used to pin log lines to a single HTTP
11
+ request to THIS service. Do NOT forward.
12
+
13
+ ``correlation_id`` (``X-Correlation-ID``) is per-business-turn — the FE /
14
+ API gateway mints it once and every downstream service forwards it
15
+ verbatim, so a search by ``correlation_id`` in Splunk / Langfuse returns
16
+ every log line + trace span for the whole turn. ALWAYS forward.
17
+
18
+ Usage::
19
+
20
+ from keystone.request_context import (
21
+ RequestContextMiddleware,
22
+ correlation_id_var,
23
+ request_id_var,
24
+ workspace_id_var,
25
+ )
26
+
27
+ app.add_middleware(RequestContextMiddleware)
28
+
29
+ # Anywhere downstream:
30
+ rid = request_id_var.get() # this hop only; log key
31
+ cid = correlation_id_var.get() # forward to upstream calls
32
+ ws = workspace_id_var.get()
33
+ """
34
+
35
+ from .caller_identity import (
36
+ LLMCallerIdentity,
37
+ gateway_identity_headers,
38
+ get_caller_identity,
39
+ get_self_agent,
40
+ reset_caller_identity,
41
+ set_caller_identity,
42
+ set_self_agent,
43
+ )
44
+ from .middleware import (
45
+ CORRELATION_ID_HEADER,
46
+ REQUEST_ID_HEADER,
47
+ RequestContextMiddleware,
48
+ normalize_correlation_id,
49
+ )
50
+ from .vars import (
51
+ agent_id_var,
52
+ correlation_id_var,
53
+ entity_id_var,
54
+ headers_var,
55
+ request_id_var,
56
+ stage_id_var,
57
+ trace_id_var,
58
+ traceparent_var,
59
+ user_id_var,
60
+ workspace_id_var,
61
+ )
62
+
63
+ __all__ = [
64
+ "CORRELATION_ID_HEADER",
65
+ "REQUEST_ID_HEADER",
66
+ "LLMCallerIdentity",
67
+ "RequestContextMiddleware",
68
+ "agent_id_var",
69
+ "correlation_id_var",
70
+ "entity_id_var",
71
+ "gateway_identity_headers",
72
+ "get_caller_identity",
73
+ "get_self_agent",
74
+ "headers_var",
75
+ "normalize_correlation_id",
76
+ "request_id_var",
77
+ "reset_caller_identity",
78
+ "set_caller_identity",
79
+ "set_self_agent",
80
+ "stage_id_var",
81
+ "trace_id_var",
82
+ "traceparent_var",
83
+ "user_id_var",
84
+ "workspace_id_var",
85
+ ]
@@ -0,0 +1,206 @@
1
+ """Per-request caller identity propagated to outbound building-block calls (llm-gateway, rag, …).
2
+
3
+ Lifted from ``services/rag`` (FDP-3118) so every keystone service AND the agent SDK share ONE
4
+ identity-forwarding contract: the gateway derives the LiteLLM credential from the caller
5
+ identity — agent (``agent_id``+``agent_version``) and/or end-user (``user_id``). FDP-1734
6
+ originally allowed EXACTLY ONE; FR-70 relaxed that so an agent serving an end user carries BOTH —
7
+ the gateway requires at least one and picks the agent's key whenever an agent is named. In an
8
+ agent pod the agent whose key must resolve is the pod itself, not the inbound caller — see
9
+ :func:`set_self_agent` (FDP-3663).
10
+
11
+ set_caller_identity(LLMCallerIdentity(...)) -> token # at the request / run / task boundary
12
+ reset_caller_identity(token) # in a finally
13
+ get_caller_identity() -> LLMCallerIdentity | None
14
+ gateway_identity_headers(auth_token=, operation=, caller_service=) -> dict[str, str]
15
+ set_self_agent(agent_id, agent_version) # once, from an agent pod's composition root
16
+
17
+ Whoever owns the boundary (a service's ``IdentityMiddleware``, or the agent runtime per run)
18
+ fills the ContextVar; outbound clients read it in their ``_headers()`` builder and forward every
19
+ field. Transparent forwarder — nothing here is service-static; ``caller_service`` names the
20
+ forwarding service (e.g. ``"rag"``, ``"agent-runtime"``).
21
+
22
+ NOTE: this bundle is a DISTINCT representation from the individual ID ContextVars in ``vars.py``
23
+ (request_id / correlation_id / entity_id / … set by ``RequestContextMiddleware`` for
24
+ logging/audit). This one is the outbound identity contract, set by the service's own
25
+ ``IdentityMiddleware`` / the agent runtime. They coexist by design; reconciling them is a later
26
+ cleanup, not this lift.
27
+ """
28
+
29
+ from __future__ import annotations
30
+
31
+ import uuid
32
+ from contextvars import ContextVar, Token
33
+ from dataclasses import dataclass, field
34
+
35
+
36
+ @dataclass(frozen=True, slots=True)
37
+ class LLMCallerIdentity:
38
+ """Per-request identity bundle forwarded to building-block services as-is.
39
+
40
+ Caller identity: agent (``agent_id`` + ``agent_version``) and/or end-user (``user_id``) —
41
+ WHO CALLED THIS PROCESS. At least one must be set — the gateway rejects a call carrying
42
+ neither. BOTH are set for the FR-70 case of an agent serving an end user, and both must then
43
+ be FORWARDED: the agent selects the credential and the agent quota tier, the user selects the
44
+ consumption tier. ``entity``/``stage`` are forwarded when present (the gateway derives them
45
+ from the workspace otherwise).
46
+
47
+ Note this is not the same question as *"whose gateway credential"* — in an agent pod that is
48
+ :func:`set_self_agent`, the agent making the outbound hop, not the one recorded here.
49
+ """
50
+
51
+ entity_id: str
52
+ stage_id: str
53
+ workspace_id: str
54
+ correlation_id: str
55
+ # who called this process — agent and/or end-user; at least one is set (FR-70 allows both),
56
+ # and a human-triggered run leaves the agent pair empty
57
+
58
+ agent_id: str = ""
59
+ agent_version: str = ""
60
+ user_id: str = ""
61
+ # FDP-3072 W5 (S3) — the turn's opaque invocation key (inv_…), minted by the gateway and
62
+ # handed to the pod per-request. A CREDENTIAL: repr=False keeps it out of any accidental
63
+ # log/str of this bundle; it is emitted outbound only as X-Invocation-Key.
64
+ invocation_key: str = field(default="", repr=False)
65
+
66
+
67
+ _caller_identity: ContextVar[LLMCallerIdentity | None] = ContextVar(
68
+ "keystone_llm_caller_identity",
69
+ default=None,
70
+ )
71
+
72
+ # The agent this PROCESS is — not who called it. A plain module global, not a ContextVar: it is
73
+ # process-static config (the operator injects it into the pod), identical for every run the process
74
+ # handles. Unset in every service that is not an agent pod, which keeps their headers unchanged.
75
+ _self_agent: tuple[str, str] | None = None
76
+
77
+
78
+ def set_self_agent(agent_id: str, agent_version: str) -> None:
79
+ """Register the agent this process IS, for outbound ``X-Agent-ID`` / ``X-Agent-Version``.
80
+
81
+ Call once from the composition root (FDP-3663). Blank ``agent_id`` clears the registration, so
82
+ a local/offline run with no injected identity behaves exactly as before.
83
+
84
+ **Why this is separate from the caller identity.** ``LLMCallerIdentity.agent_id`` answers *"who
85
+ called me"* — empty when a human triggered the run. The gateway's ``X-Agent-ID`` answers a
86
+ different question: *"whose LiteLLM key, budget, model allowlist and Langfuse project does this
87
+ call use"*, which is always the agent **making** the call. Feeding the caller's (absent) agent id
88
+ into that header sent every pro-code LLM call down the gateway's user path, so it billed the
89
+ WORKSPACE key and applied no per-agent budget, token cap or model allowlist at all.
90
+ """
91
+ global _self_agent
92
+ _self_agent = (agent_id, agent_version) if agent_id else None
93
+
94
+
95
+ def get_self_agent() -> tuple[str, str] | None:
96
+ """The ``(agent_id, agent_version)`` this process is, or ``None`` when not an agent pod."""
97
+ return _self_agent
98
+
99
+
100
+ def set_caller_identity(identity: LLMCallerIdentity) -> Token[LLMCallerIdentity | None]:
101
+ """Bind ``identity`` to the current execution context. Pass the returned token to
102
+ ``reset_caller_identity`` in a ``finally`` — failing to reset leaks identity into the next
103
+ request/run handled by the same worker process."""
104
+ return _caller_identity.set(identity)
105
+
106
+
107
+ def get_caller_identity() -> LLMCallerIdentity | None:
108
+ return _caller_identity.get()
109
+
110
+
111
+ def reset_caller_identity(token: Token[LLMCallerIdentity | None]) -> None:
112
+ _caller_identity.reset(token)
113
+
114
+
115
+ def gateway_identity_headers(*, auth_token: str, operation: str, caller_service: str) -> dict[str, str]:
116
+ """Build outbound building-block headers from the current caller-identity ContextVar.
117
+
118
+ Emits EVERY caller identity the context carries: ``X-Agent-ID`` + ``X-Agent-Version`` for an
119
+ agent/service call, ``X-User-ID`` for an end-user call, and BOTH when an agent is serving an
120
+ end user (FR-70).
121
+
122
+ This used to be exclusive — ``X-User-ID`` else the agent pair — on the pre-FR-70 rule that
123
+ the gateway 400s when both are present. FR-70 relaxed that (``auth._validate_caller_identity``
124
+ now rejects only a call carrying NEITHER), so the XOR silently dropped the agent on exactly
125
+ the case it was meant to describe: an agent run on behalf of a user reached the gateway as a
126
+ plain user call, which resolved the WORKSPACE virtual key instead of the agent's, skipped the
127
+ agent quota tier, and left per-agent cost attribution (FR-21) and the Langfuse agent blank.
128
+
129
+ ``caller_service`` → ``X-Caller-Service`` (the forwarding service). No-context fallback
130
+ (``get_caller_identity() is None``) emits empty agent headers + a random correlation id — the
131
+ gateway rejects that, which is the correct loud-failure mode for a bare/boot-time call.
132
+
133
+ When :func:`set_self_agent` has registered this process as an agent, ``X-Agent-ID`` +
134
+ ``X-Agent-Version`` are emitted for THAT agent **alongside** ``X-User-ID``. The gateway allows
135
+ both since FR-70 — *"an agent serving an end_user carries both; the agent resolves the
136
+ credential, and user_id becomes the subject of the consumption quota tier"* — and it is
137
+ explicitly "strictly more permissive, so no previously-valid call is affected" (source:
138
+ ``llm-gateway`` ``routers/middleware/auth.py::_validate_caller_identity``, which as of the
139
+ 2026-08-13 develop merge rejects only a call carrying NEITHER). Without a registered self agent
140
+ the behaviour is byte-identical to before, which is why no other service's headers change.
141
+
142
+ ``X-Agent-Version`` is REQUIRED whenever ``X-Agent-ID`` is present — the gateway 400s with
143
+ ``LLM_MISSING_CALLER_METADATA`` otherwise — so the two are always emitted as a pair.
144
+ """
145
+ caller = get_caller_identity()
146
+ if caller is None:
147
+ entity_id = stage_id = workspace_id = agent_id = agent_version = user_id = ""
148
+ correlation_id = str(uuid.uuid4())
149
+ else:
150
+ entity_id = caller.entity_id
151
+ stage_id = caller.stage_id
152
+ workspace_id = caller.workspace_id
153
+ agent_id = caller.agent_id
154
+ agent_version = caller.agent_version
155
+ user_id = caller.user_id
156
+ correlation_id = caller.correlation_id
157
+ invocation_key = "" if caller is None else caller.invocation_key
158
+ headers = {
159
+ "Content-Type": "application/json",
160
+ "X-Entity-ID": entity_id,
161
+ "X-Stage-ID": stage_id,
162
+ "X-Workspace-ID": workspace_id,
163
+ "X-Correlation-ID": correlation_id,
164
+ "X-Caller-Service": caller_service,
165
+ "X-Caller-Operation": operation,
166
+ }
167
+ # Only emit Authorization when a token is actually present. An empty token would produce the
168
+ # literal header value ``"Bearer "`` (trailing space, no credential), which httpx/h11 reject as an
169
+ # illegal header value → the request never leaves the pod. Downstreams that need no east-west
170
+ # bearer (e.g. RAG, auth-terminated at Kong and trusting the forwarded identity headers) pass
171
+ # ``auth_token=""`` and correctly get no Authorization header.
172
+ if auth_token:
173
+ headers["Authorization"] = f"Bearer {auth_token}"
174
+ # FDP-3072 W5 (S3b) — the invocation key rides its OWN header, deliberately NOT Authorization:
175
+ # that slot already carries real east-west bearers (llm-gateway), and dual-run (S5 step 1)
176
+ # requires this emission to be purely ADDITIVE — legacy identity headers keep flowing until
177
+ # Kong's strip-and-stamp flips. A separate header also dissolves the "two credentials, one
178
+ # header" routing problem Ops raised: the Kong plugin just reads X-Invocation-Key.
179
+ if invocation_key:
180
+ headers["X-Invocation-Key"] = invocation_key
181
+ self_agent = get_self_agent()
182
+ if self_agent is not None:
183
+ # Agent pod (FDP-3663): name the agent MAKING this hop — its key, budget, model allowlist
184
+ # and Langfuse project are the ones that must apply — and keep the end user alongside it so
185
+ # they stay the subject of the consumption quota tier (FR-70). Self wins over the inbound
186
+ # caller's pair: on an agent-to-agent hop the credential belongs to the agent calling now.
187
+ headers["X-Agent-ID"], headers["X-Agent-Version"] = self_agent
188
+ if user_id:
189
+ headers["X-User-ID"] = user_id
190
+ else:
191
+ # Not an agent pod (rag, memory, CLI, …): forward the inbound caller's identity as-is.
192
+ if user_id:
193
+ headers["X-User-ID"] = user_id
194
+ if agent_id or not user_id:
195
+ # The agent pair rides whenever an agent is named, alongside X-User-ID when both are
196
+ # set. The `not user_id` arm keeps the no-identity case emitting the EMPTY pair, which
197
+ # is what makes the gateway reject a context-less call loudly (see the docstring)
198
+ # instead of it leaving with no identity header at all.
199
+ #
200
+ # This is deliberately in the `else` arm: an agent pod already emitted ITS OWN pair
201
+ # above, and the inbound caller's `agent_id` must not overwrite it — that is the whole
202
+ # point of FDP-3663 (self wins on an agent-to-agent hop). Hoisting this to a bare `if`
203
+ # silently restores the workspace-key bug for exactly the A2A case.
204
+ headers["X-Agent-ID"] = agent_id
205
+ headers["X-Agent-Version"] = agent_version
206
+ return headers
@@ -0,0 +1,13 @@
1
+ Metadata-Version: 2.4
2
+ Name: keystone-request-context
3
+ Version: 0.3.0
4
+ Summary: Shared request-scoped ContextVar slots + Starlette middleware for Keystone AI services.
5
+ Requires-Python: >=3.12
6
+ Requires-Dist: starlette<2,>=1.3.1
7
+ Requires-Dist: uuid-utils<1,>=0.14
8
+ Provides-Extra: dev
9
+ Requires-Dist: pytest<9,>=8; extra == "dev"
10
+ Requires-Dist: pytest-asyncio<1,>=0.24; extra == "dev"
11
+ Requires-Dist: pytest-cov<7,>=6; extra == "dev"
12
+ Requires-Dist: httpx<1,>=0.28; extra == "dev"
13
+ Requires-Dist: ruff==0.15.15; extra == "dev"
@@ -0,0 +1,16 @@
1
+ __init__.py
2
+ caller_identity.py
3
+ middleware.py
4
+ py.typed
5
+ pyproject.toml
6
+ vars.py
7
+ ./__init__.py
8
+ ./caller_identity.py
9
+ ./middleware.py
10
+ ./py.typed
11
+ ./vars.py
12
+ keystone_request_context.egg-info/PKG-INFO
13
+ keystone_request_context.egg-info/SOURCES.txt
14
+ keystone_request_context.egg-info/dependency_links.txt
15
+ keystone_request_context.egg-info/requires.txt
16
+ keystone_request_context.egg-info/top_level.txt
@@ -0,0 +1,9 @@
1
+ starlette<2,>=1.3.1
2
+ uuid-utils<1,>=0.14
3
+
4
+ [dev]
5
+ pytest<9,>=8
6
+ pytest-asyncio<1,>=0.24
7
+ pytest-cov<7,>=6
8
+ httpx<1,>=0.28
9
+ ruff==0.15.15
@@ -0,0 +1,189 @@
1
+ """Starlette middleware that populates the request-context ``ContextVar`` slots.
2
+
3
+ Wire once at the FastAPI / Starlette app:
4
+
5
+ .. code-block:: python
6
+
7
+ from keystone.request_context import RequestContextMiddleware
8
+
9
+ app.add_middleware(RequestContextMiddleware)
10
+
11
+ Per request the middleware:
12
+
13
+ 1. Snapshots ``request.headers`` into :data:`keystone.request_context.headers_var`
14
+ with keys lowercased (Starlette already does this internally, but we
15
+ materialize a plain ``dict`` so downstream code can iterate without
16
+ pulling in Starlette types).
17
+ 2. Reads each well-known header and pushes the value into its dedicated
18
+ ``ContextVar`` slot.
19
+ 3. Auto-generates an ``uuid7`` for :data:`request_id_var` (per-API-call,
20
+ ``X-Request-ID``) when absent, and for :data:`correlation_id_var`
21
+ (per-turn, ``X-Correlation-ID``) when absent — legacy callers that
22
+ don't stamp the header still get a value to log against. Both
23
+ chosen values are echoed back on the response so callers can
24
+ correlate logs across hops.
25
+ 4. Resets every token in a ``try/finally`` block — defensive against
26
+ worker-reuse patterns (uvicorn keeps event-loop / task references
27
+ alive between requests).
28
+ """
29
+
30
+ from __future__ import annotations
31
+
32
+ from contextvars import ContextVar
33
+ from typing import TYPE_CHECKING
34
+
35
+ import uuid_utils
36
+ from starlette.middleware.base import BaseHTTPMiddleware
37
+
38
+ from .vars import (
39
+ agent_id_var,
40
+ correlation_id_var,
41
+ entity_id_var,
42
+ headers_var,
43
+ request_id_var,
44
+ stage_id_var,
45
+ trace_id_var,
46
+ traceparent_var,
47
+ user_id_var,
48
+ workspace_id_var,
49
+ )
50
+
51
+ if TYPE_CHECKING:
52
+ from collections.abc import Awaitable, Callable
53
+
54
+ from starlette.requests import Request
55
+ from starlette.responses import Response
56
+ from starlette.types import ASGIApp
57
+
58
+ REQUEST_ID_HEADER = "X-Request-ID"
59
+ CORRELATION_ID_HEADER = "X-Correlation-ID"
60
+
61
+ # Header (lowercase) → ContextVar slot. Keep alphabetical-ish by intent
62
+ # (tracing first, then identity / tenancy) to make the map easy to scan.
63
+ _HEADER_TO_VAR: tuple[tuple[str, ContextVar[str | None]], ...] = (
64
+ ("x-request-id", request_id_var),
65
+ ("x-correlation-id", correlation_id_var),
66
+ ("x-trace-id", trace_id_var),
67
+ ("traceparent", traceparent_var),
68
+ ("x-user-id", user_id_var),
69
+ ("x-workspace-id", workspace_id_var),
70
+ ("x-entity-id", entity_id_var),
71
+ ("x-stage-id", stage_id_var),
72
+ ("x-agent-id", agent_id_var),
73
+ )
74
+
75
+
76
+ def _generate_request_id() -> str:
77
+ """Return a fresh ``uuid7`` string (UUID format with dashes).
78
+
79
+ Wrapped so tests can monkey-patch a deterministic generator if needed.
80
+ """
81
+ return str(uuid_utils.uuid7())
82
+
83
+
84
+ def normalize_correlation_id(value: str | None) -> str | None:
85
+ """Strip dashes + lowercase a UUID string into the 32-hex W3C trace-id
86
+ shape so logs, downstream HTTP headers, and Langfuse trace_ids all
87
+ carry the EXACT SAME literal string.
88
+
89
+ Langfuse v3 SDK requires ``trace_id`` to be exactly 32 lowercase hex
90
+ chars (no dashes) — passing the FE-minted UUID-with-dashes raises
91
+ "is not a valid 32 lowercase hex char Langfuse trace id" and the
92
+ workflow span fails to open. Normalizing at the middleware level
93
+ means every consumer (log binding, ``X-Correlation-ID`` forward,
94
+ Langfuse start_turn) reads the already-normalized form from
95
+ ``correlation_id_var`` — no per-call-site strip.
96
+
97
+ Returns ``None`` for malformed inputs (caller falls back to
98
+ auto-minted UUID, then re-normalize).
99
+
100
+ Public since FDP-3435: agent-runtime normalizes at its own boundary and
101
+ derives the run's ``trace_id`` from stored correlation ids with the SAME
102
+ rule — one implementation, no drift.
103
+ """
104
+ if not value:
105
+ return None
106
+ cleaned = value.replace("-", "").lower()
107
+ if len(cleaned) != 32:
108
+ return None
109
+ try:
110
+ int(cleaned, 16)
111
+ except ValueError:
112
+ return None
113
+ return cleaned
114
+
115
+
116
+ # Backwards-compatible private alias (pre-FDP-3435 name).
117
+ _normalize_correlation_id = normalize_correlation_id
118
+
119
+
120
+ class RequestContextMiddleware(BaseHTTPMiddleware):
121
+ """Populate the request-context ``ContextVar`` slots for the request lifetime.
122
+
123
+ See module docstring for the full contract.
124
+ """
125
+
126
+ def __init__(self, app: ASGIApp) -> None:
127
+ super().__init__(app)
128
+
129
+ async def dispatch(
130
+ self,
131
+ request: Request,
132
+ call_next: Callable[[Request], Awaitable[Response]],
133
+ ) -> Response:
134
+ # 1. Materialize the lowercased-key header snapshot.
135
+ #
136
+ # Starlette's ``request.headers`` already supports case-insensitive
137
+ # lookup, but downstream consumers (event_bus DSL, @require) expect
138
+ # a plain dict so they aren't coupled to Starlette types.
139
+ headers_snapshot: dict[str, str] = {key.lower(): value for key, value in request.headers.items()}
140
+
141
+ # 2. Resolve / mint X-Request-ID (per-API-call) AND X-Correlation-ID
142
+ # (per-turn). Auto-gen UUID v7 for either when absent so logs always
143
+ # have a value to pivot on — semantics differ (request_id is per-hop,
144
+ # correlation_id is forwarded verbatim) but both fall back the same way.
145
+ request_id = headers_snapshot.get("x-request-id") or _generate_request_id()
146
+ # X-Correlation-ID is STORED + ECHOED as the 32-hex normalized form
147
+ # (dashes stripped, lowercased) so logs / forward / Langfuse trace_id
148
+ # all share one literal string. Malformed inbound value → re-mint a
149
+ # UUID v7 and normalize that.
150
+ correlation_raw = headers_snapshot.get("x-correlation-id")
151
+ correlation_id = _normalize_correlation_id(correlation_raw)
152
+ if correlation_id is None:
153
+ # Either no header sent, or value wasn't UUID-shaped — mint fresh.
154
+ correlation_id = _normalize_correlation_id(_generate_request_id())
155
+ # Make sure the snapshot reflects the resolved (possibly minted) values
156
+ # so DSL code reading ``headers_var`` sees a consistent view.
157
+ headers_snapshot["x-request-id"] = request_id
158
+ headers_snapshot["x-correlation-id"] = correlation_id
159
+
160
+ # 3. Set every ContextVar, retaining tokens so we can reset them
161
+ # in ``finally``. Resolved (minted-or-received) values already in
162
+ # ``headers_snapshot`` win for both request_id and correlation_id.
163
+ tokens = [headers_var.set(headers_snapshot)]
164
+ for header_name, ctx_var in _HEADER_TO_VAR:
165
+ value = headers_snapshot.get(header_name)
166
+ tokens.append(ctx_var.set(value))
167
+
168
+ try:
169
+ response = await call_next(request)
170
+ finally:
171
+ # Reset in reverse order — ContextVar.reset is strict about token
172
+ # ordering when nested calls have stacked sets on the same var.
173
+ for token in reversed(tokens):
174
+ try:
175
+ token.var.reset(token)
176
+ except (LookupError, ValueError):
177
+ # Defensive: a misbehaving downstream may have stomped
178
+ # the var with a bare ``.set()`` and dropped the token.
179
+ # We swallow because the alternative is leaking the
180
+ # exception into the response path.
181
+ pass
182
+
183
+ # 4. Echo resolved IDs so clients can correlate logs across hops.
184
+ response.headers[REQUEST_ID_HEADER] = request_id
185
+ response.headers[CORRELATION_ID_HEADER] = correlation_id
186
+ return response
187
+
188
+
189
+ __all__ = ["CORRELATION_ID_HEADER", "REQUEST_ID_HEADER", "RequestContextMiddleware"]
File without changes
@@ -0,0 +1,58 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "keystone-request-context"
7
+ version = "0.3.0"
8
+ description = "Shared request-scoped ContextVar slots + Starlette middleware for Keystone AI services."
9
+ requires-python = ">=3.12"
10
+ dependencies = [
11
+ # FDP-2360: floor raised to 1.3.1 — Starlette < 1.3.1 silently ignores
12
+ # request.form() limits for application/x-www-form-urlencoded (CVE-2026-54283 /
13
+ # GHSA-82w8-qh3p-5jfq — DoS). RequestContextMiddleware uses only stable
14
+ # BaseHTTPMiddleware / Request / Response APIs (unchanged across 0.x→1.x); the
15
+ # floor also lifts the <1 cap that pinned llm-gateway + visual-builder below the fix.
16
+ "starlette>=1.3.1,<2",
17
+ "uuid-utils>=0.14,<1",
18
+ ]
19
+
20
+ [project.optional-dependencies]
21
+ dev = [
22
+ "pytest>=8,<9",
23
+ "pytest-asyncio>=0.24,<1",
24
+ "pytest-cov>=6,<7",
25
+ "httpx>=0.28,<1",
26
+ "ruff==0.15.15",
27
+ ]
28
+
29
+ [tool.setuptools]
30
+ # Flat layout: the package source lives at the repo lib root (vars.py,
31
+ # middleware.py, __init__.py). We map the dotted import path
32
+ # ``keystone.request_context`` directly to "." so the existing
33
+ # libs/python/<name>/ pattern is preserved.
34
+ packages = ["keystone.request_context"]
35
+
36
+ [tool.setuptools.package-dir]
37
+ "keystone.request_context" = "."
38
+
39
+ [tool.setuptools.package-data]
40
+ "keystone.request_context" = ["py.typed"]
41
+
42
+ [tool.ruff]
43
+ target-version = "py312"
44
+ line-length = 120
45
+
46
+ [tool.ruff.lint]
47
+ extend-select = [
48
+ "E", # pycodestyle errors
49
+ "F", # pyflakes
50
+ "I", # isort
51
+ "UP", # pyupgrade
52
+ "B", # bugbear
53
+ ]
54
+
55
+ [tool.pytest.ini_options]
56
+ asyncio_mode = "auto"
57
+ testpaths = ["tests"]
58
+ addopts = "-ra --strict-markers"
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,119 @@
1
+ """Request-scoped ``ContextVar`` slots for the Keystone platform.
2
+
3
+ Every value is populated by :class:`keystone.request_context.RequestContextMiddleware`
4
+ at the boundary of an HTTP request, read by downstream code (services,
5
+ repositories, decorators such as ``@announce`` / ``@require``), and reset
6
+ in a ``try/finally`` block before the request returns.
7
+
8
+ Conventions
9
+ -----------
10
+
11
+ * Nine "ID" ContextVars are typed ``str | None`` with ``default=None`` —
12
+ this preserves the distinction between *absent* (no header was sent) and
13
+ *explicitly empty string* (header was sent with empty value).
14
+ * ``headers_var`` carries the full raw header snapshot (keys lowercased)
15
+ as an escape hatch for code that needs headers we haven't promoted to
16
+ first-class slots yet; defaults to an empty ``dict`` for non-request
17
+ contexts (workers, crons, tests).
18
+
19
+ The names mirror the headers documented in
20
+ ``specs/conventions/multi-tenancy.md`` and the event-bus impl plan
21
+ (``docs/architecture/event-bus/impl-plan.md`` §Phase 1).
22
+ """
23
+
24
+ from __future__ import annotations
25
+
26
+ from contextvars import ContextVar
27
+
28
+ # ---------------------------------------------------------------------------
29
+ # Tracing identifiers
30
+ # ---------------------------------------------------------------------------
31
+
32
+ request_id_var: ContextVar[str | None] = ContextVar("keystone_request_id", default=None)
33
+ """Per-API-call ID (header ``X-Request-ID``) — DO NOT forward.
34
+
35
+ Scope: a SINGLE API call. Each service mints its own value on every
36
+ inbound hop — value seen by service A is unrelated to value seen by
37
+ service B in the same business turn. Middleware auto-generates a
38
+ ``uuid7`` when the header is absent and echoes the chosen value
39
+ on the response so clients can pin log lines to one HTTP request.
40
+
41
+ → Use this to find log lines for ONE inbound request to the current
42
+ service. To stitch a multi-hop turn across services use
43
+ :data:`correlation_id_var` instead — that ID is forwarded verbatim
44
+ by every service.
45
+ """
46
+
47
+ correlation_id_var: ContextVar[str | None] = ContextVar("keystone_correlation_id", default=None)
48
+ """Per-turn correlation ID (header ``X-Correlation-ID``) — FORWARD verbatim.
49
+
50
+ Scope: a WHOLE business flow that may fan out across many services
51
+ (e.g. one user click → FE → Visual Builder → Agent Hub + LLM Gateway +
52
+ RAG + Memory + Skill + Langfuse). The FE / API gateway mints the value
53
+ ONCE (typically ``crypto.randomUUID()``) and stamps it on the first
54
+ request; every downstream service forwards the same value as-is on
55
+ every outbound hop. Middleware auto-generates a ``uuid7`` as a fallback
56
+ for legacy callers that don't stamp the header (cron jobs, internal
57
+ service-to-service calls) so logs always have something to pivot on;
58
+ the chosen value is echoed back on the response.
59
+
60
+ → Forward this on EVERY outbound HTTP / Kafka header. Search by this
61
+ single value in Splunk to see every service log for the turn; paste
62
+ into Langfuse to see the trace.
63
+ """
64
+
65
+ trace_id_var: ContextVar[str | None] = ContextVar("keystone_trace_id", default=None)
66
+ """Cross-service trace identifier (header ``X-Trace-ID``)."""
67
+
68
+ traceparent_var: ContextVar[str | None] = ContextVar("keystone_traceparent", default=None)
69
+ """W3C ``traceparent`` header verbatim (for future OTel propagation)."""
70
+
71
+ # ---------------------------------------------------------------------------
72
+ # Identity / tenancy identifiers
73
+ # ---------------------------------------------------------------------------
74
+
75
+ user_id_var: ContextVar[str | None] = ContextVar("keystone_user_id", default=None)
76
+ """Authenticated principal (header ``X-User-ID``)."""
77
+
78
+ workspace_id_var: ContextVar[str | None] = ContextVar("keystone_workspace_id", default=None)
79
+ """Active workspace scope (header ``X-Workspace-ID``)."""
80
+
81
+ entity_id_var: ContextVar[str | None] = ContextVar("keystone_entity_id", default=None)
82
+ """Active entity scope (header ``X-Entity-ID``)."""
83
+
84
+ stage_id_var: ContextVar[str | None] = ContextVar("keystone_stage_id", default=None)
85
+ """Active stage scope (header ``X-Stage-ID``, opt-in per endpoint)."""
86
+
87
+ agent_id_var: ContextVar[str | None] = ContextVar("keystone_agent_id", default=None)
88
+ """Caller agent identifier when traffic originates from another agent
89
+ (header ``X-Agent-ID``)."""
90
+
91
+ # ---------------------------------------------------------------------------
92
+ # Escape hatch — full headers snapshot
93
+ # ---------------------------------------------------------------------------
94
+
95
+ headers_var: ContextVar[dict[str, str]] = ContextVar("keystone_request_headers", default={}) # noqa: B039 — intentional read-only sentinel; see warning below
96
+ """Full inbound header snapshot, keys lowercased.
97
+
98
+ Code paths that need headers we have NOT promoted to a first-class slot
99
+ (e.g. the ``header:`` DSL in ``@announce``) read directly from here.
100
+ ``default={}`` keeps the dict truthy-comparable without ``None`` checks
101
+ in worker / no-request scenarios.
102
+
103
+ WARNING: The default value is a **shared empty dict** for ergonomic reads.
104
+ Never mutate ``headers_var.get()`` — always ``headers_var.set(new_dict)``
105
+ with a fresh dict to avoid bleeding state across requests.
106
+ """
107
+
108
+ __all__ = [
109
+ "agent_id_var",
110
+ "correlation_id_var",
111
+ "entity_id_var",
112
+ "headers_var",
113
+ "request_id_var",
114
+ "stage_id_var",
115
+ "trace_id_var",
116
+ "traceparent_var",
117
+ "user_id_var",
118
+ "workspace_id_var",
119
+ ]