Secure a Google Vertex AI agent

Identity Engine

This guide shows you how to build a Python wrapper that authenticates users with Okta, performs Okta's token exchange, and then calls a Google Vertex AI agent. In the Vertex AI, an AI agent is a Reasoning Engine (the resource type behind Vertex AI Agent Engine and Agent Builder). Your app owns the full flow: It verifies who the user is, exchanges that identity for a scoped access token, and uses that token to create a session with the Reasoning Engine. It then sends the user's prompt and polls for the AI agent's response.

The Okta authentication is a two-step token exchange that's the same for any AI agent, regardless of the platform it runs on. This guide first introduces what the integration needs to do and provides sample code functions that implement the authentication. It then shows the Google Vertex AI-specific code and the configuration that consumes it.

Note: To enable AI agent token exchange, you must first subscribe to Okta for AI Agents. Contact your Okta account team to enable the feature.

Learning outcomes

Understand what an imported AI agent must do to authenticate as a signed-in user with Okta.
Add a token exchange module to your AI agent.
Create a session with a Vertex AI Reasoning Engine, send a prompt, and poll for the AI agent's response.
Verify and test the end-to-end flow with a real Okta ID token.

What you need

An Identity Engine org with the Okta for AI Agents feature enabled
A Google Cloud project with Vertex AI enabled and an existing Reasoning Engine (Vertex AI Agent Engine) that you can call
A Google account with the Editor role or the AI Platform Editor (aiplatform.editor) role on that project
An existing Google Workspace app integration in your Okta org. The Google Vertex AI import configuration lives on this app's AI Agent Import tab.
The Google Vertex AI agent imported into Okta as an AI agent identity. See Import your AI agent from Google Vertex AI.
Python (opens new window) 3.10 or later

Overview

An AI agent has no inherent knowledge of an Okta user. To let it act for a specific user without sharing long-lived credentials, the AI agent exchanges the user's identity for a short-lived, narrowly scoped access token, and then uses that token to call protected resources.

The integration has two parts:

Okta authentication. The AI agent performs a two-step token exchange:
1. Exchange the user's id_token for an identity assertion JWT authorization grant (ID-JAG) at the org authorization server.
2. Exchange the ID-JAG for a scoped access_token at a custom authorization server.
This logic is identical for any AI agent. You add it once as a reusable module. See Add Okta authentication to your AI agent.
Platform integration (Google Vertex AI-specific). The Vertex AI Reasoning Engine API is session- and event-based, not a single synchronous call. Your wrapper creates a session, appends the user's prompt to it as an event, and polls the session's events until the AI agent responds. See Integrate the token exchange into your Vertex AI agent.

For the conceptual background on AI agent token exchange, see Set up AI agent token exchange.

Before you begin

The token exchange depends on Okta objects that you configure once per org. Confirm that the following are in place before you add any integration code. For detailed steps, see Set up imported AI agent token exchange.

An OIDC web app integration that signs users in and issues the id_token that your AI agent exchanges. Use the Authorization Code grant type and the openid profile email scopes. The id_token must have an aud claim equal to this app's client ID.
A custom authorization server. Use the built-in default server or create one.
A custom scope on the custom authorization server, such as xaa:read.
The Google Vertex AI agent imported into Okta as an AI agent identity that uses private_key_jwt client authentication, with its public key (JWK) registered. Link the OIDC web app, set the custom authorization server, include your custom scope, and activate the AI agent. See Import your AI agent from Google Vertex AI.

Note: Okta doesn't retain the AI agent's private key. Store it in a secrets manager when it's generated, because it's shown only once.
An access policy rule on the custom authorization server that enables the JWT bearer grant type (urn:ietf:params:oauth:grant-type:jwt-bearer), adds the AI agent as an allowed client, and includes the audience, the custom scope, and a user or group condition.

Import your AI agent from Google Vertex AI

Importing a Google Vertex AI agent works differently than importing an AI agent from the other supported platforms. Instead of registering the AI agent directly with its own key pair, Okta connects to your Google Cloud project through OAuth 2.0 and its Security Token Service (STS) and discovers and imports Reasoning Engines from that project. This setup lives on your org's Google Workspace app integration, not under Directory > AI Agents.

You need an existing Google Workspace app integration in your org (Applications > Applications & Resources). If you don't have one yet, create it before you continue.

In the Google Cloud Console (opens new window), select your target project and then go to APIs & Services > OAuth consent screen.

