Add exit code to process (#193)

.changeset/angry-ears-thank.md

@@ -0,0 +1,6 @@
+---
+'@e2b/python-sdk': patch
+'@e2b/sdk': patch
+---
+
+Add exit code to process

packages/js-sdk/src/session/index.ts

@@ -277,6 +277,10 @@ export class Session extends SessionConnection {
           const { promise: processExited, resolve: triggerExit } = createDeferredPromise()
 
           const output = new ProcessOutput()
+          const handleExit = (exitCode: number) => {
+            output.setExitCode(exitCode)
+            triggerExit()
+          }
 
           const handleStdout = (data: { line: string; timestamp: number }) => {
             const message = new ProcessMessage(data.line, data.timestamp, false)
@@ -304,7 +308,7 @@ export class Session extends SessionConnection {
 
           const [onExitSubID, onStdoutSubID, onStderrSubID] =
             await this.handleSubscriptions(
-              this.subscribe(processService, triggerExit, 'onExit', processID),
+              this.subscribe(processService, handleExit, 'onExit', processID),
               this.subscribe(processService, handleStdout, 'onStdout', processID),
               this.subscribe(processService, handleStderr, 'onStderr', processID),
             )

packages/js-sdk/src/session/process.ts

@@ -26,6 +26,18 @@ export class ProcessOutput {
   private readonly delimiter = '\n'
   private readonly messages: ProcessMessage[] = []
   private _error = false
+  private _exitCode?: number
+  private _finished = false
+
+  /**
+   * The exit code of the process.
+   */
+  get exitCode(): number | undefined {
+    if (!this._finished) {
+      throw new Error('Process has not finished yet')
+    }
+    return this._exitCode
+  }
 
   /**
    * Whether the process has errored.
@@ -63,6 +75,11 @@ export class ProcessOutput {
     this.insertByTimestamp(message)
   }
 
+  setExitCode(exitCode: number) {
+    this._exitCode = exitCode
+    this._finished = true
+  }
+
   private insertByTimestamp(message: ProcessMessage) {
     let i = this.messages.length - 1
     while (i >= 0 && this.messages[i].timestamp > message.timestamp) {

packages/js-sdk/test/process.test.mjs

@@ -21,6 +21,7 @@ test('process on stdout/stderr', async () => {
   expect(output.stderr).toEqual('')
   expect(stdout.map(message => message.line)).toEqual(['/tmp'])
   expect(stderr).toEqual([])
+  expect(process.output.exitCode).toEqual(0)
   await session.close()
 })
 
@@ -112,6 +113,7 @@ test('test default on stdout/stderr', async () => {
   await process.finished
   expect(onStdout).toHaveBeenCalledOnce()
   expect(onStderr).toHaveBeenCalled()
+  expect(process.output.exitCode).toEqual(1)
 
   await session.close()
 }, 10000)

packages/python-sdk/e2b/session/process.py

@@ -52,6 +52,7 @@ class ProcessOutput(BaseModel):
     messages: List[ProcessMessage] = []
 
     error: bool = False
+    exit_code: Optional[int] = None
 
     @property
     def stdout(self) -> str:
@@ -101,6 +102,15 @@ def __init__(
         self._finished = finished
         self._output = output
 
+    @property
+    def exit_code(self) -> int:
+        """
+        The exit code of the last process started by this manager.
+        """
+        if not self.finished:
+            raise ProcessException("Process has not finished yet")
+        return self.output.exit_code
+
     @property
     def output(self) -> ProcessOutput:
         """
@@ -282,6 +292,10 @@ async def start(
 
             output = ProcessOutput()
 
+            def handle_exit(exit_code: int):
+                output.exit_code = exit_code
+                future_exit(True)
+
             def handle_stdout(data: Dict[Any, Any]):
                 out = OutStdoutResponse(**data)
 
@@ -317,7 +331,7 @@ def handle_stderr(data: Dict[Any, Any]):
             try:
                 unsub_all = await self._session._handle_subscriptions(
                     self._session._subscribe(
-                        self._service_name, future_exit, "onExit", process_id
+                        self._service_name, handle_exit, "onExit", process_id
                     ),
                     self._session._subscribe(
                         self._service_name, handle_stdout, "onStdout", process_id

packages/python-sdk/example.py

@@ -2,8 +2,11 @@
 import logging
 from os import getenv
 
+from dotenv import load_dotenv
+
 from e2b import Session
 
+load_dotenv()
 id = "Python3"
 E2B_API_KEY = getenv("E2B_API_KEY")
 

packages/python-sdk/tests/test_process.py

@@ -33,6 +33,7 @@ async def test_process_on_stdout_stderr():
     assert output.stderr == ""
     assert list(map(lambda message: message.line, stdout)) == ["/tmp"]
     assert stderr == []
+    assert proc.exit_code == 0
 
     await session.close()
 
@@ -122,5 +123,6 @@ async def test_process_default_on_stdout_stderr():
 
     on_stdout.assert_called_once()
     on_stderr.assert_called()
+    assert proc.exit_code == 1
 
     await session.close()