coze-workload-identity 0.1.1

Python SDK for Coze workload identity authentication

Coze Workload Identity SDK

A Python SDK for Coze workload identity authentication using OAuth2.0 token exchange.

Features

• 🔐 Secure OAuth2.0 token exchange flow
• 🚀 Thread-safe token caching with automatic expiration
• 📝 Comprehensive error handling
• 🧪 Fully tested with high coverage
• 📦 Easy to install and use
• 🏊‍♂️ Lane support (泳道) for different environments (BOE, PPE, custom)

Installation

pip install coze-workload-identity

Quick Start

1. Set up environment variables:

export COZE_WORKLOAD_IDENTITY_CLIENT_ID="your_client_id"
export COZE_WORKLOAD_IDENTITY_CLIENT_SECRET="your_client_secret"
export COZE_WORKLOAD_IDENTITY_TOKEN_ENDPOINT="https://auth.example.com/token"
export COZE_WORKLOAD_ACCESS_TOKEN_ENDPOINT="https://auth.example.com/access-token"

2. Use the SDK:

from coze_workload_identity import Client

# Create client
client = Client()

# Get access token
token = client.get_access_token()
print(f"Access token: {token}")

# Use as context manager
with Client() as client:
    token = client.get_access_token()
    # Use the token for API calls

Configuration

The SDK requires the following environment variables:

Variable	Description	Required
COZE_WORKLOAD_IDENTITY_CLIENT_ID	Client ID for authentication	✅
COZE_WORKLOAD_IDENTITY_CLIENT_SECRET	Client secret for authentication	✅
COZE_WORKLOAD_IDENTITY_TOKEN_ENDPOINT	Token endpoint for ID token	✅
COZE_WORKLOAD_ACCESS_TOKEN_ENDPOINT	Access token endpoint for token exchange	✅
COZE_OUTBOUND_AUTH_ENDPOINT	Integration credential endpoint (required for get_integration_credential)	❌
COZE_SERVER_ENV	Lane environment (optional, default: NONE)	❌

Lane Support (泳道支持)

The SDK supports lane environments through the COZE_SERVER_ENV environment variable:

• NONE (default): No lane headers are added
• boe_*: Adds x-tt-env: boe_<lane> header
• ppe_*: Adds x-tt-env: ppe_<lane> and x-use-ppe: 1 headers
• custom: Adds x-tt-env: <custom> header

Examples:

# BOE lane
export COZE_SERVER_ENV="boe_test_lane"

# PPE lane
export COZE_SERVER_ENV="ppe_production_lane"

# Custom lane
export COZE_SERVER_ENV="my_custom_lane"

API Reference

Client

The main class for interacting with the workload identity service.

Constructor

Client()

Creates a new client instance. Configuration is loaded from environment variables.

Methods

get_access_token()

Retrieves an access token using the OAuth2.0 token exchange flow.

Returns: str - The access token

Raises:

• ConfigurationError - If required configuration is missing
• TokenRetrievalError - If ID token retrieval fails
• TokenExchangeError - If token exchange fails

Example:

try:
    token = client.get_access_token()
    # Use the token for API calls
except ConfigurationError as e:
    print(f"Configuration error: {e}")
except TokenRetrievalError as e:
    print(f"Token retrieval error: {e}")
except TokenExchangeError as e:
    print(f"Token exchange error: {e}")

close()

Closes the client and cleans up resources.

Example:

client = Client()
try:
    token = client.get_access_token()
finally:
    client.close()

get_integration_credential(integration_name)

Retrieves integration credentials for the specified integration name.

Args:

• integration_name (str): Name of the integration

Returns: str - The integration credential

Raises:

• ConfigurationError - If COZE_OUTBOUND_AUTH_ENDPOINT is not configured
• TokenRetrievalError - If access token retrieval fails
• Exception - If integration credential API returns non-200 status or invalid response

Example:

try:
    credential = client.get_integration_credential("my_integration")
    # Use the credential for integration API calls
except ConfigurationError as e:
    print(f"Configuration error: {e}")
except Exception as e:
    print(f"Integration error: {e}")

API Response Format: The integration credential endpoint should return:

{
    "code": 0,
    "msg": "success",
    "data": {
        "credential": "integration_credential_string"
    }
}

Error Handling:

• 4xx errors: Client errors (e.g., invalid integration name)
• 5xx errors: Server errors
• Non-zero code: API business logic errors

Exceptions

WorkloadIdentityError

Base exception for all workload identity SDK errors.

ConfigurationError

Raised when required configuration is missing or invalid.

TokenRetrievalError

Raised when ID token retrieval fails.

TokenExchangeError

Raised when token exchange fails.

Token Caching

The SDK automatically caches tokens to avoid unnecessary HTTP requests:

• ID tokens are cached based on their expires_in value
• Access tokens are cached based on their expires_in value
• A 1-minute buffer is applied to prevent using tokens that are about to expire
• The cache is thread-safe and supports concurrent access

Thread Safety

The SDK is designed to be thread-safe:

• Token cache operations are protected by locks
• Multiple threads can safely call get_access_token() concurrently
• The HTTP session is shared across threads

Error Handling

The SDK provides detailed error information:

from coze_workload_identity import Client, ConfigurationError, TokenRetrievalError, TokenExchangeError

client = Client()
try:
    token = client.get_access_token()
except ConfigurationError as e:
    # Handle missing or invalid configuration
    print(f"Configuration error: {e}")
except TokenRetrievalError as e:
    # Handle ID token retrieval failures
    print(f"Token retrieval error: {e}")
except TokenExchangeError as e:
    # Handle token exchange failures
    print(f"Token exchange error: {e}")
except Exception as e:
    # Handle other unexpected errors
    print(f"Unexpected error: {e}")
finally:
    client.close()

Logging

The SDK uses Python's logging module. To enable debug logging:

import logging

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger('workload_identity')
logger.setLevel(logging.DEBUG)

Development

Setting up for development

git clone <repository-url>
cd workload-identity-sdk
pip install -e .[dev]

Running tests

pytest tests/ -v --cov=workload_identity

Code formatting

black workload_identity/ tests/

Type checking

mypy workload_identity/

License

MIT License - see LICENSE file for details.

Contributing

Contributions are welcome! Please read our contributing guidelines and submit pull requests to our repository.

Support

For issues and questions:

• Create an issue in the GitHub repository
• Check existing issues for solutions
• Review the documentation and examples