imix: Separate aggregation of task and shell task outputs#1835
Merged
KCarretto merged 1 commit intofeat/noninteractive-shell-e2efrom Feb 21, 2026
Merged
Conversation
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: KCarretto <16250309+KCarretto@users.noreply.github.com>
Contributor
|
👋 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. |
github-merge-queue bot
pushed a commit
that referenced
this pull request
Feb 22, 2026
* 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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR updates the
flush_outputsmethod in theimixagent to correctly handle and aggregateshell_task_outputmessages separately from standard task outputs.Previously, aggregation was done solely based on
task_idderived from the context. Shell task outputs typically lack a task context, leading to incorrect aggregation (often under task ID 0) and potential data loss if multiple shell tasks produced output simultaneously.The new implementation:
outputinto a map keyed bytask_id.shell_task_outputinto a map keyed byshell_task_id.ReportTaskOutputRequestmessages for each aggregated entry.A regression test
test_agent_output_aggregationhas been added toimplants/imix/src/tests/agent_output_aggregation.rs.Existing tests passed (with a minor fix to
task_tests.rsto accommodate proto changes).PR created automatically by Jules for task 11423286900056386578 started by @KCarretto