@prometheus-ai/agent 0.5.3 → 0.5.8

package/src/eval/py/runner.py

@@ -51,8 +51,23 @@ from typing import Any
 # Frame writer
 # ---------------------------------------------------------------------------
 
-_RAW_STDOUT = sys.__stdout__
+# Frames travel on a private dup of the original stdout. fd 1 itself is then
+# repointed at a capture pipe: child processes spawned by user code without
+# stdout=PIPE inherit fd 1, and their output is forwarded to the host as
+# regular stdout frames by a drain thread instead of being written raw into
+# the NDJSON channel (where it would be dropped as invalid JSON — or worse,
+# spoof a frame). The wire protocol is unchanged: the host still reads NDJSON
+# frames from the subprocess stdout.
 _RAW_STDERR = sys.__stderr__
+try:
+    _FRAME_FD = os.dup(sys.__stdout__.fileno())
+    _RAW_STDOUT = os.fdopen(_FRAME_FD, "w", encoding="utf-8", errors="backslashreplace")
+    _CAPTURE_READ_FD, _capture_write_fd = os.pipe()
+    os.dup2(_capture_write_fd, sys.__stdout__.fileno())
+    os.close(_capture_write_fd)
+except (AttributeError, OSError, ValueError, io.UnsupportedOperation):
+    _RAW_STDOUT = sys.__stdout__
+    _CAPTURE_READ_FD = None
 _OUT_LOCK = threading.Lock()
 
 
@@ -78,11 +93,22 @@ def _emit(frame: dict) -> None:
 
 
 class _StreamProxy(io.TextIOBase):
-    """Emit each ``write()`` as a typed frame tied to the current request."""
+    """Emit ``write()`` data as typed frames tied to the current request.
+
+    Writes are coalesced per request: a frame is emitted once the buffer holds
+    a complete line (everything up to the last newline goes out together) or
+    grows past ``_MAX_BUFFER`` bytes, so the common ``print()`` pair of
+    ``write(text)`` + ``write("\\n")`` costs one frame instead of two. Partial
+    lines are bounded by ``flush()`` and the end-of-request flush.
+    """
+
+    _MAX_BUFFER = 8192
 
     def __init__(self, kind: str) -> None:
         super().__init__()
         self._kind = kind
+        self._lock = threading.Lock()
+        self._buffers: dict[str, str] = {}
 
     def writable(self) -> bool:  # noqa: D401 - protocol method
         return True
@@ -100,12 +126,44 @@ class _StreamProxy(io.TextIOBase):
             _RAW_STDERR.write(data)
             _RAW_STDERR.flush()
             return len(data)
-        _emit({"type": self._kind, "id": rid, "data": data})
+        emit_text = None
+        with self._lock:
+            buf = self._buffers.pop(rid, "") + data
+            if len(buf) >= self._MAX_BUFFER:
+                emit_text = buf
+            else:
+                nl = buf.rfind("\n")
+                if nl >= 0:
+                    emit_text = buf[: nl + 1]
+                    rest = buf[nl + 1 :]
+                    if rest:
+                        self._buffers[rid] = rest
+                else:
+                    self._buffers[rid] = buf
+        if emit_text:
+            _emit({"type": self._kind, "id": rid, "data": emit_text})
         return len(data)
 
     def flush(self) -> None:  # noqa: D401 - protocol method
+        rid = _CURRENT_RID.get()
+        if rid is not None:
+            self.flush_rid(rid)
         return None
 
+    def flush_rid(self, rid: str) -> None:
+        """Flush any buffered partial line for ``rid`` as its own frame."""
+        with self._lock:
+            buf = self._buffers.pop(rid, None)
+        if buf:
+            _emit({"type": self._kind, "id": rid, "data": buf})
+
+
+def _flush_stream_proxies(rid: str) -> None:
+    """Drain buffered proxy output for ``rid`` (called before its done frame)."""
+    for stream in (sys.stdout, sys.stderr):
+        if isinstance(stream, _StreamProxy):
+            stream.flush_rid(rid)
+
 
 # ---------------------------------------------------------------------------
 # Runner state