Important: Your Google account needs the Editor role or the AI Platform Editor (aiplatform.editor) role on the project to complete this configuration.
Go to the Clients tab, and then click + Create client.
Set Application type to Web application and enter a descriptive name, for example, Vertex-Agent-Import-Client.
Under Authorized redirect URIs, click Add URI, and enter {yourOktaDomain}/oauth2/v1/sts/callback.
Click Create. Note the Client ID and Client Secret that are shown in the confirmation dialog. Okta uses these credentials to discover and import your Reasoning Engines, so store them in a secrets manager.
Go to the Audience tab.
- If the consent page's User type is External, scroll to Test users, click + Add users, and add the email addresses of the accounts that complete the import consent flow.
- If the User type is Internal, no additional provisioning is required.
In the Admin Console, go to your org's Google Workspace app integration, and open its AI Agent Import tab.
Enter the client ID, client secret, your Google Cloud project ID, and the location (region) where your Reasoning Engines run, for example, us-west1. Click Test API Credentials to validate them.

Note: Okta validates these credentials against Google. If Google requires other consent, Okta returns an interaction URI. You must open that URI and complete the consent flow before validation can succeed.

Important: Only use a location value that's a genuine Google Cloud region for your project (for example, us-west1). Okta uses this value to build the request URL it calls with your Google access token attached.
Save the configuration, and then trigger the import. Select the Reasoning Engine to import from the discovered list, and complete the import.

Collect your configuration values

Your Google Vertex AI agent code reads these values as environment variables. The token exchange module consumes the first group. The second group is specific to Google Vertex AI.

Okta values (used by the token exchange):

Environment variable	Description	Where to find it
`OKTA_DOMAIN`	Okta org domain, for example `example.okta.com` (no `https://` prefix)	Admin Console > Settings > Account
`OKTA_CUSTOM_AS_ID`	Custom authorization server ID, for example `default`	Security > API
`OKTA_SCOPE`	The custom scope that the agent requests	Custom AS > Scopes
`AGENT_CLIENT_ID`	Client ID of the imported AI agent, for example `wlp9k6...`	Directory > AI Agents > (yourAgent) > Client registration
`AGENT_KEY_ID`	`kid` of the public JWK registered on the AI agent	Directory > AI Agents > (yourAgent) > Client registration > Public/private key
`AGENT_PRIVATE_KEY_JWK`	The AI agent's private JWK (single-line JSON)	Directory > AI Agents > (yourAgent) > Client registration > Public/private key. Output of Generate new key. Store the value in a secrets manager

Google Vertex AI values (used by the platform integration):

Environment variable	Description	Where to find it
`GCP_PROJECT_ID`	The Google Cloud project that hosts your Reasoning Engine, for example, `imports-testing`	Google Cloud Console > project selector
`GCP_LOCATION`	The region where your Reasoning Engine runs, for example, `us-west1`	Google Cloud Console > Vertex AI > Agent Engine
`GCP_REASONING_ENGINE_ID`	The numeric ID of the Reasoning Engine to call	The trailing segment of the engine's `name` field, for example, `8411791718067732480` in `projects/{project}/locations/{location}/reasoningEngines/8411791718067732480`

Note: These values identify the Reasoning Engine that you're calling. They're separate from the client ID and client secret that you created in Import your AI agent from Google Vertex AI, which Okta uses only to discover and import AI agents, not at runtime.

Add Okta authentication to your AI agent

The following example token_exchange.py module that you create here has no dependency on Google Cloud or Vertex AI.

Install the token exchange dependencies

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 the token exchange module

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 resource

A few details that this module encodes:

The client assertion function is invoked twice. 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.
The audience parameter in Step 1 is the custom authorization server's issuer URL (https://{yourOktaDomain}/oauth2/{custom-as-id}), not its token endpoint.
Step 1 requires the Okta imported AI Agent client. An OIDC app client can't perform this exchange.

Note: For production workloads, cache the ID-JAG and access token in process until their exp claim expires. This avoids a fresh two-step exchange on every user request.

Integrate the token exchange into your Vertex AI agent

This section is specific to Google Vertex AI. Here you call get_id_jag and get_access_token from your AI agent, then use the resulting access token to create a session with the Reasoning Engine, send the user's prompt, and poll for a response.

Note: The following examples pass the Okta-issued access_token directly to the Vertex AI REST API as a bearer credential. Confirm this against your own Reasoning Engine setup. If your org instead requires exchanging the Okta token for a Google Cloud-native token (for example, through Workforce Identity Federation), add that exchange before the calls that are shown here.

Add the Vertex AI dependencies

Add these to the same requirements.txt, alongside the token exchange dependencies:

requests
python-dotenv

Install the complete set of dependencies:

pip install -r requirements.txt

Create a session

Start a session with the Reasoning Engine. The response includes a name field. The last path segment is the session ID that you use in the following steps.

import os
import requests

GCP_PROJECT_ID = os.environ["GCP_PROJECT_ID"]
GCP_LOCATION = os.environ["GCP_LOCATION"]
GCP_REASONING_ENGINE_ID = os.environ["GCP_REASONING_ENGINE_ID"]

BASE_URL = (
    f"https://{GCP_LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{GCP_PROJECT_ID}"
    f"/locations/{GCP_LOCATION}/reasoningEngines/{GCP_REASONING_ENGINE_ID}"
)


