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)

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

1. MCPSessionMiddleware (outermost):
2. Extracts Mcp-Session-Id header from requests to /mcp endpoint
3. Stores session ID in ASGI scope for downstream access

4. Enables session correlation across tool calls

5. AuthMiddleware (inner):

6. Validates Bearer tokens on /mcp and /object-transfer/* endpoints using Entra ID
7. Extracts user identity (oid/sub) and token claims
8. Stores authenticated context (token, claims, user_identity)
9. Bypasses /health and /.well-known/* endpoints (no auth required)
10. 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:

1. AssetResolutionMiddleware (added during init):
2. Resolves tagged asset references like id to local cache paths
3. 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:

1. Auto-Extraction (execute_code_tool):
2. Handle IDs (h_xxxxxxxxxxxx) and asset tags (id) embedded as string literals in code are automatically detected via AST analysis.
3. Matched literals are replaced with synthetic variables that hold the resolved object (for handles) or a Path to the cached file (for assets).

4. 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.

5. Middleware-Based Injection (domain tools):

6. Assets embedded in natural parameter values: grid_file="xyz"
7. Middleware extracts, validates, and resolves before Pydantic validation
8. Use case: LLM-driven tool calls where assets are encoded in argument values
9. 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

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,
):
    """
    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.
    """
    super().__init__()
    self.server_config = server_config
    self.tool_registry = tool_registry

    # 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_proxies_injected: set[str] = set()
    self._tool_search_backends: list[Any] = []
    self._custom_tool_search_backend = tool_search_backend
    self._publishers: "list[AssetPublisher]" = list(publishers or [])
    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())
        LOGGER.info("Created default SessionManager.")
    else:
        _session_manager = session_manager

    self.session_manager = _session_manager

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

    # Entra client/tenant IDs for RFC 9728 OAuth protected-resource metadata.
    # Resolution order: auth_config validator → ServerConfig → environment variable.
    self.entra_client_id: Optional[str] = None
    self.entra_tenant_id: Optional[str] = None
    if hasattr(auth_config.token_validator, "_client_id"):
        self.entra_client_id = auth_config.token_validator._client_id
        self.entra_tenant_id = auth_config.token_validator._tenant_id
    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.max_timeout = server_config.max_timeout
    self.default_timeout = server_config.default_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="tools-py")
    LOGGER.info(f"✓ Environment '{self.server_config.name}' is warm and ready.")

main(*, default_host='0.0.0.0', 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).

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

def main(self, *, default_host: str = "0.0.0.0", 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).
    """
    import argparse

    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).")

    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).",
    )
    args = parser.parse_args()

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

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_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)