Sessions

Session Container

agora_workbench.code_execution.sessions.session.Session(session_id, data, session_type, user_identity, user_token, token_claims, metadata=None, data_manager=None)

Bases: Generic[T]

Generic session container managing stateful data across MCP server interactions.

Sessions provide lifecycle management, ownership tracking, and metadata storage for persistent resources like code execution environments, database connections, or computation state. Each session is owned by a specific user (identified via JWT token claims) and tracks access patterns, status transitions, and cleanup requirements.

Attributes:

Name	Type	Description
session_id	str	Unique identifier for the session.
data	T	The session's payload data, type-parameterized for type safety.
session_type	str	Categorizes session type (e.g., "python", "database").
user_identity	str	Owner's composite identifier from JWT token (oid@tid).
user_token	str	User's bearer token for authentication.
metadata	Dict	Optional key-value metadata for session configuration.
token_claims	Dict	Optional cached JWT token claims for session authorization. These claims are used to restore authentication context without re-validating the JWT token. Intentionally excluded from get_info() for security.
created_at	datetime	Timestamp when session was created.
last_accessed	datetime	Timestamp of most recent session access.
status	str	Current session state (e.g., "created", "active", "error").
data_manager	DataLakeDataManager	Manager for DataLake asset access. Owned by the session — :meth:cleanup tears it down, so an injected manager must not be shared between sessions.

Example

from .session import Session session = Session( ... session_id="sess_123", ... data={"counter": 0}, ... session_type="demo", ... user_identity="user-oid-xyz", ... user_token="eyJ...", ... metadata={"version": "1.0"}, ... token_claims={"oid": "user-oid-xyz", "exp": 1234567890}, ... ) session.touch() # Update last accessed time session.update_status("active") info = session.get_info()

Initialize a session.

Parameters:

Name	Type	Description	Default
session_id	str	Unique identifier for the session.	required
data	T	The session's payload data.	required
session_type	str	Categorizes session type (e.g. "python").	required
user_identity	str	Owner's composite identifier from JWT token (oid@tid).	required
user_token	str	User's bearer token for authentication.	required
token_claims	Dict	Cached JWT claims for the user token.	required
metadata	Optional[Dict]	Optional key-value metadata for session configuration.	None
data_manager	Optional[DataLakeDataManager]	Optional pre-built data manager for DataLake asset access. When omitted, a default :class:DataLakeDataManager is constructed. The session takes ownership of whichever manager it ends up with — :meth:cleanup calls cleanup() on it — so a caller supplying one must pass a fresh instance per session rather than a shared singleton.	None

Source code in src/agora_workbench/code_execution/sessions/session.py

def __init__(
    self,
    session_id: str,
    data: T,
    session_type: str,
    user_identity: str,
    user_token: str,
    token_claims: Dict,
    metadata: Optional[Dict] = None,
    data_manager: Optional["DataLakeDataManager"] = None,
):
    """
    Initialize a session.

    Args:
        session_id: Unique identifier for the session.
        data: The session's payload data.
        session_type: Categorizes session type (e.g. ``"python"``).
        user_identity: Owner's composite identifier from JWT token (``oid@tid``).
        user_token: User's bearer token for authentication.
        token_claims: Cached JWT claims for the user token.
        metadata: Optional key-value metadata for session configuration.
        data_manager: Optional pre-built data manager for DataLake asset
            access. When omitted, a default :class:`DataLakeDataManager` is
            constructed. The session takes ownership of whichever manager it
            ends up with — :meth:`cleanup` calls ``cleanup()`` on it — so a
            caller supplying one must pass a fresh instance per session
            rather than a shared singleton.
    """
    self.session_id = session_id
    self.data = data
    self.created_at = datetime.now()
    self.last_accessed = datetime.now()
    self.metadata = metadata or {}
    self.user_identity = user_identity
    self.user_token = user_token
    self.token_claims = token_claims
    self.status = "created"
    self.session_type = session_type
    self._asset_counter: int = 0
    self._status_history = [("created", datetime.now())]

    # Initialize data manager for DataLake asset access. Constructing the
    # default lazily matters: DataLakeDataManager.__init__ eagerly allocates
    # a temp cache dir, so building one only to discard it would leak it.
    if data_manager is None:
        from ..data_access.manager import DataLakeDataManager

        data_manager = DataLakeDataManager()

    self.data_manager = data_manager

    # Initialize object store for asset objects
    self.object_store = ObjectStore()

touch()

Update last accessed timestamp.

Source code in src/agora_workbench/code_execution/sessions/session.py

def touch(self):
    """Update last accessed timestamp."""
    self.last_accessed = datetime.now()

update_status(new_status)

Update session status with history tracking.

Source code in src/agora_workbench/code_execution/sessions/session.py

def update_status(self, new_status: str):
    """Update session status with history tracking."""
    self.status = new_status
    self._status_history.append((new_status, datetime.now()))
    self.touch()

get_info()

Return session information.

Source code in src/agora_workbench/code_execution/sessions/session.py

def get_info(self) -> Dict[str, Any]:
    """Return session information."""
    age_seconds = (datetime.now() - self.created_at).total_seconds()
    idle_seconds = (datetime.now() - self.last_accessed).total_seconds()

    return {
        "session_id": self.session_id,
        "session_type": self.session_type,
        "status": self.status,
        "created_at": self.created_at.isoformat(),
        "last_accessed": self.last_accessed.isoformat(),
        "age_seconds": age_seconds,
        "idle_seconds": idle_seconds,
        "metadata": self.metadata,
        "user_identity": self.user_identity,
        "status_history": [{"status": s, "timestamp": t.isoformat()} for s, t in self._status_history],
    }

