Reuse StreamPrinter implementation in imix/shell#1837
Conversation
Moved `StreamPrinter` from `imix/src/task.rs` to `imix/src/printer.rs` and updated it to distinguish between `Stdout` and `Stderr` using an `OutputKind` enum. Updated `imix/src/shell/repl.rs` to reuse `StreamPrinter` instead of the custom `ShellPrinter`. Refactored `imix/src/task.rs` to use the shared `StreamPrinter` and properly report stderr to `TaskError` field. Updated `imix/src/lib.rs` and `imix/src/main.rs` to include the new `printer` module. This change unifies output handling across the agent and improves maintainability. Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com>
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
* Implement Agent Shell Logic and Synchronization (#1832) * Implement agent-side shell features and synchronization - Update `ShellTask` schema with execution timestamps. - Update `c2` and `portal` protocol buffers to support `ShellTask` and `ShellTaskOutput`. - Implement `ClaimTasks` and `ReportTaskOutput` logic for shell tasks in Go backend. - Implement `ShellManager` in Rust agent (`imix`) to manage `Eldritch` interpreters per shell. - Integrate `ShellManager` with `ImixAgent` and `portal` streams for input/output routing. - Ensure shell tasks are claimed and output is buffered/reported correctly. Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * Implement agent-side shell features and synchronization - Update `ShellTask` schema with execution timestamps. - Update `c2` and `portal` protocol buffers to support `ShellTask` and `ShellTaskOutput`. - Implement `ClaimTasks` and `ReportTaskOutput` logic for shell tasks in Go backend. - Implement `ShellManager` in Rust agent (`imix`) to manage `Eldritch` interpreters per shell. - Integrate `ShellManager` with `ImixAgent` and `portal` streams for input/output routing. - Ensure shell tasks are claimed and output is buffered/reported correctly. - Fixed `tavern/internal/http/shell/handler.go` by removing invalid `StreamId` initialization. Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * Fix ShellPayload compilation error in tavern - Removed invalid `StreamId` field initialization from `portalpb.ShellPayload` struct literal in `tavern/internal/http/shell/handler.go`. - This field was removed from the protobuf definition in a previous commit. Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * Fix Rust compilation errors in implants/imix and stdlib - Added missing `shell_task_output: None` to `ReportTaskOutputRequest` struct literals in `implants/imix/src/tests/agent_trait_tests.rs` and `implants/lib/eldritch/stdlib/eldritch-libagent/src/std/report_task_output_impl.rs` to match the updated protobuf definition. Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * Run cargo fmt to ensure consistent code style Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * Fix ShellPrinter output sequence ID handling - Updated `c2.proto` to include `sequence_id` and `stream_id` in `ShellTask`. - Updated `tavern` to populate these fields when claiming tasks. - Updated `imix` `ShellManager` to propagate `sequence_id` and `stream_id` from input tasks/motes to output motes. - Fixed an issue where output was being dropped in interactive sessions due to mismatched sequence IDs (output motes were using `seq_id: 0`). Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * delete weird stuff * Separate aggregation of task and shell task outputs in imix agent (#1835) Updated `implants/imix/src/agent.rs` to aggregate `ReportTaskOutputRequest` messages based on their content type: - Standard task outputs (`output` field) are aggregated by `task_id`. - Shell task outputs (`shell_task_output` field) are aggregated by `shell_task_id`. This ensures that shell task outputs are correctly reported and not merged with standard task outputs (or dropped/overwritten if they share the default task context). Also added a new test `test_agent_output_aggregation` in `implants/imix/src/tests/agent_output_aggregation.rs` to verify this behavior. Fixed a missing field `shell_tasks` in `ClaimTasksResponse` construction in `implants/imix/src/tests/task_tests.rs`. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * fix newline prints * Update shell manager to send shell task output with the result of interpret() (#1843) Refactored `implants/imix/src/shell/manager.rs` to: - Capture the result of `state.interpreter.interpret(&task.input)`. - Use helper functions `send_shell_output` and `send_shell_error` to send `ReportTaskOutputRequest` and `Mote` messages. - Send the `Debug` representation of the result (if not `Value::None`) as task output, consistent with REPL behavior. - Ensure correct `Arc` cloning to avoid borrow checker errors. - Handle interpretation errors by reporting them as task errors. This ensures that return values of shell commands (e.g. `process.list()`) are displayed to the user even if they are not explicitly printed. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * Refactor imix/shell to reuse StreamPrinter implementation (#1837) Moved `StreamPrinter` from `imix/src/task.rs` to `imix/src/printer.rs` and updated it to distinguish between `Stdout` and `Stderr` using an `OutputKind` enum. Updated `imix/src/shell/repl.rs` to reuse `StreamPrinter` instead of the custom `ShellPrinter`. Refactored `imix/src/task.rs` to use the shared `StreamPrinter` and properly report stderr to `TaskError` field. Updated `imix/src/lib.rs` and `imix/src/main.rs` to include the new `printer` module. This change unifies output handling across the agent and improves maintainability. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * Fix imix agent stream printing output separation (#1848) - Modified `StreamPrinter` in `implants/imix/src/printer.rs` to use separate channels for stdout and stderr. - Updated `implants/imix/src/task.rs` to instantiate `StreamPrinter` with two channels and consume them concurrently using `tokio::select!`. - Updated `implants/imix/src/shell/repl.rs` to consume both stdout and stderr streams in the REPL loop. - Verified that `eprint` output is correctly routed to the `TaskError` field via existing tests. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * fixed a panic in eldritch * feat(imix): Implement panic resilience in ShellManager (#1850) * Fix imix shell manager async runtime panic (#1851) * Refactor imix shell manager to run interpreter in blocking task Fixes a panic "Cannot start a runtime from within a runtime" when the eldritch interpreter calls async agent functions (which use `block_on`) from within the imix shell manager's async runtime context. Refactored `ShellManager::handle_task` and `handle_portal_payload` to offload the synchronous `interpreter.interpret` call to `tokio::task::spawn_blocking`. Introduced `run_interpreter_task` helper to manage state transfer between async and blocking contexts. Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * Refactor imix shell manager to run interpreter in dedicated blocking thread Refactored `ShellManager` to execute the `eldritch` interpreter in a dedicated, long-running blocking thread (`tokio::task::spawn_blocking`) for each shell session. This solves a panic ("Cannot start a runtime from within a runtime") caused by the interpreter calling `block_on` from within an async runtime context. Introduced `InterpreterCommand` enum and `InterpreterHandle` to manage communication between the async manager and the blocking interpreter thread via `tokio::sync::mpsc` channels. This ensures tasks are processed sequentially and prevents state race conditions. Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * some small fixes * small cleanup * Refactor shell_manager_tx in portals (#1854) * Refactor shell_manager_tx in portals to use mpsc::Sender directly This change simplifies the passing of `shell_manager_tx` to portal functions by removing the `Arc<Mutex<Option<...>>>` wrapper. The sender is now passed directly as `mpsc::Sender<ShellManagerMessage>`. - Updated `run_create_portal` in `implants/imix/src/portal/mod.rs` to accept `mpsc::Sender`. - Updated `run` and `handle_incoming_mote` in `implants/imix/src/portal/run.rs` to accept `mpsc::Sender` and remove unnecessary locking. - Updated `create_portal` in `implants/imix/src/agent.rs` to extract and clone the sender before spawning the portal task. - `handle_incoming_mote` now uses `try_send` to avoid blocking the portal loop on shell manager backpressure. Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * Refactor shell_manager_tx in ImixAgent to be a direct mpsc::Sender - Updated `ImixAgent` struct: `shell_manager_tx` is now `tokio::sync::mpsc::Sender<ShellManagerMessage>`. - Introduced `shell_manager_rx: Arc<Mutex<Option<...>>>` to handle the receiver side in `ImixAgent::new`. - Updated `start_shell_manager` to claim the receiver from `shell_manager_rx`. - Simplified usages of `shell_manager_tx` in `process_job_request` and `create_portal` (removed unnecessary locking/cloning). - This aligns with the requirement that `shell_manager_tx` should not be an Option and should be readily available. Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * Refactor shell_manager_tx in ImixAgent to be a direct mpsc::Sender - Updated `ImixAgent` struct: `shell_manager_tx` is now `tokio::sync::mpsc::Sender<ShellManagerMessage>`. - Updated `ImixAgent::new` to accept `shell_manager_tx` as an argument, implementing dependency injection. - Updated `start_shell_manager` to accept a `ShellManager` instance instead of creating the channel internally. - Refactored `run_agent` in `run.rs` to create the channel and `ShellManager` explicitly. - Updated all relevant tests to match the new API. - Simplified usages of `shell_manager_tx` in `process_job_request` and `create_portal` (removed unnecessary locking/cloning/unwrapping). Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> * minor cleanup * remove useless comments --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com>
Refactor imix/shell to reuse StreamPrinter implementation
Moved
StreamPrinterfromimix/src/task.rstoimix/src/printer.rsand updated it to distinguish betweenStdoutandStderrusing anOutputKindenum.Updated
imix/src/shell/repl.rsto reuseStreamPrinterinstead of the customShellPrinter.Refactored
imix/src/task.rsto use the sharedStreamPrinterand properly report stderr toTaskErrorfield.Updated
imix/src/lib.rsandimix/src/main.rsto include the newprintermodule.This change unifies output handling across the agent and improves maintainability.
PR created automatically by Jules for task 10663336750935921147 started by @KCarretto