Allow closing the stdin stream

This allows executing programs that read until output is complete.

Author: shepmaster

compiler/base/orchestrator/src/coordinator.rs

@@ -777,10 +777,22 @@ impl Container {
 
         let task = tokio::spawn({
             async move {
+                let mut stdin_open = true;
+
                 loop {
                     select! {
-                        Some(stdin) = stdin_rx.recv() => {
-                            let msg = CoordinatorMessage::StdinPacket(stdin);
+                        stdin = stdin_rx.recv(), if stdin_open => {
+                            let msg = match stdin {
+                                Some(stdin) => {
+                                    CoordinatorMessage::StdinPacket(stdin)
+                                }
+
+                                None => {
+                                    stdin_open = false;
+                                    CoordinatorMessage::StdinClose
+                                }
+                            };
+
                             trace!("processing {msg:?}");
                             to_worker_tx.send(msg).await.context(StdinSnafu)?;
                         },
@@ -1802,6 +1814,62 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    #[snafu::report]
+    async fn execute_stdin_close() -> Result<()> {
+        let coordinator = new_coordinator().await;
+
+        let request = ExecuteRequest {
+            code: r#"
+                fn main() {
+                    let mut input = String::new();
+                    while let Ok(n) = std::io::stdin().read_line(&mut input) {
+                        if n == 0 {
+                            break;
+                        }
+                        println!("You entered >>>{input:?}<<<");
+                        input.clear();
+                    }
+                }
+            "#
+            .into(),
+            ..ARBITRARY_EXECUTE_REQUEST
+        };
+
+        let ActiveExecution {
+            task,
+            stdin_tx,
+            stdout_rx,
+            stderr_rx,
+        } = coordinator.begin_execute(request).await.unwrap();
+
+        for i in 0..3 {
+            stdin_tx.send(format!("line {i}\n")).await.unwrap();
+        }
+
+        stdin_tx.send("no newline".into()).await.unwrap();
+        drop(stdin_tx); // Close the stdin handle
+
+        let WithOutput {
+            response,
+            stdout,
+            stderr,
+        } = WithOutput::try_absorb(task, stdout_rx, stderr_rx)
+            .with_timeout()
+            .await
+            .unwrap();
+
+        assert!(response.success, "{stderr}");
+        assert_contains!(stdout, r#">>>"line 0\n"<<<"#);
+        assert_contains!(stdout, r#">>>"line 1\n"<<<"#);
+        assert_contains!(stdout, r#">>>"line 2\n"<<<"#);
+        assert_contains!(stdout, r#">>>"no newline"<<<"#);
+
+        coordinator.shutdown().await?;
+
+        Ok(())
+    }
+
     const HELLO_WORLD_CODE: &str = r#"fn main() { println!("Hello World!"); }"#;
 
     const ARBITRARY_COMPILE_REQUEST: CompileRequest = CompileRequest {

compiler/base/orchestrator/src/message.rs

@@ -26,6 +26,7 @@ pub enum CoordinatorMessage {
     ReadFile(ReadFileRequest),
     ExecuteCommand(ExecuteCommandRequest),
     StdinPacket(String),
+    StdinClose,
 }
 
 impl_narrow_to_broad!(

compiler/base/orchestrator/src/worker.rs

@@ -186,6 +186,14 @@ async fn handle_coordinator_message(
                             .drop_error_details()
                             .context(UnableToSendStdinPacketSnafu)?;
                     }
+
+                    CoordinatorMessage::StdinClose => {
+                        process_tx
+                            .send(Multiplexed(job_id, ProcessCommand::StdinClose))
+                            .await
+                            .drop_error_details()
+                            .context(UnableToSendStdinCloseSnafu)?;
+                    }
                 }
             }
 
@@ -216,6 +224,9 @@ pub enum HandleCoordinatorMessageError {
     #[snafu(display("Failed to send stdin packet to the command task"))]
     UnableToSendStdinPacket { source: mpsc::error::SendError<()> },
 
+    #[snafu(display("Failed to send stdin close request to the command task"))]
+    UnableToSendStdinClose { source: mpsc::error::SendError<()> },
+
     #[snafu(display("A coordinator command handler background task panicked"))]
     TaskPanicked { source: tokio::task::JoinError },
 }
@@ -371,6 +382,7 @@ fn parse_working_dir(cwd: Option<String>, project_path: impl Into<PathBuf>) -> P
 enum ProcessCommand {
     Start(ExecuteCommandRequest, MultiplexingSender),
     Stdin(String),
+    StdinClose,
 }
 
 struct ProcessState {
@@ -478,6 +490,8 @@ async fn manage_processes(
                     ProcessCommand::Start(req, worker_msg_tx) => state.start(job_id, req, worker_msg_tx).await?,
 
                     ProcessCommand::Stdin(packet) => state.stdin(job_id, packet).await?,
+
+                    ProcessCommand::StdinClose => state.stdin_close(job_id),
                 }
             }
 

tests/spec/features/streaming_spec.rb

@@ -63,6 +63,30 @@
     end
   end
 
+  scenario "input can be closed" do
+    editor.set <<~EOF
+      fn main() {
+          let mut input = String::new();
+          while std::io::stdin().read_line(&mut input).unwrap() != 0 {
+              println!("You entered >>>{input:?}<<<");
+              input.clear();
+          }
+          println!("All done");
+      }
+    EOF
+
+    click_on("Run")
+
+    within(:stdin) do
+      click_on 'Execution control'
+      click_on 'Close stdin'
+    end
+
+    within(:output, :stdout) do
+      expect(page).to have_content 'All done'
+    end
+  end
+
   def editor
     Editor.new(page)
   end

ui/frontend/Stdin.module.css

@@ -25,3 +25,22 @@
   visibility: hidden;
   white-space: pre-wrap;
 }
+
+.buttons {
+  font-size: 11px;
+}
+
+.menu {
+  list-style: none;
+  margin: 0;
+  padding: 0;
+}
+
+.button {
+  composes: -menuItemFullButton from './shared.module.css';
+  padding: 1em;
+
+  &:hover {
+    color: var(--header-tint);
+  }
+}

ui/frontend/Stdin.tsx

@@ -1,7 +1,9 @@
 import React, { ChangeEvent, FormEvent, KeyboardEvent, useCallback, useRef, useState } from 'react';
 import { useDispatch, useSelector } from 'react-redux';
 
-import { wsExecuteStdin } from './reducers/output/execute';
+import { Button, ButtonSet, IconButton } from './ButtonSet';
+import PopButton, { ButtonProps, MenuProps } from './PopButton';
+import { wsExecuteStdin, wsExecuteStdinClose } from './reducers/output/execute';
 import { enableStdinSelector } from './selectors';
 
 import styles from './Stdin.module.css';
@@ -20,6 +22,10 @@ const Stdin: React.FC = () => {
         e.preventDefault();
         form.current?.dispatchEvent(new Event('submit', { bubbles: true, cancelable: true }));
       }
+
+      if (e.key === 'd' && e.ctrlKey && content.length === 0) {
+        dispatch(wsExecuteStdinClose());
+      }
     },
     [dispatch, form, content],
   );
@@ -42,26 +48,73 @@ const Stdin: React.FC = () => {
     [dispatch, content, setContent],
   );
 
+  const menuContainer = useRef<HTMLDivElement | null>(null);
+
+  return (
+    <div data-test-id="stdin">
+      <form onSubmit={handleSubmit} className={styles.form} ref={form}>
+        <div className={styles.multiLine}>
+          <textarea
+            rows={1}
+            onKeyDown={handleKeyDown}
+            onChange={handleChange}
+            name="content"
+            autoComplete="off"
+            spellCheck="false"
+            className={styles.text}
+            value={content}
+            disabled={disabled}
+          ></textarea>
+          <p className={styles.sizer}>{content} </p>
+        </div>
+
+        <ButtonSet className={styles.buttons}>
+          <Button isPrimary isSmall type="submit" disabled={disabled} iconRight={() => '⏎'}>
+            Send
+          </Button>
+
+          <PopButton Button={MoreButton} Menu={MoreMenu} menuContainer={menuContainer} />
+        </ButtonSet>
+      </form>
+      <div ref={menuContainer} />
+    </div>
+  );
+};
+
+const MoreButton = React.forwardRef<HTMLButtonElement, ButtonProps>(({ toggle }, ref) => {
+  const disabled = !useSelector(enableStdinSelector);
+
+  return (
+    <IconButton
+      isSmall
+      type="button"
+      ref={ref}
+      title="Execution control"
+      onClick={toggle}
+      disabled={disabled}
+    >
+      ⋮
+    </IconButton>
+  );
+});
+MoreButton.displayName = 'MoreButton';
+
+const MoreMenu: React.FC<MenuProps> = ({ close }) => {
+  const dispatch = useDispatch();
+
+  const stdinClose = useCallback(() => {
+    dispatch(wsExecuteStdinClose());
+    close();
+  }, [dispatch, close]);
+
   return (
-    <form onSubmit={handleSubmit} className={styles.form} data-test-id="stdin" ref={form}>
-      <div className={styles.multiLine}>
-        <textarea
-          rows={1}
-          onKeyDown={handleKeyDown}
-          onChange={handleChange}
-          name="content"
-          autoComplete="off"
-          spellCheck="false"
-          className={styles.text}
-          value={content}
-          disabled={disabled}
-        ></textarea>
-        <p className={styles.sizer}>{content} </p>
-      </div>
-      <button type="submit" disabled={disabled}>
-        Send
-      </button>
-    </form>
+    <ul className={styles.menu}>
+      <li>
+        <button type="button" className={styles.button} onClick={stdinClose}>
+          Close stdin
+        </button>
+      </li>
+    </ul>
   );
 };
 

ui/frontend/reducers/output/execute.ts

@@ -1,4 +1,4 @@
-import { Draft, createAsyncThunk, createSlice } from '@reduxjs/toolkit';
+import { AnyAction, Draft, createAsyncThunk, createSlice } from '@reduxjs/toolkit';
 import * as z from 'zod';
 
 import { SimpleThunkAction, adaptFetchError, jsonPost, routes } from '../../actions';
@@ -78,6 +78,14 @@ export const performExecute = createAsyncThunk(sliceName, async (payload: Execut
   adaptFetchError(() => jsonPost<ExecuteResponseBody>(routes.execute, payload)),
 );
 
+const prepareWithCurrentSequenceNumber = <P>(payload: P, sequenceNumber: number) => ({
+  payload,
+  meta: {
+    websocket: true,
+    sequenceNumber,
+  },
+});
+
 const sequenceNumberMatches =
   <P>(whenMatch: (state: Draft<State>, payload: P) => void) =>
   (state: Draft<State>, action: WsPayloadAction<P>) => {
@@ -111,13 +119,12 @@ const slice = createSlice({
     wsExecuteStdin: {
       reducer: () => {},
 
-      prepare: (payload: string, sequenceNumber: number) => ({
-        payload,
-        meta: {
-          websocket: true,
-          sequenceNumber,
-        },
-      }),
+      prepare: prepareWithCurrentSequenceNumber,
+    },
+    wsExecuteStdinClose: {
+      reducer: () => {},
+
+      prepare: prepareWithCurrentSequenceNumber,
     },
   },
   extraReducers: (builder) => {
@@ -192,17 +199,27 @@ export const performCommonExecute =
     }
   };
 
-export const wsExecuteStdin =
-  (payload: string): SimpleThunkAction =>
+const dispatchWhenSequenceNumber =
+  <A extends AnyAction>(cb: (sequenceNumber: number) => A): SimpleThunkAction =>
   (dispatch, getState) => {
     const state = getState();
-    const { requestsInProgress, sequenceNumber } = state.output.execute;
-    if (requestsInProgress === 0 || !sequenceNumber) {
-      return;
+    const { sequenceNumber } = state.output.execute;
+    if (sequenceNumber) {
+      const action = cb(sequenceNumber);
+      dispatch(action);
     }
-    dispatch(slice.actions.wsExecuteStdin(payload, sequenceNumber));
   };
 
+export const wsExecuteStdin = (payload: string): SimpleThunkAction =>
+  dispatchWhenSequenceNumber((sequenceNumber) =>
+    slice.actions.wsExecuteStdin(payload, sequenceNumber),
+  );
+
+export const wsExecuteStdinClose = (): SimpleThunkAction =>
+  dispatchWhenSequenceNumber((sequenceNumber) =>
+    slice.actions.wsExecuteStdinClose(undefined, sequenceNumber),
+  );
+
 export { wsExecuteBeginSchema, wsExecuteStdoutSchema, wsExecuteStderrSchema, wsExecuteEndSchema };
 
 export default slice.reducer;