cleanup()

Cleanup session resources including session files.

Raises:

Type	Description
Exception	If cleanup fails, to allow calling code to handle the failure

Source code in src/agora_workbench/code_execution/sessions/session.py

def cleanup(self):
    """
    Cleanup session resources including session files.

    Raises:
        Exception: If cleanup fails, to allow calling code to handle the failure
    """
    # Clean up data manager cache directory
    self.data_manager.cleanup()

    # Call cleanup on data if it has a cleanup method
    if hasattr(self.data, "cleanup"):
        self.data.cleanup()  # pyright: ignore reportAttributeAccessIssue

    # Clean up session file if it exists
    if isinstance(self.data, dict) and "session_file" in self.data:
        session_file = Path(self.data["session_file"])
        if session_file.exists():
            # Remove the session file
            session_file.unlink()

            # Try to remove the parent directory if it's a temp directory for this session
            session_dir = session_file.parent
            if f"session_{self.session_id}" in str(session_dir):
                try:
                    shutil.rmtree(session_dir)
                except OSError:
                    # Directory might not be empty or already deleted, that's ok
                    pass

Session Manager

agora_workbench.code_execution.sessions.manager.SessionManager(config=None, kernel_name='tools-py')

Manages the lifecycle of code-execution sessions and their Jupyter kernels.

Responsibilities include session creation, retrieval, timeout-based cleanup, kernel provisioning (one AsyncKernelManager per session), background-job tracking, and artifact registration for the /artifacts download endpoint.

Thread-safety is provided by an internal RLock for session-lifecycle mutations and per-session asyncio.Lock instances for kernel access.

Source code in src/agora_workbench/code_execution/sessions/manager.py

def __init__(self, config: Optional[SessionConfig] = None, kernel_name: str = "tools-py"):
    self.config = config or SessionConfig()
    self.kernel_name = kernel_name
    self.storage = self.config.storage_backend
    self._last_cleanup = datetime.now()
    self._session_lifecycle_lock = RLock()
    timeout_seconds = self.config.timeout.total_seconds()
    self.execution_session_keepalive_seconds = max(0.5, min(timeout_seconds / 10.0, 60.0))

    # session_id -> (KernelManager, KernelClient)
    self._kernels: dict[str, Tuple[AsyncKernelManager, "AsyncKernelClient"]] = {}
    self._kernel_last_used: dict[str, float] = {}  # session_id -> timestamp
    self._kernel_tokens: dict[str, Optional[str]] = {}  # session_id -> last injected user token
    # Per-session lock that serializes execute_code_for_session calls so the
    # shared Jupyter kernel client (single iopub queue) cannot be raced by
    # concurrent callers (e.g. four parallel push_object MCP calls).
    self._kernel_execute_locks: dict[str, asyncio.Lock] = {}
    self._background_jobs: dict[str, _BackgroundJob] = {}
    self._session_running_jobs: dict[str, str] = {}
    # Artifact pipeline state: the token -> record map used by the HTTP
    # download endpoint.
    self._session_artifacts: dict[str, dict[str, _ArtifactRecord]] = {}

    # Kernel bootstrap tracking.  Every kernel process gets a globally
    # unique, monotonically increasing *generation* id when it is started.
    # One-time bootstrap work (the AGORA_OUTPUT_DIR preamble, tool proxy
    # injection, ...) is recorded against that generation rather than
    # against the session id, so a session whose kernel is rebuilt is
    # re-bootstrapped automatically.  Because generations are never
    # reused, correctness does not depend on any teardown path
    # remembering to clear this state -- the entries below are dropped in
    # :meth:`_shutdown_kernel` purely to bound memory.
    self._kernel_generation_seq: int = 0
    self._kernel_generations: dict[str, int] = {}  # session_id -> generation
    self._kernel_bootstrap_state: dict[int, set[str]] = {}  # generation -> completed keys

    # In-flight kernel teardowns, so concurrent closes coalesce onto one
    # task instead of racing, callers can await teardown, and the task is
    # strongly referenced for its lifetime (an unreferenced task may be
    # garbage-collected mid-flight).
    self._kernel_shutdown_tasks: dict[str, "asyncio.Task[None]"] = {}

    LOGGER.info(
        f"Initialized SessionManager: max_sessions={self.config.max_sessions}, "
        f"timeout={self.config.timeout.total_seconds() / 60}min"
    )

create_session(data, user_identity, user_token, token_claims, metadata=None, session_id=None)

Create a new session and return its ID.

Parameters:

Name	Type	Description	Default
data	Any	The data/object to store in the session	required
user_identity	str	User identity (Entra ID, email, etc.)	required
user_token	str	User's bearer token for authentication	required
metadata	Optional[dict]	Optional metadata dict	None
session_id	Optional[str]	Optional custom session ID (default: UUID)	None
token_claims	dict	Cached JWT claims for the user token. Stored on the session so that background tasks can restore the full auth context without re-validating the token.	required

Returns:

Type	Description
str	Session ID string

Note

When SessionConfig.data_manager_factory is configured it is invoked here, once per session, and the resulting manager is passed to the Session. The session owns it and cleans it up.

Raises:

Type	Description
MaxSessionsReachedError	If the session limit has been reached.
TypeError	If a configured data_manager_factory returns something that is not a usable data manager.

Source code in src/agora_workbench/code_execution/sessions/manager.py