@@ -125,13 +183,53 @@ class _RunnerState:
         self.last_install_marker: int = 0
         self.loop: asyncio.AbstractEventLoop | None = None
         self.active_executions: int = 0
+        # Best-effort attribution target for captured fd-1 bytes (child
+        # processes inheriting stdout). With overlapping requests the most
+        # recently started one wins — strictly better than dropping the bytes.
+        self.capture_rid: str | None = None
 
 
-_CURRENT_RID: contextvars.ContextVar[str | None] = contextvars.ContextVar("prometheus_current_rid", default=None)
+_CURRENT_RID: contextvars.ContextVar[str | None] = contextvars.ContextVar("omp_current_rid", default=None)
 
 _STATE = _RunnerState()
 
 
+def _drain_captured_stdout() -> None:
+    """Forward bytes written to the captured fd 1 as stdout frames.
+
+    Runs on a daemon thread for the life of the process. Child processes that
+    inherit fd 1 (any ``subprocess`` call without ``stdout=PIPE``) land here.
+    """
+    if _CAPTURE_READ_FD is None:
+        return
+    import codecs
+
+    decoder = codecs.getincrementaldecoder("utf-8")("replace")
+    while True:
+        try:
+            chunk = os.read(_CAPTURE_READ_FD, 65536)
+        except OSError:
+            return
+        if not chunk:
+            return
+        text = decoder.decode(chunk)
+        if not text:
+            continue
+        rid = _STATE.capture_rid
+        if rid is None:
+            _RAW_STDERR.write(text)
+            _RAW_STDERR.flush()
+        else:
+            _emit({"type": "stdout", "id": rid, "data": text})
+
+
+def _start_capture_drain() -> None:
+    if _CAPTURE_READ_FD is None:
+        return
+    thread = threading.Thread(target=_drain_captured_stdout, name="prometheus-fd1-capture", daemon=True)
+    thread.start()
+
+
 # ---------------------------------------------------------------------------
 # Magic source transformer
 # ---------------------------------------------------------------------------
@@ -170,11 +268,11 @@ def transform_cell(source: str) -> str:
 
     Rules
     -----
-    * ``%name args``              -> ``__prometheus_magic("name", "args")``
-    * ``var = %name args``        -> ``var = __prometheus_magic("name", "args")``
-    * ``!cmd``                    -> ``__prometheus_shell("cmd")``
-    * ``var = !cmd``              -> ``var = __prometheus_shell("cmd")``
-    * ``%%name args\\n<body>``    -> ``__prometheus_magic_cell("name", "args", "<body>")``
+    * ``%name args``              -> ``__omp_magic("name", "args")``
+    * ``var = %name args``        -> ``var = __omp_magic("name", "args")``
+    * ``!cmd``                    -> ``__omp_shell("cmd")``
+    * ``var = !cmd``              -> ``var = __omp_shell("cmd")``
+    * ``%%name args\\n<body>``    -> ``__omp_magic_cell("name", "args", "<body>")``
       (cell magic must be the first non-whitespace token of a top-level line and
       consumes the remainder of the cell)
 
@@ -202,7 +300,7 @@ def transform_cell(source: str) -> str:
             body_lines = lines[i + 1 :]
             body = "\n".join(body_lines)
             out.append(
-                f"{indent}__prometheus_magic_cell({_quote_arg(name)}, {_quote_arg(args)}, {_quote_arg(body)})"
+                f"{indent}__omp_magic_cell({_quote_arg(name)}, {_quote_arg(args)}, {_quote_arg(body)})"
             )
             return "\n".join(out)
 
@@ -213,7 +311,7 @@ def transform_cell(source: str) -> str:
             indent = folded[: len(folded) - len(stripped_folded)]
             head, _ = _split_magic_head(stripped_folded[1:])
             name, args = head
-            out.append(f"{indent}__prometheus_magic({_quote_arg(name)}, {_quote_arg(args)})")
+            out.append(f"{indent}__omp_magic({_quote_arg(name)}, {_quote_arg(args)})")
             i += consumed
             continue
 
@@ -222,7 +320,7 @@ def transform_cell(source: str) -> str:
             stripped_folded = folded.lstrip()
             indent = folded[: len(folded) - len(stripped_folded)]
             cmd = stripped_folded[1:].strip()
