SemianiakaVY/openpencil
1
0
Fork 0
Code Issues Pull requests Actions 1 Packages Projects Releases Wiki Activity

feat(orchestrator): concurrent executor — semaphore + serialized replay

Browse source 
Task B2 of S3b-2: implements `run_concurrent` — the concurrent executor
that drives screen-group workers via a tokio Semaphore (RAII permits,
FIFO-fair, cap = effective_concurrency), collects per-worker buffered
EditorCommands, replays them into the real DocSink in subtask-plan-index
order, and fan-ins Progress events via mpsc channel.

Also splits concurrent.rs test block into concurrent_tests.rs (STEP 0)
to keep both files under the 800-line ceiling, wired via
`#[path = "concurrent_tests.rs"] mod tests` (same pattern as plan_repair).

Updated run_screen_group_worker signature: now takes Arc<Semaphore> +
mpsc::UnboundedSender<Progress> instead of &mut dyn FnMut(Progress).

Added CountingLlm to test_support for semaphore-cap verification.
Added tokio (sync + rt + macros) to op-orchestrator Cargo.toml.

208 tests pass; clippy -D warnings clean; fmt clean.
concurrent.rs: 432 lines; concurrent_tests.rs: 727 lines.
This commit is contained in:
Fini 2026-05-23 04:06:51 +08:00
parent 4def1dd62b
commit 7f5152ac07
6 changed files with 1234 additions and 487 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
Cargo.lock generated
View file

@ -2707,6 +2707,7 @@ dependencies = [
"op-editor-core",
"serde",
"serde_json",
"tokio",
"tracing",
]

1
crates/op-orchestrator/Cargo.toml
View file

@ -15,6 +15,7 @@ serde = { workspace = true }
serde_json = { workspace = true }
futures = "0.3"
tracing = { workspace = true }
tokio = { version = "1", features = ["sync", "rt", "macros"] }
[dev-dependencies]
serde = { workspace = true, features = ["derive"] }

701
crates/op-orchestrator/src/concurrent.rs
View file