def create_session(
    self,
    data: Any,
    user_identity: str,
    user_token: str,
    token_claims: dict,
    metadata: Optional[dict] = None,
    session_id: Optional[str] = None,
) -> str:
    """
    Create a new session and return its ID.

    Args:
        data: The data/object to store in the session
        user_identity: User identity (Entra ID, email, etc.)
        user_token: User's bearer token for authentication
        metadata: Optional metadata dict
        session_id: Optional custom session ID (default: UUID)
        token_claims: Cached JWT claims for the user token.
            Stored on the session so that background tasks can restore
            the full auth context without re-validating the token.

    Returns:
        Session ID string

    Note:
        When ``SessionConfig.data_manager_factory`` is configured it is
        invoked here, once per session, and the resulting manager is passed
        to the ``Session``. The session owns it and cleans it up.

    Raises:
        MaxSessionsReachedError: If the session limit has been reached.
        TypeError: If a configured ``data_manager_factory`` returns
            something that is not a usable data manager.
    """
    # Run periodic cleanup
    self._maybe_cleanup()

    with self._session_lifecycle_lock:
        # Enforce max sessions limit
        self._enforce_max_sessions()

        # Generate session ID
        if session_id is None:
            session_id = str(uuid.uuid4())

        # Build a customized data manager when a factory is configured, so
        # the Session never constructs (and immediately discards) a default
        # one — DataLakeDataManager allocates a temp cache dir eagerly.
        data_manager = None
        if self.config.data_manager_factory is not None:
            data_manager = self.config.data_manager_factory(
                SessionContext(
                    session_id=session_id,
                    user_identity=user_identity,
                    user_token=user_token,
                    token_claims=token_claims,
                    session_type="default",
                    metadata=metadata or {},
                )
            )
            # Validate eagerly. A factory returning None is the dangerous
            # case: Session would fall back to building a default manager,
            # silently discarding the customization, so the mistake would
            # surface later as unexplained default behaviour instead of an
            # error at the point of the bug.
            if data_manager is None or not callable(getattr(data_manager, "cleanup", None)):
                raise TypeError(
                    "SessionConfig.data_manager_factory must return a data manager instance with a "
                    f"cleanup() method, but it returned {type(data_manager).__name__}. Returning None "
                    "would silently fall back to a default DataLakeDataManager and discard the "
                    "customization the factory exists to provide."
                )

        # Create session
        session = Session(
            session_id=session_id,
            data=data,
            session_type="default",
            user_identity=user_identity,
            user_token=user_token,
            token_claims=token_claims,
            metadata=metadata,
            data_manager=data_manager,
        )

        # Store
        self.storage.store(session_id, session)

        # Create the per-session outputs directory.  Done eagerly so the
        # kernel can write to it on the very first execute.  Failure to
        # create is non-fatal: artifact discovery just becomes a no-op for
        # this session.
        try:
            self._get_outputs_dir(session_id).mkdir(parents=True, exist_ok=True)
        except OSError:
            LOGGER.warning(
                "Could not create outputs dir for session %s under %s; "
                "artifact download will be disabled for this session.",
                session_id,
                _OUTPUTS_BASE_DIR,
                exc_info=True,
            )

    LOGGER.info(f"Created session {session_id} (total={self.storage.count()})")

    return session_id

get_session(session_id)

Get a session by ID, updating its access time.

Returns:

Type	Description
Session	Session object

Raises:

Type	Description
ValueError	If session not found or expired

Source code in src/agora_workbench/code_execution/sessions/manager.py

def get_session(self, session_id: str) -> Session:
    """
    Get a session by ID, updating its access time.

    Returns:
        Session object

    Raises:
        ValueError: If session not found or expired
    """
    self._maybe_cleanup()

    session = self.storage.retrieve(session_id)

    if session is None:
        raise ValueError(
            f"Session {session_id} not found. It may have expired or been "
            f"cleaned up. Active sessions: {self.storage.count()}"
        )

    # Update access time
    session.touch()
    self.storage.store(session_id, session)

    return session

update_session(session_id, session)

Update an existing session.

Source code in src/agora_workbench/code_execution/sessions/manager.py

def update_session(self, session_id: str, session: Session) -> None:
    """Update an existing session."""
    if self.storage.retrieve(session_id) is None:
        raise ValueError(f"Session {session_id} not found")

    session.touch()
    self.storage.store(session_id, session)

update_status(session_id, status)

Update the status of a session.

Source code in src/agora_workbench/code_execution/sessions/manager.py

def update_status(self, session_id: str, status: str) -> None:
    """Update the status of a session."""
    session = self.get_session(session_id)
    session.update_status(status)
    self.storage.store(session_id, session)

close_session(session_id)

Explicitly close a session.

Attempts to clean up resources first. If cleanup fails, the session is still deleted to prevent session accumulation, but the error is logged.

Kernel teardown is asynchronous. This method schedules it and returns the task, so the session is removed from storage immediately but its kernel process may still be running when the call returns. Callers that need the kernel's resources actually released — GPU memory, for instance — should use :meth:aclose_session instead, or await the returned task.

Parameters:

Name	Type	Description	Default
session_id	str	ID of the session to close	required

Returns:

Type	Description
Optional[Task[None]]	The teardown task, or None when there was no kernel to tear
Optional[Task[None]]	down or no running event loop to schedule it on.

Source code in src/agora_workbench/code_execution/sessions/manager.py