-            out.append(f"{indent}__prometheus_shell({_quote_arg(cmd)})")
+            out.append(f"{indent}__omp_shell({_quote_arg(cmd)})")
             i += consumed
             continue
 
@@ -232,14 +330,14 @@ def transform_cell(source: str) -> str:
             rhs = m.group("rhs").strip()
             if rhs.startswith("!"):
                 cmd = rhs[1:].strip()
-                out.append(f"{m.group('indent')}{m.group('lhs').rstrip()} = __prometheus_shell({_quote_arg(cmd)})")
+                out.append(f"{m.group('indent')}{m.group('lhs').rstrip()} = __omp_shell({_quote_arg(cmd)})")
                 i += 1
                 continue
             if rhs.startswith("%") and not rhs.startswith("%%"):
                 head, _ = _split_magic_head(rhs[1:])
                 name, args = head
                 out.append(
-                    f"{m.group('indent')}{m.group('lhs').rstrip()} = __prometheus_magic({_quote_arg(name)}, {_quote_arg(args)})"
+                    f"{m.group('indent')}{m.group('lhs').rstrip()} = __omp_magic({_quote_arg(name)}, {_quote_arg(args)})"
                 )
                 i += 1
                 continue
@@ -512,14 +610,14 @@ def _run_shell_body(body: str, *, shell_arg: str) -> int:
     return proc.returncode
 
 
-def __prometheus_magic(name: str, args: str) -> Any:
+def __omp_magic(name: str, args: str) -> Any:
     fn = _LINE_MAGICS.get(name)
     if fn is None:
         raise NameError(f"UsageError: Line magic function '%{name}' not found.")
     return fn(args)
 
 
-def __prometheus_magic_cell(name: str, args: str, body: str) -> Any:
+def __omp_magic_cell(name: str, args: str, body: str) -> Any:
     fn = _CELL_MAGICS.get(name)
     if fn is None:
         raise NameError(f"UsageError: Cell magic function '%%{name}' not found.")
@@ -542,7 +640,7 @@ class _ShellResult(list):
         return " ".join(self)
 
 
