Correctly display native REPL execution status (#23797)

Resolves: #23739

python_files/python_server.py

@@ -13,28 +13,38 @@
 USER_GLOBALS = {}
 
 
-def send_message(msg: str):
+def _send_message(msg: str):
     length_msg = len(msg)
     STDOUT.buffer.write(f"Content-Length: {length_msg}\r\n\r\n{msg}".encode())
     STDOUT.buffer.flush()
 
 
+def send_message(**kwargs):
+    _send_message(json.dumps({"jsonrpc": "2.0", **kwargs}))
+
+
 def print_log(msg: str):
-    send_message(json.dumps({"jsonrpc": "2.0", "method": "log", "params": msg}))
+    send_message(method="log", params=msg)
 
 
-def send_response(response: str, response_id: int):
-    send_message(json.dumps({"jsonrpc": "2.0", "id": response_id, "result": response}))
+def send_response(
+    response: str,
+    response_id: int,
+    execution_status: bool = True,  # noqa: FBT001, FBT002
+):
+    send_message(
+        id=response_id,
+        result={"status": execution_status, "output": response},
+    )
 
 
 def send_request(params: Optional[Union[List, Dict]] = None):
     request_id = uuid.uuid4().hex
     if params is None:
-        send_message(json.dumps({"jsonrpc": "2.0", "id": request_id, "method": "input"}))
+        send_message(id=request_id, method="input")
     else:
-        send_message(
-            json.dumps({"jsonrpc": "2.0", "id": request_id, "method": "input", "params": params})
-        )
+        send_message(id=request_id, method="input", params=params)
+
     return request_id
 
 
@@ -105,24 +115,28 @@ def execute(request, user_globals):
         original_stdin = sys.stdin
         try:
             sys.stdin = str_input
-            exec_user_input(request["params"], user_globals)
+            execution_status = exec_user_input(request["params"], user_globals)
         finally:
             sys.stdin = original_stdin
-    send_response(str_output.get_value(), request["id"])
+
+    send_response(str_output.get_value(), request["id"], execution_status)
 
 
-def exec_user_input(user_input, user_globals):
+def exec_user_input(user_input, user_globals) -> bool:
     user_input = user_input[0] if isinstance(user_input, list) else user_input
 
     try:
         callable_ = exec_function(user_input)
         retval = callable_(user_input, user_globals)
         if retval is not None:
             print(retval)
+        return True
     except KeyboardInterrupt:
         print(traceback.format_exc())
+        return False
     except Exception:
         print(traceback.format_exc())
+        return False
 
 
 class CustomIO(io.TextIOWrapper):

src/client/repl/pythonServer.ts

@@ -9,9 +9,13 @@ import { EventName } from '../telemetry/constants';
 
 const SERVER_PATH = path.join(EXTENSION_ROOT_DIR, 'python_files', 'python_server.py');
 let serverInstance: PythonServer | undefined;
+export interface ExecutionResult {
+    status: boolean;
+    output: string;
+}
 
 export interface PythonServer extends Disposable {
-    execute(code: string): Promise<string>;
+    execute(code: string): Promise<ExecutionResult | undefined>;
     interrupt(): void;
     input(): void;
     checkValidCommand(code: string): Promise<boolean>;
@@ -52,8 +56,15 @@ class PythonServerImpl implements Disposable {
     }
 
     @captureTelemetry(EventName.EXECUTION_CODE, { scope: 'selection' }, false)
-    public execute(code: string): Promise<string> {
-        return this.connection.sendRequest('execute', code);
+    public async execute(code: string): Promise<ExecutionResult | undefined> {
+        try {
+            const result = await this.connection.sendRequest('execute', code);
+            return result as ExecutionResult;
+        } catch (err) {
+            const error = err as Error;
+            traceError(`Error getting response from REPL server:`, error);
+        }
+        return undefined;
     }
 
     public interrupt(): void {

src/client/repl/replController.ts

@@ -22,27 +22,17 @@ export function createReplController(
         for (const cell of cells) {
             const exec = controller.createNotebookCellExecution(cell);
             exec.start(Date.now());
-            try {
-                const result = await server.execute(cell.document.getText());
-                if (result !== '') {
-                    exec.replaceOutput([
-                        new vscode.NotebookCellOutput([vscode.NotebookCellOutputItem.text(result, 'text/plain')]),
-                    ]);
-                }
-                exec.end(true);
-            } catch (err) {
-                const error = err as Error;
+
+            const result = await server.execute(cell.document.getText());
+
+            if (result?.output) {
                 exec.replaceOutput([
-                    new vscode.NotebookCellOutput([
-                        vscode.NotebookCellOutputItem.error({
-                            name: error.name,
-                            message: error.message,
-                            stack: error.stack,
-                        }),
-                    ]),
+                    new vscode.NotebookCellOutput([vscode.NotebookCellOutputItem.text(result.output, 'text/plain')]),
                 ]);
-                exec.end(false);
+                // TODO: Properly update via NotebookCellOutputItem.error later.
             }
+
+            exec.end(result?.status);
         }
     };
     disposables.push(controller);