def close_session(self, session_id: str) -> "Optional[asyncio.Task[None]]":
    """
    Explicitly close a session.

    Attempts to clean up resources first. If cleanup fails, the session
    is still deleted to prevent session accumulation, but the error is logged.

    Kernel teardown is asynchronous. This method schedules it and returns
    the task, so the session is removed from storage immediately but its
    kernel process may still be running when the call returns. Callers that
    need the kernel's resources actually released — GPU memory, for
    instance — should use :meth:`aclose_session` instead, or await the
    returned task.

    Args:
        session_id: ID of the session to close

    Returns:
        The teardown task, or ``None`` when there was no kernel to tear
        down or no running event loop to schedule it on.
    """
    running_job_id = self._get_running_job_for_session(session_id)
    if running_job_id:
        job = self._background_jobs.get(running_job_id)
        if job:
            job.success = False
            job.error = f"Session {session_id} was closed while job {running_job_id} was running"
            job.status = "failed"
            if job.task and not job.task.done():
                job.task.cancel()

    shutdown_task = self._schedule_kernel_shutdown(session_id, caller="close_session()")

    session = self.storage.retrieve(session_id)

    if session:
        # Run cleanup
        cleanup_failed = False
        try:
            session.cleanup()
        except Exception as e:
            cleanup_failed = True
            LOGGER.error(
                f"Error during cleanup of session {session_id}: {e}. "
                f"Session will still be removed, but resources may be leaked."
            )

        # Remove from storage
        self.storage.delete(session_id)
        if cleanup_failed:
            LOGGER.warning(f"Closed session {session_id} with failed cleanup (remaining={self.storage.count()})")
        else:
            LOGGER.info(f"Closed session {session_id} (remaining={self.storage.count()})")

    return shutdown_task

aclose_session(session_id) async

Close a session and wait for its kernel to actually shut down.

The awaitable counterpart to :meth:close_session. Prefer this wherever the kernel's resources must be released before proceeding — freeing GPU memory, tearing down a batch's child sessions, or reclaiming capacity before starting new work.

Source code in src/agora_workbench/code_execution/sessions/manager.py

async def aclose_session(self, session_id: str) -> None:
    """Close a session and wait for its kernel to actually shut down.

    The awaitable counterpart to :meth:`close_session`. Prefer this
    wherever the kernel's resources must be released before proceeding —
    freeing GPU memory, tearing down a batch's child sessions, or
    reclaiming capacity before starting new work.
    """
    self.close_session(session_id)
    await self.await_kernel_shutdown(session_id)

get_kernel_generation(session_id)

Return the generation id of the session's current kernel.

Generation ids are globally unique and monotonically increasing across the lifetime of this manager: restarting a session's kernel always yields a strictly greater id, and an id is never reused. Callers can therefore cache per-kernel state keyed on this value and detect a rebuilt kernel by comparing against the value they captured.

Returns None when the session has no live kernel.

Source code in src/agora_workbench/code_execution/sessions/manager.py

def get_kernel_generation(self, session_id: str) -> Optional[int]:
    """Return the generation id of the session's current kernel.

    Generation ids are globally unique and monotonically increasing across
    the lifetime of this manager: restarting a session's kernel always
    yields a strictly greater id, and an id is never reused. Callers can
    therefore cache per-kernel state keyed on this value and detect a
    rebuilt kernel by comparing against the value they captured.

    Returns ``None`` when the session has no live kernel.
    """
    return self._kernel_generations.get(session_id)

is_kernel_bootstrapped(session_id, key)

Whether key has been completed against the session's current kernel.

Returns False when the session has no kernel, or when its kernel has been rebuilt since the bootstrap step ran.

Source code in src/agora_workbench/code_execution/sessions/manager.py

def is_kernel_bootstrapped(self, session_id: str, key: str) -> bool:
    """Whether ``key`` has been completed against the session's *current* kernel.

    Returns ``False`` when the session has no kernel, or when its kernel
    has been rebuilt since the bootstrap step ran.
    """
    generation = self._kernel_generations.get(session_id)
    if generation is None:
        return False
    return key in self._kernel_bootstrap_state.get(generation, frozenset())

mark_kernel_bootstrapped(session_id, key)

Record that key has been completed against the current kernel.

Returns False (and records nothing) when the session has no live kernel, so the caller's work will be retried against the next one.

Source code in src/agora_workbench/code_execution/sessions/manager.py

def mark_kernel_bootstrapped(self, session_id: str, key: str) -> bool:
    """Record that ``key`` has been completed against the current kernel.

    Returns ``False`` (and records nothing) when the session has no live
    kernel, so the caller's work will be retried against the next one.
    """
    generation = self._kernel_generations.get(session_id)
    if generation is None:
        return False
    self._kernel_bootstrap_state.setdefault(generation, set()).add(key)
    return True

get_artifact_record(session_id, token)

Look up an artifact for the download endpoint.

Returns None for unknown sessions/tokens and for tokens whose backing file has been removed — both surface as 404 at the HTTP layer.

Source code in src/agora_workbench/code_execution/sessions/manager.py

def get_artifact_record(self, session_id: str, token: str) -> Optional[_ArtifactRecord]:
    """Look up an artifact for the download endpoint.

    Returns ``None`` for unknown sessions/tokens *and* for tokens whose
    backing file has been removed — both surface as 404 at the HTTP
    layer.
    """
    record = self._session_artifacts.get(session_id, {}).get(token)
    if record is None:
        return None
    if not record.path.is_file():
        return None
    return record

find_artifact_by_name(session_id, artifact_name)

Find a registered artifact by its relative name within a session.

Looks up a previously registered artifact (i.e. one that was discovered during a snapshot-diff after an execute) by its relative filename. This is used by the publish pipeline to resolve the local path of an artifact before pushing it to remote storage.

Parameters:

Name	Type	Description	Default
session_id	str	The session that owns the artifact.	required
artifact_name	str	Relative filename as registered (e.g. "results.csv" or "subdir/report.pdf").	required

Returns:

Name	Type	Description
The	Optional[_ArtifactRecord]	class:_ArtifactRecord if found and the backing file still
	Optional[_ArtifactRecord]	exists, otherwise None.

