error: {request.args.get('error')}\n"
+ f"description: {request.args.get('error_description')}",
+ 400,
+ )
+
+ # Validate state to protect against CSRF.
+ expected_state = session.pop("oauth_state", None)
+ if not expected_state or request.args.get("state") != expected_state:
+ return "error: state mismatch", 400 + + code = request.args.get("code") + if not code: + return "
error: no authorization code returned", 400 + + resp = requests.post(f"{OKTA_DOMAIN}/oauth2/v1/token", data={ + "grant_type": "authorization_code", + "code": code, + "redirect_uri": REDIRECT_URI, + "client_id": OIDC_CLIENT_ID, + "client_secret": OIDC_CLIENT_SECRET, + }) + if not resp.ok: + return f"
token endpoint error ({resp.status_code}):\n{resp.text}", 400
+
+ id_token = resp.json().get("id_token", "")
+ return f"""\
+
+
+
+
+ | Token | Value |
|---|---|
| ID token | +{id_token} | +
error: {request.args.get('error')}\n"
+ f"description: {request.args.get('error_description')}",
+ 400,
+ )
+
+ # Validate state to protect against CSRF.
+ expected_state = session.pop("oauth_state", None)
+ if not expected_state or request.args.get("state") != expected_state:
+ return "error: state mismatch", 400 + + code = request.args.get("code") + if not code: + return "
error: no authorization code returned", 400 + + resp = requests.post(f"{OKTA_DOMAIN}/oauth2/v1/token", data={ + "grant_type": "authorization_code", + "code": code, + "redirect_uri": REDIRECT_URI, + "client_id": OIDC_CLIENT_ID, + "client_secret": OIDC_CLIENT_SECRET, + }) + resp.raise_for_status() + id_token = resp.json()["id_token"] + + id_jag = get_id_jag(id_token) + access_token = get_access_token(id_jag) + + return f"""\ + + + + +
| Token | Value |
|---|---|
| ID token | +{id_token} | +
| ID-JAG | +{id_jag} | +
| Access token | +{access_token} | +
+ The module needs only a JWT library and an HTTP client. Add these to your project's requirements.txt:
+
PyJWT[crypto]>=2.8.0
+requests>=2.31.0
+ Create a file named token_exchange.py. It reads the Okta values from the environment, signs the client assertion, and exposes two functions, get_id_jag and get_access_token, that your agent calls in order.
+
"""Okta token exchange for AI agents.
+
+Turns a signed-in user's id_token into a scoped access_token:
+ id_token -> ID-JAG (org AS) -> access_token (custom AS)
+
+Exposes get_id_jag() and get_access_token(). No platform dependencies.
+"""
+
+import json, os, time, uuid
+import jwt
+import requests
+from jwt.algorithms import RSAAlgorithm
+
+# --- Okta configuration (from environment) ---
+OKTA_DOMAIN = os.environ["OKTA_DOMAIN"] # for example, example.okta.com
+CUSTOM_AS_ID = os.environ.get("OKTA_CUSTOM_AS_ID", "default")
+REQUESTED_SCOPE = os.environ.get("OKTA_SCOPE", "xaa:read")
+AGENT_CLIENT_ID = os.environ["AGENT_CLIENT_ID"]
+AGENT_KEY_ID = os.environ["AGENT_KEY_ID"]
+AGENT_PRIVATE_KEY_JWK = json.loads(os.environ["AGENT_PRIVATE_KEY_JWK"])
+
+ORG_TOKEN_URL = f"https://{OKTA_DOMAIN}/oauth2/v1/token"
+CUSTOM_AS_TOKEN_URL = f"https://{OKTA_DOMAIN}/oauth2/{CUSTOM_AS_ID}/v1/token"
+CUSTOM_AS_AUDIENCE = f"https://{OKTA_DOMAIN}/oauth2/{CUSTOM_AS_ID}"
+
+
+def build_client_assertion(audience: str) -> str:
+ """Sign a short-lived client assertion JWT for the given token endpoint."""
+ private_key = RSAAlgorithm.from_jwk(json.dumps(AGENT_PRIVATE_KEY_JWK))
+ now = int(time.time())
+ return jwt.encode(
+ {
+ "iss": AGENT_CLIENT_ID,
+ "sub": AGENT_CLIENT_ID,
+ "aud": audience, # must match the endpoint this assertion is sent to
+ "iat": now,
+ "exp": now + 300, # valid for 5 minutes
+ "jti": str(uuid.uuid4()),
+ },
+ private_key,
+ algorithm="RS256",
+ headers={"kid": AGENT_KEY_ID},
+ )
+
+
+def get_id_jag(id_token: str) -> str:
+ """Step 1: exchange the user's id_token for an ID-JAG at the org AS."""
+ r = requests.post(ORG_TOKEN_URL, data={
+ "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+ "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
+ "client_assertion": build_client_assertion(ORG_TOKEN_URL),
+ "subject_token": id_token,
+ "subject_token_type": "urn:ietf:params:oauth:token-type:id_token",
+ "requested_token_type": "urn:ietf:params:oauth:token-type:id-jag",
+ "scope": REQUESTED_SCOPE,
+ "audience": CUSTOM_AS_AUDIENCE,
+ }, timeout=10)
+ r.raise_for_status()
+ return r.json()["access_token"] # the ID-JAG
+
+
+def get_access_token(id_jag: str) -> str:
+ """Step 2: exchange the ID-JAG for a scoped access token at the custom AS."""
+ r = requests.post(CUSTOM_AS_TOKEN_URL, data={
+ "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+ "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
+ "client_assertion": build_client_assertion(CUSTOM_AS_TOKEN_URL),
+ "assertion": id_jag,
+ }, timeout=10)
+ r.raise_for_status()
+ return r.json()["access_token"] # scoped access token for the resourceA few details that this module encodes:
+ +build_client_assertion is called once per step, each time with the aud set to the token endpoint it targets: the org token URL for Step 1, and the custom authorization server token URL for Step 2. The kid header must match the public JWK registered on the agent.
+ audience parameter in Step 1 is the custom authorization server's issuer URL (https://{yourOktaDomain}/oauth2/{custom-as-id}), not its token endpoint.
+ +++ Note: For production workloads, cache the ID-JAG and access token in process until their
+expclaim expires. This avoids a fresh two-step exchange on every user request. +