-def __prometheus_shell(cmd: str) -> _ShellResult:
+def __omp_shell(cmd: str) -> _ShellResult:
     proc = subprocess.run(
         cmd,
         shell=True,
@@ -632,7 +730,7 @@ def _emit_display(bundle: dict, *, kind: str = "display") -> None:
     _emit({"type": kind, "id": rid, "bundle": bundle})
 
 
-def __prometheus_display(value: Any, *, raw: bool = False, kind: str = "display") -> None:
+def __omp_display(value: Any, *, raw: bool = False, kind: str = "display") -> None:
     if raw:
         if not isinstance(value, dict):
             raise TypeError("display(..., raw=True) requires a MIME bundle dict")
@@ -680,12 +778,12 @@ os.environ.setdefault("MPLBACKEND", "Agg")
 
 
 def _install_builtins(ns: dict) -> None:
-    ns["display"] = __prometheus_display
-    ns["__prometheus_display"] = __prometheus_display
-    ns["__prometheus_magic"] = __prometheus_magic
-    ns["__prometheus_magic_cell"] = __prometheus_magic_cell
-    ns["__prometheus_shell"] = __prometheus_shell
-    ns["__prometheus_current_run_id__"] = lambda: _CURRENT_RID.get()
+    ns["display"] = __omp_display
+    ns["__omp_display"] = __omp_display
+    ns["__omp_magic"] = __omp_magic
+    ns["__omp_magic_cell"] = __omp_magic_cell
+    ns["__omp_shell"] = __omp_shell
+    ns["__omp_current_run_id__"] = lambda: _CURRENT_RID.get()
 
 
 _install_builtins(_STATE.user_ns)
@@ -763,12 +861,12 @@ def _exec_source(source: str, ns: dict) -> None:
     if has_expr and expr_code is not None:
         value = _run_compiled_sync(expr_code, ns, want_value=True)
         if value is not None:
-            __prometheus_display(value, kind="result")
+            __omp_display(value, kind="result")
 
 
 async def _exec_source_async(source: str, ns: dict) -> None:
     """Compile + execute ``source``; if the last node is an expression, route
-    its value through ``__prometheus_display`` so dataframes/figures render rich.
+    its value through ``__omp_display`` so dataframes/figures render rich.
     Top-level ``await`` / ``async for`` / ``async with`` is permitted; awaited
     regions yield to other requests in the runner's persistent event loop."""
     body_code, expr_code, has_expr = _compile_source(source)
@@ -778,7 +876,7 @@ async def _exec_source_async(source: str, ns: dict) -> None:
     if has_expr and expr_code is not None:
         value = await _run_compiled_async(expr_code, ns, want_value=True)
         if value is not None:
-            __prometheus_display(value, kind="result")
+            __omp_display(value, kind="result")
 
 
 # ---------------------------------------------------------------------------
@@ -819,6 +917,7 @@ _MANAGED_ENV_KEYS = (
     "PROMETHEUS_TOOL_BRIDGE_URL",
     "PROMETHEUS_TOOL_BRIDGE_TOKEN",
     "PROMETHEUS_TOOL_BRIDGE_SESSION",
+    "PROMETHEUS_EVAL_LOCAL_ROOTS",
 )
 
 
@@ -879,7 +978,8 @@ def _start_parent_watchdog() -> None:
 async def _handle_request_async(req: dict) -> None:
     rid = str(req.get("id"))
     token = _CURRENT_RID.set(rid)
-    _STATE.user_ns["__prometheus_run_id__"] = rid
+    _STATE.capture_rid = rid
+    _STATE.user_ns["__omp_run_id__"] = rid
     _STATE.cancel_requested = False
     _STATE.execution_count += 1
     execution_count = _STATE.execution_count
@@ -933,6 +1033,7 @@ async def _handle_request_async(req: dict) -> None:
             except Exception:
                 pass
 
+        _flush_stream_proxies(rid)
         _emit({
             "type": "done",
             "id": rid,
@@ -941,6 +1042,9 @@ async def _handle_request_async(req: dict) -> None:
             "cancelled": cancelled,
         })
     finally:
+        if _STATE.capture_rid == rid:
+            _STATE.capture_rid = None
+        _flush_stream_proxies(rid)
         _CURRENT_RID.reset(token)
 
 
@@ -985,6 +1089,7 @@ async def _main_async() -> None:
     sys.stderr = _StreamProxy("stderr")
     _install_idle_sigint()
     _start_parent_watchdog()
+    _start_capture_drain()
 
     stdin = sys.__stdin__
     if stdin is None:

package/src/eval/py/runtime.ts

@@ -5,6 +5,7 @@
  * for both the shared gateway and local kernel spawning.
  */
 import * as fs from "node:fs";
+import * as os from "node:os";
 import * as path from "node:path";
 import { $env, $which, getPythonEnvDir } from "@prometheus-ai/utils";
 
@@ -79,7 +80,7 @@ const DEFAULT_ENV_DENYLIST = new Set([
 	"MISTRAL_API_KEY",
 ]);
 
-const DEFAULT_ENV_ALLOW_PREFIXES = ["LC_", "XDG_", "PROMETHEUS_"];
+const DEFAULT_ENV_ALLOW_PREFIXES = ["LC_", "XDG_", "PROMETHEUS_", "PI_"];
 
 const CASE_INSENSITIVE_ENV = process.platform === "win32";
 const BASE_ENV_ALLOWLIST = new Set([...DEFAULT_ENV_ALLOWLIST, ...WINDOWS_ENV_ALLOWLIST]);
@@ -182,6 +183,42 @@ function venvBinDir(venvPath: string): string {
 	return process.platform === "win32" ? path.join(venvPath, "Scripts") : path.join(venvPath, "bin");
 }
 
+function detectExplicitVenv(pythonPath: string): { venvPath: string; binDir: string } | undefined {
+	const binDir = path.dirname(pythonPath);
+	const venvPath = path.dirname(binDir);
+	if (fs.existsSync(path.join(venvPath, "pyvenv.cfg"))) {
+		return { venvPath, binDir };
+	}
+	return undefined;
+}
+
+/**
+ * Resolve an explicitly configured interpreter (`python.interpreter`) into a
+ * runtime, bypassing discovery. Does not probe or validate the executable —
+ * callers must check it actually runs. `~` expands to the home directory and
+ * relative paths resolve against `cwd`. When the interpreter sits inside a
+ * virtualenv (a `pyvenv.cfg` above its bin dir), the venv activation env is
+ * applied so subprocesses and `pip` resolve consistently.
+ */
+export function resolveExplicitPythonRuntime(
+	interpreter: string,
+	cwd: string,
+	baseEnv: Record<string, string | undefined>,
+): PythonRuntime {
+	const expanded =
+		interpreter === "~"
+			? os.homedir()
+			: interpreter.startsWith("~/")
+				? path.join(os.homedir(), interpreter.slice(2))
+				: interpreter;
+	const pythonPath = path.isAbsolute(expanded) ? expanded : path.resolve(cwd, expanded);
+	const venv = detectExplicitVenv(pythonPath);
+	if (venv) {
+		return { pythonPath, env: applyVenvEnv(baseEnv, venv.venvPath, venv.binDir), venvPath: venv.venvPath };
+	}
+	return { pythonPath, env: { ...baseEnv } };
+}
+
 /**
  * Enumerate candidate Python runtimes in priority order: an active/project venv,
  * the managed `~/.prometheus/python-env`, then the system interpreter on PATH. Every

package/src/exa/index.ts

@@ -1,27 +1,2 @@
-/**
- * Exa MCP Tools
- *
- * 22 tools for Exa's MCP servers:
- * - 4 search tools (search, deep, code, crawl)
- * - 1 LinkedIn search tool
- * - 1 company research tool
- * - 2 researcher tools (start, poll)
- * - 14 websets tools (CRUD, items, search, enrichment, monitor)
- */
-import type { CustomTool } from "../extensibility/custom-tools/types";
-import { researcherTools } from "./researcher";
-import { searchTools } from "./search";
-import type { ExaRenderDetails } from "./types";
-import { websetsTools } from "./websets";
-
-/** All Exa tools (22 total) - static export for backward compatibility */
-export const exaTools: CustomTool<any, ExaRenderDetails>[] = [...searchTools, ...researcherTools, ...websetsTools];
-
 export * from "./mcp-client";
-export { renderExaCall, renderExaResult } from "./render";
-export { researcherTools } from "./researcher";
-// Re-export individual modules for selective importing
-export { searchTools } from "./search";
-// Re-export types and utilities
-export type { ExaRenderDetails, ExaSearchResponse, ExaSearchResult, MCPToolWrapperConfig } from "./types";
-export { websetsTools } from "./websets";
+export type { ExaSearchResponse, MCPCallResponse, MCPTool, MCPToolsResponse, MCPToolWrapperConfig } from "./types";

package/src/exa/mcp-client.ts

@@ -2,14 +2,14 @@ import type { TSchema } from "@prometheus-ai/ai";
 import { $env, logger } from "@prometheus-ai/utils";
 import type { CustomTool, CustomToolResult } from "../extensibility/custom-tools/types";
 import { type CallMcpOptions, callMCP } from "../mcp/json-rpc";
-import type {
-	ExaRenderDetails,
-	ExaSearchResponse,
-	MCPCallResponse,
-	MCPTool,
-	MCPToolsResponse,
-	MCPToolWrapperConfig,
-} from "./types";
+import type { ExaSearchResponse, MCPCallResponse, MCPTool, MCPToolsResponse, MCPToolWrapperConfig } from "./types";
+
+type MCPWrappedToolDetails = {
+	response?: ExaSearchResponse;
+	error?: string;
+	toolName?: string;
+	raw?: unknown;
+};
 
 /** Find EXA_API_KEY from Bun.env or .env files */
 export function findApiKey(): string | null {
@@ -296,7 +296,7 @@ export async function fetchMCPToolSchema(
  * This allows tools to be generated from MCP server schemas without hardcoding,
  * reducing drift when MCP servers add new parameters.
  */
-export class MCPWrappedTool implements CustomTool<TSchema, ExaRenderDetails> {
+export class MCPWrappedTool implements CustomTool<TSchema, MCPWrappedToolDetails> {
 	readonly name: string;
 	readonly label: string;
 
@@ -315,7 +315,7 @@ export class MCPWrappedTool implements CustomTool<TSchema, ExaRenderDetails> {
 		_onUpdate?: unknown,
 		_ctx?: unknown,
 		_signal?: AbortSignal,
-	): Promise<CustomToolResult<ExaRenderDetails>> {
+	): Promise<CustomToolResult<MCPWrappedToolDetails>> {
 		try {
 			const apiKey = findApiKey();
 			// Websets tools require an API key; basic Exa MCP tools work without one

package/src/exa/types.ts

@@ -5,10 +5,6 @@
  */
 import type { TSchema } from "@prometheus-ai/ai";
 
-/** MCP endpoint URLs */
-export const EXA_MCP_URL = "https://mcp.exa.ai/mcp";
-export const WEBSETS_MCP_URL = "https://websetsmcp.exa.ai/mcp";
-
 /** MCP tool definition from server */
 export interface MCPTool {
 	name: string;
@@ -71,96 +67,3 @@ export interface ExaSearchResponse {
 	searchTime?: number;
 	requestId?: string;
 }
-
-/** Researcher task status */
-export interface ResearcherStatus {
-	id: string;
-	status: "pending" | "running" | "completed" | "failed";
-	result?: string;
-	error?: string;
-}
-
-/** Webset definition */
-export interface Webset {
-	id: string;
-	name: string;
-	description?: string;
-	createdAt?: string;
-	updatedAt?: string;
-}
-
-/** Webset item */
-export interface WebsetItem {
-	id: string;
-	websetId: string;
-	url: string;
-	title?: string;
-	content?: string;
-	metadata?: Record<string, unknown>;
-}
-
-/** Webset search */
-export interface WebsetSearch {
-	id: string;
-	websetId: string;
-	query: string;
-	status: "pending" | "running" | "completed" | "cancelled";
-	resultCount?: number;
-}
-
-/** Webset enrichment */
-export interface WebsetEnrichment {
-	id: string;
-	websetId: string;
-	name: string;
-	prompt: string;
-	status: "pending" | "running" | "completed" | "cancelled";
-}
-
-/** Tool name mappings: MCP name -> our tool name */
-export const EXA_TOOL_MAPPINGS = {
-	// Search tools
-	web_search_exa: "exa_search",
-	get_code_context_exa: "exa_search_code",
-	crawling_exa: "exa_crawl",
-	// LinkedIn
-	linkedin_search_exa: "exa_linkedin",
-	// Company
-	company_research_exa: "exa_company",
-	// Researcher
-	deep_researcher_start: "exa_researcher_start",
-	deep_researcher_check: "exa_researcher_poll",
-} as const;
-
-export const WEBSETS_TOOL_MAPPINGS = {
-	create_webset: "webset_create",
-	list_websets: "webset_list",
-	get_webset: "webset_get",
-	update_webset: "webset_update",
-	delete_webset: "webset_delete",
-	list_webset_items: "webset_items_list",
-	get_item: "webset_item_get",
-	create_search: "webset_search_create",
-	get_search: "webset_search_get",
-	cancel_search: "webset_search_cancel",
-	create_enrichment: "webset_enrichment_create",
-	get_enrichment: "webset_enrichment_get",
-	update_enrichment: "webset_enrichment_update",
-	delete_enrichment: "webset_enrichment_delete",
-	cancel_enrichment: "webset_enrichment_cancel",
-	create_monitor: "webset_monitor_create",
-} as const;
-
-export type ExaMcpToolName = keyof typeof EXA_TOOL_MAPPINGS;
-export type WebsetsMcpToolName = keyof typeof WEBSETS_TOOL_MAPPINGS;
-export type ExaToolName = (typeof EXA_TOOL_MAPPINGS)[ExaMcpToolName];
-export type WebsetsToolName = (typeof WEBSETS_TOOL_MAPPINGS)[WebsetsMcpToolName];
-
-/** Render details for TUI */
-export interface ExaRenderDetails {
-	response?: ExaSearchResponse;
-	error?: string;
-	toolName?: string;
-	/** Raw result for non-search responses */
-	raw?: unknown;
-}