Source code in src/agora_workbench/code_execution/sessions/manager.py

def find_artifact_by_name(self, session_id: str, artifact_name: str) -> Optional[_ArtifactRecord]:
    """Find a registered artifact by its relative name within a session.

    Looks up a previously registered artifact (i.e. one that was discovered
    during a snapshot-diff after an execute) by its relative filename.  This
    is used by the publish pipeline to resolve the local path of an artifact
    before pushing it to remote storage.

    Args:
        session_id: The session that owns the artifact.
        artifact_name: Relative filename as registered (e.g. ``"results.csv"``
            or ``"subdir/report.pdf"``).

    Returns:
        The :class:`_ArtifactRecord` if found and the backing file still
        exists, otherwise ``None``.
    """
    registry = self._session_artifacts.get(session_id, {})
    for record in registry.values():
        if record.name == artifact_name:
            if record.path.is_file():
                return record
            return None
    return None

start_background_execution_for_session(session_id, code, timeout, working_dir=None) async

Start execution and return immediately with a background job id.

Source code in src/agora_workbench/code_execution/sessions/manager.py

async def start_background_execution_for_session(
    self, session_id: str, code: str, timeout: float, working_dir: Optional[str] = None
) -> dict[str, str]:
    """Start execution and return immediately with a background job id."""
    try:
        session = self.get_session(session_id)
        user_token = session.user_token
        user_identity = session.user_identity
    except ValueError:
        raise ValueError(
            f"Session {session_id} is no longer available (expired or cleaned up). "
            "Please create a new session and retry."
        ) from None

    running_job_id = self._get_running_job_for_session(session_id)
    if running_job_id:
        raise RuntimeError(f"Session busy — job {running_job_id} is still running")

    km, kc = await self._get_or_create_kernel(
        session_id, working_dir, user_token=user_token, user_identity=user_identity
    )
    # Same preamble shape as the foreground path: outputs preamble
    # must run before the token preamble so AGORA_OUTPUT_DIR is
    # populated even when this background call is the kernel's
    # first execute.  Snapshot before kc.execute so the terminal
    # handler can diff for artifacts.
    code = self._prepare_outputs_preamble(session_id) + self._prepare_code_with_token_preamble(
        session_id, code, user_token
    )
    outputs_before = self._snapshot_outputs_dir(session_id)
    msg_id = kc.execute(code)

    job_id = f"j_{uuid.uuid4().hex[:12]}"
    job = _BackgroundJob(
        job_id=job_id,
        session_id=session_id,
        msg_id=msg_id,
        timeout=timeout,
        start_time=time.monotonic(),
        user_identity=user_identity,
        outputs_before=outputs_before,
    )
    job.task = asyncio.create_task(self._collect_background_job(job, km, kc))
    self._background_jobs[job_id] = job
    self._session_running_jobs[session_id] = job_id

    return {"job_id": job_id, "status": job.status, "session_id": session_id}

start_promoted_execution_for_session(session_id, code, timeout, promotion_threshold_s, working_dir=None) async

Execute code synchronously, promoting to a background job if it exceeds the threshold.

Starts executing like the normal foreground path. If the kernel reaches idle within promotion_threshold_s seconds the result is returned as a (stdout, stderr, success, displays, artifacts) tuple — identical to :meth:execute_code_for_session. If the threshold expires while the kernel is still busy, the in-flight execution is registered as a :class:_BackgroundJob and a job-handle dict is returned instead.

Parameters:

Name	Type	Description	Default
session_id	str	Session identifier.	required
code	str	Python code to execute. This method applies the outputs and token preambles internally (same as the foreground path), so callers should pass the raw user code — not pre-preambled.	required
timeout	float	Total execution timeout in seconds.	required
promotion_threshold_s	float	Seconds to wait before promoting.	required
working_dir	Optional[str]	Optional working directory for the kernel.	None

Returns:

Type	Description
Tuple[str, str, bool, list[dict], list[dict]] | dict[str, Any]	Either the 5-tuple (stdout, stderr, success, displays, artifacts)
Tuple[str, str, bool, list[dict], list[dict]] | dict[str, Any]	when the execution completes within the threshold, or a dict with
Tuple[str, str, bool, list[dict], list[dict]] | dict[str, Any]	job_id, status, session_id, and promoted=True.

Source code in src/agora_workbench/code_execution/sessions/manager.py

