src/managers/execution/namespace_kernel.py

"""
Shared namespace-exec core (ADR-0007 Phase 1).

`PythonExecutor` (in-process, the default) and the sandbox exec-kernel
(`sandbox/kernel.py`, runs code inside an isolated per-session container) both
need the *same* semantics: a persistent Jupyter-kernel-style namespace, seeded
with pandas/numpy + file-writing helpers, into which code is `exec`'d with stdout
captured and capped. Rather than duplicate that logic across process boundaries,
it lives here once as pure functions + a tiny `NamespaceKernel` holder.

`PythonExecutor` is a thin in-process wrapper over this; the sandbox kernel wraps
the SAME core behind an HTTP surface. Keeping one implementation means the two
executors can never drift in their exec/capture/truncation behavior.

Nothing here imports Docker, requests, or MCP — it is the innermost, dependency
-light layer (pandas/numpy are imported lazily inside `seed_namespace`, exactly
as `PythonExecutor.reset_environment` did, so importing this module never fails
when those heavy deps are absent).
"""

from __future__ import annotations

import io
import sys
from typing import Any

__all__ = [
    "MAX_OUTPUT_CHARS",
    "fresh_namespace",
    "seed_namespace",
    "exec_capture",
    "NamespaceKernel",
]

# Cap on stdout captured per execution. Large prints (e.g. full dataframes) get
# re-sent to the LLM on every subsequent step, so an uncapped print can multiply
# token cost across the rest of the run. (Moved verbatim from PythonExecutor.)
MAX_OUTPUT_CHARS = 4000


def fresh_namespace() -> dict[str, Any]:
    """Return a bare namespace dict with builtins available."""
    return {"__builtins__": __builtins__}


def seed_namespace(namespace: dict[str, Any]) -> None:
    """Seed a namespace with the standard libs + file-writing helpers.

    Mirrors `PythonExecutor.reset_environment` exactly. pandas/numpy are imported
    LAZILY so a namespace can still be seeded (minus those names) where the heavy
    deps are absent — the same tolerant posture the original had.
    """
    try:
        import os
        from pathlib import Path

        import numpy as np
        import pandas as pd

        namespace.update({"pd": pd, "np": np, "os": os, "Path": Path})

        def write_text_file(filepath, content):
            """Write text content to a file."""
            with open(filepath, "w") as f:
                f.write(content)
            return f"File written to {filepath}"

        def write_dataframe_to_csv(df, filepath, **kwargs):
            """Write pandas DataFrame to CSV file."""
            df.to_csv(filepath, **kwargs)
            return f"DataFrame written to {filepath}"

        def write_dataframe_to_tsv(df, filepath, **kwargs):
            """Write pandas DataFrame to TSV file."""
            df.to_csv(filepath, sep="\t", **kwargs)
            return f"DataFrame written to {filepath}"

        def create_directory(dirpath):
            """Create directory if it doesn't exist."""
            Path(dirpath).mkdir(parents=True, exist_ok=True)
            return f"Directory created: {dirpath}"

        namespace.update(
            {
                "write_text_file": write_text_file,
                "write_dataframe_to_csv": write_dataframe_to_csv,
                "write_dataframe_to_tsv": write_dataframe_to_tsv,
                "create_directory": create_directory,
            }
        )
    except ImportError as e:
        print(f"Warning: Could not import library: {e}")


def exec_capture(code: str, namespace: dict[str, Any]) -> str:
    """`exec` `code` into `namespace`, returning captured (and capped) stdout.

    Same contract as `PythonExecutor.execute`: empty output → the
    "Code executed successfully" sentinel; oversized output truncated with a
    self-describing tail; any exception → an "Error: …" string (never raised).
    """
    try:
        old_stdout = sys.stdout
        sys.stdout = captured_output = io.StringIO()
        try:
            # Executing LLM-generated code is this agent's core purpose. In-process
            # this runs with the container's privileges (the accepted prototype
            # risk); the ADR-0007 sandbox reuses THIS same function inside an
            # isolated per-session container so the untrusted code is confined.
            exec(code, namespace)  # nosec B102
            result = captured_output.getvalue().strip()
            if not result:
                return "Code executed successfully"
            if len(result) > MAX_OUTPUT_CHARS:
                omitted = len(result) - MAX_OUTPUT_CHARS
                result = (
                    result[:MAX_OUTPUT_CHARS]
                    + f"\n...[output truncated, {omitted} more characters omitted. "
                    f"Print a smaller summary (e.g. .head(), .describe(), or specific columns) "
                    f"if you need to inspect this further.]"
                )
            return result
        finally:
            sys.stdout = old_stdout
    except Exception as e:
        return f"Error: {str(e)}"


class NamespaceKernel:
    """A persistent, seeded namespace with `exec`-and-capture semantics.

    The shared engine behind BOTH executors: `PythonExecutor` delegates to an
    instance of this in-process; the sandbox `kernel.py` HTTP server holds one
    per session and drives it over HTTP. State (variables, injected functions)
    persists across `execute()` calls, matching the Jupyter-kernel model the
    agent relies on.
    """

    def __init__(self) -> None:
        self.namespace: dict[str, Any] = {}
        self.reset()

    def reset(self) -> None:
        """Reset to a clean, freshly-seeded namespace."""
        self.namespace = fresh_namespace()
        seed_namespace(self.namespace)

    def send_functions(self, functions: dict[str, Any]) -> None:
        """Inject callables (or any names) into the namespace."""
        if functions:
            self.namespace.update(functions)

    def send_variables(self, variables: dict[str, Any]) -> None:
        """Inject variables into the namespace."""
        if variables:
            self.namespace.update(variables)

    def execute(self, code: str) -> str:
        """Execute `code` in the persistent namespace; return captured stdout."""
        return exec_capture(code, self.namespace)