@ -1,7 +1,9 @@
//! S3b-2 Tasks A1 + B1: concurrency decision + screen grouping + worker.
//! S3b-2 Tasks A1 + B1 + B2: concurrency decision + screen grouping + worker
//! + concurrent executor.
//!
//! Port of `orchestrator.ts:780-810` (Task A1) and
//! `orchestrator-sub-agent.ts:248-312` (Task B1 worker body).
//! Port of `orchestrator.ts:780-810` (Task A1),
//! `orchestrator-sub-agent.ts:248-312` (Task B1 worker body), and
//! `orchestrator-sub-agent.ts:227-327` (Task B2 concurrent executor).
//!
//! ## Task A1
//! - `DesignRequest.concurrency` (added in `types.rs`)
@ -21,14 +23,35 @@
//! in order using the 2-attempt concurrent retry ladder (vs the sequential
//! path's 3-attempt tier-gated ladder in `run.rs`).
//!
//! Callers land in later S3b-2 tasks; scaffolding symbols are allowed to be unused.
//! ## Task B2
//! - `run_concurrent` — the concurrent executor:
//! - a shared `Arc<tokio::sync::Semaphore>` (permits = `effective_concurrency`)
//! caps in-flight subtask LLM calls; RAII permit drop replaces TS `releaseSlot`
//! - N worker futures (one per screen group, each owning its own `BufferDocSink`)
//! driven together by `futures::future::join_all` on a SINGLE task (no
//! `tokio::spawn` — avoids `Send` bounds; LLM I/O is offloaded in
//! `LlmClient::call`). `join_all` polls every worker on each wake, so
//! different groups' LLM calls genuinely overlap — the concurrent equivalent
//! of JS `Promise.all(workers)`.
//! - progress fan-in: workers send `Progress` on an mpsc channel; driver drains
//! + forwards to `on_progress` after `join_all`
//! - serialized replay: after `join_all`, replays every worker's buffered
//! commands into the real `&mut DocSink` in subtask-plan-index order
//! (deterministic)
//!
//! Callers for C1/C2 land in later S3b-2 tasks; scaffolding symbols are allowed
//! to be unused until then.
#![allow(dead_code)]
use crate::plan::{OrchestratorPlan, Subtask};
use crate::retry::is_non_retryable;
use crate::subagent::run_subtask;
use crate::types::{AbortFlag, DesignRequest, DocSink, LlmClient, Progress, SubtaskOutcome};
use futures::future::join_all;
use op_editor_core::{EditorCommand, EditorState};
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::sync::Semaphore;
/// A group of subtask indices that share the same screen.
/// `screen` is the screen name; `indices` are into the plan's `subtasks` slice.
@ -121,7 +144,7 @@ pub(crate) fn effective_concurrency(concurrency: u32, screen_group_count: usize)
/// be taken just before the concurrent phase starts (after the N-root scaffold
/// is applied to the real sink) — that is Task B2's responsibility.
///
/// After `join_all` the orchestrator replays `commands` into the real
/// After all workers finish the orchestrator replays `commands` into the real
/// `DocSink` in subtask-plan-index order (Task B2: serialized replay).
pub(crate) struct BufferDocSink {
/// Snapshot of `EditorState` taken before the concurrent phase.
@ -167,13 +190,31 @@ impl DocSink for BufferDocSink {
}
}
// ── Task B1: concurrent screen-group worker fn ────────────────────────────────
// ── Task B1/B2: concurrent screen-group worker fn ─────────────────────────────
/// The output of one screen-group worker: its private command buffer plus the
/// per-subtask outcomes (each tagged with its plan index for serialized replay).
pub(crate) struct WorkerResult {
/// The worker's private buffer — replayed into the real sink after `join_all`.
pub buffer: BufferDocSink,
/// `(plan_index, outcome)` for every subtask the worker ran, in `indices` order.
pub outcomes: Vec<(usize, SubtaskOutcome)>,
}
/// Runs all subtasks for one screen group using the **2-attempt concurrent
/// retry ladder** (vs the sequential path's 3-attempt tier-gated ladder).
///
/// Port of `orchestrator-sub-agent.ts:248-312`.
///
/// ## Owned buffer (genuine concurrency)
/// The worker **owns** its [`BufferDocSink`] by value (passed in, returned in
/// [`WorkerResult`]). `run_subtask` writes into this stack-local buffer via
/// `&mut` — the `&mut` is to the worker's *own* buffer, never a shared `Vec`
/// slot. This is what lets [`run_concurrent`] drive N worker futures with
/// `futures::future::join_all` without the borrow checker rejecting aliased
/// `&mut` borrows: each future owns its buffer, so there is no shared mutable
/// state.
///
/// ## Retry ladder
/// | Attempt | `reduced_complexity` | `minimal_skills` |
/// |---------|----------------------|------------------|
@ -182,30 +223,30 @@ impl DocSink for BufferDocSink {
///
/// Skips the sequential path's tier-gated middle attempt.
///
/// ## Semaphore
/// Each subtask acquires one permit from `semaphore` before the LLM call.
/// RAII permit drop (when the block completes) replaces TS `releaseSlot`.
///
/// ## Retryable-failure gate (identical to sequential path)
/// `error.is_some() && node_count == 0 && !abort.is_set() && !is_non_retryable(&err)`
///
/// Partial results (`node_count > 0`) are **never** retried.
///
/// ## Progress
/// `progress_fn` is called with `SubtaskStarted` before each subtask and
/// `SubtaskDone` / `SubtaskFailed` after. In Task B2 this will be wired to
/// an `mpsc::UnboundedSender` so multiple workers can fan-in; for B1 a plain
/// `&mut dyn FnMut(Progress)` suffices.
///
/// ## Return value
/// A `Vec<SubtaskOutcome>` with one entry per subtask index in `group.indices`,
/// in the same order as `group.indices` (i.e. plan order within this screen).
/// Progress events are sent via `progress_tx` (mpsc channel).
/// The `run_concurrent` driver drains the receiver and forwards to `on_progress`.
#[allow(clippy::too_many_arguments)]
pub(crate) async fn run_screen_group_worker(
group: &ScreenGroup,
plan: &OrchestratorPlan,
request: &DesignRequest,
llm: &dyn LlmClient,
abort: &AbortFlag,
sink: &mut BufferDocSink,
progress_fn: &mut dyn FnMut(Progress),
) -> Vec<SubtaskOutcome> {
let mut outcomes: Vec<SubtaskOutcome> = Vec::with_capacity(group.indices.len());
mut buffer: BufferDocSink,
semaphore: Arc<Semaphore>,
progress_tx: mpsc::UnboundedSender<Progress>,
) -> WorkerResult {
let mut outcomes: Vec<(usize, SubtaskOutcome)> = Vec::with_capacity(group.indices.len());
for &idx in &group.indices {
// Abort check at the top of every iteration (spec §4.4).
@ -215,13 +256,30 @@ pub(crate) async fn run_screen_group_worker(
let subtask = &plan.subtasks[idx];
progress_fn(Progress::SubtaskStarted {
let _ = progress_tx.send(Progress::SubtaskStarted {
id: subtask.id.clone(),
label: subtask.label.clone(),
});
// Acquire a semaphore permit before the LLM call.
// RAII: permit drops when the block completes (= releaseSlot in TS).
let _permit = semaphore
.acquire()
.await
.expect("semaphore should not be closed");
// --- Attempt 1: full complexity ---
let outcome1 = run_subtask(subtask, plan, request, llm, sink, abort, false, false).await;
let outcome1 = run_subtask(
subtask,
plan,
request,
llm,
&mut buffer,
abort,
false,
false,
)
.await;
// Evaluate the non-retryable predicate once from attempt-1's error
// (matches TS `isNonRetryable` computed before the retry chain).
@ -244,7 +302,7 @@ pub(crate) async fn run_screen_group_worker(
error = outcome1.error.as_deref().unwrap_or(""),
"concurrent worker: subtask empty, retrying with minimal skills (attempt 2)"
);
run_subtask(subtask, plan, request, llm, sink, abort, true, true).await
run_subtask(subtask, plan, request, llm, &mut buffer, abort, true, true).await
} else {
outcome1
};
@ -255,482 +313,153 @@ pub(crate) async fn run_screen_group_worker(
let err_msg = final_outcome.error.clone();
if zero {
progress_fn(Progress::SubtaskFailed {
let _ = progress_tx.send(Progress::SubtaskFailed {
id: subtask.id.clone(),
error: err_msg.unwrap_or_default(),
});
} else {
progress_fn(Progress::SubtaskDone {
let _ = progress_tx.send(Progress::SubtaskDone {
id: subtask.id.clone(),
node_count,
});
}
outcomes.push(final_outcome);
outcomes.push((idx, final_outcome));
}
outcomes
WorkerResult { buffer, outcomes }
}
// ── Task B2: run_concurrent ────────────────────────────────────────────────────
/// Concurrent executor — semaphore-gated screen-group workers + serialized replay.
///
/// Port of `orchestrator-sub-agent.ts:227-327` (semaphore + workers + collection).
///
/// ## Genuine concurrency via `join_all`
/// One **worker future per screen group** is built; all N are driven together
/// with [`futures::future::join_all`] on the SINGLE current task (no
/// `tokio::spawn` — avoids `Send` bounds; the LLM I/O is already offloaded
/// inside `LlmClient::call`). `join_all` polls every worker on each wake: when
/// worker A awaits an LLM stream chunk, `join_all` polls worker B, which
/// proceeds with *its* LLM call — the calls genuinely overlap in wall-clock
/// time. This is the concurrent equivalent of JS `Promise.all(workers)`.
///
/// The borrow checker is satisfied because **each worker future owns its own
/// `BufferDocSink`** (constructed here, moved in, returned in [`WorkerResult`])
/// — there is no shared `&mut`. The genuinely-shared values are all shared
/// `&` (`&dyn LlmClient`, `&AbortFlag`, `&plan`, `&request`, `Arc<Semaphore>`);
/// the `mpsc` sender is cloned per worker.
///
/// ## Semaphore
/// A shared `Arc<Semaphore>` (permits = `effective_concurrency(concurrency,
/// groups.len())`) caps the total in-flight subtask LLM calls across all
/// genuinely-overlapping workers.
///
/// ## Progress fan-in
/// Workers send `Progress` events on cloned `mpsc::UnboundedSender<Progress>`s;
/// after `join_all` the driver drains the receiver and forwards to
/// `on_progress`. The `on_progress` consumer must tolerate interleaved
/// `SubtaskStarted` events (multiple started before corresponding done).
///
/// ## Serialized replay (spec §5)
/// `join_all` yields `Vec<WorkerResult>`; the driver then replays every
/// worker's buffered `EditorCommand`s into the real `&mut DocSink` in
/// ascending plan-index order — deterministic, matches the TS indexed-results
/// array, never torn.
///
/// Sort-by-min-index: within each group subtasks already ran in ascending plan
/// index order, so the buffer contains commands in ascending plan-index order.
/// Cross-group: sorting groups by min(indices) ensures the group with the
/// lower-numbered subtask replays first.
///
/// ## Return value
/// `Vec<Option<SubtaskOutcome>>` indexed by plan order (0..plan.subtasks.len()).
/// A slot is `None` if the worker was aborted before reaching that subtask.
/// The all-zero-nodes failure verdict (Task C1) is NOT checked here.
#[allow(clippy::too_many_arguments)]
pub(crate) async fn run_concurrent(
groups: &[ScreenGroup],
plan: &OrchestratorPlan,
request: &DesignRequest,
llm: &dyn LlmClient,
abort: &AbortFlag,
snapshot: EditorState,
real_sink: &mut dyn DocSink,
on_progress: &mut dyn FnMut(Progress),
) -> Vec<Option<SubtaskOutcome>> {
let n = plan.subtasks.len();
let ec = effective_concurrency(request.concurrency, groups.len());
let semaphore = Arc::new(Semaphore::new(ec as usize));
// Progress channel: workers send; we drain after `join_all`.
// The master sender is dropped after building the per-worker clones so the
// channel closes once every worker (and thus every clone) has finished.
let (progress_tx, mut progress_rx) = mpsc::unbounded_channel::<Progress>();
// Build one worker future per screen group. Each future OWNS its own
// `BufferDocSink` (moved in) — no shared `&mut`, so `join_all` is legal.
let worker_futures: Vec<_> = groups
.iter()
.map(|group| {
let buffer = BufferDocSink::new(snapshot.clone());
let sem = Arc::clone(&semaphore);
let ptx = progress_tx.clone();
run_screen_group_worker(group, plan, request, llm, abort, buffer, sem, ptx)
})
.collect();
// Drop the master sender so the channel closes when all worker clones drop.
drop(progress_tx);
// Drive ALL workers concurrently on the current task. `join_all` polls
// each worker on every wake — LLM calls from different groups genuinely
// overlap. Returns one `WorkerResult` per group, in `groups` order.
let worker_results: Vec<WorkerResult> = join_all(worker_futures).await;
// Drain the progress channel and forward to `on_progress`.
while let Ok(event) = progress_rx.try_recv() {
on_progress(event);
}
// Collect per-subtask outcomes indexed by plan order.
let mut per_subtask: Vec<Option<SubtaskOutcome>> = vec![None; n];
for result in &worker_results {
for (plan_idx, outcome) in &result.outcomes {
per_subtask[*plan_idx] = Some(outcome.clone());
}
}
// Serialized replay in plan-index order.
// `worker_results` is in `groups` order; sort it by each group's minimum
// plan index so the group containing the lowest-indexed subtask replays
// first. Within a group subtasks already ran in ascending index order.
let mut replay_order: Vec<usize> = (0..groups.len()).collect();
replay_order.sort_by_key(|&g| {
groups[g]
.indices
.iter()
.copied()
.min()
.unwrap_or(usize::MAX)
});
let mut worker_results = worker_results;
for g_idx in replay_order {
for cmd in worker_results[g_idx].buffer.commands.drain(..) {
real_sink.apply(cmd);
}
}
per_subtask
}
// Tests are split across two sibling files to honor the 800-line cap:
// `concurrent_tests.rs` — Task A1 + B1 (grouping / decision / worker / sink)
// `concurrent_tests_b2.rs` — Task B2 (`run_concurrent` concurrent executor)
#[cfg(test)]
#[path = "concurrent_tests.rs"]
mod tests;
#[cfg(test)]
mod tests {
use super::*;
use crate::plan::{Region, Subtask};
fn subtask_with_screen(id: &str, screen: Option<&str>) -> Subtask {
Subtask {
id: id.into(),
label: id.into(),
id_prefix: id.into(),
region: Region {
width: 1200.0,
height: 400.0,
},
parent_frame_id: None,
elements: None,
screen: screen.map(|s| s.to_string()),
}
}
// ── effective_concurrency ──────────────────────────────────────────────────
/// (concurrency=1, 3 screens) → 1 (single-threaded even with many screens)
#[test]
fn effective_concurrency_one_concurrency_three_screens_gives_one() {
assert_eq!(effective_concurrency(1, 3), 1);
}
/// (concurrency=4, 1 screen) → 1 (only one group → sequential)
#[test]
fn effective_concurrency_four_concurrency_one_screen_gives_one() {
assert_eq!(effective_concurrency(4, 1), 1);
}
/// (concurrency=4, 3 screens) → 4
#[test]
fn effective_concurrency_four_concurrency_three_screens_gives_four() {
assert_eq!(effective_concurrency(4, 3), 4);
}
/// Clamp: (concurrency=99, 3 screens) → 6
#[test]
fn effective_concurrency_clamps_to_six() {
assert_eq!(effective_concurrency(99, 3), 6);
}
// ── group_subtasks_by_screen ───────────────────────────────────────────────
/// Basic grouping: [login, home, login] → 2 groups {login:[0,2], home:[1]}
#[test]
fn group_subtasks_three_entries_two_screens() {
let subtasks = vec![
subtask_with_screen("a", Some("login")),
subtask_with_screen("b", Some("home")),
subtask_with_screen("c", Some("login")),
];
let groups = group_subtasks_by_screen(&subtasks);
assert_eq!(groups.len(), 2);
assert_eq!(groups[0].screen, "login");
assert_eq!(groups[0].indices, vec![0, 2]);
assert_eq!(groups[1].screen, "home");
assert_eq!(groups[1].indices, vec![1]);
}
/// A subtask with no screen falls back to first_screen.
#[test]
fn group_subtasks_no_screen_falls_back_to_first_screen() {
let subtasks = vec![
subtask_with_screen("a", Some("login")),
subtask_with_screen("b", None), // no screen → "login"
subtask_with_screen("c", Some("home")),
];
let groups = group_subtasks_by_screen(&subtasks);
// "login" gets index 0 and 1 (b bucketed under first_screen="login")
assert_eq!(groups.len(), 2);
assert_eq!(groups[0].screen, "login");
assert_eq!(groups[0].indices, vec![0, 1]);
assert_eq!(groups[1].screen, "home");
assert_eq!(groups[1].indices, vec![2]);
}
/// All subtasks have no screen → empty result (no groups).
#[test]
fn group_subtasks_all_no_screen_returns_empty() {
let subtasks = vec![
subtask_with_screen("a", None),
subtask_with_screen("b", None),
];
let groups = group_subtasks_by_screen(&subtasks);
assert!(groups.is_empty());
}
/// Empty subtask slice → empty result.
#[test]
fn group_subtasks_empty_slice_returns_empty() {
let groups = group_subtasks_by_screen(&[]);
assert!(groups.is_empty());
}
/// Group order is first-seen order of distinct screen values.
#[test]
fn group_subtasks_preserves_first_seen_order() {
let subtasks = vec![
subtask_with_screen("a", Some("profile")),
subtask_with_screen("b", Some("settings")),
subtask_with_screen("c", Some("profile")),
subtask_with_screen("d", Some("dashboard")),
];
let groups = group_subtasks_by_screen(&subtasks);
assert_eq!(groups.len(), 3);
assert_eq!(groups[0].screen, "profile");
assert_eq!(groups[1].screen, "settings");
assert_eq!(groups[2].screen, "dashboard");
}
/// first_screen fallback is "page" when no subtask has a screen
/// (but `has_any_screen` is false, so returns empty — tested separately).
/// Here we test that a None screen at the START falls back to first_screen
/// from a LATER subtask — i.e., first_screen is derived from the first
/// subtask that HAS a screen.
#[test]
fn group_subtasks_no_screen_at_start_uses_later_first_screen() {
let subtasks = vec![
subtask_with_screen("a", None), // no screen
subtask_with_screen("b", Some("home")), // first with screen → first_screen="home"
subtask_with_screen("c", Some("settings")),
];
let groups = group_subtasks_by_screen(&subtasks);
// "home" first (b has screen), but "a" has no screen → bucketed under first_screen="home"
assert_eq!(groups.len(), 2);
assert_eq!(groups[0].screen, "home");
// "a" (index 0) bucketed under "home" — it's first because it's processed first
// and maps to first_screen="home" which is first seen when processing "a".
assert!(groups[0].indices.contains(&0));
assert!(groups[0].indices.contains(&1));
assert_eq!(groups[1].screen, "settings");
assert_eq!(groups[1].indices, vec![2]);
}
// ── Task B1: BufferDocSink ────────────────────────────────────────────────
use crate::test_support::{ScriptResponse, ScriptedLlm};
use crate::types::LlmError;
use futures::executor::block_on;
use op_editor_core::EditorState;
fn make_plan_with_subtasks(subtask_ids: &[&str]) -> crate::plan::OrchestratorPlan {
use crate::plan::{OrchestratorPlan, Region, RootFrameSpec};
OrchestratorPlan {
root_frame: RootFrameSpec {
id: "root".into(),
name: "Page".into(),
width: 1200.0,
height: 800.0,
layout: None,
gap: None,
padding: None,
fill: None,
},
subtasks: subtask_ids
.iter()
.map(|id| crate::plan::Subtask {
id: id.to_string(),
label: id.to_string(),
region: Region {
width: 1200.0,
height: 400.0,
},
id_prefix: id.to_string(),
parent_frame_id: Some("root".into()),
elements: None,
screen: None,
})
.collect(),
style_guide_name: None,
}
}
fn make_req() -> crate::types::DesignRequest {
crate::types::DesignRequest {
prompt: "test".into(),
model: None,
provider: None,
design_md: None,
concurrency: 2,
}
}
const NODE_JSON: &str = r#"[{"type":"frame","id":"h1","name":"H","x":0,"y":0,"width":100,"height":100,"children":[]}]"#;
/// `BufferDocSink` collects commands without modifying a real doc.
#[test]
fn buffer_doc_sink_collects_commands() {
let mut sink = BufferDocSink::new(EditorState::new());
// apply() should always return true and buffer the command.
let applied = sink.apply(EditorCommand::InsertSubtree {
nodes: vec![],
parent_id: op_editor_core::NodeId::NONE,
});
assert!(applied, "BufferDocSink.apply must always return true");
assert_eq!(sink.commands.len(), 1);
}
/// `state()` on `BufferDocSink` returns the snapshot passed at construction.
#[test]
fn buffer_doc_sink_state_returns_snapshot() {
let state = EditorState::new();
let sink = BufferDocSink::new(state.clone());
// The returned reference should be valid (same content).
let _ = sink.state();
}
/// `BufferDocSink` tracks undo-batch depth correctly.
#[test]
fn buffer_doc_sink_undo_batch_depth() {
let mut sink = BufferDocSink::new(EditorState::new());
assert_eq!(sink.batch_depth, 0);
sink.begin_undo_batch();
assert_eq!(sink.batch_depth, 1);
sink.end_undo_batch();
assert_eq!(sink.batch_depth, 0);
}
// ── Task B1: run_screen_group_worker ─────────────────────────────────────
/// Happy path: both subtasks succeed on attempt 1.
#[test]
fn worker_happy_path_both_succeed() {
let plan = make_plan_with_subtasks(&["s0", "s1"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![
ScriptResponse::Text(NODE_JSON.into()), // s0 attempt 1
ScriptResponse::Text(NODE_JSON.into()), // s1 attempt 1
]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0, 1],
};
let abort = AbortFlag::new();
let mut sink = BufferDocSink::new(EditorState::new());
let mut events: Vec<Progress> = Vec::new();
let outcomes = block_on(run_screen_group_worker(
&group,
&plan,
&req,
&llm,
&abort,
&mut sink,
&mut |p| events.push(p),
));
assert_eq!(outcomes.len(), 2);
assert_eq!(outcomes[0].node_count, 1);
assert_eq!(outcomes[1].node_count, 1);
// Buffer collected 2 InsertSubtree commands (one per subtask).
assert_eq!(sink.commands.len(), 2);
}
/// Attempt-1 zero-node + retryable → attempt 2 runs with (true, true).
/// We verify this by scripting: attempt 1 → error, attempt 2 → nodes.
#[test]
fn worker_attempt1_zero_retryable_triggers_attempt2() {
let plan = make_plan_with_subtasks(&["s0"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![
// Attempt 1 — LLM error → zero nodes, retryable
ScriptResponse::Fail(LlmError {
message: "timeout".into(),
aborted: false,
}),
// Attempt 2 — succeeds
ScriptResponse::Text(NODE_JSON.into()),
]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0],
};
let abort = AbortFlag::new();
let mut sink = BufferDocSink::new(EditorState::new());
let mut events: Vec<Progress> = Vec::new();
let outcomes = block_on(run_screen_group_worker(
&group,
&plan,
&req,
&llm,
&abort,
&mut sink,
&mut |p| events.push(p),
));
assert_eq!(outcomes.len(), 1);
// Final outcome should show nodes from attempt 2.
assert_eq!(
outcomes[0].node_count, 1,
"attempt 2 should have produced a node"
);
assert_eq!(sink.commands.len(), 1, "one InsertSubtree from attempt 2");
}
/// Non-retryable error on attempt 1 → NO attempt 2.
#[test]
fn worker_non_retryable_error_no_retry() {
let plan = make_plan_with_subtasks(&["s0"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![
// Attempt 1 — non-retryable HTTP 401
ScriptResponse::Fail(LlmError {
message: "HTTP 401 Unauthorized".into(),
aborted: false,
}),
// Attempt 2 — should NOT be called; if it is, the test fails by
// observing the scripted LLM's exhausted error.
]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0],
};
let abort = AbortFlag::new();
let mut sink = BufferDocSink::new(EditorState::new());
let mut events: Vec<Progress> = Vec::new();
let outcomes = block_on(run_screen_group_worker(
&group,
&plan,
&req,
&llm,
&abort,
&mut sink,
&mut |p| events.push(p),
));
assert_eq!(outcomes.len(), 1);
assert_eq!(outcomes[0].node_count, 0, "non-retryable must stay zero");
// Verify the error message is the original 401, not "scripted LLM exhausted"
assert!(
outcomes[0].error.as_deref().unwrap_or("").contains("401"),
"error should carry original 401 message"
);
// No InsertSubtree in the buffer (nothing was applied).
assert_eq!(sink.commands.len(), 0);
}
/// Partial result (node_count > 0) on attempt 1 → never retried.
#[test]
fn worker_partial_result_never_retried() {
let plan = make_plan_with_subtasks(&["s0"]);
let req = make_req();
// Attempt 1 produces nodes — the scripted LLM has only one response.
let llm = ScriptedLlm::new(vec![ScriptResponse::Text(NODE_JSON.into())]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0],
};
let abort = AbortFlag::new();
let mut sink = BufferDocSink::new(EditorState::new());
let outcomes = block_on(run_screen_group_worker(
&group,
&plan,
&req,
&llm,
&abort,
&mut sink,
&mut |_| {},
));
assert_eq!(outcomes.len(), 1);
assert_eq!(outcomes[0].node_count, 1);
// Only 1 command buffered (attempt 1 only).
assert_eq!(sink.commands.len(), 1);
}
/// Abort flag set before the loop → worker returns immediately with no outcomes.
#[test]
fn worker_aborts_at_loop_head() {
let plan = make_plan_with_subtasks(&["s0", "s1"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![]); // should not be called
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0, 1],
};
let abort = AbortFlag::new();
abort.set(); // set before the worker runs
let mut sink = BufferDocSink::new(EditorState::new());
let outcomes = block_on(run_screen_group_worker(
&group,
&plan,
&req,
&llm,
&abort,
&mut sink,
&mut |_| {},
));
assert!(
outcomes.is_empty(),
"aborted before loop: no outcomes expected"
);
assert_eq!(sink.commands.len(), 0);
}
/// Progress events: SubtaskStarted is emitted before each subtask;
/// SubtaskDone emitted for successful subtasks.
#[test]
fn worker_emits_started_and_done_progress() {
let plan = make_plan_with_subtasks(&["s0", "s1"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![
ScriptResponse::Text(NODE_JSON.into()),
ScriptResponse::Text(NODE_JSON.into()),
]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0, 1],
};
let abort = AbortFlag::new();
let mut sink = BufferDocSink::new(EditorState::new());
let mut events: Vec<String> = Vec::new();
block_on(run_screen_group_worker(
&group,
&plan,
&req,
&llm,
&abort,
&mut sink,
&mut |p| match &p {
Progress::SubtaskStarted { id, .. } => events.push(format!("start:{id}")),
Progress::SubtaskDone { id, .. } => events.push(format!("done:{id}")),
Progress::SubtaskFailed { id, .. } => events.push(format!("fail:{id}")),
_ => {}
},
));
assert_eq!(events, vec!["start:s0", "done:s0", "start:s1", "done:s1"],);
}
/// SubtaskFailed progress emitted when final outcome is zero nodes.
#[test]
fn worker_emits_failed_progress_on_zero_nodes() {
let plan = make_plan_with_subtasks(&["s0"]);
let req = make_req();
// Both attempts produce garbage → zero nodes.
let llm = ScriptedLlm::new(vec![
ScriptResponse::Text("not json".into()),
ScriptResponse::Text("also not json".into()),
]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0],
};
let abort = AbortFlag::new();
let mut sink = BufferDocSink::new(EditorState::new());
let mut events: Vec<String> = Vec::new();
block_on(run_screen_group_worker(
&group,
&plan,
&req,
&llm,
&abort,
&mut sink,
&mut |p| {
if let Progress::SubtaskFailed { id, .. } = &p {
events.push(format!("fail:{id}"));
}
},
));
assert_eq!(events, vec!["fail:s0"]);
}
}
#[path = "concurrent_tests_b2.rs"]
mod tests_b2;

457
crates/op-orchestrator/src/concurrent_tests.rs Normal file
View file

@ -0,0 +1,457 @@
//! `concurrent` tests, part 1 — Task A1 (`effective_concurrency`,
//! `group_subtasks_by_screen`) + Task B1 (`BufferDocSink`,
//! `run_screen_group_worker`).
//!
//! Task B2 (`run_concurrent`) tests live in the sibling `concurrent_tests_b2.rs`
//! — split so both test files stay under the 800-line cap.
//!
//! Wired as `#[path = "concurrent_tests.rs"] mod tests;` inside `concurrent.rs`,
//! so this stays a child module of `concurrent` and `use super::*` resolves to
//! `concurrent`.
use super::*;
use crate::plan::{Region, Subtask};
use crate::test_support::{ScriptResponse, ScriptedLlm};
use crate::types::LlmError;
use futures::executor::block_on;
use op_editor_core::EditorState;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::sync::Semaphore;
// ── Helpers ────────────────────────────────────────────────────────────────
fn subtask_with_screen(id: &str, screen: Option<&str>) -> Subtask {
Subtask {
id: id.into(),
label: id.into(),
id_prefix: id.into(),
region: Region {
width: 1200.0,
height: 400.0,
},
parent_frame_id: None,
elements: None,
screen: screen.map(|s| s.to_string()),
}
}
fn make_plan_with_subtasks(subtask_ids: &[&str]) -> crate::plan::OrchestratorPlan {
use crate::plan::{OrchestratorPlan, Region, RootFrameSpec};
OrchestratorPlan {
root_frame: RootFrameSpec {
id: "root".into(),
name: "Page".into(),
width: 1200.0,
height: 800.0,
layout: None,
gap: None,
padding: None,
fill: None,
},
subtasks: subtask_ids
.iter()
.map(|id| crate::plan::Subtask {
id: id.to_string(),
label: id.to_string(),
region: Region {
width: 1200.0,
height: 400.0,
},
id_prefix: id.to_string(),
parent_frame_id: Some("root".into()),
elements: None,
screen: None,
})
.collect(),
style_guide_name: None,
}
}
fn make_req() -> crate::types::DesignRequest {
crate::types::DesignRequest {
prompt: "test".into(),
model: None,
provider: None,
design_md: None,
concurrency: 2,
}
}
const NODE_JSON: &str =
r#"[{"type":"frame","id":"h1","name":"H","x":0,"y":0,"width":100,"height":100,"children":[]}]"#;
/// Build a default semaphore (2 permits) and an mpsc sender for B1 worker tests.
/// Returns `(semaphore, sender, receiver)`.
fn make_worker_channel() -> (
Arc<Semaphore>,
mpsc::UnboundedSender<Progress>,
mpsc::UnboundedReceiver<Progress>,
) {
let sem = Arc::new(Semaphore::new(2));
let (tx, rx) = mpsc::unbounded_channel::<Progress>();
(sem, tx, rx)
}
/// Drain all progress events from a receiver into a Vec.
fn drain_events(mut rx: mpsc::UnboundedReceiver<Progress>) -> Vec<Progress> {
let mut events = Vec::new();
while let Ok(e) = rx.try_recv() {
events.push(e);
}
events
}
// ── effective_concurrency ──────────────────────────────────────────────────
/// (concurrency=1, 3 screens) → 1 (single-threaded even with many screens)
#[test]
fn effective_concurrency_one_concurrency_three_screens_gives_one() {
assert_eq!(effective_concurrency(1, 3), 1);
}
/// (concurrency=4, 1 screen) → 1 (only one group → sequential)
#[test]
fn effective_concurrency_four_concurrency_one_screen_gives_one() {
assert_eq!(effective_concurrency(4, 1), 1);
}
/// (concurrency=4, 3 screens) → 4
#[test]
fn effective_concurrency_four_concurrency_three_screens_gives_four() {
assert_eq!(effective_concurrency(4, 3), 4);
}
/// Clamp: (concurrency=99, 3 screens) → 6
#[test]
fn effective_concurrency_clamps_to_six() {
assert_eq!(effective_concurrency(99, 3), 6);
}
// ── group_subtasks_by_screen ───────────────────────────────────────────────
/// Basic grouping: [login, home, login] → 2 groups {login:[0,2], home:[1]}
#[test]
fn group_subtasks_three_entries_two_screens() {
let subtasks = vec![
subtask_with_screen("a", Some("login")),
subtask_with_screen("b", Some("home")),
subtask_with_screen("c", Some("login")),
];
let groups = group_subtasks_by_screen(&subtasks);
assert_eq!(groups.len(), 2);
assert_eq!(groups[0].screen, "login");
assert_eq!(groups[0].indices, vec![0, 2]);
assert_eq!(groups[1].screen, "home");
assert_eq!(groups[1].indices, vec![1]);
}
/// A subtask with no screen falls back to first_screen.
#[test]
fn group_subtasks_no_screen_falls_back_to_first_screen() {
let subtasks = vec![
subtask_with_screen("a", Some("login")),
subtask_with_screen("b", None), // no screen → "login"
subtask_with_screen("c", Some("home")),
];
let groups = group_subtasks_by_screen(&subtasks);
assert_eq!(groups.len(), 2);
assert_eq!(groups[0].screen, "login");
assert_eq!(groups[0].indices, vec![0, 1]);
assert_eq!(groups[1].screen, "home");
assert_eq!(groups[1].indices, vec![2]);
}
/// All subtasks have no screen → empty result (no groups).
#[test]
fn group_subtasks_all_no_screen_returns_empty() {
let subtasks = vec![
subtask_with_screen("a", None),
subtask_with_screen("b", None),
];
let groups = group_subtasks_by_screen(&subtasks);
assert!(groups.is_empty());
}
/// Empty subtask slice → empty result.
#[test]
fn group_subtasks_empty_slice_returns_empty() {
let groups = group_subtasks_by_screen(&[]);
assert!(groups.is_empty());
}
/// Group order is first-seen order of distinct screen values.
#[test]
fn group_subtasks_preserves_first_seen_order() {
let subtasks = vec![
subtask_with_screen("a", Some("profile")),
subtask_with_screen("b", Some("settings")),
subtask_with_screen("c", Some("profile")),
subtask_with_screen("d", Some("dashboard")),
];
let groups = group_subtasks_by_screen(&subtasks);
assert_eq!(groups.len(), 3);
assert_eq!(groups[0].screen, "profile");
assert_eq!(groups[1].screen, "settings");
assert_eq!(groups[2].screen, "dashboard");
}
/// A None screen at the START falls back to first_screen from a LATER subtask.
#[test]
fn group_subtasks_no_screen_at_start_uses_later_first_screen() {
let subtasks = vec![
subtask_with_screen("a", None), // no screen
subtask_with_screen("b", Some("home")), // first with screen → first_screen="home"
subtask_with_screen("c", Some("settings")),
];
let groups = group_subtasks_by_screen(&subtasks);
assert_eq!(groups.len(), 2);
assert_eq!(groups[0].screen, "home");
assert!(groups[0].indices.contains(&0));
assert!(groups[0].indices.contains(&1));
assert_eq!(groups[1].screen, "settings");
assert_eq!(groups[1].indices, vec![2]);
}
// ── Task B1: BufferDocSink ────────────────────────────────────────────────
/// `BufferDocSink` collects commands without modifying a real doc.
#[test]
fn buffer_doc_sink_collects_commands() {
let mut sink = BufferDocSink::new(EditorState::new());
let applied = sink.apply(EditorCommand::InsertSubtree {
nodes: vec![],
parent_id: op_editor_core::NodeId::NONE,
});
assert!(applied, "BufferDocSink.apply must always return true");
assert_eq!(sink.commands.len(), 1);
}
/// `state()` on `BufferDocSink` returns the snapshot passed at construction.
#[test]
fn buffer_doc_sink_state_returns_snapshot() {
let state = EditorState::new();
let sink = BufferDocSink::new(state.clone());
let _ = sink.state();
}
/// `BufferDocSink` tracks undo-batch depth correctly.
#[test]
fn buffer_doc_sink_undo_batch_depth() {
let mut sink = BufferDocSink::new(EditorState::new());
assert_eq!(sink.batch_depth, 0);
sink.begin_undo_batch();
assert_eq!(sink.batch_depth, 1);
sink.end_undo_batch();
assert_eq!(sink.batch_depth, 0);
}
// ── Task B1/B2: run_screen_group_worker ──────────────────────────────────
/// Happy path: both subtasks succeed on attempt 1.
#[test]
fn worker_happy_path_both_succeed() {
let plan = make_plan_with_subtasks(&["s0", "s1"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![
ScriptResponse::Text(NODE_JSON.into()), // s0 attempt 1
ScriptResponse::Text(NODE_JSON.into()), // s1 attempt 1
]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0, 1],
};
let abort = AbortFlag::new();
let buffer = BufferDocSink::new(EditorState::new());
let (sem, tx, rx) = make_worker_channel();
let result = block_on(run_screen_group_worker(
&group, &plan, &req, &llm, &abort, buffer, sem, tx,
));
let events = drain_events(rx);
assert_eq!(result.outcomes.len(), 2);
assert_eq!(result.outcomes[0].1.node_count, 1);
assert_eq!(result.outcomes[1].1.node_count, 1);
// Buffer collected 2 InsertSubtree commands (one per subtask).
assert_eq!(result.buffer.commands.len(), 2);
// Progress events: started + done for each subtask.
assert_eq!(events.len(), 4);
}
/// Attempt-1 zero-node + retryable → attempt 2 runs with (true, true).
#[test]
fn worker_attempt1_zero_retryable_triggers_attempt2() {
let plan = make_plan_with_subtasks(&["s0"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![
ScriptResponse::Fail(LlmError {
message: "timeout".into(),
aborted: false,
}),
ScriptResponse::Text(NODE_JSON.into()),
]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0],
};
let abort = AbortFlag::new();
let buffer = BufferDocSink::new(EditorState::new());
let (sem, tx, _rx) = make_worker_channel();
let result = block_on(run_screen_group_worker(
&group, &plan, &req, &llm, &abort, buffer, sem, tx,
));
assert_eq!(result.outcomes.len(), 1);
assert_eq!(
result.outcomes[0].1.node_count, 1,
"attempt 2 should have produced a node"
);
assert_eq!(
result.buffer.commands.len(),
1,
"one InsertSubtree from attempt 2"
);
}
/// Non-retryable error on attempt 1 → NO attempt 2.
#[test]
fn worker_non_retryable_error_no_retry() {
let plan = make_plan_with_subtasks(&["s0"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![ScriptResponse::Fail(LlmError {
message: "HTTP 401 Unauthorized".into(),
aborted: false,
})]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0],
};
let abort = AbortFlag::new();
let buffer = BufferDocSink::new(EditorState::new());
let (sem, tx, _rx) = make_worker_channel();
let result = block_on(run_screen_group_worker(
&group, &plan, &req, &llm, &abort, buffer, sem, tx,
));
assert_eq!(result.outcomes.len(), 1);
assert_eq!(
result.outcomes[0].1.node_count, 0,
"non-retryable must stay zero"
);
assert!(
result.outcomes[0]
.1
.error
.as_deref()
.unwrap_or("")
.contains("401"),
"error should carry original 401 message"
);
assert_eq!(result.buffer.commands.len(), 0);
}
/// Partial result (node_count > 0) on attempt 1 → never retried.
#[test]
fn worker_partial_result_never_retried() {
let plan = make_plan_with_subtasks(&["s0"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![ScriptResponse::Text(NODE_JSON.into())]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0],
};
let abort = AbortFlag::new();
let buffer = BufferDocSink::new(EditorState::new());
let (sem, tx, _rx) = make_worker_channel();
let result = block_on(run_screen_group_worker(
&group, &plan, &req, &llm, &abort, buffer, sem, tx,
));
assert_eq!(result.outcomes.len(), 1);
assert_eq!(result.outcomes[0].1.node_count, 1);
assert_eq!(result.buffer.commands.len(), 1);
}
/// Abort flag set before the loop → worker returns immediately with no outcomes.
#[test]
fn worker_aborts_at_loop_head() {
let plan = make_plan_with_subtasks(&["s0", "s1"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0, 1],
};
let abort = AbortFlag::new();
abort.set();
let buffer = BufferDocSink::new(EditorState::new());
let (sem, tx, _rx) = make_worker_channel();
let result = block_on(run_screen_group_worker(
&group, &plan, &req, &llm, &abort, buffer, sem, tx,
));
assert!(
result.outcomes.is_empty(),
"aborted before loop: no outcomes expected"
);
assert_eq!(result.buffer.commands.len(), 0);
}
/// Progress events: SubtaskStarted emitted before each subtask;
/// SubtaskDone emitted for successful subtasks.
#[test]
fn worker_emits_started_and_done_progress() {
let plan = make_plan_with_subtasks(&["s0", "s1"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![
ScriptResponse::Text(NODE_JSON.into()),
ScriptResponse::Text(NODE_JSON.into()),
]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0, 1],
};
let abort = AbortFlag::new();
let buffer = BufferDocSink::new(EditorState::new());
let (sem, tx, rx) = make_worker_channel();
block_on(run_screen_group_worker(
&group, &plan, &req, &llm, &abort, buffer, sem, tx,
));
let events: Vec<String> = drain_events(rx)
.into_iter()
.filter_map(|p| match &p {
Progress::SubtaskStarted { id, .. } => Some(format!("start:{id}")),
Progress::SubtaskDone { id, .. } => Some(format!("done:{id}")),
Progress::SubtaskFailed { id, .. } => Some(format!("fail:{id}")),
_ => None,
})
.collect();
assert_eq!(events, vec!["start:s0", "done:s0", "start:s1", "done:s1"]);
}
/// SubtaskFailed progress emitted when final outcome is zero nodes.
#[test]
fn worker_emits_failed_progress_on_zero_nodes() {
let plan = make_plan_with_subtasks(&["s0"]);
let req = make_req();
let llm = ScriptedLlm::new(vec![
ScriptResponse::Text("not json".into()),
ScriptResponse::Text("also not json".into()),
]);
let group = ScreenGroup {
screen: "home".into(),
indices: vec![0],
};
let abort = AbortFlag::new();
let buffer = BufferDocSink::new(EditorState::new());
let (sem, tx, rx) = make_worker_channel();
block_on(run_screen_group_worker(
&group, &plan, &req, &llm, &abort, buffer, sem, tx,
));
let fail_events: Vec<String> = drain_events(rx)
.into_iter()
.filter_map(|p| {
if let Progress::SubtaskFailed { id, .. } = &p {
Some(format!("fail:{id}"))
} else {
None
}
})
.collect();
assert_eq!(fail_events, vec!["fail:s0"]);
}

428
crates/op-orchestrator/src/concurrent_tests_b2.rs Normal file
View file

@ -0,0 +1,428 @@
//! `concurrent` Task B2 tests — `run_concurrent` (concurrent executor).
//!
//! Split from `concurrent_tests.rs` to keep both test files under the
//! 800-line cap. Wired as a second `#[path = ...] mod tests_b2;` inside
//! `concurrent.rs`; like `concurrent_tests.rs` this stays a child module of
//! `concurrent`, so `use super::*` resolves to `concurrent`.
use super::*;
use crate::test_support::{CountingLlm, ScriptResponse, ScriptedLlm};
use futures::executor::block_on;
use op_editor_core::{EditorState, PenNodeExt};
const NODE_JSON: &str =
r#"[{"type":"frame","id":"h1","name":"H","x":0,"y":0,"width":100,"height":100,"children":[]}]"#;
/// A default `DesignRequest` (concurrency overridden per-test as needed).
fn make_req() -> crate::types::DesignRequest {
crate::types::DesignRequest {
prompt: "test".into(),
model: None,
provider: None,
design_md: None,
concurrency: 2,
}
}
/// Build a multi-screen plan: s0=login (plan index 0), s1=home (plan index 1).
fn make_two_screen_plan() -> crate::plan::OrchestratorPlan {
use crate::plan::{OrchestratorPlan, Region, RootFrameSpec};
OrchestratorPlan {
root_frame: RootFrameSpec {
id: "root".into(),
name: "Page".into(),
width: 1200.0,
height: 800.0,
layout: None,
gap: None,
padding: None,
fill: None,
},
subtasks: vec![
crate::plan::Subtask {
id: "s0".into(),
label: "Login".into(),
region: Region {
width: 1200.0,
height: 400.0,
},
id_prefix: "s0".into(),
parent_frame_id: Some("root".into()),
elements: None,
screen: Some("login".into()),
},
crate::plan::Subtask {
id: "s1".into(),
label: "Home".into(),
region: Region {
width: 1200.0,
height: 400.0,
},
id_prefix: "s1".into(),
parent_frame_id: Some("root".into()),
elements: None,
screen: Some("home".into()),
},
],
style_guide_name: None,
}
}
/// Build a 3-screen plan (s0=a, s1=b, s2=c) — used for the genuine-interleaving
/// test where all 3 groups should overlap.
fn make_three_screen_plan() -> crate::plan::OrchestratorPlan {
use crate::plan::{OrchestratorPlan, Region, RootFrameSpec};
let mk = |id: &str, screen: &str| crate::plan::Subtask {
id: id.into(),
label: id.into(),
region: Region {
width: 1200.0,
height: 400.0,
},
id_prefix: id.into(),
parent_frame_id: Some("root".into()),
elements: None,
screen: Some(screen.into()),
};
OrchestratorPlan {
root_frame: RootFrameSpec {
id: "root".into(),
name: "Page".into(),
width: 1200.0,
height: 800.0,
layout: None,
gap: None,
padding: None,
fill: None,
},
subtasks: vec![mk("s0", "a"), mk("s1", "b"), mk("s2", "c")],
style_guide_name: None,
}
}
/// Happy path: 2 screen groups, both succeed.
/// After `run_concurrent`, the real sink receives 2 commands (one per subtask)
/// in plan order, and all 4 progress events arrive.
#[test]
fn concurrent_two_groups_both_succeed_replays_in_plan_order() {
let plan = make_two_screen_plan();
let req = make_req();
let groups = group_subtasks_by_screen(&plan.subtasks);
assert_eq!(groups.len(), 2);
let llm = ScriptedLlm::new(vec![
ScriptResponse::Text(NODE_JSON.into()), // group[0] s0
ScriptResponse::Text(NODE_JSON.into()), // group[1] s1
]);
let abort = AbortFlag::new();
let snapshot = EditorState::new();
let mut real_sink = BufferDocSink::new(snapshot.clone());
let mut progress_events: Vec<String> = Vec::new();
let outcomes = block_on(run_concurrent(
&groups,
&plan,
&req,
&llm,
&abort,
snapshot,
&mut real_sink,
&mut |p| match &p {
Progress::SubtaskStarted { id, .. } => progress_events.push(format!("start:{id}")),
Progress::SubtaskDone { id, .. } => progress_events.push(format!("done:{id}")),
Progress::SubtaskFailed { id, .. } => progress_events.push(format!("fail:{id}")),
_ => {}
},
));
// Both subtasks produced 1 node each.
assert_eq!(outcomes.len(), 2);
assert_eq!(outcomes[0].as_ref().map(|o| o.node_count), Some(1));
assert_eq!(outcomes[1].as_ref().map(|o| o.node_count), Some(1));
// Real sink got 2 commands (one per subtask) — replayed in plan order.
assert_eq!(real_sink.commands.len(), 2);
// All 4 progress events arrived.
assert!(
progress_events.contains(&"start:s0".to_string()),
"missing start:s0, events: {progress_events:?}"
);
assert!(
progress_events.contains(&"done:s0".to_string()),
"missing done:s0"
);
assert!(
progress_events.contains(&"start:s1".to_string()),
"missing start:s1"
);
assert!(
progress_events.contains(&"done:s1".to_string()),
"missing done:s1"
);
}
/// **Genuine interleaving** — with `concurrency=3` and 3 screen groups, the
/// 3 workers' LLM calls overlap in wall-clock time. `CountingLlm`'s stream
/// yields once (Poll::Pending) before producing its chunk, so under `join_all`
/// worker A goes pending → worker B's `call` runs → worker C's `call` runs →
/// all 3 are in-flight simultaneously. `max_concurrent` must reach 3.
///
/// This test would FAIL with a sequential `for worker.await` loop: a sequential
/// loop awaits worker A to FULL completion before starting B, so `current`
/// never exceeds 1 — `max_concurrent` would be 1, not 3. (Verified by
/// temporarily reverting `run_concurrent` to a sequential loop — this test then
/// fails with `left: 1, right: 3`.)
#[test]
fn concurrent_workers_genuinely_interleave() {
let plan = make_three_screen_plan();
let mut req = make_req();
req.concurrency = 3; // 3 groups → effective_concurrency = 3 (no cap pressure)
let groups = group_subtasks_by_screen(&plan.subtasks);
assert_eq!(groups.len(), 3);
let llm = CountingLlm::new(vec![NODE_JSON.into(), NODE_JSON.into(), NODE_JSON.into()]);
let abort = AbortFlag::new();
let snapshot = EditorState::new();
let mut real_sink = BufferDocSink::new(snapshot.clone());
let _outcomes = block_on(run_concurrent(
&groups,
&plan,
&req,
&llm,
&abort,
snapshot,
&mut real_sink,
&mut |_| {},
));
// All 3 workers' LLM calls genuinely overlapped.
assert_eq!(
llm.max_concurrent(),
3,
"with concurrency=3 + 3 groups, all 3 workers must be in-flight at once \
(a sequential loop would yield 1)"
);
}
/// Semaphore cap: with `concurrency=1` and 2 groups, the semaphore caps
/// in-flight LLM calls at 1 — even though `join_all` drives both workers
/// concurrently, the single permit serializes their LLM calls.
/// `max_concurrent` must stay ≤ 1.
#[test]
fn concurrent_semaphore_caps_in_flight_calls() {
let plan = make_two_screen_plan();
let mut req = make_req();
req.concurrency = 1; // 2 groups → effective_concurrency = 1 (hard cap)
let groups = group_subtasks_by_screen(&plan.subtasks);
let llm = CountingLlm::new(vec![NODE_JSON.into(), NODE_JSON.into()]);
let abort = AbortFlag::new();
let snapshot = EditorState::new();
let mut real_sink = BufferDocSink::new(snapshot.clone());
let _outcomes = block_on(run_concurrent(
&groups,
&plan,
&req,
&llm,
&abort,
snapshot,
&mut real_sink,
&mut |_| {},
));
// The cap-1 semaphore must serialize the LLM calls despite join_all driving
// both workers — at most 1 call in-flight at a time.
assert_eq!(
llm.max_concurrent(),
1,
"cap-1 semaphore must keep in-flight ≤ 1, got {}",
llm.max_concurrent()
);
}
/// Semaphore cap honored under pressure: with `concurrency=2` and 3 groups,
/// the LLM calls overlap (peak > 1) but never exceed the cap of 2.
#[test]
fn concurrent_semaphore_cap_two_with_three_groups() {
let plan = make_three_screen_plan();
let mut req = make_req();
req.concurrency = 2; // 3 groups → effective_concurrency = 2 (cap < group count)
let groups = group_subtasks_by_screen(&plan.subtasks);
let llm = CountingLlm::new(vec![NODE_JSON.into(), NODE_JSON.into(), NODE_JSON.into()]);
let abort = AbortFlag::new();
let snapshot = EditorState::new();
let mut real_sink = BufferDocSink::new(snapshot.clone());
let _outcomes = block_on(run_concurrent(
&groups,
&plan,
&req,
&llm,
&abort,
snapshot,
&mut real_sink,
&mut |_| {},
));
let peak = llm.max_concurrent();
assert!(
peak >= 2,
"concurrency=2 should let 2 workers overlap, got peak {peak}"
);
assert!(
peak <= 2,
"concurrency=2 cap must not be exceeded, got peak {peak}"
);
}
/// Replay order: real sink receives commands in plan-index order.
/// We use distinct JSON payloads so we can identify which subtask each command
/// came from by looking at the first node's id in the InsertSubtree.
#[test]
fn concurrent_replay_is_in_plan_index_order() {
let plan = make_two_screen_plan();
let req = make_req();
let groups = group_subtasks_by_screen(&plan.subtasks);
let s0_json = r#"[{"type":"frame","id":"s0-node","name":"S0","x":0,"y":0,"width":10,"height":10,"children":[]}]"#;
let s1_json = r#"[{"type":"frame","id":"s1-node","name":"S1","x":0,"y":0,"width":20,"height":20,"children":[]}]"#;
let llm = ScriptedLlm::new(vec![
ScriptResponse::Text(s0_json.into()),
ScriptResponse::Text(s1_json.into()),
]);
let abort = AbortFlag::new();
let snapshot = EditorState::new();
let mut real_sink = BufferDocSink::new(snapshot.clone());
let _outcomes = block_on(run_concurrent(
&groups,
&plan,
&req,
&llm,
&abort,
snapshot,
&mut real_sink,
&mut |_| {},
));
assert_eq!(real_sink.commands.len(), 2);
// First command's nodes come from s0 (plan index 0).
match &real_sink.commands[0] {
EditorCommand::InsertSubtree { nodes, .. } => {
if !nodes.is_empty() {
let first_id = nodes[0].id_str();
assert!(
first_id.starts_with("s0"),
"first command should be from s0, got id={first_id}"
);
}
}
other => panic!("expected InsertSubtree, got {other:?}"),
}
// Second command's nodes come from s1 (plan index 1).
match &real_sink.commands[1] {
EditorCommand::InsertSubtree { nodes, .. } => {
if !nodes.is_empty() {
let first_id = nodes[0].id_str();
assert!(
first_id.starts_with("s1"),
"second command should be from s1, got id={first_id}"
);
}
}
other => panic!("expected InsertSubtree, got {other:?}"),
}
}
/// Abort flag set before run_concurrent → all outcomes are None, no LLM calls.
#[test]
fn concurrent_aborts_before_any_work() {
let plan = make_two_screen_plan();
let req = make_req();
let groups = group_subtasks_by_screen(&plan.subtasks);
let llm = ScriptedLlm::new(vec![]);
let abort = AbortFlag::new();
abort.set();
let snapshot = EditorState::new();
let mut real_sink = BufferDocSink::new(snapshot.clone());
let mut started_ids: Vec<String> = Vec::new();
let outcomes = block_on(run_concurrent(
&groups,
&plan,
&req,
&llm,
&abort,
snapshot,
&mut real_sink,
&mut |p| {
if let Progress::SubtaskStarted { id, .. } = &p {
started_ids.push(id.clone());
}
},
));
// All outcomes are None (aborted before any subtask ran).
assert!(
outcomes.iter().all(|o| o.is_none()),
"aborted run should produce only None outcomes"
);
// No commands replayed.
assert_eq!(real_sink.commands.len(), 0);
}
/// Progress fan-in: all 4 events (2 started + 2 done) arrive after run_concurrent.
#[test]
fn concurrent_progress_fanin_all_events_arrive() {
let plan = make_two_screen_plan();
let req = make_req();
let groups = group_subtasks_by_screen(&plan.subtasks);
let llm = ScriptedLlm::new(vec![
ScriptResponse::Text(NODE_JSON.into()),
ScriptResponse::Text(NODE_JSON.into()),
]);
let abort = AbortFlag::new();
let snapshot = EditorState::new();
let mut real_sink = BufferDocSink::new(snapshot.clone());
let mut events: Vec<String> = Vec::new();
block_on(run_concurrent(
&groups,
&plan,
&req,
&llm,
&abort,
snapshot,
&mut real_sink,
&mut |p| match &p {
Progress::SubtaskStarted { id, .. } => events.push(format!("start:{id}")),
Progress::SubtaskDone { id, .. } => events.push(format!("done:{id}")),
_ => {}
},
));
assert_eq!(
events.len(),
4,
"expected 4 progress events, got {events:?}"
);
assert!(events.contains(&"start:s0".to_string()));
assert!(events.contains(&"done:s0".to_string()));
assert!(events.contains(&"start:s1".to_string()));
assert!(events.contains(&"done:s1".to_string()));
}

133
crates/op-orchestrator/src/test_support.rs
View file

@ -2,9 +2,13 @@
use crate::types::{CallRequest, DocSink, LlmChunk, LlmClient, LlmError};
use futures::stream::BoxStream;
use futures::Stream;
use op_editor_core::{EditorCommand, EditorState};
use std::collections::VecDeque;
use std::sync::Mutex;
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
/// 录命令 + 持内存 `EditorState` 的 `DocSink` 实现。
pub(crate) struct VecDocSink {
@ -74,3 +78,130 @@ impl LlmClient for ScriptedLlm {
Box::pin(futures::stream::iter(items))
}
}
/// Shared in-flight counters for [`CountingLlm`].
struct CallCounters {
/// Current number of in-flight calls (a call is "in-flight" from the
/// moment `call` returns until its stream yields the text chunk).
current: AtomicUsize,
/// Maximum observed concurrent in-flight count.
max_concurrent: AtomicUsize,
}
impl CallCounters {
/// Record one new in-flight call: bump `current` and update `max_concurrent`.
fn enter(&self) {
let new_val = self.current.fetch_add(1, Ordering::SeqCst) + 1;
let mut cur_max = self.max_concurrent.load(Ordering::SeqCst);
while new_val > cur_max {
match self.max_concurrent.compare_exchange(
cur_max,
new_val,
Ordering::SeqCst,
Ordering::SeqCst,
) {
Ok(_) => break,
Err(actual) => cur_max = actual,
}
}
}
/// Record one call completing: decrement `current`.
fn leave(&self) {
self.current.fetch_sub(1, Ordering::SeqCst);
}
}
/// A `Stream` that **yields `Poll::Pending` once** before producing its single
/// text chunk. The initial pending poll is what lets `futures::future::join_all`
/// poll a *different* worker — so multiple `CountingLlm` calls genuinely overlap
/// (their `current` counter is > 1 simultaneously) instead of running strictly
/// one-at-a-time.
///
/// Lifecycle vs the counter:
/// - The call's `enter()` is recorded by `CountingLlm::call` BEFORE the stream
/// is built, so the call counts as in-flight immediately.
/// - The stream's first poll registers a self-wake and returns `Pending` — the
/// executor is now free to poll another worker, raising `current`.
/// - The stream's second poll calls `leave()` (the call completes) and returns
/// the text chunk.
/// - The third poll returns `None` (stream end).
struct YieldingChunkStream {
yielded_once: bool,
/// `Some` until the text chunk has been emitted.
text: Option<String>,
counters: Arc<CallCounters>,
}
impl Stream for YieldingChunkStream {
type Item = Result<LlmChunk, LlmError>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
if !this.yielded_once {
// First poll: yield control so join_all can poll another worker.
this.yielded_once = true;
cx.waker().wake_by_ref();
return Poll::Pending;
}
match this.text.take() {
Some(t) => {
// Second poll: the call completes — decrement the in-flight count.
this.counters.leave();
Poll::Ready(Some(Ok(LlmChunk::Text(t))))
}
None => Poll::Ready(None),
}
}
}
/// An `LlmClient` that tracks the maximum number of concurrent in-flight calls.
///
/// Used in Task B2 tests to verify the semaphore genuinely caps concurrency.
/// Each scripted response is a `String` (node JSON text — always succeeds).
///
/// Crucially the returned stream **yields once** before producing its chunk
/// (see [`YieldingChunkStream`]), so under `join_all` multiple workers' calls
/// genuinely overlap — `max_concurrent` reflects real wall-clock overlap, not a
/// vacuous always-1.
pub(crate) struct CountingLlm {
responses: Mutex<VecDeque<String>>,
counters: Arc<CallCounters>,
}
impl CountingLlm {
pub(crate) fn new(responses: Vec<String>) -> Self {
Self {
responses: Mutex::new(responses.into()),
counters: Arc::new(CallCounters {
current: AtomicUsize::new(0),
max_concurrent: AtomicUsize::new(0),
}),
}
}
/// Returns the maximum number of concurrent in-flight calls observed.
pub(crate) fn max_concurrent(&self) -> usize {
self.counters.max_concurrent.load(Ordering::SeqCst)
}
}
impl LlmClient for CountingLlm {
fn call(&self, _req: CallRequest) -> BoxStream<'static, Result<LlmChunk, LlmError>> {
// Record the call as in-flight immediately (before the stream is polled).
self.counters.enter();
let text = self
.responses
.lock()
.unwrap()
.pop_front()
.unwrap_or_default();
Box::pin(YieldingChunkStream {
yielded_once: false,
text: Some(text),
counters: Arc::clone(&self.counters),
})
}
}