async def start_promoted_execution_for_session(
    self,
    session_id: str,
    code: str,
    timeout: float,
    promotion_threshold_s: float,
    working_dir: Optional[str] = None,
) -> "Tuple[str, str, bool, list[dict], list[dict]] | dict[str, Any]":
    """Execute code synchronously, promoting to a background job if it exceeds the threshold.

    Starts executing like the normal foreground path. If the kernel reaches
    ``idle`` within *promotion_threshold_s* seconds the result is returned
    as a ``(stdout, stderr, success, displays, artifacts)`` tuple — identical
    to :meth:`execute_code_for_session`.  If the threshold expires while the
    kernel is still busy, the in-flight execution is registered as a
    :class:`_BackgroundJob` and a job-handle dict is returned instead.

    Args:
        session_id: Session identifier.
        code: Python code to execute.  This method applies the outputs and
            token preambles internally (same as the foreground path), so
            callers should pass the raw user code — **not** pre-preambled.
        timeout: Total execution timeout in seconds.
        promotion_threshold_s: Seconds to wait before promoting.
        working_dir: Optional working directory for the kernel.

    Returns:
        Either the 5-tuple ``(stdout, stderr, success, displays, artifacts)``
        when the execution completes within the threshold, or a dict with
        ``job_id``, ``status``, ``session_id``, and ``promoted=True``.
    """
    if timeout <= promotion_threshold_s:
        raise ValueError(
            f"timeout ({timeout}s) must be greater than promotion_threshold_s ({promotion_threshold_s}s)"
        )

    try:
        session = self.get_session(session_id)
        user_token = session.user_token
        user_identity = session.user_identity
    except ValueError:
        return (
            "",
            (
                f"Session {session_id} is no longer available (expired or cleaned up). "
                "Please create a new session and retry."
            ),
            False,
            [],
            [],
        )

    running_job_id = self._get_running_job_for_session(session_id)
    if running_job_id:
        return "", f"Session busy — job {running_job_id} is still running", False, [], []

    async with self._get_kernel_execute_lock(session_id):
        km, kc = await self._get_or_create_kernel(
            session_id, working_dir, user_token=user_token, user_identity=user_identity
        )
        code = self._prepare_outputs_preamble(session_id) + self._prepare_code_with_token_preamble(
            session_id, code, user_token
        )
        outputs_before = self._snapshot_outputs_dir(session_id)
        msg_id = kc.execute(code)

        # --- Collect iopub messages for up to promotion_threshold_s ---
        stdout_parts: list[str] = []
        stderr_parts: list[str] = []
        displays: list[dict] = []
        displays_bytes = 0
        success = True
        completed = False

        def _maybe_capture_display(content_data: dict, content_meta: dict) -> None:
            nonlocal displays_bytes
            entry = _extract_display(content_data, content_meta)
            if entry is None:
                return
            size = len(entry["data"]) if isinstance(entry["data"], (str, bytes)) else 0
            if displays_bytes + size > _MAX_DISPLAY_BYTES_PER_EXECUTE:
                return
            displays_bytes += size
            displays.append(entry)

        start_time = time.monotonic()
        deadline = start_time + promotion_threshold_s

        while True:
            now = time.monotonic()
            remaining = deadline - now
            if remaining <= 0:
                break  # threshold reached, promote

            try:
                msg = await asyncio.wait_for(
                    kc.get_iopub_msg(timeout=min(remaining, 1.0)),
                    timeout=remaining + 0.5,
                )
            except (queue.Empty, asyncio.TimeoutError):
                if time.monotonic() >= deadline:
                    break
                continue

            msg_type = msg["msg_type"]
            content = msg.get("content", {})
            parent_id = msg.get("parent_header", {}).get("msg_id")
            if parent_id != msg_id:
                continue

            if msg_type == "stream":
                stream_name = content.get("name", "stdout")
                text = content.get("text", "")
                if stream_name == "stdout":
                    stdout_parts.append(text)
                elif stream_name == "stderr":
                    stderr_parts.append(text)
            elif msg_type == "error":
                traceback = "\n".join(content.get("traceback", []))
                stderr_parts.append(_strip_ansi(traceback))
                success = False
            elif msg_type == "execute_result":
                data = content.get("data", {})
                text_result = data.get("text/plain", "")
                if text_result:
                    stdout_parts.append(text_result)
                _maybe_capture_display(data, content.get("metadata", {}))
            elif msg_type in ("display_data", "update_display_data"):
                data = content.get("data", {})
                text_result = data.get("text/plain", "")
                if text_result:
                    stdout_parts.append(text_result)
                _maybe_capture_display(data, content.get("metadata", {}))
            elif msg_type == "status" and content.get("execution_state") == "idle":
                # Drain trailing messages
                while True:
                    try:
                        trailing_msg = await asyncio.wait_for(kc.get_iopub_msg(timeout=0.2), timeout=0.5)
                    except (queue.Empty, asyncio.TimeoutError, Exception):
                        break
                    r_type = trailing_msg["msg_type"]
                    r_content = trailing_msg.get("content", {})
                    r_parent = trailing_msg.get("parent_header", {}).get("msg_id")
                    if r_parent != msg_id:
                        continue
                    if r_type == "stream":
                        txt = r_content.get("text", "")
                        if r_content.get("name") == "stdout":
                            stdout_parts.append(txt)
                        elif r_content.get("name") == "stderr":
                            stderr_parts.append(txt)
                    elif r_type == "execute_result":
                        data = r_content.get("data", {})
                        t = data.get("text/plain", "")
                        if t:
                            stdout_parts.append(t)
                        _maybe_capture_display(data, r_content.get("metadata", {}))
                    elif r_type in ("display_data", "update_display_data"):
                        data = r_content.get("data", {})
                        t = data.get("text/plain", "")
                        if t:
                            stdout_parts.append(t)
                        _maybe_capture_display(data, r_content.get("metadata", {}))
                completed = True
                break

        if completed:
            # Finished within threshold — return as normal foreground result
            stdout = "".join(stdout_parts)
            stderr = "".join(stderr_parts)
            self._kernel_last_used[session_id] = time.time()
            try:
                outputs_after = self._snapshot_outputs_dir(session_id)
                artifacts = self._register_artifacts_from_diff(session_id, outputs_before, outputs_after)
            except Exception:
                LOGGER.warning("Artifact discovery failed for session %s", session_id, exc_info=True)
                artifacts = []
            return stdout, stderr, success, displays, artifacts

        # --- Promote to background job ---
        job_id = f"j_{uuid.uuid4().hex[:12]}"
        job = _BackgroundJob(
            job_id=job_id,
            session_id=session_id,
            msg_id=msg_id,
            timeout=timeout,
            start_time=start_time,
            user_identity=user_identity or "",
            stdout_parts=stdout_parts,
            stderr_parts=stderr_parts,
            success=success,
            outputs_before=outputs_before,
        )
        job.task = asyncio.create_task(self._collect_background_job(job, km, kc))
        self._background_jobs[job_id] = job
        self._session_running_jobs[session_id] = job_id

        return {
            "job_id": job_id,
            "status": "running",
            "session_id": session_id,
            "promoted": True,
        }

