Add on exit callback

Author: jakubno

js/src/messaging.ts

@@ -311,7 +311,8 @@ export async function parseOutput(
   line: string,
   onStdout?: (output: OutputMessage) => Promise<any> | any,
   onStderr?: (output: OutputMessage) => Promise<any> | any,
-  onResult?: (data: Result) => Promise<any> | any
+  onResult?: (data: Result) => Promise<any> | any,
+  onError?: (error: ExecutionError) => Promise<any> | any
 ) {
   const msg = JSON.parse(line)
 
@@ -348,6 +349,9 @@ export async function parseOutput(
       break
     case 'error':
       execution.error = new ExecutionError(msg.name, msg.value, msg.traceback)
+      if (onError) {
+        await onError(execution.error)
+      }
       break
     case 'number_of_executions':
       execution.executionCount = msg.execution_count

js/src/sandbox.ts

@@ -1,6 +1,6 @@
 import { Sandbox as BaseSandbox, InvalidArgumentError } from 'e2b'
 
-import { Result, Execution, OutputMessage, parseOutput, extractError } from './messaging'
+import {Result, Execution, OutputMessage, parseOutput, extractError, ExecutionError} from './messaging'
 import { formatExecutionTimeoutError, formatRequestTimeoutError, readLines } from "./utils";
 import { JUPYTER_PORT, DEFAULT_TIMEOUT_MS } from './consts'
 export type Context = {
@@ -25,6 +25,7 @@ export class Sandbox extends BaseSandbox {
    * @param opts.onStdout Callback for handling stdout messages
    * @param opts.onStderr Callback for handling stderr messages
    * @param opts.onResult Callback for handling the final result
+   * @param opts.onError Callback for handling the `ExecutionError` object
    * @param opts.envs Environment variables to set for the execution
    * @param opts.timeoutMs Max time to wait for the execution to finish
    * @param opts.requestTimeoutMs Max time to wait for the request to finish
@@ -37,6 +38,7 @@ async runCode(
       onStdout?: (output: OutputMessage) => (Promise<any> | any),
       onStderr?: (output: OutputMessage) => (Promise<any> | any),
       onResult?: (data: Result) => (Promise<any> | any),
+      onError?: (error: ExecutionError) => (Promise<any> | any),
       envs?: Record<string, string>,
       timeoutMs?: number,
       requestTimeoutMs?: number,
@@ -52,6 +54,7 @@ async runCode(
    * @param opts.onStdout Callback for handling stdout messages
    * @param opts.onStderr Callback for handling stderr messages
    * @param opts.onResult Callback for handling the final result
+   * @param opts.onError Callback for handling the `ExecutionError` object
    * @param opts.envs Environment variables to set for the execution
    * @param opts.timeoutMs Max time to wait for the execution to finish
    * @param opts.requestTimeoutMs Max time to wait for the request to finish
@@ -64,6 +67,7 @@ async runCode(
       onStdout?: (output: OutputMessage) => (Promise<any> | any),
       onStderr?: (output: OutputMessage) => (Promise<any> | any),
       onResult?: (data: Result) => (Promise<any> | any),
+      onError?: (error: ExecutionError) => (Promise<any> | any),
       envs?: Record<string, string>,
       timeoutMs?: number,
       requestTimeoutMs?: number,
@@ -77,6 +81,7 @@ async runCode(
       onStdout?: (output: OutputMessage) => (Promise<any> | any),
       onStderr?: (output: OutputMessage) => (Promise<any> | any),
       onResult?: (data: Result) => (Promise<any> | any),
+      onError?: (error: ExecutionError) => (Promise<any> | any),
       envs?: Record<string, string>,
       timeoutMs?: number,
       requestTimeoutMs?: number,
@@ -135,7 +140,7 @@ async runCode(
 
       try {
         for await (const chunk of readLines(res.body)) {
-          await parseOutput(execution, chunk, opts?.onStdout, opts?.onStderr, opts?.onResult)
+          await parseOutput(execution, chunk, opts?.onStdout, opts?.onStderr, opts?.onResult, opts?.onError)
         }
       } catch (error) {
         throw formatExecutionTimeoutError(error)

js/tests/callbacks.test.ts

@@ -0,0 +1,44 @@
+import { expect } from 'vitest'
+
+import { sandboxTest } from './setup'
+
+sandboxTest('callback results', async ({ sandbox }) => {
+  const results = []
+  const result = await sandbox.runCode('x =1; x', {
+    onResult: (result) => results.push(result),
+  })
+
+  expect(results.length).toBe(1)
+  expect(result.results[0].text).toBe('1')
+})
+
+
+sandboxTest('callback error', async ({ sandbox }) => {
+  const errors = []
+  const result = await sandbox.runCode('xyz', {
+    onError: (error) => errors.push(error),
+  })
+
+  expect(errors.length).toBe(1)
+  expect(result.error.name).toBe('NameError')
+})
+
+sandboxTest('callback stdout', async ({ sandbox }) => {
+  const stdout = []
+  const result = await sandbox.runCode('print("hello")', {
+    onStdout: (out) => stdout.push(out),
+  })
+
+  expect(stdout.length).toBe(1)
+  expect(result.logs.stdout).toEqual(['hello\n'])
+})
+
+sandboxTest('callback stderr', async ({ sandbox }) => {
+  const stderr = []
+  const result = await sandbox.runCode('import sys;print("This is an error message", file=sys.stderr)', {
+    onStderr: (err) => stderr.push(err),
+  })
+
+  expect(stderr.length).toBe(1)
+  expect(result.logs.stderr).toEqual(['This is an error message\n'])
+})

python/e2b_code_interpreter/code_interpreter_async.py

@@ -17,6 +17,7 @@
 )
 from e2b_code_interpreter.models import (
     Execution,
+    ExecutionError,
     Context,
     Result,
     aextract_exception,
@@ -54,6 +55,7 @@ async def run_code(
         on_stdout: Optional[OutputHandler[OutputMessage]] = None,
         on_stderr: Optional[OutputHandler[OutputMessage]] = None,
         on_result: Optional[OutputHandler[Result]] = None,
+        on_error: Optional[OutputHandler[ExecutionError]] = None,
         envs: Optional[Dict[str, str]] = None,
         timeout: Optional[float] = None,
         request_timeout: Optional[float] = None,
@@ -67,6 +69,7 @@ async def run_code(
         on_stdout: Optional[OutputHandler[OutputMessage]] = None,
         on_stderr: Optional[OutputHandler[OutputMessage]] = None,
         on_result: Optional[OutputHandler[Result]] = None,
+        on_error: Optional[OutputHandler[ExecutionError]] = None,
         envs: Optional[Dict[str, str]] = None,
         timeout: Optional[float] = None,
         request_timeout: Optional[float] = None,
@@ -80,6 +83,7 @@ async def run_code(
         on_stdout: Optional[OutputHandler[OutputMessage]] = None,
         on_stderr: Optional[OutputHandler[OutputMessage]] = None,
         on_result: Optional[OutputHandler[Result]] = None,
+        on_error: Optional[OutputHandler[ExecutionError]] = None,
         envs: Optional[Dict[str, str]] = None,
         timeout: Optional[float] = None,
         request_timeout: Optional[float] = None,
@@ -94,6 +98,7 @@ async def run_code(
         :param on_stdout: Callback for stdout messages
         :param on_stderr: Callback for stderr messages
         :param on_result: Callback for the `Result` object
+        :param on_error: Callback for the `ExecutionError` object
         :param envs: Environment variables
         :param timeout: Max time to wait for the execution to finish
         :param request_timeout: Max time to wait for the request to finish
@@ -136,6 +141,7 @@ async def run_code(
                         on_stdout=on_stdout,
                         on_stderr=on_stderr,
                         on_result=on_result,
+                        on_error=on_error,
                     )
 
                 return execution

python/e2b_code_interpreter/code_interpreter_sync.py

@@ -11,6 +11,7 @@
     DEFAULT_TIMEOUT,
 )
 from e2b_code_interpreter.models import (
+    ExecutionError,
     Execution,
     Context,
     Result,
@@ -46,6 +47,7 @@ def run_code(
         on_stdout: Optional[OutputHandler[OutputMessage]] = None,
         on_stderr: Optional[OutputHandler[OutputMessage]] = None,
         on_result: Optional[OutputHandler[Result]] = None,
+        on_error: Optional[OutputHandler[ExecutionError]] = None,
         envs: Optional[Dict[str, str]] = None,
         timeout: Optional[float] = None,
         request_timeout: Optional[float] = None,
@@ -59,6 +61,7 @@ def run_code(
         on_stdout: Optional[OutputHandler[OutputMessage]] = None,
         on_stderr: Optional[OutputHandler[OutputMessage]] = None,
         on_result: Optional[OutputHandler[Result]] = None,
+        on_error: Optional[OutputHandler[ExecutionError]] = None,
         envs: Optional[Dict[str, str]] = None,
         timeout: Optional[float] = None,
         request_timeout: Optional[float] = None,
@@ -72,6 +75,7 @@ def run_code(
         on_stdout: Optional[OutputHandler[OutputMessage]] = None,
         on_stderr: Optional[OutputHandler[OutputMessage]] = None,
         on_result: Optional[OutputHandler[Result]] = None,
+        on_error: Optional[OutputHandler[ExecutionError]] = None,
         envs: Optional[Dict[str, str]] = None,
         timeout: Optional[float] = None,
         request_timeout: Optional[float] = None,
@@ -86,6 +90,7 @@ def run_code(
         :param on_stdout: Callback for stdout messages
         :param on_stderr: Callback for stderr messages
         :param on_result: Callback for the `Result` object
+        :param on_error: Callback for the `ExecutionError` object
         :param envs: Environment variables
         :param timeout: Max time to wait for the execution to finish
         :param request_timeout: Max time to wait for the request to finish
@@ -128,6 +133,7 @@ def run_code(
                         on_stdout=on_stdout,
                         on_stderr=on_stderr,
                         on_result=on_result,
+                        on_error=on_error,
                     )
 
                 return execution

python/e2b_code_interpreter/models.py

@@ -394,6 +394,7 @@ def parse_output(
     on_stdout: Optional[OutputHandler[OutputMessage]] = None,
     on_stderr: Optional[OutputHandler[OutputMessage]] = None,
     on_result: Optional[OutputHandler[Result]] = None,
+    on_error: Optional[OutputHandler[ExecutionError]] = None,
 ):
     data = json.loads(output)
     data_type = data.pop("type")
@@ -413,6 +414,8 @@ def parse_output(
             on_stderr(OutputMessage(data["text"], data["timestamp"], True))
     elif data_type == "error":
         execution.error = ExecutionError(data["name"], data["value"], data["traceback"])
+        if on_error:
+            on_error(execution.error)
     elif data_type == "number_of_executions":
         execution.execution_count = data["execution_count"]
 

python/tests/async/test_async_callbacks.py

@@ -0,0 +1,38 @@
+from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
+
+
+async def test_resuls(async_sandbox: AsyncSandbox):
+    results = []
+    execution = await async_sandbox.run_code(
+        "x = 1;x", on_result=lambda result: results.append(result)
+    )
+    assert len(results) == 1
+    assert execution.results[0].text == "1"
+
+
+async def test_error(async_sandbox: AsyncSandbox):
+    errors = []
+    execution = await async_sandbox.run_code(
+        "xyz", on_error=lambda error: errors.append(error)
+    )
+    assert len(errors) == 1
+    assert execution.error.name == "NameError"
+
+
+async def test_stdout(async_sandbox: AsyncSandbox):
+    stdout = []
+    execution = await async_sandbox.run_code(
+        "print('Hello from e2b')", on_stdout=lambda out: stdout.append(out)
+    )
+    assert len(stdout) == 1
+    assert execution.logs.stdout == ["Hello from e2b\n"]
+
+
+async def test_stderr(async_sandbox: AsyncSandbox):
+    stderr = []
+    execution = await async_sandbox.run_code(
+        'import sys;print("This is an error message", file=sys.stderr)',
+        on_stderr=lambda err: stderr.append(err),
+    )
+    assert len(stderr) == 1
+    assert execution.logs.stderr == ["This is an error message\n"]

python/tests/sync/test_callbacks.py

@@ -0,0 +1,36 @@
+from e2b_code_interpreter.code_interpreter_sync import Sandbox
+
+
+def test_resuls(sandbox: Sandbox):
+    results = []
+    execution = sandbox.run_code(
+        "x = 1;x", on_result=lambda result: results.append(result)
+    )
+    assert len(results) == 1
+    assert execution.results[0].text == "1"
+
+
+def test_error(sandbox: Sandbox):
+    errors = []
+    execution = sandbox.run_code("xyz", on_error=lambda error: errors.append(error))
+    assert len(errors) == 1
+    assert execution.error.name == "NameError"
+
+
+def test_stdout(sandbox: Sandbox):
+    stdout = []
+    execution = sandbox.run_code(
+        "print('Hello from e2b')", on_stdout=lambda out: stdout.append(out)
+    )
+    assert len(stdout) == 1
+    assert execution.logs.stdout == ["Hello from e2b\n"]
+
+
+def test_stderr(sandbox: Sandbox):
+    stderr = []
+    execution = sandbox.run_code(
+        'import sys;print("This is an error message", file=sys.stderr)',
+        on_stderr=lambda err: stderr.append(err),
+    )
+    assert len(stderr) == 1
+    assert execution.logs.stderr == ["This is an error message\n"]