agent-identity-python-sdk 0.1.1__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.
Files changed (20) hide show
  1. agent_identity_python_sdk-0.1.1/PKG-INFO +192 -0
  2. agent_identity_python_sdk-0.1.1/README.md +171 -0
  3. agent_identity_python_sdk-0.1.1/setup.cfg +4 -0
  4. agent_identity_python_sdk-0.1.1/setup.py +45 -0
  5. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/__init__.py +16 -0
  6. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/context/__init__.py +6 -0
  7. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/context/context.py +106 -0
  8. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/core/__init__.py +6 -0
  9. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/core/decorators.py +270 -0
  10. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/core/identity.py +400 -0
  11. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/model/__init__.py +11 -0
  12. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/model/stscredential.py +12 -0
  13. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/utils/__init__.py +0 -0
  14. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/utils/cache.py +67 -0
  15. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk/utils/config.py +70 -0
  16. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk.egg-info/PKG-INFO +192 -0
  17. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk.egg-info/SOURCES.txt +18 -0
  18. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk.egg-info/dependency_links.txt +1 -0
  19. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk.egg-info/requires.txt +11 -0
  20. agent_identity_python_sdk-0.1.1/src/agent_identity_python_sdk.egg-info/top_level.txt +1 -0
@@ -0,0 +1,192 @@
1
+ Metadata-Version: 2.1
2
+ Name: agent_identity_python_sdk
3
+ Version: 0.1.1
4
+ Summary: Python SDK for using Agent Identity Service
5
+ Home-page: https://github.com/aliyun/agent-identity-dev-kit
6
+ Author: shaoheng
7
+ Author-email: liuyuhao.lyh@alibaba-inc.com
8
+ License: Apache-2.0
9
+ Classifier: Development Status :: 4 - Beta
10
+ Classifier: Intended Audience :: Developers
11
+ Classifier: License :: OSI Approved :: Apache Software License
12
+ Classifier: Operating System :: OS Independent
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
+ Requires-Python: >=3.10
19
+ Description-Content-Type: text/markdown
20
+ Provides-Extra: dev
21
+
22
+ # Agent Identity Python SDK
23
+
24
+ The Agent Identity Python SDK is a Python development toolkit for accessing Agent Identity services. This SDK provides identity authentication, token management, API key acquisition, and more, supporting both synchronous and asynchronous invocation modes.
25
+
26
+ ## Features
27
+
28
+ - **OAuth2 Access Token Acquisition**: Supports multiple OAuth2 flows for access token retrieval
29
+ - **API Key Acquisition**: Programmatic API key retrieval
30
+ - **STS Credential Acquisition**: Obtain temporary security token service credentials
31
+ - **Context Management**: Thread-safe context variable management
32
+ - **Caching Mechanism**: Built-in credential caching for improved performance
33
+ - **Decorator Support**: Simplified authentication flow integration via decorators
34
+ - **Concurrency Safety**: Supports multi-threading and asynchronous environments
35
+
36
+ ## Installation
37
+
38
+ ```bash
39
+ pip install agent-identity-python-sdk
40
+ ```
41
+
42
+ ## Quick Start
43
+
44
+ ### Basic Configuration
45
+
46
+ Before using the SDK, ensure you have set the correct environment variables:
47
+
48
+ ```bash
49
+ export AGENT_IDENTITY_REGION_ID="cn-beijing" # Optional, set your Region ID, default is cn-beijing
50
+ export AGENT_IDENTITY_WORKLOAD_ACCESS_TOKEN="<your-workload-access-token>" # Optional, set your workload access token, if not specified, the Agent Identity service will be automatically called to obtain it
51
+ export AGENT_IDENTITY_WORKLOAD_IDENTITY_NAME="<your-workload-identity-name>" # Optional, set your workload identity name, if specified, the workload identity will be used as the agent's identity, otherwise a random workload identity will be generated
52
+ export AGENT_IDENTITY_USE_STS="true/false" # Optional, set whether to use the agent identity associated role for resource credential acquisition, default is true
53
+ ```
54
+
55
+ ### Using Decorators to Automatically Obtain Tokens
56
+
57
+ ```python
58
+ from agent_identity_python_sdk.core import requires_access_token
59
+
60
+ @requires_access_token(
61
+ credential_provider_name="your-provider-name",
62
+ inject_param_name="access_token",
63
+ auth_flow="USER_FEDERATION",
64
+ on_auth_url= lambda url: print(f"Please visit {url} to authenticate."),
65
+ scopes=["openid", "profile", "email"],
66
+ force_authentication=False,
67
+ callback_url="http://localhost:8080",
68
+ custom_parameters={
69
+ "custom_param_1": "value_1",
70
+ "custom_param_2": "value_2"
71
+ }
72
+ )
73
+ def my_function(access_token: str):
74
+ # Use access_token here
75
+ print(f"Access token: {access_token}")
76
+ # Your business logic
77
+
78
+ # Call the function
79
+ my_function()
80
+ ```
81
+
82
+ ### Using Decorators to Obtain API Keys
83
+
84
+ ```python
85
+ from agent_identity_python_sdk.core.decorators import requires_api_key
86
+
87
+ @requires_api_key(credential_provider_name="your-provider-name", inject_param_name="api_key")
88
+ def my_function(api_key: str):
89
+ # Use api_key here
90
+ print(f"API key: {api_key}")
91
+ # Your business logic
92
+
93
+ # Call the function
94
+ my_function()
95
+ ```
96
+
97
+ ### Using Decorators to Obtain STS Credentials
98
+
99
+ ```python
100
+ from agent_identity_python_sdk.core.decorators import requires_sts_token
101
+ from agent_identity_python_sdk.model.stscredential import STSCredential
102
+
103
+ @requires_sts_token(inject_param_name="sts_credential")
104
+ def my_function(sts_credential: STSCredential):
105
+ # Use sts_credential here
106
+ print(f"STS Access Key ID: {sts_credential.access_key_id}")
107
+ # Your business logic
108
+
109
+ # Call the function
110
+ my_function()
111
+ ```
112
+
113
+ ## Core Modules
114
+
115
+ ### IdentityClient
116
+
117
+ IdentityClient is the core identity management client that provides methods for creating and managing identities and acquiring various types of credentials.
118
+
119
+ ```python
120
+ from agent_identity_python_sdk.core.identity import IdentityClient
121
+
122
+ client = IdentityClient(region_id="cn-beijing")
123
+
124
+ # Create workload identity
125
+ workload_identity_name = client.create_workload_identity(
126
+ workload_identity_name="my-workload",
127
+ allowed_resource_oauth2_return_urls=["https://example.com/callback"],
128
+ role_arn="acs:ram::12****:role/example-role",
129
+ )
130
+
131
+ # Get workload access token
132
+ token = client.get_workload_access_token(
133
+ workload_name=workload_identity_name,
134
+ user_token="ejwyJ9***",
135
+ user_id="example-user"
136
+ ) # Prioritizes using user_token to obtain workload access token; if not available, uses user_id to obtain workload access token; if both are absent, obtains workload access token without end-user information
137
+ ```
138
+
139
+ ### Context Management
140
+
141
+ The SDK provides context managers for storing thread/async task isolated data:
142
+
143
+ #### AgentIdentityContext
144
+
145
+ Used to manage workload access tokens, user ID, user tokens, session ID, etc. The SDK will read the current thread's context variables to obtain the workload access token when acquiring workload access tokens.
146
+
147
+ ```python
148
+ from agent_identity_python_sdk.context.context import AgentIdentityContext
149
+
150
+ # Set workload access token
151
+ AgentIdentityContext.set_workload_access_token("your-token")
152
+
153
+ # Get workload access token
154
+ token = AgentIdentityContext.get_workload_access_token()
155
+
156
+ # Set user token
157
+ AgentIdentityContext.set_user_token("your-token")
158
+
159
+ # Set user ID
160
+ AgentIdentityContext.set_user_id("user-123")
161
+
162
+ # Set custom state
163
+ AgentIdentityContext.set_custom_state("your-state")
164
+
165
+ # Clear current thread context, needs to be actively cleared at the end of a single session, otherwise permission leakage may occur due to thread sharing
166
+ AgentIdentityContext.clear()
167
+ ```
168
+
169
+ If a workload access token is set in the context, it will be prioritized when acquiring workload access tokens from the current thread context variables.
170
+
171
+ If no workload access token is set, when the SDK automatically acquires a workload access token, it will retrieve user ID/user token information from the current thread context variables, following the **user token/user ID/none** priority to acquire the workload access token.
172
+
173
+ If a custom state is set, during OAuth2 authorization, the custom state will be passed along. User applications can use the custom state to handle authorization callbacks and perform verification. It is recommended that applications use custom state for session verification to prevent malicious sharing of authorization links to obtain other users' permissions.
174
+
175
+ ⚠️ **Note**: After the current workflow execution is completed, you need to actively clear the current thread context, otherwise permission leakage may occur due to thread sharing.
176
+
177
+ ## Environment Variables Configuration
178
+
179
+ | Variable Name | Description | Default Value |
180
+ |---------------|-------------|---------------|
181
+ | AGENT_IDENTITY_REGION_ID | Region identifier | cn-beijing |
182
+ | AGENT_IDENTITY_WORKLOAD_ACCESS_TOKEN | Workload identity token | None |
183
+ | AGENT_IDENTITY_WORKLOAD_IDENTITY_NAME | Workload identity name | None |
184
+ | AGENT_IDENTITY_USE_STS | Whether to use STS for resource credential acquisition | true |
185
+
186
+ ## Contributing
187
+
188
+ Issues and Pull Requests are welcome to help improve this SDK.
189
+
190
+ ## License
191
+
192
+ This project is licensed under the Apache-2.0 License. See the [LICENSE](../LICENSE) file for details.
@@ -0,0 +1,171 @@
1
+ # Agent Identity Python SDK
2
+
3
+ The Agent Identity Python SDK is a Python development toolkit for accessing Agent Identity services. This SDK provides identity authentication, token management, API key acquisition, and more, supporting both synchronous and asynchronous invocation modes.
4
+
5
+ ## Features
6
+
7
+ - **OAuth2 Access Token Acquisition**: Supports multiple OAuth2 flows for access token retrieval
8
+ - **API Key Acquisition**: Programmatic API key retrieval
9
+ - **STS Credential Acquisition**: Obtain temporary security token service credentials
10
+ - **Context Management**: Thread-safe context variable management
11
+ - **Caching Mechanism**: Built-in credential caching for improved performance
12
+ - **Decorator Support**: Simplified authentication flow integration via decorators
13
+ - **Concurrency Safety**: Supports multi-threading and asynchronous environments
14
+
15
+ ## Installation
16
+
17
+ ```bash
18
+ pip install agent-identity-python-sdk
19
+ ```
20
+
21
+ ## Quick Start
22
+
23
+ ### Basic Configuration
24
+
25
+ Before using the SDK, ensure you have set the correct environment variables:
26
+
27
+ ```bash
28
+ export AGENT_IDENTITY_REGION_ID="cn-beijing" # Optional, set your Region ID, default is cn-beijing
29
+ export AGENT_IDENTITY_WORKLOAD_ACCESS_TOKEN="<your-workload-access-token>" # Optional, set your workload access token, if not specified, the Agent Identity service will be automatically called to obtain it
30
+ export AGENT_IDENTITY_WORKLOAD_IDENTITY_NAME="<your-workload-identity-name>" # Optional, set your workload identity name, if specified, the workload identity will be used as the agent's identity, otherwise a random workload identity will be generated
31
+ export AGENT_IDENTITY_USE_STS="true/false" # Optional, set whether to use the agent identity associated role for resource credential acquisition, default is true
32
+ ```
33
+
34
+ ### Using Decorators to Automatically Obtain Tokens
35
+
36
+ ```python
37
+ from agent_identity_python_sdk.core import requires_access_token
38
+
39
+ @requires_access_token(
40
+ credential_provider_name="your-provider-name",
41
+ inject_param_name="access_token",
42
+ auth_flow="USER_FEDERATION",
43
+ on_auth_url= lambda url: print(f"Please visit {url} to authenticate."),
44
+ scopes=["openid", "profile", "email"],
45
+ force_authentication=False,
46
+ callback_url="http://localhost:8080",
47
+ custom_parameters={
48
+ "custom_param_1": "value_1",
49
+ "custom_param_2": "value_2"
50
+ }
51
+ )
52
+ def my_function(access_token: str):
53
+ # Use access_token here
54
+ print(f"Access token: {access_token}")
55
+ # Your business logic
56
+
57
+ # Call the function
58
+ my_function()
59
+ ```
60
+
61
+ ### Using Decorators to Obtain API Keys
62
+
63
+ ```python
64
+ from agent_identity_python_sdk.core.decorators import requires_api_key
65
+
66
+ @requires_api_key(credential_provider_name="your-provider-name", inject_param_name="api_key")
67
+ def my_function(api_key: str):
68
+ # Use api_key here
69
+ print(f"API key: {api_key}")
70
+ # Your business logic
71
+
72
+ # Call the function
73
+ my_function()
74
+ ```
75
+
76
+ ### Using Decorators to Obtain STS Credentials
77
+
78
+ ```python
79
+ from agent_identity_python_sdk.core.decorators import requires_sts_token
80
+ from agent_identity_python_sdk.model.stscredential import STSCredential
81
+
82
+ @requires_sts_token(inject_param_name="sts_credential")
83
+ def my_function(sts_credential: STSCredential):
84
+ # Use sts_credential here
85
+ print(f"STS Access Key ID: {sts_credential.access_key_id}")
86
+ # Your business logic
87
+
88
+ # Call the function
89
+ my_function()
90
+ ```
91
+
92
+ ## Core Modules
93
+
94
+ ### IdentityClient
95
+
96
+ IdentityClient is the core identity management client that provides methods for creating and managing identities and acquiring various types of credentials.
97
+
98
+ ```python
99
+ from agent_identity_python_sdk.core.identity import IdentityClient
100
+
101
+ client = IdentityClient(region_id="cn-beijing")
102
+
103
+ # Create workload identity
104
+ workload_identity_name = client.create_workload_identity(
105
+ workload_identity_name="my-workload",
106
+ allowed_resource_oauth2_return_urls=["https://example.com/callback"],
107
+ role_arn="acs:ram::12****:role/example-role",
108
+ )
109
+
110
+ # Get workload access token
111
+ token = client.get_workload_access_token(
112
+ workload_name=workload_identity_name,
113
+ user_token="ejwyJ9***",
114
+ user_id="example-user"
115
+ ) # Prioritizes using user_token to obtain workload access token; if not available, uses user_id to obtain workload access token; if both are absent, obtains workload access token without end-user information
116
+ ```
117
+
118
+ ### Context Management
119
+
120
+ The SDK provides context managers for storing thread/async task isolated data:
121
+
122
+ #### AgentIdentityContext
123
+
124
+ Used to manage workload access tokens, user ID, user tokens, session ID, etc. The SDK will read the current thread's context variables to obtain the workload access token when acquiring workload access tokens.
125
+
126
+ ```python
127
+ from agent_identity_python_sdk.context.context import AgentIdentityContext
128
+
129
+ # Set workload access token
130
+ AgentIdentityContext.set_workload_access_token("your-token")
131
+
132
+ # Get workload access token
133
+ token = AgentIdentityContext.get_workload_access_token()
134
+
135
+ # Set user token
136
+ AgentIdentityContext.set_user_token("your-token")
137
+
138
+ # Set user ID
139
+ AgentIdentityContext.set_user_id("user-123")
140
+
141
+ # Set custom state
142
+ AgentIdentityContext.set_custom_state("your-state")
143
+
144
+ # Clear current thread context, needs to be actively cleared at the end of a single session, otherwise permission leakage may occur due to thread sharing
145
+ AgentIdentityContext.clear()
146
+ ```
147
+
148
+ If a workload access token is set in the context, it will be prioritized when acquiring workload access tokens from the current thread context variables.
149
+
150
+ If no workload access token is set, when the SDK automatically acquires a workload access token, it will retrieve user ID/user token information from the current thread context variables, following the **user token/user ID/none** priority to acquire the workload access token.
151
+
152
+ If a custom state is set, during OAuth2 authorization, the custom state will be passed along. User applications can use the custom state to handle authorization callbacks and perform verification. It is recommended that applications use custom state for session verification to prevent malicious sharing of authorization links to obtain other users' permissions.
153
+
154
+ ⚠️ **Note**: After the current workflow execution is completed, you need to actively clear the current thread context, otherwise permission leakage may occur due to thread sharing.
155
+
156
+ ## Environment Variables Configuration
157
+
158
+ | Variable Name | Description | Default Value |
159
+ |---------------|-------------|---------------|
160
+ | AGENT_IDENTITY_REGION_ID | Region identifier | cn-beijing |
161
+ | AGENT_IDENTITY_WORKLOAD_ACCESS_TOKEN | Workload identity token | None |
162
+ | AGENT_IDENTITY_WORKLOAD_IDENTITY_NAME | Workload identity name | None |
163
+ | AGENT_IDENTITY_USE_STS | Whether to use STS for resource credential acquisition | true |
164
+
165
+ ## Contributing
166
+
167
+ Issues and Pull Requests are welcome to help improve this SDK.
168
+
169
+ ## License
170
+
171
+ This project is licensed under the Apache-2.0 License. See the [LICENSE](../LICENSE) file for details.
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,45 @@
1
+ from setuptools import setup, find_packages
2
+
3
+ with open("README.md", encoding="utf-8") as f:
4
+ long_description = f.read()
5
+
6
+ setup(
7
+ name="agent_identity_python_sdk",
8
+ version="0.1.1",
9
+ description="Python SDK for using Agent Identity Service",
10
+ long_description=long_description,
11
+ long_description_content_type="text/markdown",
12
+ license="Apache-2.0",
13
+ author="shaoheng",
14
+ author_email="liuyuhao.lyh@alibaba-inc.com",
15
+ url="https://github.com/aliyun/agent-identity-dev-kit",
16
+ packages=find_packages(where="src"),
17
+ package_dir={"": "src"},
18
+ python_requires=">=3.10",
19
+ install_requires=[
20
+ "alibabacloud-agentidentity20250901>=1.0.1",
21
+ "alibabacloud-agentidentitydata20251127>=1.0.2",
22
+ "setuptools",
23
+ "pydantic==2.11.7",
24
+ "urllib3==2.3.0",
25
+ "utils==1.0.2",
26
+ ],
27
+ extras_require={
28
+ "dev": [
29
+ "pytest>=8.4.1",
30
+ "pytest-asyncio>=0.24.0",
31
+ "pytest-cov>=6.0.0",
32
+ ],
33
+ },
34
+ classifiers=[
35
+ "Development Status :: 4 - Beta",
36
+ "Intended Audience :: Developers",
37
+ "License :: OSI Approved :: Apache Software License",
38
+ "Operating System :: OS Independent",
39
+ "Programming Language :: Python :: 3",
40
+ "Programming Language :: Python :: 3.10",
41
+ "Programming Language :: Python :: 3.11",
42
+ "Programming Language :: Python :: 3.12",
43
+ "Programming Language :: Python :: 3.13",
44
+ ],
45
+ )
@@ -0,0 +1,16 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ __version__ = "0.1.0"
4
+
5
+ from .context import AgentIdentityContext
6
+ from .core import requires_access_token, requires_sts_token, requires_api_key
7
+ from .core import IdentityClient
8
+
9
+
10
+ __all__ = [
11
+ "IdentityClient",
12
+ "requires_access_token",
13
+ "requires_sts_token",
14
+ "requires_api_key",
15
+ "AgentIdentityContext"
16
+ ]
@@ -0,0 +1,6 @@
1
+ """Agent Identity Runtime Package.
2
+ """
3
+
4
+ from .context import AgentIdentityContext
5
+
6
+ __all__ = ["AgentIdentityContext"]
@@ -0,0 +1,106 @@
1
+ import os
2
+ from contextvars import ContextVar
3
+ from typing import Optional
4
+
5
+ class AgentIdentityContext:
6
+ """
7
+ AgentIdentityContext is a context management class used to store user-related information in concurrent environments.
8
+ It utilizes contextvars to ensure data isolation between threads/asynchronous tasks.
9
+
10
+ This class is primarily used to store the following pieces of information:
11
+ 1. user_id: Unique identifier for the user, used to identify the current operating user
12
+ 2. user_token: User access token, used for authentication and authorization
13
+ 3. custom_state: Custom state parameter, used in OAuth2 flow to prevent CSRF attacks and verify request sources during callback
14
+ 4. workload_access_token: Token for accessing workload resources, which can be retrieved from context or environment variable
15
+ 5. session_id: Unique identifier for the session, used to track and manage user sessions
16
+
17
+ These pieces of information are isolated within threads, allowing safe usage in asynchronous operations or multi-threaded environments
18
+ without risk of data confusion.
19
+ """
20
+
21
+ # Session ID context variable, used to pass session identifier within request chain
22
+ _session_id: ContextVar[Optional[str]] = ContextVar("session_id")
23
+
24
+ # User ID context variable, used to pass user id within request chain
25
+ _user_id: ContextVar[Optional[str]] = ContextVar("user_id")
26
+
27
+ # User token context variable, used to pass user token within request chain
28
+ _user_token: ContextVar[Optional[str]] = ContextVar("user_token")
29
+
30
+ # Custom state context variable, used to pass custom state within request chain.
31
+ # The custom state will be carried when OAuth2 authorization is successful and redirected to the user application.
32
+ # It is recommended that the user application perform ownership verification of the callback login status and state
33
+ # to prevent the link from being maliciously disseminated and causing unauthorized access.
34
+ _custom_state: ContextVar[Optional[str]] = ContextVar("custom_state")
35
+
36
+ # Workload access token context variable, used to pass workload authentication token within request chain
37
+ # The workload access token will be retrieved from this context first, if not present,
38
+ # it will read from environment variable. This allows the platform to preset the workload
39
+ # access token to the agent execution environment in certain scenarios. If neither exists,
40
+ # the Agent Identity SDK will automatically acquire it for the client based on the current context.
41
+ _workload_access_token: ContextVar[Optional[str]] = ContextVar("workload_access_token")
42
+
43
+ @classmethod
44
+ def set_user_id(cls, user_id: str):
45
+ # Set the user ID in the context
46
+ cls._user_id.set(user_id)
47
+
48
+ @classmethod
49
+ def get_user_id(cls) -> Optional[str]:
50
+ # Get the user ID from context
51
+ try:
52
+ return cls._user_id.get()
53
+ except LookupError:
54
+ return None
55
+
56
+ @classmethod
57
+ def set_user_token(cls, token: str):
58
+ # Set the user token in the context
59
+ cls._user_token.set(token)
60
+
61
+ @classmethod
62
+ def get_user_token(cls) -> Optional[str]:
63
+ # Get the user token from context
64
+ try:
65
+ return cls._user_token.get()
66
+ except LookupError:
67
+ return None
68
+
69
+ @classmethod
70
+ def set_custom_state(cls, state: str):
71
+ # Set the custom state in the context
72
+ cls._custom_state.set(state)
73
+
74
+ @classmethod
75
+ def get_custom_state(cls) -> Optional[str]:
76
+ # Get the custom state from context
77
+ try:
78
+ return cls._custom_state.get()
79
+ except LookupError:
80
+ return None
81
+
82
+ @classmethod
83
+ def set_workload_access_token(cls, token: str):
84
+ # Set the workload access token in the context
85
+ cls._workload_access_token.set(token)
86
+
87
+ @classmethod
88
+ def get_workload_access_token(cls) -> Optional[str]:
89
+ # Get the workload access token from context or environment variable
90
+ try:
91
+ workload_access_token = cls._workload_access_token.get()
92
+ if workload_access_token is None:
93
+ return os.environ.get("AGENT_IDENTITY_WORKLOAD_ACCESS_TOKEN", None)
94
+ return workload_access_token # Return the context value if it's not None
95
+ except LookupError:
96
+ return None
97
+
98
+ @classmethod
99
+ def clear(cls):
100
+ # Clear all context variables
101
+ cls._session_id.set(None)
102
+ cls._user_id.set(None)
103
+ cls._user_token.set(None)
104
+ cls._custom_state.set(None)
105
+ cls._workload_access_token.set(None)
106
+
@@ -0,0 +1,6 @@
1
+ """Agent identity core package."""
2
+
3
+ from .decorators import requires_access_token, requires_api_key, requires_sts_token
4
+ from .identity import IdentityClient
5
+
6
+ __all__ = ["requires_access_token", "requires_api_key", "requires_sts_token", "IdentityClient"]