Code Execution Server¶

`agora_workbench.code_execution.server` ¶

Base Code Execution Server for MCP.

This module provides a base class for creating MCP servers that execute Python code in isolated environments with domain-specific packages.

`CodeExecutionServer(server_config, tool_registry=None, session_manager=None, auth_config=None, working_dir=None, tool_search_backend=None, publishers=None, skills=None, states=None, catalog=None)` ¶

Bases: BaseMCPServer

Base class for MCP code execution servers.

Subclass this to create execution environments with specific Python interpreters and package sets. Each server exposes an MCP tool for executing code in its isolated environment.

Security features: - Timeout enforcement - Separate process execution - Optional working directory isolation - Configurable resource limits (override in subclass) - Pluggable authentication (Entra ID, no-op/dev, or custom) - Managed identity for downstream Azure resource access

Middleware Architecture: The server uses two layers of middleware at different levels of the stack:

Starlette Middleware (ASGI level, outermost layer): Applied to the entire HTTP application before FastMCP processing. Added in _create_middleware() and registered via app.add_middleware() during serve().

MCPSessionMiddleware (outermost):
Extracts Mcp-Session-Id header from requests to /mcp endpoint
Stores session ID in ASGI scope for downstream access
Enables session correlation across tool calls
AuthMiddleware (inner):
Validates Bearer tokens on /mcp and /object-transfer/* endpoints using Entra ID
Extracts user identity (oid/sub) and token claims
Stores authenticated context (token, claims, user_identity)
Bypasses /health and /.well-known/* endpoints (no auth required)
Returns RFC 9728 WWW-Authenticate header on 401 for OAuth discovery

FastMCP Middleware (tool call level, innermost layer): Applied to MCP tool calls after HTTP auth but before Pydantic validation. Added directly to the FastMCP instance:

AssetResolutionMiddleware (added during init):
Resolves tagged asset references like id to local cache paths
Runs BEFORE Pydantic validation so Path parameters receive proper values

Full execution flow: HTTP Request → MCPSessionMiddleware → AuthMiddleware → FastMCP routing → AssetResolutionMiddleware → Pydantic validation → Tool callback

Constraints: - Asset resolution happens first, so asset parameters are properly typed

Asset/Handle Injection Patterns: The server supports two different patterns for injecting assets and handles:

Auto-Extraction (execute_code_tool):
Handle IDs (h_xxxxxxxxxxxx) and asset tags (id) embedded as string literals in code are automatically detected via AST analysis.
Matched literals are replaced with synthetic variables that hold the resolved object (for handles) or a Path to the cached file (for assets).
Validation: References must appear as complete string literals in assignments, function arguments, return statements, or container literals (list/dict/tuple/set). Unresolvable references fail fast.
Middleware-Based Injection (domain tools):
Assets embedded in natural parameter values: grid_file="xyz"
Middleware extracts, validates, and resolves before Pydantic validation
Use case: LLM-driven tool calls where assets are encoded in argument values
Validation: AssetResolutionMiddleware

Initialize the code execution server.

Parameters:

Name	Type	Description	Default
`server_config`	`ServerConfig`	Server configuration (environment, assets, execution policy, features)	required
`tool_registry`	`Optional[ToolRegistry]`	Optional ToolRegistry containing domain-specific tools	`None`
`session_manager`	`Optional[SessionManager]`	Optional SessionManager for stateful tool support (auto-created with defaults if None)	`None`
`auth_config`	`Optional[AuthConfig]`	Authentication configuration providing token validation, identity extraction, and credential provisioning.	`None`
`working_dir`	`Optional[Path]`	Working directory for code execution (None = temp dir per execution)	`None`
`tool_search_backend`	`Optional[ToolSearchBackend]`	Optional pre-configured ToolSearchBackend instance. If provided, this backend is used instead of creating one from config. Enables custom search backends (e.g. vector DB, Elasticsearch) without modifying the built-in factory.	`None`
`publishers`	`Optional[list[AssetPublisher]]`	Optional list of :class:`~.data_access.AssetPublisher` instances that the `{name}_send` MCP tool will dispatch to. Publishers are routed by `destination_name`; the first match wins. A :class:`~.data_access.GuiPublisher` is always prepended (see `__init__`) so the send tool is registered unconditionally and the `"user"` destination works even when this list is `None`.	`None`
`skills`	`Optional[list[Skill]]`	Optional list of :class:`~.skills.Skill` objects for workflow planning and skill loading. Use :func:`~.skills.discover_skills` to scan a directory.	`None`
`states`	`Optional[list[State]]`	Optional list of :class:`~.tool_registry.State` objects defining the domain's state vocabulary with descriptions and affordances.	`None`
`catalog`	`Optional[CatalogIntegration]`	Optional policy-aware data catalog integration. When omitted, catalog lifecycle, tools, and session behavior remain unchanged.	`None`

Source code in src/agora_workbench/code_execution/server.py

def __init__(
    self,
    server_config: ServerConfig,
    tool_registry: Optional["ToolRegistry"] = None,
    session_manager: Optional["SessionManager"] = None,
    auth_config: Optional["AuthConfig"] = None,
    working_dir: Optional[Path] = None,
    tool_search_backend: Optional["ToolSearchBackend"] = None,
    publishers: "Optional[list[AssetPublisher]]" = None,
    skills: "Optional[list[Skill]]" = None,
    states: "Optional[list[State]]" = None,
    catalog: "Optional[CatalogIntegration]" = None,
):
    """
    Initialize the code execution server.

    Args:
        server_config: Server configuration (environment, assets, execution policy, features)
        tool_registry: Optional ToolRegistry containing domain-specific tools
        session_manager: Optional SessionManager for stateful tool support (auto-created with defaults if None)
        auth_config: Authentication configuration providing token validation,
            identity extraction, and credential provisioning.
        working_dir: Working directory for code execution (None = temp dir per execution)
        tool_search_backend: Optional pre-configured ToolSearchBackend instance.
            If provided, this backend is used instead of creating one from config.
            Enables custom search backends (e.g. vector DB, Elasticsearch) without
            modifying the built-in factory.
        publishers: Optional list of :class:`~.data_access.AssetPublisher` instances
            that the ``{name}_send`` MCP tool will dispatch to.
            Publishers are routed by ``destination_name``; the first match wins.
            A :class:`~.data_access.GuiPublisher` is always prepended (see
            ``__init__``) so the send tool is registered unconditionally and
            the ``"user"`` destination works even when this list is ``None``.
        skills: Optional list of :class:`~.skills.Skill` objects for workflow planning
            and skill loading. Use :func:`~.skills.discover_skills` to scan a directory.
        states: Optional list of :class:`~.tool_registry.State` objects defining the
            domain's state vocabulary with descriptions and affordances.
        catalog: Optional policy-aware data catalog integration. When
            omitted, catalog lifecycle, tools, and session behavior remain
            unchanged.
    """
    super().__init__()
    self.server_config = server_config
    self.kernel_name = f"tools-py-{server_config.name}"
    self.tool_registry = tool_registry
    self.catalog = catalog

    # Sidecar processes (e.g. a shared model service). Lazily started in
    # _startup and stopped in _shutdown; no-op when none are declared.
    self._sidecar_manager = SidecarManager(server_config)

    self.skills: list["Skill"] = list(skills or [])
    self.states: list["State"] = list(states or [])
    self._state_affordances: dict[str, list[str]] = {s.token: s.affordances for s in self.states if s.affordances}
    self._tool_search_backends: list[Any] = []
    self._closed_tool_search_backends: set[int] = set()
    self._custom_tool_search_backend = tool_search_backend
    self._publishers: "list[AssetPublisher]" = list(publishers or [])
    self._closed_publishers: set[int] = set()
    self._activity_publisher_stopped = False
    self._parallel_jobs: dict[str, dict[str, Any]] = {}
    self._parallel_batches: dict[str, dict[str, Any]] = {}
    self._parallel_job_by_session: dict[str, str] = {}
    self._parallel_state_lock = asyncio.Lock()

    # GuiPublisher is always available so agents can use <gui>name</gui>
    # to make outputs downloadable without requiring external storage.
    from .data_access.publishers import GuiPublisher

    self._gui_publisher = GuiPublisher(public_url_fn=self.public_url)
    self._publishers.insert(0, self._gui_publisher)

    # Shared peer registry: name → base URL of peer servers reachable via
    # {name}_send(to=<peer>). Resolved once so the send tool can construct a
    # ServerPublisher on demand instead of requiring one pre-registered per peer.
    self._peer_registry: dict[str, str] = self._load_peer_registry()
    # Resolution: ServerConfig overrides env var; env var provides deployment default.
    if server_config.parallel_max_concurrency is not None:
        parallel_execute_max_concurrency = server_config.parallel_max_concurrency
    else:
        parallel_execute_max_concurrency_raw = os.getenv("PARALLEL_EXECUTE_MAX_CONCURRENCY", "0").strip()
        try:
            parallel_execute_max_concurrency = int(parallel_execute_max_concurrency_raw)
        except ValueError:
            LOGGER.warning(
                "Invalid PARALLEL_EXECUTE_MAX_CONCURRENCY value %r; using default 0.",
                parallel_execute_max_concurrency_raw,
            )
            parallel_execute_max_concurrency = 0
    self.parallel_max_concurrency = max(0, parallel_execute_max_concurrency)
    self._parallel_semaphore: Optional[asyncio.Semaphore] = (
        asyncio.Semaphore(self.parallel_max_concurrency) if self.parallel_max_concurrency > 0 else None
    )

    # Auto-create session manager with defaults if not provided
    if session_manager is None:
        _session_manager = SessionManager(
            SessionConfig(kernel_network_mode=server_config.kernel_network_mode or "inherit")
        )
        LOGGER.info("Created default SessionManager.")
    else:
        _session_manager = session_manager
        if (
            server_config.kernel_network_mode is not None
            and _session_manager.config.kernel_network_mode != server_config.kernel_network_mode
        ):
            LOGGER.warning(
                "ServerConfig.kernel_network_mode=%s overrides the custom SessionManager setting %s.",
                server_config.kernel_network_mode,
                _session_manager.config.kernel_network_mode,
            )
            _session_manager.set_kernel_network_mode(server_config.kernel_network_mode)
    if _session_manager.config.kernel_network_mode in {"isolated", "seccomp"} and server_config.sidecars:
        raise ValueError("restricted kernel networking cannot be combined with loopback HTTP sidecars")

    self.session_manager = _session_manager
    self.session_manager.kernel_name = self.kernel_name

    # --- Authentication configuration ---
    if auth_config is None:
        raise ValueError(
            "auth_config is required. Use create_entra_auth_config() for Entra ID "
            "or create_noop_auth_config() for development."
        )
    self.auth_config = auth_config
    if self.catalog is not None:
        self._configure_catalog_sessions()

    # Entra client/tenant IDs for RFC 9728 OAuth protected-resource metadata.
    # Preferred source is auth_config.protected_resource_metadata; these attributes
    # remain for back-compat and for deployments that configure Entra IDs directly.
    # Resolution order: auth_config validator → ServerConfig → environment variable.
    #
    # Reading private validator attributes is a deprecated legacy convention, kept
    # for third-party validators written against it. Probed symmetrically so a
    # validator exposing only one of the two cannot raise or be partially ignored.
    self.entra_client_id: Optional[str] = getattr(auth_config.token_validator, "_client_id", None)
    self.entra_tenant_id: Optional[str] = getattr(auth_config.token_validator, "_tenant_id", None)
    if not self.entra_client_id:
        self.entra_client_id = server_config.entra_client_id or os.getenv("ENTRA_CLIENT_ID")
    if not self.entra_tenant_id:
        self.entra_tenant_id = server_config.entra_tenant_id or os.getenv("ENTRA_TENANT_ID")

    self._warn_if_oauth_metadata_unresolvable()

    self.max_timeout = server_config.max_timeout
    self.default_timeout = min(server_config.default_timeout, self.max_timeout)

    # Resolution: ServerConfig overrides env var; env var provides deployment default.
    if server_config.output_truncation_threshold is not None:
        self.output_truncation_threshold = server_config.output_truncation_threshold
    else:
        env_threshold = os.getenv("CODE_OUTPUT_TRUNCATION_THRESHOLD", "50000")
        normalized = env_threshold.strip().replace("_", "")
        try:
            parsed = int(normalized)
            if parsed < 0:
                raise ValueError("threshold must be >= 0")
            self.output_truncation_threshold = parsed
        except ValueError:
            LOGGER.warning(
                "Invalid CODE_OUTPUT_TRUNCATION_THRESHOLD=%r; using default 50000",
                env_threshold,
            )
            self.output_truncation_threshold = 50_000

    self.working_dir = working_dir
    self._python_executable: Optional[Path] = None
    self._environment_ready = False

    # Best-effort activity publisher (silent no-op when ACTIVITY_UI_URL is unset).
    from .activity_publisher import ActivityPublisher

    self.activity_publisher = ActivityPublisher(server_name=server_config.name)

    self.mcp = FastMCP(
        f"{server_config.name}-executor",
        instructions=server_config.server_description or server_config.description,
    )

    # AssetResolutionMiddleware resolves tagged asset references before Pydantic validation.
    self.mcp.add_middleware(AssetResolutionMiddleware(self))

    self._setup_tools()

`warm()` `async` ¶

Pre-initialize the execution environment without serving requests.

Call this during Docker builds or process startup to avoid cold-start latency. It prepares the Python environment, provisions assets according to the ServerConfig (e.g. when auto_provision is enabled), and registers the execution kernel.

Source code in src/agora_workbench/code_execution/server.py

async def warm(self) -> None:
    """Pre-initialize the execution environment without serving requests.

    Call this during Docker builds or process startup to avoid cold-start
    latency. It prepares the Python environment, provisions assets according
    to the ServerConfig (e.g. when auto_provision is enabled), and registers
    the execution kernel.
    """
    LOGGER.info(f"Warming environment: {self.server_config.name}")
    await self._ensure_environment()
    await self._register_kernel(kernel_name=self.kernel_name)
    LOGGER.info(f"✓ Environment '{self.server_config.name}' is warm and ready.")

`main(*, default_host='127.0.0.1', default_port=8000)` ¶

CLI entrypoint that handles --warm, --host, and --port.

Call this from your server's if __name__ == "__main__" block to get standard flag handling without manual sys.argv parsing::

server = MyDomainServer(...)
server.main()

Flags

--warm Pre-initialize the environment and exit (no HTTP server). --host HOST Bind address (default: default_host, or HOST env var). --port PORT Bind port (default: default_port, or PORT env var). --allow-unauthenticated-remote Permit open/no-op auth on a non-loopback bind. --no-allow-unauthenticated-remote Disable an acknowledgement inherited from the environment.

Source code in src/agora_workbench/code_execution/server.py

def main(self, *, default_host: str = "127.0.0.1", default_port: int = 8000) -> None:
    """CLI entrypoint that handles ``--warm``, ``--host``, and ``--port``.

    Call this from your server's ``if __name__ == "__main__"`` block to get
    standard flag handling without manual ``sys.argv`` parsing::

        server = MyDomainServer(...)
        server.main()

    Flags:
        --warm          Pre-initialize the environment and exit (no HTTP server).
        --host HOST     Bind address (default: default_host, or HOST env var).
        --port PORT     Bind port (default: default_port, or PORT env var).
        --allow-unauthenticated-remote
                        Permit open/no-op auth on a non-loopback bind.
        --no-allow-unauthenticated-remote
                        Disable an acknowledgement inherited from the environment.
    """
    import argparse

    from agora_workbench.base import ALLOW_UNAUTHENTICATED_REMOTE_ENV_VAR

    parser = argparse.ArgumentParser(
        description=f"{self.server_config.name} — CodeExecutionServer",
    )
    parser.add_argument(
        "--warm",
        action="store_true",
        help="Pre-initialize the environment and exit without starting the server.",
    )
    env_host = os.getenv("HOST")
    env_port = os.getenv("PORT")
    try:
        port_default = int(env_port) if env_port else default_port
    except ValueError:
        parser.error(f"Invalid PORT env var: {env_port!r} (must be an integer).")
    allow_remote_raw = os.getenv(ALLOW_UNAUTHENTICATED_REMOTE_ENV_VAR, "")
    normalized_allow_remote = allow_remote_raw.strip().lower()
    if normalized_allow_remote in {"", "0", "false", "no", "off"}:
        allow_remote_default = False
    elif normalized_allow_remote in {"1", "true", "yes", "on"}:
        allow_remote_default = True
    else:
        parser.error(
            f"Invalid {ALLOW_UNAUTHENTICATED_REMOTE_ENV_VAR}={allow_remote_raw!r} "
            "(expected one of: 1, true, yes, on, 0, false, no, off)."
        )

    parser.add_argument(
        "--host",
        default=env_host or default_host,
        help=f"Bind address (default: {default_host}, or HOST env var).",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=port_default,
        help=f"Bind port (default: {default_port}, or PORT env var).",
    )
    parser.add_argument(
        "--allow-unauthenticated-remote",
        action=argparse.BooleanOptionalAction,
        default=allow_remote_default,
        help=(
            "Allow no-op/open authentication on a non-loopback bind. "
            f"Defaults from {ALLOW_UNAUTHENTICATED_REMOTE_ENV_VAR}; use the --no- form to override it."
        ),
    )
    args = parser.parse_args()

    if args.warm:
        asyncio.run(self.warm())
    else:
        asyncio.run(
            self.run_http(
                host=args.host,
                port=args.port,
                allow_unauthenticated_remote=args.allow_unauthenticated_remote,
            )
        )

`get_tool_name()` ¶

Get the MCP tool name. Override to customize.

Source code in src/agora_workbench/code_execution/server.py

def get_tool_name(self) -> str:
    """Get the MCP tool name. Override to customize."""
    return f"execute_{self.server_config.name}_code"

`public_url()` ¶

Best-effort base URL for outbound references (e.g. download links).

Falls back to http://localhost:<port> when called before :meth:run_http has stashed the bind host/port. Callers that need the real externally-visible URL (containerized deployments behind a proxy / Ingress) should set SERVER_PUBLIC_URL instead — that env var takes precedence wherever this method is consulted.

Source code in src/agora_workbench/code_execution/server.py

def public_url(self) -> str:
    """Best-effort base URL for outbound references (e.g. download links).

    Falls back to ``http://localhost:<port>`` when called before
    :meth:`run_http` has stashed the bind host/port.  Callers that need
    the real externally-visible URL (containerized deployments behind a
    proxy / Ingress) should set ``SERVER_PUBLIC_URL`` instead — that env
    var takes precedence wherever this method is consulted.
    """
    host = getattr(self, "_bind_host", None) or "localhost"
    port = getattr(self, "_bind_port", None) or 8000
    # 0.0.0.0 is a bind address, not a reachable one.  Map it back to
    # localhost for outward-facing URLs; for non-localhost deployments
    # the operator should set SERVER_PUBLIC_URL explicitly.
    if host == "0.0.0.0":
        host = "localhost"
    return f"http://{host}:{port}"

`preprocess_code(code)` ¶

Preprocess code before execution.

Override to add imports, setup code, or modify the input.

Source code in src/agora_workbench/code_execution/server.py

def preprocess_code(self, code: str) -> str:
    """
    Preprocess code before execution.

    Override to add imports, setup code, or modify the input.
    """
    return code

`validate_code(code)` ¶

Validate user-provided code before execution. This is a thin instance-method wrapper around the default implementation in execution_defaults.validate_code, which expects self as its first argument.

Override to adjust behavior.

Source code in src/agora_workbench/code_execution/server.py

def validate_code(self, code: str) -> tuple[bool, "Optional[str]"]:
    """
    Validate user-provided code before execution.
    This is a thin instance-method wrapper around the default implementation
    in `execution_defaults.validate_code`, which expects `self` as its first
    argument.

    Override to adjust behavior.
    """
    return execution_defaults.validate_code(self, code)

`postprocess_result(result)` ¶

Postprocess execution result.

Override to filter output, add metadata, or modify the result.

Source code in src/agora_workbench/code_execution/server.py

def postprocess_result(self, result: CodeExecutionResult) -> CodeExecutionResult:
    """
    Postprocess execution result.

    Override to filter output, add metadata, or modify the result.
    """
    return result

`get_data_lake_capabilities(session)` `async` ¶

Return caller-effective catalog capabilities for application adapters.

Source code in src/agora_workbench/code_execution/server.py

async def get_data_lake_capabilities(self, session: "Session") -> tuple[Any, ...]:
    """Return caller-effective catalog capabilities for application adapters."""
    if self.catalog is None:
        return ()
    binding: "CatalogSessionBinding | None" = session.extensions.get("catalog")
    if binding is None:
        return ()
    snapshot = binding.snapshot()
    try:
        return await self.catalog.capabilities(snapshot)
    finally:
        snapshot.close()

`get_python_executable()` `async` ¶

Return the path to the Python interpreter for this environment.

This will automatically build the environment if needed and build_environment=True.

Source code in src/agora_workbench/code_execution/server.py

async def get_python_executable(self) -> str:
    """
    Return the path to the Python interpreter for this environment.

    This will automatically build the environment if needed and build_environment=True.
    """
    if self._python_executable and self._environment_ready:
        return str(self._python_executable)

    await self._ensure_environment()

    if not self._python_executable:
        raise RuntimeError("Failed to determine Python executable path")

    return str(self._python_executable)

Code Execution Server¶

agora_workbench.code_execution.server ¶