Sessions

Session Container

agora_workbench.code_execution.sessions.session.Session(session_id, data, session_type, user_identity, user_token, token_claims, metadata=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.

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

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,
):
    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
    from ..data_access.manager import DataLakeDataManager

    self.data_manager = DataLakeDataManager()

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

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):
    self.config = config or SessionConfig()
    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.  ``_session_artifacts`` is the
    # token -> record map used by the HTTP download endpoint;
    # ``_kernel_outputs_initialized`` tracks which kernels have already
    # received the AGORA_OUTPUT_DIR preamble so we only inject it once.
    self._session_artifacts: dict[str, dict[str, _ArtifactRecord]] = {}
    self._kernel_outputs_initialized: set[str] = set()

    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

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
    """
    # 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())

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

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

Parameters:

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

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

def close_session(self, session_id: str) -> 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.

    Args:
        session_id: ID of the session to close
    """
    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()

    if session_id in self._kernels:
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                asyncio.create_task(self._shutdown_kernel(session_id))
            else:
                loop.run_until_complete(self._shutdown_kernel(session_id))
        except Exception as e:
            LOGGER.error(f"Failed to shutdown kernel for session {session_id}: {e}")

    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()})")

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

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)

Configuration for session manager.

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,
):
    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()