def create_session(access_token: str, user_id: str) -> str:
    resp = requests.post(
        f"{BASE_URL}/sessions",
        headers={"Authorization": f"Bearer {access_token}"},
        json={"userId": user_id},
    )
    resp.raise_for_status()
    # "name" looks like ".../reasoningEngines/{id}/sessions/{session_id}"
    return resp.json()["name"].rsplit("/", 1)[-1]

Note: The userId identifies the session owner to the Reasoning Engine. Use a stable identifier for the signed-in user, such as the sub claim from their id_token.

Send the prompt

Append the user's prompt to the session as a USER event.

def send_prompt(access_token: str, session_id: str, prompt: str) -> None:
    resp = requests.post(
        f"{BASE_URL}/sessions/{session_id}:appendEvent",
        headers={"Authorization": f"Bearer {access_token}"},
        json={
            "author": "USER",
            "rawEvent": {"text": prompt},
        },
    )
    resp.raise_for_status()

Poll for the response

The Reasoning Engine responds asynchronously. Poll the session's events and scan them for the first one that the AI agent authors.

import time


def get_agent_response(access_token: str, session_id: str, timeout_seconds: int = 30) -> str:
    deadline = time.time() + timeout_seconds
    while time.time() < deadline:
        resp = requests.get(
            f"{BASE_URL}/sessions/{session_id}/events",
            headers={"Authorization": f"Bearer {access_token}"},
        )
        resp.raise_for_status()
        events = resp.json().get("sessionEvents", [])
        for event in reversed(events):
            if event.get("author") in ("AGENT", "ASSISTANT"):
                return event.get("rawEvent", {}).get("text", "")
        time.sleep(1)
    raise TimeoutError("Timed out waiting for a response from the Reasoning Engine.")

Note: Confirm the exact shape of an AI agent-authored event (whether the response text is always under rawEvent.text) against a real response from your Reasoning Engine. Also confirm the sessionEvents wrapper key on the list-events response. Current Google documentation describes this operation through the Python SDK's events.list() method and doesn't surface a wrapper field name to check against.

Wire it into the entry point

In your app's entry point, call the two token exchange functions in order, and then create the session, send the prompt, and poll for the response. The following example main_vertex.py imports the reusable token exchange module and adds only the Vertex AI-specific wiring:

import jwt

from token_exchange import get_id_jag, get_access_token
# create_session, send_prompt, get_agent_response from the previous steps


def ask_vertex_agent(id_token: str, prompt: str) -> str:
    # Okta authentication
    id_jag = get_id_jag(id_token)
    access_token = get_access_token(id_jag)

    # The id_token was already verified by the org authorization server in
    # Step 1. Decoding it here only reads the sub claim to identify the
    # session owner. It isn't used to make an authorization decision.
    user_claims = jwt.decode(id_token, options={"verify_signature": False})

    # Platform integration (Google Vertex AI)
    session_id = create_session(access_token, user_claims["sub"])
    send_prompt(access_token, session_id, prompt)
    return get_agent_response(access_token, session_id)


if __name__ == "__main__":
    id_token = input("Paste a test Okta id_token: ").strip()
    prompt = input("💬 Enter custom prompt for vertex ai: ").strip()
    print(ask_vertex_agent(id_token, prompt))

Verify the configuration

After you add the code, verify the Okta-side configuration:

Go to Directory > AI Agents and confirm that the agent appears with Status: Active and the expected owners, connections, and user app.
(Optional) Go to Identity Governance > Access Certifications to confirm that the agent's user sign-on app is visible for future certification campaigns.

Obtain a test ID token

To exercise the flow, you need an ID token from the OIDC app that's linked to the AI agent. Complete an OIDC sign-in against that app to obtain one. For an example Authorization Code sign-in helper that authenticates with private_key_jwt, see Create an app to obtain a test ID token.

Note: Add the helper's callback URL (for example, http://localhost:5000/callback) to the linked OIDC app's Sign-in redirect URIs before you run it, and remove it after verification is complete.

Run an end-to-end invocation

Run the entry point locally, passing a test ID token to confirm the full id_token → ID-JAG → access_token → Vertex AI round trip:

python3 main_vertex.py

A successful run looks like this, and confirms the full round trip:

Paste a test Okta id_token: eyJraWQiOiJoZnpMS3...
💬 Enter custom prompt for vertex ai: Hello
Hello! How can I help you today?

Troubleshoot your integration

The following errors come from the Okta token exchange and are covered in Set up imported AI agent token exchange: Troubleshooting:

invalid_scope: openid not allowed
invalid_client: JWKSet not configured
invalid_client: kid is invalid
access_denied: no_matching_policy
Only service apps can use client_credentials

Next steps

Your AI agent can now authenticate as a user and call Okta-protected resources on their behalf. To define the resources and scopes you permit the AI agent to reach, see Set up AI agent token exchange and the Okta for AI Agents documentation on governing access to AI agents.

On this page