check_background_job(job_id, caller_identity=None)

Return current status/output for a background job.

Parameters:

Name	Type	Description	Default
job_id	str	Background job identifier.	required
caller_identity	Optional[str]	When provided, the caller's user identity is compared to the identity that submitted the job. Both a missing job and an identity mismatch raise ValueError("Job {job_id} not found") so that callers cannot use differing error messages to probe whether a job ID exists.	None

Source code in src/agora_workbench/code_execution/sessions/manager.py

def check_background_job(self, job_id: str, caller_identity: Optional[str] = None) -> dict[str, Any]:
    """Return current status/output for a background job.

    Args:
        job_id: Background job identifier.
        caller_identity: When provided, the caller's user identity is compared to
            the identity that submitted the job. Both a missing job and an identity
            mismatch raise ``ValueError("Job {job_id} not found")`` so that callers
            cannot use differing error messages to probe whether a job ID exists.
    """
    job = self._background_jobs.get(job_id)
    if not job:
        raise ValueError(f"Job {job_id} not found")

    if caller_identity and job.user_identity and caller_identity != job.user_identity:
        raise ValueError(f"Job {job_id} not found")

    now = time.monotonic()
    elapsed_seconds = (job.completed_at or now) - job.start_time
    result: dict[str, Any] = {
        "job_id": job.job_id,
        "session_id": job.session_id,
        "status": job.status,
        "elapsed_seconds": round(elapsed_seconds, 3),
    }

    if job.status == "running":
        result["stdout"] = "".join(job.stdout_parts)
        result["stderr"] = "".join(job.stderr_parts)
        return result

    result["stdout"] = "".join(job.stdout_parts)
    result["stderr"] = "".join(job.stderr_parts)
    result["success"] = job.success
    if job.error:
        result["error"] = job.error
    # Artifacts are populated by _finalize_background_artifacts when the
    # job reaches a terminal state; carry them here without download URLs
    # — URL composition happens in the MCP server layer where the public
    # base URL is known (matches the foreground path's contract).
    result["artifacts"] = list(job.artifacts)
    return result

await_background_job(job_id) async

Wait for a background job's task to reach a terminal state, then return its status.

Returns None if the job was never registered or has already been purged. Used by the activity publisher to emit job_finished events; callers must treat exceptions on the underlying task as terminal (a failed task still leaves _BackgroundJob.status set by _collect_background_job).

Source code in src/agora_workbench/code_execution/sessions/manager.py

async def await_background_job(self, job_id: str) -> Optional[dict[str, Any]]:
    """Wait for a background job's task to reach a terminal state, then return its status.

    Returns ``None`` if the job was never registered or has already been purged.
    Used by the activity publisher to emit ``job_finished`` events; callers must
    treat exceptions on the underlying task as terminal (a failed task still
    leaves ``_BackgroundJob.status`` set by ``_collect_background_job``).
    """
    job = self._background_jobs.get(job_id)
    if job is None or job.task is None:
        return None
    try:
        await job.task
    except Exception:
        LOGGER.debug("Background job task raised for job_id=%s", job_id, exc_info=True)
    try:
        return self.check_background_job(job_id)
    except ValueError:
        return None

execute_code_for_session(session_id, code, timeout, working_dir=None) async

Execute code in the session's Jupyter kernel.

Concurrent calls for the same session_id are serialized by a per-session asyncio.Lock so they cannot race on the shared Jupyter KernelClient iopub stream.

Parameters:

Name	Type	Description	Default
session_id	str	Session identifier	required
code	str	Python code to execute	required
timeout	float	Execution timeout in seconds	required
working_dir	Optional[str]	Optional working directory for kernel	None

Returns:

Type	Description
str	Tuple of (stdout, stderr, success, displays, artifacts).
str	displays contains rich-output payloads emitted by the kernel,
bool	with MIME type, data, and metadata fields. artifacts contains
list[dict]	metadata for each new or modified file under the session output
list[dict]	directory, including its name, size, MIME type, modification time,
Tuple[str, str, bool, list[dict], list[dict]]	and download token.

Source code in src/agora_workbench/code_execution/sessions/manager.py

async def execute_code_for_session(
    self, session_id: str, code: str, timeout: float, working_dir: Optional[str] = None
) -> Tuple[str, str, bool, list[dict], list[dict]]:
    """
    Execute code in the session's Jupyter kernel.

    Concurrent calls for the same ``session_id`` are serialized by a
    per-session asyncio.Lock so they cannot race on the shared
    Jupyter ``KernelClient`` iopub stream.

    Args:
        session_id: Session identifier
        code: Python code to execute
        timeout: Execution timeout in seconds
        working_dir: Optional working directory for kernel

    Returns:
        Tuple of ``(stdout, stderr, success, displays, artifacts)``.
        ``displays`` contains rich-output payloads emitted by the kernel,
        with MIME type, data, and metadata fields. ``artifacts`` contains
        metadata for each new or modified file under the session output
        directory, including its name, size, MIME type, modification time,
        and download token.
    """
    # Look up session to get current user credentials, ensuring session
    # access goes through the manager (cleanup, expiry check, touch).
    try:
        session = self.get_session(session_id)
        user_token = session.user_token
        user_identity = session.user_identity
    except ValueError:
        return (
            "",
            (
                f"Session {session_id} is no longer available (expired or cleaned up). "
                "Please create a new session and retry."
            ),
            False,
            [],
            [],
        )

    running_job_id = self._get_running_job_for_session(session_id)
    if running_job_id:
        return "", f"Session busy — job {running_job_id} is still running", False, [], []

    async with self._get_kernel_execute_lock(session_id):
        return await self._execute_code_locked(
            session_id=session_id,
            code=code,
            timeout=timeout,
            working_dir=working_dir,
            user_token=user_token,
            user_identity=user_identity,
        )

