Merge pull request #11 from scaleapi/anishxyz/authn

[Auth] Authorized Registration

Author: anishxyz

src/agentex/lib/cli/handlers/deploy_handlers.py

@@ -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
@@ -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()

src/agentex/lib/cli/handlers/run_handlers.py

@@ -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
 
@@ -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

src/agentex/lib/cli/utils/__init__.py

@@ -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')

src/agentex/lib/cli/utils/auth_utils.py

@@ -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):
@@ -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:

src/agentex/lib/environment_variables.py

@@ -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
@@ -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:
         """

src/agentex/lib/sdk/config/agent_manifest.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Dict
 
 from pydantic import Field
 
@@ -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"""
 

src/agentex/lib/sdk/config/deployment_config.py

@@ -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
@@ -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
@@ -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: