Skip to content

Commit bdb2f95

Browse files
feat(keycardai-temporal): Keycard model provider for the OpenAI Agents plugin
ECO-395. temporalio's OpenAIAgentsPlugin builds its OpenAI client before any activity exists, so neither @grant nor the interceptor can supply the model credential and every consumer hand-rolled a boot-time mint. KeycardOpenAIProvider (keycardai.temporal.openai_agents, optional extra keycardai-temporal[openai-agents]) fills the plugin's model_provider seat: the OpenAI client carries an async key callable that mints the vaulted key with a client-credentials grant on first use, reuses it for a refresh window bounded by the grant's expires_in, and re-mints after, so rotation in Keycard propagates without a worker restart. Permanent grant failures fail the model activity with the same non-retryable KeycardAccessDenied a @grant activity raises; transient ones follow the model activity's retry policy. Base import stays free of the optional dependencies. The motivating consumer (mongodb-developer/mdb-temporal-keycard#1) deletes its inline boot mint on release.
1 parent 6d255b2 commit bdb2f95

6 files changed

Lines changed: 695 additions & 23 deletions

File tree

‎packages/temporal/README.md‎

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,12 +137,70 @@ interceptor = KeycardInterceptor(
137137
)
138138
```
139139

140+
## The OpenAI Agents plugin
141+
142+
`temporalio.contrib.openai_agents.OpenAIAgentsPlugin` runs every model call as an activity of its own, registered by the plugin before any of your activities exist, so neither `@grant` nor the interceptor can hand it a credential. `keycardai.temporal.openai_agents.KeycardOpenAIProvider` fills that seat: a `ModelProvider` whose OpenAI client resolves its key through Keycard on each request instead of reading `OPENAI_API_KEY` once at startup.
143+
144+
```bash
145+
pip install "keycardai-temporal[openai-agents]" "temporalio[openai-agents]"
146+
```
147+
148+
The intended end state for a worker: the interceptor once, `@grant` on each activity, and the plugin configured through Keycard for the model.
149+
150+
```python
151+
from datetime import timedelta
152+
153+
from temporalio.client import Client
154+
from temporalio.contrib.openai_agents import ModelActivityParameters, OpenAIAgentsPlugin
155+
from temporalio.worker import Worker
156+
157+
from keycardai.temporal import KeycardInterceptor, access, grant
158+
from keycardai.temporal.openai_agents import KeycardOpenAIProvider
159+
160+
ZONE = "https://your-zone.keycard.cloud"
161+
162+
163+
@grant("urn:mongodb:atlas:orders")
164+
async def load_orders(customer_id: str) -> list[dict]:
165+
return await atlas_client(access().access_token).find(customer_id)
166+
167+
168+
async def main() -> None:
169+
client = await Client.connect(
170+
"localhost:7233",
171+
plugins=[
172+
OpenAIAgentsPlugin(
173+
model_params=ModelActivityParameters(
174+
start_to_close_timeout=timedelta(seconds=60)
175+
),
176+
model_provider=KeycardOpenAIProvider(
177+
ZONE, "urn:vault:openai-api-key", refresh=timedelta(minutes=5)
178+
),
179+
)
180+
],
181+
)
182+
worker = Worker(
183+
client,
184+
task_queue="orders",
185+
workflows=[OrderWorkflow],
186+
activities=[load_orders],
187+
interceptors=[KeycardInterceptor(ZONE)],
188+
)
189+
await worker.run()
190+
```
191+
192+
`KeycardOpenAIProvider(zone_url, resource, credential=None, *, refresh=timedelta(minutes=5), base_url=None, use_responses=None)` takes the zone URL and the identifier of the resource whose vault holds the OpenAI key. The credential is the same as the interceptor's: an explicit `ApplicationCredential`, or the one `discover_credential()` finds in the environment when omitted. It must be a `ClientSecret`, as for any client-credentials `@grant`; other credential types raise `GrantConfigurationError` at construction.
193+
194+
Refresh: the key is minted with a client-credentials grant on the first model call, not at worker startup, and reused for `refresh` (or the grant's `expires_in`, whichever is shorter). The next model call after the window mints again, so a key rotated in Keycard reaches the worker within one refresh window, with no restart. Concurrent model calls on an expired cache share one mint. A permanent grant failure fails the model activity with the same non-retryable `KeycardAccessDenied` a `@grant` activity would raise; a transient one is left to the model activity's retry policy. The plugin requires an explicit `start_to_close_timeout` or `schedule_to_close_timeout` whenever a custom provider is set.
195+
196+
The optional dependencies (`openai-agents`, `openai`) are imported only inside `keycardai.temporal.openai_agents`, and only when a provider is built; `import keycardai.temporal` stays free of them. Without the extra installed, building a provider raises an `ImportError` naming it. `temporalio[openai-agents]` itself stays yours to install, since its version is pinned by your worker, not by this package.
197+
140198
## Design notes
141199

142200
- Tokens never touch durable state. Workflow history is replayable and permanent, so unlike header-based context-propagation interceptors, nothing is written to activity headers, arguments, or return values. The token exists only inside one execution's context. This is also why on-behalf-of activities receive an identity reference instead of a token: the session lookup and exchange happen at the edge, inside the execution, so a session revoked mid-workflow is never replayed from state.
143201
- A mint failure for any declared resource raises before the activity body runs, and there is no fallback path. Tokens live in the SDK's shared `AccessContext` (the same container `keycardai-mcp` uses), but where that idiom is non-throwing, the interceptor converts recorded errors into raises on purpose: in Temporal, raising is the error channel.
144202
- Transient mint failures are retryable; permanent failures are not. The classification is `keycardai-oauth`'s: every typed error carries a `retryable` property, and `exchange_tokens_for_resources` records it in the error dict it stores on the `AccessContext`. Network failures, 5xx and 429 responses, and unclassified errors let the activity retry policy govern what happens next, while permanent failures (`access_denied`, `insufficient_authorization`, `invalid_client`, other 4xx responses, configuration and authentication errors) raise `ApplicationError(type="KeycardAccessDenied", non_retryable=True)` immediately. Misdeclarations surface as `GrantConfigurationError`, retryable by default so a worker redeploy with the fix lets the next retry succeed; list `"GrantConfigurationError"` in the retry policy's `non_retryable_error_types` to give up sooner.
145-
- No token caching, per Keycard's credential rules; per-call mint is the contract. One OAuth client is created per worker and reused; only the tokens are fresh.
203+
- No token caching, per Keycard's credential rules; per-call mint is the contract. One OAuth client is created per worker and reused; only the tokens are fresh. The one exception is the OpenAI Agents provider's model key, which is a vaulted third-party secret rather than a Keycard token: it is reused for a short `refresh` window rather than minted per request, and re-minted after it, so rotation still propagates without a restart.
146204
- The package wraps its own `keycardai` imports in `workflow.unsafe.imports_passed_through()`, the idiom from Temporal's sentry sample, so consumers import it normally even in files that define workflows.
147205
- Works with async activities and with sync activities on the thread-pool executor (the Temporal SDK copies contextvars into the thread). Sync activities on a process-pool executor are not supported because contextvars do not cross processes.
148206

@@ -152,4 +210,4 @@ interceptor = KeycardInterceptor(
152210
cd packages/temporal && uv run --extra test pytest tests/ -v
153211
```
154212

155-
`tests/test_interceptor.py` drives the interceptor chain directly with the OAuth client stubbed. `tests/test_history_hygiene.py` runs a real workflow against a local Temporal dev server (downloaded by `temporalio` on first use), then scans the recorded history, including base64-decoded payloads, and asserts the minted token appears nowhere. Neither needs a Keycard zone.
213+
`tests/test_interceptor.py` drives the interceptor chain directly with the OAuth client stubbed. `tests/test_openai_agents.py` covers the OpenAI Agents provider with the OAuth client stubbed and the clock faked; two tests that need the Agents SDK and the real plugin skip unless the `openai-agents` extra (and `temporalio[openai-agents]`) is installed. `tests/test_history_hygiene.py` runs a real workflow against a local Temporal dev server (downloaded by `temporalio` on first use), then scans the recorded history, including base64-decoded payloads, and asserts the minted token appears nowhere. Neither needs a Keycard zone.

‎packages/temporal/pyproject.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@ classifiers = [
2626
]
2727

2828
[project.optional-dependencies]
29+
openai-agents = [
30+
"openai-agents>=0.20.0",
31+
"openai>=2.45.0",
32+
]
2933
test = [
3034
"pytest>=8.4.1",
3135
"pytest-asyncio>=1.1.0",

‎packages/temporal/src/keycardai/temporal/__init__.py‎

Lines changed: 33 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
import inspect
4242
from collections.abc import Awaitable, Callable
4343
from dataclasses import dataclass, fields
44-
from typing import Annotated, Any, get_args, get_origin, get_type_hints
44+
from typing import Annotated, Any, NoReturn, get_args, get_origin, get_type_hints
4545

4646
from temporalio import workflow
4747
from temporalio.exceptions import ApplicationError
@@ -337,6 +337,20 @@ def access(resource: str | None = None) -> TokenResponse:
337337
return ctx.access(resource)
338338

339339

340+
def _worker_credential(
341+
credential: ApplicationCredential | None,
342+
) -> ApplicationCredential:
343+
"""The explicit credential, or the one the environment describes."""
344+
if credential is not None:
345+
return credential
346+
try:
347+
return discover_credential()
348+
except CredentialDiscoveryError as e:
349+
# Same taxonomy as every other worker-configuration problem in this
350+
# package, so retry policies see one error class.
351+
raise GrantConfigurationError(str(e)) from e
352+
353+
340354
class KeycardInterceptor(Interceptor):
341355
"""Worker interceptor that mints fresh Keycard tokens per activity execution.
342356
@@ -371,13 +385,7 @@ def __init__(
371385
credential: ApplicationCredential | None = None,
372386
subject_token_provider: SubjectTokenProvider | None = None,
373387
) -> None:
374-
if credential is None:
375-
try:
376-
credential = discover_credential()
377-
except CredentialDiscoveryError as e:
378-
# Same taxonomy as every other worker-configuration problem
379-
# in this package, so retry policies see one error class.
380-
raise GrantConfigurationError(str(e)) from e
388+
credential = _worker_credential(credential)
381389
self._credential = credential
382390
# One client for the worker's lifetime: endpoint discovery runs once
383391
# and is cached on the instance. Tokens are still minted per call;
@@ -428,15 +436,7 @@ async def _mint(self, grant: _Grant, input: ExecuteActivityInput) -> AccessConte
428436
resource=resource
429437
)
430438
except Exception as e:
431-
if not error_retryable(e):
432-
code = getattr(e, "error", None) or type(e).__name__
433-
raise ApplicationError(
434-
f"Keycard grant for {resource} failed permanently: "
435-
f"{code}: {e}",
436-
type="KeycardAccessDenied",
437-
non_retryable=True,
438-
) from e
439-
raise # transient: the activity retry policy governs it
439+
_raise_for_grant_failure(resource, e)
440440
ctx.set_token(resource, resp)
441441
elif grant.impersonate:
442442
# client.impersonate() authenticates only at the HTTP layer, and
@@ -497,6 +497,22 @@ async def execute_activity(self, input: ExecuteActivityInput) -> Any:
497497
_ctx.reset(token)
498498

499499

500+
def _raise_for_grant_failure(resource: str, e: Exception) -> NoReturn:
501+
"""Classify a failed client-credentials grant the way the interceptor does.
502+
503+
Permanent failures become the non-retryable ``KeycardAccessDenied``;
504+
transient ones re-raise as-is so the activity retry policy governs them.
505+
"""
506+
if not error_retryable(e):
507+
code = getattr(e, "error", None) or type(e).__name__
508+
raise ApplicationError(
509+
f"Keycard grant for {resource} failed permanently: {code}: {e}",
510+
type="KeycardAccessDenied",
511+
non_retryable=True,
512+
) from e
513+
raise e
514+
515+
500516
def _raise_on_mint_error(ctx: AccessContext, *resources: str) -> None:
501517
"""Raise for the first failed resource, permanent failures first.
502518
Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
"""Keycard-minted OpenAI credential for temporalio's OpenAI Agents plugin.
2+
3+
``OpenAIAgentsPlugin`` runs every model call as a Temporal activity of its
4+
own, and that activity asks its ``ModelProvider`` for a model on each call.
5+
Neither ``@grant`` nor :class:`keycardai.temporal.KeycardInterceptor` can
6+
reach that activity: the plugin registers it and builds its OpenAI client
7+
before any consumer activity exists. :class:`KeycardOpenAIProvider` closes
8+
the gap from the provider side. It hands the OpenAI client an async key
9+
callback instead of a key string, and the client awaits that callback before
10+
every request it sends. The callback mints the OpenAI key from a vaulted
11+
Keycard resource with a client-credentials grant and caches it for a short
12+
``refresh`` window, so a key rotated in Keycard is picked up by the next
13+
model call after the window closes, with no worker restart.
14+
15+
Optional dependencies (``openai``, ``agents``) are imported only here, and
16+
only when a provider is built, so ``import keycardai.temporal`` stays clean
17+
without the ``openai-agents`` extra.
18+
"""
19+
20+
from __future__ import annotations
21+
22+
import asyncio
23+
import time
24+
from datetime import timedelta
25+
from typing import TYPE_CHECKING, Any
26+
27+
from temporalio import workflow
28+
29+
with workflow.unsafe.imports_passed_through():
30+
from keycardai.oauth import AsyncClient
31+
from keycardai.oauth.server import ApplicationCredential, ClientSecret
32+
33+
from keycardai.temporal import (
34+
GrantConfigurationError,
35+
_raise_for_grant_failure,
36+
_worker_credential,
37+
)
38+
39+
if TYPE_CHECKING:
40+
from agents import Model, ModelProvider
41+
42+
_ProviderBase = ModelProvider
43+
else:
44+
# Duck-typed at runtime so this module imports without ``agents``; the
45+
# plugin only ever calls get_model() and aclose().
46+
_ProviderBase = object
47+
48+
__all__ = ["KeycardOpenAIKey", "KeycardOpenAIProvider"]
49+
50+
_MISSING_EXTRA = (
51+
"keycardai.temporal.openai_agents needs the OpenAI Agents SDK and the openai "
52+
"client: install keycardai-temporal[openai-agents] next to "
53+
"temporalio[openai-agents]."
54+
)
55+
56+
57+
def _openai_deps() -> tuple[Any, Any]:
58+
"""Import the optional dependencies on first use: (OpenAIProvider, AsyncOpenAI)."""
59+
try:
60+
from agents.models.openai_provider import OpenAIProvider
61+
from openai import AsyncOpenAI
62+
except ImportError as e:
63+
raise ImportError(_MISSING_EXTRA) from e
64+
return OpenAIProvider, AsyncOpenAI
65+
66+
67+
class KeycardOpenAIKey:
68+
"""Async callable returning the OpenAI key vaulted behind a Keycard resource.
69+
70+
Shaped for ``openai.AsyncOpenAI(api_key=...)``, which awaits it before
71+
each request. The key is minted with a client-credentials grant and
72+
cached for ``refresh`` (bounded by the grant's ``expires_in`` when the
73+
zone reports a shorter one); concurrent callers on a cold or expired
74+
cache share one mint. A permanent grant failure raises the same
75+
non-retryable ``KeycardAccessDenied`` ``ApplicationError`` the
76+
interceptor raises, so the plugin's model activity fails the same way a
77+
``@grant`` activity would; a transient failure propagates as-is and the
78+
model activity's retry policy governs it.
79+
80+
Args:
81+
zone_url: Keycard zone issuer URL.
82+
resource: Identifier of the vaulted OpenAI key resource in the zone.
83+
credential: How this application authenticates to the zone; when
84+
omitted, discovered from the environment exactly as
85+
:class:`keycardai.temporal.KeycardInterceptor` does. Must be a
86+
``ClientSecret`` (the client-credentials path has no bridge for
87+
assertion credentials yet).
88+
refresh: How long a minted key is reused before the next call mints
89+
again. Default five minutes.
90+
91+
Raises:
92+
GrantConfigurationError: the credential could not be discovered, or
93+
is not a ``ClientSecret``.
94+
"""
95+
96+
def __init__(
97+
self,
98+
zone_url: str,
99+
resource: str,
100+
credential: ApplicationCredential | None = None,
101+
*,
102+
refresh: timedelta = timedelta(minutes=5),
103+
) -> None:
104+
credential = _worker_credential(credential)
105+
if not isinstance(credential, ClientSecret):
106+
raise GrantConfigurationError(
107+
f"{type(self).__name__} for {resource} mints with client "
108+
"credentials, which requires a ClientSecret credential; the "
109+
f"worker has {type(credential).__name__}."
110+
)
111+
if refresh <= timedelta(0):
112+
raise GrantConfigurationError(
113+
f"{type(self).__name__}: refresh must be positive, got {refresh}."
114+
)
115+
self.resource = resource
116+
self.refresh = refresh
117+
self._client = AsyncClient(zone_url, auth=credential.get_http_client_auth())
118+
self._key: str | None = None
119+
self._fresh_until = 0.0
120+
self._lock = asyncio.Lock()
121+
122+
def _cached(self) -> str | None:
123+
if self._key is not None and time.monotonic() < self._fresh_until:
124+
return self._key
125+
return None
126+
127+
async def __call__(self) -> str:
128+
key = self._cached()
129+
if key is not None:
130+
return key
131+
async with self._lock:
132+
key = self._cached()
133+
if key is not None:
134+
return key
135+
try:
136+
resp = await self._client.client_credentials_grant(
137+
resource=self.resource
138+
)
139+
except Exception as e:
140+
_raise_for_grant_failure(self.resource, e)
141+
window = self.refresh.total_seconds()
142+
if resp.expires_in is not None:
143+
window = min(window, float(resp.expires_in))
144+
self._key = resp.access_token
145+
self._fresh_until = time.monotonic() + window
146+
return self._key
147+
148+
async def aclose(self) -> None:
149+
"""Close the Keycard client; the next call opens a fresh one."""
150+
self._key = None
151+
self._fresh_until = 0.0
152+
await self._client.aclose()
153+
154+
155+
class KeycardOpenAIProvider(_ProviderBase):
156+
"""``agents.ModelProvider`` whose OpenAI key comes from Keycard per call.
157+
158+
Pass it as ``OpenAIAgentsPlugin(model_provider=...)``. Models come from
159+
the SDK's own ``OpenAIProvider`` over one ``AsyncOpenAI`` client that
160+
carries a :class:`KeycardOpenAIKey` in place of a key string, so every
161+
request the plugin's model activity sends first resolves (and, past the
162+
``refresh`` window, re-mints) the key. ``max_retries`` on that client
163+
is 0, as in the plugin's default provider, so the model activity's
164+
retry policy is the only retry loop.
165+
166+
The plugin requires an explicit ``start_to_close_timeout`` or
167+
``schedule_to_close_timeout`` in ``ModelActivityParameters`` whenever a
168+
custom provider is set; it only defaults the timeout for its own.
169+
170+
Args:
171+
zone_url, resource, credential, refresh: as :class:`KeycardOpenAIKey`.
172+
base_url: OpenAI-compatible base URL, when not api.openai.com.
173+
use_responses: Forwarded to ``OpenAIProvider`` (Responses API vs
174+
Chat Completions); ``None`` keeps the SDK default.
175+
176+
Raises:
177+
GrantConfigurationError: as :class:`KeycardOpenAIKey`.
178+
ImportError: the ``openai-agents`` extra is not installed.
179+
"""
180+
181+
def __init__(
182+
self,
183+
zone_url: str,
184+
resource: str,
185+
credential: ApplicationCredential | None = None,
186+
*,
187+
refresh: timedelta = timedelta(minutes=5),
188+
base_url: str | None = None,
189+
use_responses: bool | None = None,
190+
) -> None:
191+
OpenAIProvider, AsyncOpenAI = _openai_deps()
192+
self.api_key = KeycardOpenAIKey(zone_url, resource, credential, refresh=refresh)
193+
self._openai = AsyncOpenAI(
194+
api_key=self.api_key, base_url=base_url, max_retries=0
195+
)
196+
self._provider = OpenAIProvider(
197+
openai_client=self._openai, use_responses=use_responses
198+
)
199+
200+
def get_model(self, model_name: str | None) -> Model:
201+
return self._provider.get_model(model_name)
202+
203+
async def aclose(self) -> None:
204+
"""Release the inner provider, the OpenAI client, and the Keycard client."""
205+
await self._provider.aclose()
206+
await self._openai.close()
207+
await self.api_key.aclose()

0 commit comments

Comments
 (0)