await_kernel_shutdown(session_id) async

Wait for any in-flight teardown of this session's kernel to finish.

No-op when none is running. Shielded, so a cancelled caller does not cancel the teardown itself; failures are already reported by the done-callback, so awaiting is purely for sequencing.

Source code in src/agora_workbench/code_execution/sessions/manager.py

async def await_kernel_shutdown(self, session_id: str) -> None:
    """Wait for any in-flight teardown of this session's kernel to finish.

    No-op when none is running. Shielded, so a cancelled caller does not
    cancel the teardown itself; failures are already reported by the
    done-callback, so awaiting is purely for sequencing.
    """
    task = self._kernel_shutdown_tasks.get(session_id)
    if task is None or task.done():
        return
    try:
        await asyncio.shield(task)
    except asyncio.CancelledError:
        raise
    except Exception:
        LOGGER.debug("Awaited kernel shutdown for session %s failed", session_id, exc_info=True)

cleanup_idle_kernels(max_idle_time=3600.0) async

Cleanup kernels that have been idle for too long.

Source code in src/agora_workbench/code_execution/sessions/manager.py

async def cleanup_idle_kernels(self, max_idle_time: float = 3600.0):
    """Cleanup kernels that have been idle for too long."""
    now = time.time()
    idle_sessions = [sid for sid, last_used in self._kernel_last_used.items() if now - last_used > max_idle_time]

    for session_id in idle_sessions:
        LOGGER.info(f"Cleaning up idle kernel for session {session_id}")
        await self._shutdown_kernel(session_id)

list_sessions()

List all active sessions with metadata.

Returns:

Type	Description
list[dict[str, Any]]	List of session info dicts

Source code in src/agora_workbench/code_execution/sessions/manager.py

def list_sessions(self) -> list[dict[str, Any]]:
    """
    List all active sessions with metadata.

    Returns:
        List of session info dicts
    """
    self._maybe_cleanup()

    sessions = self.storage.list_all()

    result = []
    for session in sessions.values():
        result.append(session.get_info())

    return result

Session Configuration

agora_workbench.code_execution.sessions.manager.SessionConfig(max_sessions=100, timeout_minutes=30, cleanup_interval_seconds=300, storage_backend=None, data_manager_factory=None)

Configuration for session manager.

Initialize session manager configuration.

Parameters:

Name	Type	Description	Default
max_sessions	int	Maximum number of concurrent sessions.	100
timeout_minutes	int	Idle time after which a session is cleaned up.	30
cleanup_interval_seconds	int	Minimum interval between cleanup sweeps.	300
storage_backend	Optional[SessionStorageBackend]	Optional session storage backend.	None
data_manager_factory	Optional[Callable[[SessionContext], DataLakeDataManager]]	Optional callable invoked once per session to build its :class:DataLakeDataManager, receiving a :class:SessionContext describing the session being created. Use it to supply a customized manager — a different credential, a custom artifact resolver, extra fetchers, or configuration derived from user_identity / user_token. The factory must return a fresh instance per call. The session takes ownership of the manager and calls cleanup() on it when the session ends, so returning a shared singleton would let the first session torn down destroy a manager still in use by the others. When omitted, each session builds a default DataLakeDataManager(), matching previous behavior. A factory that returns None — or any object without a cleanup() method — raises TypeError from create_session, rather than silently falling back to the default manager.	None

Source code in src/agora_workbench/code_execution/sessions/manager.py

def __init__(
    self,
    max_sessions: int = 100,
    timeout_minutes: int = 30,
    cleanup_interval_seconds: int = 300,  # 5 minutes
    storage_backend: Optional[SessionStorageBackend] = None,
    data_manager_factory: Optional[Callable[[SessionContext], "DataLakeDataManager"]] = None,
):
    """
    Initialize session manager configuration.

    Args:
        max_sessions: Maximum number of concurrent sessions.
        timeout_minutes: Idle time after which a session is cleaned up.
        cleanup_interval_seconds: Minimum interval between cleanup sweeps.
        storage_backend: Optional session storage backend.
        data_manager_factory: Optional callable invoked once per session to
            build its :class:`DataLakeDataManager`, receiving a
            :class:`SessionContext` describing the session being created.
            Use it to supply a customized manager — a different credential,
            a custom artifact resolver, extra fetchers, or configuration
            derived from ``user_identity`` / ``user_token``.

            The factory **must return a fresh instance per call**. The
            session takes ownership of the manager and calls ``cleanup()``
            on it when the session ends, so returning a shared singleton
            would let the first session torn down destroy a manager still
            in use by the others.

            When omitted, each session builds a default
            ``DataLakeDataManager()``, matching previous behavior.

            A factory that returns ``None`` — or any object without a
            ``cleanup()`` method — raises ``TypeError`` from
            ``create_session``, rather than silently falling back to the
            default manager.
    """
    self.max_sessions = max_sessions
    self.timeout = timedelta(minutes=timeout_minutes)
    self.cleanup_interval = timedelta(seconds=cleanup_interval_seconds)
    self.storage_backend = storage_backend or InMemoryStorage()
    self.data_manager_factory = data_manager_factory