Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/agentex/lib/cli/handlers/deploy_handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,10 @@
from pydantic import BaseModel, Field
from rich.console import Console

from agentex.lib.cli.utils.auth_utils import _encode_principal_context
from agentex.lib.cli.utils.exceptions import DeploymentError, HelmError
from agentex.lib.cli.utils.kubectl_utils import check_and_switch_cluster_context
from agentex.lib.environment_variables import EnvVarKeys
from agentex.lib.sdk.config.agent_config import AgentConfig
from agentex.lib.sdk.config.agent_manifest import AgentManifest
from agentex.lib.sdk.config.deployment_config import ClusterConfig
Expand Down Expand Up @@ -168,6 +170,10 @@ def merge_deployment_configs(
if TEMPORAL_WORKER_KEY in helm_values:
helm_values[TEMPORAL_WORKER_KEY]["env"] = agent_config.env

encoded_principal = _encode_principal_context(manifest)
if encoded_principal:
helm_values["env"][EnvVarKeys.AUTH_PRINCIPAL_B64] = encoded_principal

if manifest.deployment and manifest.deployment.imagePullSecrets:
pull_secrets = [
pull_secret.to_dict()
Expand Down
7 changes: 7 additions & 0 deletions src/agentex/lib/cli/handlers/run_handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,12 @@
from rich.console import Console
from rich.panel import Panel

from agentex.lib.cli.utils.auth_utils import _encode_principal_context
from agentex.lib.cli.handlers.cleanup_handlers import (
cleanup_agent_workflows,
should_cleanup_on_restart
)
from agentex.lib.environment_variables import EnvVarKeys
from agentex.lib.sdk.config.agent_manifest import AgentManifest
from agentex.lib.utils.logging import make_logger

Expand Down Expand Up @@ -427,6 +429,11 @@ def create_agent_environment(manifest: AgentManifest) -> dict[str, str]:
"ACP_PORT": str(manifest.local_development.agent.port),
}

# Add authorization principal if set
encoded_principal = _encode_principal_context(manifest)
if encoded_principal:
env_vars[EnvVarKeys.AUTH_PRINCIPAL_B64] = encoded_principal

# Add description if available
if manifest.agent.description:
env_vars["AGENT_DESCRIPTION"] = manifest.agent.description
Expand Down
Empty file.
18 changes: 18 additions & 0 deletions src/agentex/lib/cli/utils/auth_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
import base64
import json

from agentex.lib.sdk.config.agent_manifest import AgentManifest


# Base 64 encode principal dictionary
def _encode_principal_context(manifest: AgentManifest):
if manifest.auth is None:
return None

principal = manifest.auth.principal
if principal is None:
return None

json_str = json.dumps(principal, separators=(',', ':'))
encoded_bytes = base64.b64encode(json_str.encode('utf-8'))
return encoded_bytes.decode('utf-8')
5 changes: 4 additions & 1 deletion src/agentex/lib/environment_variables.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,11 @@ class EnvVarKeys(str, Enum):
ACP_URL = "ACP_URL"
ACP_PORT = "ACP_PORT"
ACP_TYPE = "ACP_TYPE"
# Workflow Configuraiton
# Workflow Configuration
WORKFLOW_NAME = "WORKFLOW_NAME"
WORKFLOW_TASK_QUEUE = "WORKFLOW_TASK_QUEUE"
# Auth Configuration
AUTH_PRINCIPAL_B64 = "AUTH_PRINCIPAL_B64"


class Environment(str, Enum):
Expand Down Expand Up @@ -54,6 +56,7 @@ class EnvironmentVariables(BaseModel):
# Workflow Configuration
WORKFLOW_TASK_QUEUE: str | None = None
WORKFLOW_NAME: str | None = None
AUTH_PRINCIPAL_B64: str | None = None

@classmethod
def refresh(cls) -> EnvironmentVariables | None:
Expand Down
3 changes: 2 additions & 1 deletion src/agentex/lib/sdk/config/agent_manifest.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@

from agentex.lib.sdk.config.agent_config import AgentConfig
from agentex.lib.sdk.config.build_config import BuildConfig
from agentex.lib.sdk.config.deployment_config import DeploymentConfig
from agentex.lib.sdk.config.deployment_config import DeploymentConfig, AuthenticationConfig
from agentex.lib.sdk.config.local_development_config import LocalDevelopmentConfig
from agentex.lib.utils.logging import make_logger
from agentex.lib.utils.model_utils import BaseModel
Expand All @@ -36,6 +36,7 @@ class AgentManifest(BaseModel):
deployment: DeploymentConfig | None = Field(
default=None, description="Deployment configuration for the agent"
)
auth: AuthenticationConfig | None = Field(default=None, description="Authentication configuration")

def context_manager(self, build_context_root: Path) -> BuildContextManager:
"""
Expand Down
6 changes: 5 additions & 1 deletion src/agentex/lib/sdk/config/deployment_config.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Any
from typing import Any, Dict

from pydantic import Field

Expand Down Expand Up @@ -90,6 +90,10 @@ class ClusterConfig(BaseModel):
)


class AuthenticationConfig(BaseModel):
principal: Dict[str, Any] = Field(description="Principal used for authorization on registration")


class InjectedImagePullSecretValues(BaseModel):
"""Values for image pull secrets"""

Expand Down
14 changes: 14 additions & 0 deletions src/agentex/lib/sdk/fastacp/base/base_acp_server.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import asyncio
import base64
import inspect
import json
from collections.abc import AsyncGenerator, Awaitable, Callable
from contextlib import asynccontextmanager
from typing import Any
Expand Down Expand Up @@ -341,6 +343,16 @@ def run(self, host: str = "0.0.0.0", port: int = 8000, **kwargs):
"""Start the Uvicorn server for async handlers."""
uvicorn.run(self, host=host, port=port, **kwargs)

def _get_auth_principal(self, env_vars: EnvironmentVariables):
if not env_vars.AUTH_PRINCIPAL_B64:
return None

try:
decoded_str = base64.b64decode(env_vars.AUTH_PRINCIPAL_B64).decode('utf-8')
return json.loads(decoded_str)
except Exception:
return None

async def _register_agent(self, env_vars: EnvironmentVariables):
"""Register this agent with the Agentex server"""
# Build the agent's own URL
Expand All @@ -350,12 +362,14 @@ async def _register_agent(self, env_vars: EnvironmentVariables):
env_vars.AGENT_DESCRIPTION
or f"Generic description for agent: {env_vars.AGENT_NAME}"
)

# Prepare registration data
registration_data = {
"name": env_vars.AGENT_NAME,
"description": description,
"acp_url": full_acp_url,
"acp_type": env_vars.ACP_TYPE,
"principal_context": self._get_auth_principal(env_vars)
}

if env_vars.AGENT_ID:
Expand Down
Loading
Loading