diff --git a/crates/tui/src/core/engine.rs b/crates/tui/src/core/engine.rs index 0505944615..8de15bd0f4 100644 --- a/crates/tui/src/core/engine.rs +++ b/crates/tui/src/core/engine.rs @@ -9079,6 +9079,7 @@ use context::{ use context::{context_input_budget_for_provider, effective_max_output_tokens}; mod dispatch; mod lsp_hooks; +mod partial_input; pub(crate) mod reviewer; mod streaming; mod token_estimate_cache; diff --git a/crates/tui/src/core/engine/partial_input.rs b/crates/tui/src/core/engine/partial_input.rs new file mode 100644 index 0000000000..8465a37451 --- /dev/null +++ b/crates/tui/src/core/engine/partial_input.rs @@ -0,0 +1,150 @@ +//! Suffix-only proof for the existing per-delta tool-input parser. +//! +//! Replaces repeated full-prefix evaluation only when an unfinished string +//! cannot be repaired and its selected raw JSON segment has not changed. +//! All other input retains dispatch::parse_tool_input semantics. +//! +//! Strict parsing and repair cannot close a valid unfinished string. In that +//! case only the raw balanced-segment fallback can produce a value. Reuse the +//! prior Some/None disposition only while that segment's byte range is stable. +//! This state is owned by one append-only tool block, never by a tool ID. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Segment { + Root(usize), + Array(usize, usize), +} +#[derive(Debug, Clone, Default)] +pub(super) struct PartialInputProof { + bytes: usize, + first_seen: bool, + starts_object: bool, + disabled: bool, + in_string: bool, + escape: bool, + unicode: u8, + root_start: Option, + root_depth: i32, + root_end: Option, + array_start: Option, + array_depth: i32, + array_end: Option, + certified: bool, + previous: Option, + last_some: bool, +} +impl PartialInputProof { + /// The owner must append exactly `delta` and invalidate before any other + /// buffer or mirror mutation. Length/suffix checks reject stale hints, but + /// cannot certify ownership after an equal-length prefix replacement. + pub(super) fn reuse(&mut self, buffer: &str, delta: &str) -> Option { + if self.disabled { + return None; + } + if buffer.len().checked_sub(delta.len()) != Some(self.bytes) || !buffer.ends_with(delta) { + self.invalidate(); + return None; + } + self.feed(delta).then_some(self.last_some) + } + pub(super) fn record(&mut self, some: bool) { + self.last_some = some; + } + + pub(super) fn invalidate(&mut self) { + self.disabled = true; + self.certified = false; + } + fn feed(&mut self, text: &str) -> bool { + for (offset, b) in text.bytes().enumerate() { + let pos = self.bytes + offset; + if !self.first_seen && !b.is_ascii_whitespace() { + self.first_seen = true; + self.starts_object = b == b'{'; + if !self.starts_object { + // Unsupported top-level forms always take the parser path; + // avoid scanning their later suffixes for a certificate. + self.disabled = true; + return false; + } + } + // Match extract_balanced_segment's raw (not quote-aware) selection. + if self.root_start.is_none() && b == b'{' { + self.root_start = Some(pos); + self.root_depth = 1; + } else if self.root_start.is_some() && self.root_end.is_none() { + if b == b'{' { + self.root_depth += 1; + } else if b == b'}' { + self.root_depth -= 1; + } + if self.root_depth == 0 { + self.root_end = Some(pos + 1); + } + } + if self.array_start.is_none() && b == b'[' { + self.array_start = Some(pos); + self.array_depth = 1; + } else if self.array_start.is_some() && self.array_end.is_none() { + if b == b'[' { + self.array_depth += 1; + } else if b == b']' { + self.array_depth -= 1; + } + if self.array_depth == 0 { + self.array_end = Some(pos + 1); + } + } + // No fence stripping, invalid escapes or raw control removal may + // alter quote/escape parity under the certificate. + if b == b'`' { + self.disabled = true; + } + if self.in_string { + if self.unicode > 0 { + if !b.is_ascii_hexdigit() { + self.disabled = true; + } + self.unicode -= 1; + } else if self.escape { + self.escape = false; + match b { + b'u' => self.unicode = 4, + b'"' | b'\\' | b'/' | b'b' | b'f' | b'n' | b'r' | b't' => {} + _ => self.disabled = true, + } + } else { + match b { + b'"' => self.in_string = false, + b'\\' => self.escape = true, + 0..=31 => self.disabled = true, + _ => {} + } + } + } else { + match b { + b'"' => self.in_string = true, + b'\\' => self.disabled = true, + 0..=31 if !b.is_ascii_whitespace() => self.disabled = true, + _ => {} + } + } + } + self.bytes += text.len(); + let selected = if let Some(end) = self.root_end { + Some(Segment::Root(end)) + } else if let (Some(start), Some(end)) = (self.array_start, self.array_end) { + Some(Segment::Array(start, end)) + } else { + None + }; + let eligible = self.starts_object + && !self.disabled + && self.in_string + && !self.escape + && self.unicode == 0; + let reuse = eligible && self.certified && selected == self.previous; + self.certified = eligible; + self.previous = selected; + reuse + } +} diff --git a/crates/tui/src/core/engine/partial_input_differential_tests.rs b/crates/tui/src/core/engine/partial_input_differential_tests.rs new file mode 100644 index 0000000000..157cdf5c11 --- /dev/null +++ b/crates/tui/src/core/engine/partial_input_differential_tests.rs @@ -0,0 +1,256 @@ +//! Compare the production mirror update with the authoritative parser at every +//! delivered prefix, including reads between deltas and terminal mutations. +use super::*; + +fn chunks(text: &str, width: usize) -> Vec<&str> { + let mut parts = Vec::new(); + let mut start = 0; + while start < text.len() { + let mut end = (start + width).min(text.len()); + while !text.is_char_boundary(end) { + end += 1; + } + parts.push(&text[start..end]); + start = end; + } + parts +} + +fn assert_mirror(baseline: &ToolUseState, candidate: &ToolUseState) { + assert_eq!(candidate.input_buffer, baseline.input_buffer); + assert_eq!( + candidate.input, baseline.input, + "prefix: {:?}", + baseline.input_buffer + ); + assert_eq!(candidate.input_parse_error, baseline.input_parse_error); +} + +fn compare_delta(baseline: &mut ToolUseState, candidate: &mut ToolUseState, delta: &str) { + baseline.input_buffer.push_str(delta); + let expected_some = if let Some(value) = parse_tool_input(&baseline.input_buffer) { + baseline.input = value; + true + } else { + false + }; + candidate.input_buffer.push_str(delta); + assert_eq!( + candidate.refresh_partial_input(delta), + expected_some, + "Some/None disposition at prefix {:?}", + baseline.input_buffer + ); + assert_mirror(baseline, candidate); +} + +fn compare_trace(text: &str, width: usize) { + let mut baseline = tool_state(json!({"initial": "retained until parse succeeds"}), ""); + let mut candidate = baseline.clone(); + for delta in chunks(text, width) { + compare_delta(&mut baseline, &mut candidate, delta); + } + assert_eq!(final_tool_input(&candidate), final_tool_input(&baseline)); +} + +#[test] +fn partial_input_differential_prefixes_preserve_parser_disposition_and_mirror() { + let fixtures = [ + r#"{"a":1,"b":"unfinished"#, + r#"{"s":"aaa[1,2]bbbbbbbbbbbb"}"#, + // Completing a raw root segment changes precedence over the array, + // even though the root is malformed and the old mirror remains Some. + r#"{"s":"[1] more text } trailing"#, + r#"{"s":"} extra [1]more"}"#, + "```json\n{\"a\":1}\n```", + r#""{\"a\":1}""#, + "{not json", + r#"{"a":1} trailing"x"#, + r#"{"a":"\u1234\n\\\"text"}"#, + r#"{"a":"\uZZZZmore"}"#, + r#"{"a":"\q invalid escape more"}"#, + "{\"a\":\"raw\u{0001}control\"}", + "{\"a\":\"raw\nline\tcontrol\"}", + r#"{"a":"foo`bar"}"#, + // Fence stripping after a closed raw root can change quote parity. + "{\"a\": \"}\"}\n```z\" tail", + // Quote-unaware closer removal can merge escape runs. The invalid + // escape must disable reuse before that repair can change the value. + r#"{"k": "\]\" tail"#, + r#"{"δΈ­ζ–‡πŸ™‚":"textπŸ™‚","nested":[{"a":true}],}"#, + r#"[1,"text"]"#, + " ", + "", + ]; + for fixture in fixtures { + for width in [1, 2, 7, 64] { + compare_trace(fixture, width); + } + } + + // Fixed-seed generated inputs cover both serialized values and malformed + // quote/escape/segment combinations without requiring a fuzz dependency. + let mut seed = 37_u64; + let alphabet = b"{}[]\"\\`,:abcdef0123 \n\t"; + for case in 0..64 { + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + let content = [ + "δΈ­ζ–‡πŸ™‚ plain", + "\\\"\n[]{}", + "[1,2] after", + "} then [3]", + "`fence", + ][seed as usize % 5] + .repeat(4); + compare_trace( + &json!({"content": content, "number": case, "nested": [{"a":true}]}).to_string(), + 1 + case % 31, + ); + let mut malformed = String::from("{"); + for _ in 0..64 { + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + malformed.push(alphabet[(seed >> 24) as usize % alphabet.len()] as char); + } + compare_trace(&malformed, 1 + case % 13); + } + // Exercise repair-size boundaries with coarse deltas to keep this focused + // regression cheap; the small fixtures above cover every UTF-8 prefix. + let overhead = json!({"content":"", "meta":{"a":[1,2]}}).to_string().len(); + for size in [1024 * 1024 - 1, 1024 * 1024, 1024 * 1024 + 1] { + let text = json!({"content": "x".repeat(size - overhead), "meta": {"a":[1,2]}}).to_string(); + assert_eq!(text.len(), size); + compare_trace(&text, 65536); + } +} + +#[test] +fn partial_input_differential_repair_parity_at_adversarial_delta_boundaries() { + for deltas in [ + vec!["{\"a\": \"}", "\"}\n```z\""], + vec![r#"{"k": ""#, "\\", "]", "\\\"", " tail"], + ] { + let mut baseline = tool_state(json!({}), ""); + let mut candidate = baseline.clone(); + for delta in deltas { + compare_delta(&mut baseline, &mut candidate, delta); + } + } +} + +#[test] +fn partial_input_proof_reuses_dispositions_and_rechecks_changed_segments() { + use super::partial_input::PartialInputProof; + + for (prefix, expected_some) in [(r#"{"content":""#, false), (r#"{"s":"[1]"#, true)] { + let mut proof = PartialInputProof::default(); + let mut buffer = prefix.to_owned(); + assert_eq!(proof.reuse(&buffer, prefix), None); + assert_eq!(parse_tool_input(&buffer).is_some(), expected_some); + proof.record(expected_some); + buffer.push_str(" more"); + assert_eq!(proof.reuse(&buffer, " more"), Some(expected_some)); + assert_eq!(parse_tool_input(&buffer).is_some(), expected_some); + + // A newly completed raw root takes precedence over an extracted array. + buffer.push_str("}"); + assert_eq!(proof.reuse(&buffer, "}"), None); + let some = parse_tool_input(&buffer).is_some(); + proof.record(some); + buffer.push_str(" tail"); + assert_eq!(proof.reuse(&buffer, " tail"), Some(some)); + assert_eq!(parse_tool_input(&buffer).is_some(), some); + } +} + +#[test] +fn partial_input_proof_rejects_stale_buffers_and_explicit_invalidation() { + use super::partial_input::PartialInputProof; + + let prefix = r#"{"content":""#; + let mut certified = PartialInputProof::default(); + assert_eq!(certified.reuse(prefix, prefix), None); + certified.record(false); + assert_eq!( + certified.clone().reuse(&format!("{prefix}x"), "x"), + Some(false) + ); + + // Wrong previous length and wrong suffix are independent stale hints. + for (buffer, delta) in [(format!("{prefix}xx"), "x"), (format!("{prefix}x"), "y")] { + let mut proof = certified.clone(); + assert_eq!(proof.reuse(&buffer, delta), None); + } + let mut invalidated = certified; + invalidated.invalidate(); + assert_eq!(invalidated.reuse(&format!("{prefix}x"), "x"), None); +} + +#[test] +fn partial_input_differential_interleaving_continuation_and_new_attempt() { + let mut baseline = [ + tool_state(json!({}), ""), + tool_state(json!({"initial":2}), ""), + ]; + let mut candidate = baseline.clone(); + for (index, delta) in [ + (0, r#"{"content":""#), + (1, r#"{"s":"[1]"#), + (0, "δΈ­ζ–‡πŸ™‚ text"), + (1, " more"), + (0, " continuation"), + (1, "} selected segment changes"), + (0, r#"","meta":{"a":1}}"#), + ] { + compare_delta(&mut baseline[index], &mut candidate[index], delta); + for index in 0..2 { + assert_mirror(&baseline[index], &candidate[index]); + } + } + // Reused provider ID must acquire fresh state with a different initial + // value, rather than carrying the previous attempt's certificate/mirror. + baseline[0] = tool_state(json!({"retry":true}), ""); + candidate[0] = baseline[0].clone(); + for delta in [r#"{"retry":""#, "fresh", " value", r#""}"#] { + compare_delta(&mut baseline[0], &mut candidate[0], delta); + } +} + +#[test] +fn partial_input_differential_invalidation_preserves_terminal_mutations_and_late_delta() { + let mut baseline = tool_state(json!({}), ""); + let mut candidate = baseline.clone(); + for delta in [r#"{"s":"[1]"#, " unfinished"] { + compare_delta(&mut baseline, &mut candidate, delta); + } + // These writes model block-stop error tagging or schema normalization. + // They must survive when the next full evaluation returns None, and be + // replaced when the authoritative parser returns Some (the extracted array). + candidate.invalidate_input_parse_hint(); + baseline.input = json!({"terminal":"mutation"}); + baseline.input_parse_error = Some("terminal parse error".into()); + candidate.input = baseline.input.clone(); + candidate.input_parse_error = baseline.input_parse_error.clone(); + compare_delta(&mut baseline, &mut candidate, " late text"); + compare_delta(&mut baseline, &mut candidate, "} changes selected segment"); + + baseline = tool_state(json!({}), ""); + candidate = baseline.clone(); + for delta in [r#"{"content":""#, "unfinished"] { + compare_delta(&mut baseline, &mut candidate, delta); + } + candidate.invalidate_input_parse_hint(); + baseline.input = json!({"raw_arguments":baseline.input_buffer}); + baseline.input_parse_error = Some("malformed final input".into()); + candidate.input = baseline.input.clone(); + candidate.input_parse_error = baseline.input_parse_error.clone(); + compare_delta(&mut baseline, &mut candidate, " late unfinished text"); + assert!(candidate.input.get("raw_arguments").is_some()); + + // A snapshot/buffer replacement uses a fresh proof; initial-input-only + // finalization remains authoritative when no delta ever arrives. + baseline = tool_state(json!({"initial_only":true}), ""); + candidate = baseline.clone(); + assert_eq!(final_tool_input(&candidate), json!({"initial_only":true})); + compare_delta(&mut baseline, &mut candidate, r#"{"new":""#); + compare_delta(&mut baseline, &mut candidate, "content"); +} diff --git a/crates/tui/src/core/engine/streaming.rs b/crates/tui/src/core/engine/streaming.rs index 17a3b12514..cfdc5b4041 100644 --- a/crates/tui/src/core/engine/streaming.rs +++ b/crates/tui/src/core/engine/streaming.rs @@ -4,6 +4,7 @@ //! content block kind tracking, streamed tool-use buffers, transparent retry //! policy, and scrubbers for text that looks like a forged tool-call wrapper. +use super::{dispatch::parse_tool_input, partial_input::PartialInputProof}; use crate::models::ToolCaller; use std::time::Duration; @@ -25,6 +26,28 @@ pub(super) struct ToolUseState { pub(super) thought_signature: Option, pub(super) input_buffer: String, pub(super) input_parse_error: Option, + pub(super) input_parse_hint: PartialInputProof, +} + +impl ToolUseState { + /// Called after appending/logging the raw delta, preserving parsed-log order. + pub(super) fn refresh_partial_input(&mut self, delta: &str) -> bool { + if let Some(some) = self.input_parse_hint.reuse(&self.input_buffer, delta) { + return some; + } + let some = if let Some(value) = parse_tool_input(&self.input_buffer) { + self.input = value; + true + } else { + false + }; + self.input_parse_hint.record(some); + some + } + + pub(super) fn invalidate_input_parse_hint(&mut self) { + self.input_parse_hint.invalidate(); + } } /// Maximum total bytes of text/thinking content before aborting the stream. diff --git a/crates/tui/src/core/engine/tests.rs b/crates/tui/src/core/engine/tests.rs index dc866cbf5e..e7dda49cf3 100644 --- a/crates/tui/src/core/engine/tests.rs +++ b/crates/tui/src/core/engine/tests.rs @@ -22858,6 +22858,7 @@ fn tool_state(initial: serde_json::Value, buffer: &str) -> ToolUseState { thought_signature: None, input_buffer: buffer.into(), input_parse_error: None, + input_parse_hint: Default::default(), } } @@ -29264,6 +29265,7 @@ async fn hydration_receipt_audit_payload_carries_the_executed_marker() { thought_signature: None, input_buffer: String::new(), input_parse_error: None, + input_parse_hint: Default::default(), }]; let mut active = HashSet::new(); let mut budget = crate::tools::tool_call_budget::ToolCallBudget::new(surface.max_tool_calls); @@ -29583,3 +29585,6 @@ fn composer_bang_ask_rule_site_forwards_session_workspace_roots() { "the engine.rs composer-bang call site must forward the session's workspace_roots" ); } + +#[path = "partial_input_differential_tests.rs"] +mod partial_input_differential_tests; diff --git a/crates/tui/src/core/engine/turn_loop.rs b/crates/tui/src/core/engine/turn_loop.rs index 5919c766f7..9a1b3ab323 100644 --- a/crates/tui/src/core/engine/turn_loop.rs +++ b/crates/tui/src/core/engine/turn_loop.rs @@ -2734,11 +2734,13 @@ impl Engine { thought_signature: None, input_buffer: String::new(), input_parse_error: None, + input_parse_hint: Default::default(), }); } } for tool in &mut tool_uses { + tool.invalidate_input_parse_hint(); let Some(schema) = tool_catalog .iter() .find(|candidate| candidate.name == tool.name) @@ -5999,6 +6001,7 @@ impl Engine { thought_signature, input_buffer: String::new(), input_parse_error: None, + input_parse_hint: Default::default(), }); } ContentBlockStart::ServerToolUse { id, name, input } => { @@ -6021,6 +6024,7 @@ impl Engine { thought_signature: None, input_buffer: String::new(), input_parse_error: None, + input_parse_hint: Default::default(), }); } }, @@ -6109,8 +6113,7 @@ impl Engine { tool_state.name, partial_json, tool_state.input_buffer )); } - if let Some(value) = parse_tool_input(&tool_state.input_buffer) { - tool_state.input = value.clone(); + if tool_state.refresh_partial_input(&partial_json) { if self.active_turn_tool_security.is_some() { crate::logging::info( "Restricted tool input parsed (details redacted)", @@ -6118,7 +6121,7 @@ impl Engine { } else if crate::logging::is_verbose() { crate::logging::info(format!( "Tool '{}' input parsed: {:?}", - tool_state.name, value + tool_state.name, tool_state.input )); } } @@ -6174,6 +6177,7 @@ impl Engine { if let Some(tool_idx) = current_tool_indices.remove(&index) && let Some(tool_state) = tool_uses.get_mut(tool_idx) { + tool_state.invalidate_input_parse_hint(); if self.active_turn_tool_security.is_some() { crate::logging::info("Restricted tool block stop (details redacted)"); } else { diff --git a/crates/tui/tests/README.md b/crates/tui/tests/README.md index bb46d5f69d..7123cbf436 100644 --- a/crates/tui/tests/README.md +++ b/crates/tui/tests/README.md @@ -16,7 +16,7 @@ the `codewhale-tui` bin unit tests): | harness | binary | what lives there | why it stays separate | |---|---|---|---| -| `tests/integration/main.rs` | `integration` | 15 plain `#[test]`/`#[tokio::test]` suites: `adaptive_evidence_acceptance`, `diagnostic_read_only`, `dotenv_authority`, `eval_harness`, `exec_persistent_service`, `exec_stream_drop_acceptance`, `exec_turn_usage`, `integration_mock_llm`, `lifecycle_outbox_exec`, `palette_audit`, `protocol_recovery`, `skill_cli`, `telemetry_contract`, `verifiers_harness_contract`, `workflow_tool_stream_acceptance` | All are process-level but require no PTY or Gherkin runner; they share `wiremock`/`tempfile` and link the TUI once instead of 15 times. `crate::` for `eval`/`models`/`llm_client`/`palette`/`network_policy`/`config`/`install` is satisfied by `integration/main.rs` re-exporting those `#[path]` modules at the harness crate root so `crate::config` etc. resolve. | +| `tests/integration/main.rs` | `integration` | 16 plain `#[test]`/`#[tokio::test]` suites: `adaptive_evidence_acceptance`, `diagnostic_read_only`, `dotenv_authority`, `eval_harness`, `exec_persistent_service`, `exec_stream_drop_acceptance`, `exec_turn_usage`, `integration_mock_llm`, `lifecycle_outbox_exec`, `palette_audit`, `partial_input_engine`, `protocol_recovery`, `skill_cli`, `telemetry_contract`, `verifiers_harness_contract`, `workflow_tool_stream_acceptance` | All are process-level but require no PTY or Gherkin runner; they share `wiremock`/`tempfile` and link the TUI once instead of 16 times. `crate::` for `eval`/`models`/`llm_client`/`palette`/`network_policy`/`config`/`install` is satisfied by `integration/main.rs` re-exporting those `#[path]` modules at the harness crate root so `crate::config` etc. resolve. | | `tests/cucumber/main.rs` | `cucumber` | 5 Gherkin runners: `core_session_command_extraction`, `directory_listing_acceptance`, `epic_acceptance_harness`, `eval_smoke_acceptance`, `tool_lifecycle_acceptance` | Each defines a distinct `cucumber::World`; steps are registered per-World via inventory, so merging is safe and cuts the `cucumber` link jobs. `plugin_e2e_acceptance` is no longer a Gherkin runner: its production-mirror scanner World/steps were removed with the worthless-test sweep (the real `scan_plugin_dir` contract stays covered by unit tests in `src/tools/plugin.rs`), leaving the plain binary smoke + `#[cfg(all(unix, feature="long-running-tests"))]` PTY acceptance, dormant in the default run. | The former `tests/pty` harness was removed. It accumulated full-screen copy, color, timing, and geometry assertions that were expensive to link and made the @@ -33,6 +33,7 @@ Filtering still works via the module path: cargo test -p codewhale-tui --tests -- --list | grep adaptive_evidence cargo test -p codewhale-tui --test integration adaptive_evidence_acceptance -- --nocapture cargo test -p codewhale-tui --test cucumber tool_lifecycle -- --nocapture +cargo test -p codewhale-tui --test integration partial_input_engine -- --nocapture ``` The shared helpers in `crates/tui/tests/support/` (`qa_harness`, `llm_client`) and fixtures in `crates/tui/tests/fixtures/` are untouched β€” harnesses reach them via `../support`. diff --git a/crates/tui/tests/integration/main.rs b/crates/tui/tests/integration/main.rs index 395b440137..29086c6ab8 100644 --- a/crates/tui/tests/integration/main.rs +++ b/crates/tui/tests/integration/main.rs @@ -102,6 +102,7 @@ mod exec_turn_usage; mod integration_mock_llm; mod lifecycle_outbox_exec; mod palette_audit; +mod partial_input_engine; mod protocol_recovery; mod skill_cli; mod telemetry_contract; diff --git a/crates/tui/tests/integration/partial_input_engine.rs b/crates/tui/tests/integration/partial_input_engine.rs new file mode 100644 index 0000000000..8c9aab6fc6 --- /dev/null +++ b/crates/tui/tests/integration/partial_input_engine.rs @@ -0,0 +1,805 @@ +//! Public-boundary integration evidence: real Engine, synthetic local tools only. +use codewhale_tui::{ + AppMode, ApprovalMode, + config::Config, + core::{ + engine::{Engine, EngineConfig, EngineHandle}, + events::{Event, TurnOutcomeStatus}, + model_client::{ModelClient, SharedModelClient}, + ops::{Op, UserInputProvenance}, + }, + models::{ContentBlock, MessageRequest, MessageResponse, StreamEvent}, +}; +use serde_json::json; +use std::{ + collections::VecDeque, + fs, + path::Path, + sync::{Arc, Mutex}, + time::Duration, +}; +use tempfile::tempdir; + +type EventStream = + std::pin::Pin> + Send>>; + +fn model_turn_event_timeout() -> Duration { + Duration::from_secs(30) +} +fn deterministic_engine_config(workspace: &Path) -> EngineConfig { + EngineConfig { + workspace: workspace.to_path_buf(), + snapshots_enabled: false, + subagents_enabled: false, + skills_dir: workspace.join("skills"), + explicit_skills_root_only: true, + mcp_config_path: workspace.join("no-mcp.json"), + notes_path: workspace.join("notes.md"), + memory_path: workspace.join("memory.md"), + ..EngineConfig::default() + } +} +fn external_user_message_op(content: &str, mode: AppMode, config: &Config) -> Op { + Op::SendMessage { + content: content.into(), + mode, + route: Box::new( + codewhale_tui::route_runtime::resolve_runtime_route( + config, + config.api_provider(), + Some("local-model"), + ) + .expect("test route"), + ), + compaction: Default::default(), + goal_objective: None, + goal_token_budget: None, + goal_status: Default::default(), + reasoning_effort: None, + reasoning_effort_auto: false, + auto_model: false, + allow_shell: false, + trust_mode: false, + auto_approve: false, + approval_mode: ApprovalMode::Suggest, + translation_enabled: false, + allowed_tools: Some(vec!["File".into()]), + dynamic_tools: Vec::new(), + hook_executor: None, + verbosity: None, + provenance: UserInputProvenance::ExternalUser, + turn_tool_security: None, + submission_id: None, + } +} + +struct MockLlmClient { + turns: Mutex>>>, + requests: Mutex>, + initial_network_error: AtomicBool, +} +impl MockLlmClient { + fn new(turns: Vec>) -> Self { + Self { + turns: Mutex::new(turns.into_iter().map(Ok).collect()), + requests: Mutex::new(Vec::new()), + initial_network_error: AtomicBool::new(false), + } + } + fn call_count(&self) -> usize { + self.requests.lock().unwrap().len() + } + fn captured_requests(&self) -> Vec { + self.requests.lock().unwrap().clone() + } +} +#[async_trait::async_trait] +impl ModelClient for MockLlmClient { + fn provider_name(&self) -> &str { + "partial-tool-test" + } + fn model(&self) -> &str { + "local-model" + } + async fn create_message(&self, _: MessageRequest) -> anyhow::Result { + anyhow::bail!("streaming required") + } + async fn create_message_stream(&self, request: MessageRequest) -> anyhow::Result { + self.requests.lock().unwrap().push(request); + if self.initial_network_error.swap(false, Ordering::SeqCst) { + return Ok(Box::pin(futures_util::stream::iter([Err( + anyhow::anyhow!("Stream read error: error decoding response body"), + )]))); + } + let events = self + .turns + .lock() + .unwrap() + .pop_front() + .expect("unexpected extra model request")?; + Ok(Box::pin(futures_util::stream::iter( + events.into_iter().map(Ok), + ))) + } + async fn health_check(&self) -> anyhow::Result { + Ok(true) + } +} + +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +fn argument_chunks(text: &str) -> Vec<&str> { + let mut parts = Vec::new(); + let mut start = 0; + while start < text.len() { + let mut end = (start + 64).min(text.len()); + while !text.is_char_boundary(end) { + end += 1; + } + parts.push(&text[start..end]); + start = end; + } + parts +} + +async fn terminal_events(handle: &EngineHandle) -> Vec { + let mut events = Vec::new(); + loop { + let event = tokio::time::timeout(model_turn_event_timeout(), async { + handle.rx_event.write().await.recv().await + }) + .await + .expect("partial-input Engine timeout") + .expect("Engine event stream ended before terminal event"); + let done = matches!(event, Event::TurnComplete { .. }); + events.push(event); + if done { + return events; + } + } +} + +fn auto_write_op(config: &Config) -> Op { + let mut op = external_user_message_op( + "Exercise synthetic workspace tools.", + AppMode::Agent, + config, + ); + if let Op::SendMessage { auto_approve, .. } = &mut op { + *auto_approve = true; + } + op +} + +#[tokio::test] +async fn partial_input_engine_interleaved_tools_keep_complete_arguments_and_history() { + let workspace = tempdir().expect("workspace"); + fs::write(workspace.path().join("fixture.txt"), "read fixture").expect("fixture"); + let content = "fn sample() { println!(\"δΈ­ζ–‡πŸ™‚\"); }\n".repeat(384); + let write = json!({"action":"write","path":"written.txt","content":content}); + let read = json!({"action":"read","path":"fixture.txt"}); + let write_json = serde_json::to_string(&write).expect("serialize write"); + let read_json = serde_json::to_string(&read).expect("serialize read"); + let mut turn = vec![ + canned::message_start("partial-input-interleaved"), + canned::tool_use_block_start(0, "write-call", "File"), + canned::tool_use_block_start(1, "read-call", "File"), + ]; + let writes = argument_chunks(&write_json); + let reads = argument_chunks(&read_json); + for index in 0..writes.len().max(reads.len()) { + if let Some(part) = writes.get(index) { + turn.push(canned::tool_input_delta(0, part)); + } + if let Some(part) = reads.get(index) { + turn.push(canned::tool_input_delta(1, part)); + } + } + turn.extend([ + canned::block_stop(1), + canned::block_stop(0), + canned::message_delta("tool_use", None), + canned::message_stop(), + ]); + let mock = Arc::new(MockLlmClient::new(vec![ + turn, + canned::simple_text_turn("done"), + ])); + let client: SharedModelClient = mock.clone(); + let config = Config::default(); + let (engine, handle) = Engine::new_with_model_client( + deterministic_engine_config(workspace.path()), + &config, + client, + ); + let task = tokio::spawn(engine.run()); + handle.send(auto_write_op(&config)).await.expect("send"); + let events = terminal_events(&handle).await; + let mut starts = std::collections::HashMap::new(); + let mut successes = 0; + for event in &events { + match event { + Event::ToolCallStarted { id, input, .. } => { + assert!( + starts.insert(id.as_str(), input).is_none(), + "duplicate tool start" + ); + } + Event::ToolCallComplete { result, .. } => { + assert!(result.as_ref().expect("tool result").success); + successes += 1; + } + Event::TurnComplete { status, error, .. } => { + assert_eq!(*status, TurnOutcomeStatus::Completed, "{error:?}") + } + _ => {} + } + } + assert_eq!(starts.len(), 2); + assert_eq!(starts["write-call"], &write); + assert_eq!(starts["read-call"], &read); + assert_eq!(successes, 2); + assert_eq!( + fs::read_to_string(workspace.path().join("written.txt")).expect("written file"), + content + ); + let requests = mock.captured_requests(); + assert_eq!(requests.len(), 2); + let history_input = requests[1] + .messages + .iter() + .flat_map(|message| &message.content) + .find_map(|block| { + if let ContentBlock::ToolUse { id, input, .. } = block { + (id == "write-call").then_some(input) + } else { + None + } + }) + .expect("tool arguments in follow-up history"); + assert_eq!(history_input["path"], write["path"]); + assert_eq!(history_input["content"], write["content"]); + handle.send(Op::Shutdown).await.expect("shutdown"); + task.await.expect("Engine task"); +} + +#[tokio::test] +async fn partial_input_engine_malformed_final_preserves_raw_arguments_and_does_not_write() { + let workspace = tempdir().expect("workspace"); + let raw = format!( + "{{\"action\":\"write\",\"path\":\"never.txt\",\"content\":\"{}", + "unfinished ".repeat(256) + ); + let mut turn = vec![ + canned::message_start("malformed-partial"), + canned::tool_use_block_start(0, "bad-call", "File"), + ]; + turn.extend( + argument_chunks(&raw) + .into_iter() + .map(|part| canned::tool_input_delta(0, part)), + ); + turn.extend([ + canned::block_stop(0), + canned::message_delta("tool_use", None), + canned::message_stop(), + ]); + let mock = Arc::new(MockLlmClient::new(vec![ + turn, + canned::simple_text_turn("recovered"), + ])); + let client: SharedModelClient = mock.clone(); + let config = Config::default(); + let (engine, handle) = Engine::new_with_model_client( + deterministic_engine_config(workspace.path()), + &config, + client, + ); + let task = tokio::spawn(engine.run()); + handle.send(auto_write_op(&config)).await.expect("send"); + let events = terminal_events(&handle).await; + let input = events + .iter() + .find_map(|event| match event { + Event::ToolCallStarted { id, input, .. } if id == "bad-call" => Some(input), + _ => None, + }) + .expect("malformed tool lifecycle"); + assert_eq!(input, &json!({"raw_arguments":raw})); + assert!(!workspace.path().join("never.txt").exists()); + assert!(events.iter().any(|event| matches!(event, Event::ToolCallComplete { result, .. } if result.as_ref().map_or(true, |value| !value.success)))); + assert_eq!( + mock.call_count(), + 2, + "validation feedback must reach the next model request" + ); + handle.send(Op::Shutdown).await.expect("shutdown"); + task.await.expect("Engine task"); +} + +struct PartialToolModel { + requests: Mutex>, + calls: AtomicUsize, + entered: Arc, + dropped: Arc, + network_error: bool, +} + +struct StreamDropFlag(Arc); +impl Drop for StreamDropFlag { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } +} + +#[async_trait::async_trait] +impl ModelClient for PartialToolModel { + fn provider_name(&self) -> &str { + "partial-tool-test" + } + fn model(&self) -> &str { + "local-model" + } + async fn create_message( + &self, + _: codewhale_tui::models::MessageRequest, + ) -> anyhow::Result { + anyhow::bail!("partial tool test expects streaming") + } + async fn create_message_stream( + &self, + request: codewhale_tui::models::MessageRequest, + ) -> anyhow::Result { + self.requests.lock().unwrap().push(request); + let call = self.calls.fetch_add(1, Ordering::SeqCst); + if call > 0 { + let events = if call == 1 { + canned::tool_call_turn( + "reused-call", + "File", + r#"{"action":"read","path":"fixture.txt"}"#, + ) + } else { + canned::simple_text_turn("retry completed") + }; + return Ok(Box::pin(futures_util::stream::iter( + events.into_iter().map(Ok), + ))); + } + let entered = self.entered.clone(); + let dropped = self.dropped.clone(); + let network_error = self.network_error; + Ok(Box::pin(async_stream::stream! { + let _drop_flag = StreamDropFlag(dropped); + yield Ok(canned::message_start("partial-tool-start")); + yield Ok(canned::tool_use_block_start(0, "reused-call", "File")); + yield Ok(canned::tool_input_delta(0, r#"{"action":"write","path":"never.txt","#)); + yield Ok(canned::tool_input_delta(0, r#""content":"unfinished"#)); + for _ in 0..16 { yield Ok(canned::tool_input_delta(0, " longer unfinished text")); } + entered.notify_one(); + if network_error { + yield Err(anyhow::anyhow!("Stream read error: error decoding response body")); + } else { + std::future::pending::<()>().await; + } + })) + } + async fn health_check(&self) -> anyhow::Result { + Ok(true) + } +} + +#[tokio::test] +async fn partial_input_engine_cancel_then_new_attempt_resets_reused_tool_id() { + let workspace = tempdir().expect("workspace"); + fs::write(workspace.path().join("fixture.txt"), "fixture").expect("fixture"); + let model = Arc::new(PartialToolModel { + requests: Mutex::new(Vec::new()), + calls: AtomicUsize::new(0), + entered: Arc::new(tokio::sync::Notify::new()), + dropped: Arc::new(AtomicBool::new(false)), + network_error: false, + }); + let client: SharedModelClient = model.clone(); + let config = Config::default(); + let (engine, handle) = Engine::new_with_model_client( + deterministic_engine_config(workspace.path()), + &config, + client, + ); + let task = tokio::spawn(engine.run()); + handle.send(auto_write_op(&config)).await.expect("send"); + tokio::time::timeout(model_turn_event_timeout(), model.entered.notified()) + .await + .expect("partial deltas not consumed"); + handle.cancel(); + let cancelled = terminal_events(&handle).await; + assert!(cancelled.iter().any(|event| matches!( + event, + Event::TurnComplete { + status: TurnOutcomeStatus::Interrupted, + .. + } + ))); + assert!( + !cancelled + .iter() + .any(|event| matches!(event, Event::ToolCallStarted { .. })) + ); + assert!(model.dropped.load(Ordering::SeqCst)); + assert_eq!( + model.calls.load(Ordering::SeqCst), + 1, + "cancel must not retry" + ); + assert!(!workspace.path().join("never.txt").exists()); + handle + .send(external_user_message_op( + "Read fixture after cancellation.", + AppMode::Agent, + &config, + )) + .await + .expect("retry send"); + let retried = terminal_events(&handle).await; + let input = retried + .iter() + .find_map(|event| match event { + Event::ToolCallStarted { id, input, .. } if id == "reused-call" => Some(input), + _ => None, + }) + .expect("new tool start"); + assert_eq!(input, &json!({"action":"read","path":"fixture.txt"})); + assert!(retried.iter().any(|event| matches!( + event, + Event::ToolCallComplete { id, result: Ok(result), .. } + if id == "reused-call" && result.success + ))); + assert!(retried.iter().any(|event| matches!( + event, + Event::TurnComplete { + status: TurnOutcomeStatus::Completed, + .. + } + ))); + assert_eq!(model.calls.load(Ordering::SeqCst), 3); + assert!(!workspace.path().join("never.txt").exists()); + handle.send(Op::Shutdown).await.expect("shutdown"); + task.await.expect("Engine task"); +} + +#[tokio::test] +async fn partial_input_engine_network_error_preserves_partial_tool_feedback() { + let workspace = tempdir().expect("workspace"); + fs::write(workspace.path().join("fixture.txt"), "fixture").expect("fixture"); + let model = Arc::new(PartialToolModel { + requests: Mutex::new(Vec::new()), + calls: AtomicUsize::new(0), + entered: Arc::new(tokio::sync::Notify::new()), + dropped: Arc::new(AtomicBool::new(false)), + network_error: true, + }); + let config = Config::default(); + let (engine, handle) = Engine::new_with_model_client( + deterministic_engine_config(workspace.path()), + &config, + model.clone(), + ); + let task = tokio::spawn(engine.run()); + handle.send(auto_write_op(&config)).await.expect("send"); + let events = terminal_events(&handle).await; + assert!(events.iter().any(|event| matches!(event, + Event::TurnComplete {status: TurnOutcomeStatus::Failed,error:Some(error),..} + if error.contains("response body")))); + // The current interactive path retains the partial mirror and reports its + // validation failure through the ordinary tool-feedback loop. It then lets + // the model finish that loop before surfacing the original stream error. + // This pins compatibility, not a general guarantee against partial dispatch. + let completions: Vec<_> = events + .iter() + .filter_map(|event| match event { + Event::ToolCallComplete { id, result, .. } if id == "reused-call" => Some(result), + _ => None, + }) + .collect(); + assert_eq!(completions.len(), 2); + assert!( + completions[0] + .as_ref() + .expect_err("partial writer validation failure") + .to_string() + .contains("missing required field 'content'") + ); + assert!( + completions[1] + .as_ref() + .expect("complete reader result") + .success + ); + let starts: Vec<_> = events + .iter() + .filter_map(|event| match event { + Event::ToolCallStarted { id, input, .. } => Some((id, input)), + _ => None, + }) + .collect(); + assert_eq!( + starts.len(), + 1, + "only the complete follow-up reader publishes a start" + ); + assert_eq!(starts[0].0, "reused-call"); + assert_eq!(starts[0].1, &json!({"action":"read","path":"fixture.txt"})); + assert_eq!( + model.calls.load(Ordering::SeqCst), + 3, + "tool-error feedback and reader result feedback" + ); + { + let requests = model.requests.lock().unwrap(); + let partial = requests[1] + .messages + .iter() + .flat_map(|message| &message.content) + .find_map(|block| match block { + ContentBlock::ToolUse { id, input, .. } if id == "reused-call" => Some(input), + _ => None, + }) + .expect("partial mirror in tool-error feedback history"); + assert_eq!(partial["path"], "never.txt"); + assert!(partial.get("content").is_none()); + } + assert!(model.dropped.load(Ordering::SeqCst)); + assert!(!workspace.path().join("never.txt").exists()); + handle.send(Op::Shutdown).await.expect("shutdown"); + task.await.expect("Engine task"); +} + +mod canned { + use codewhale_tui::models::{ + ContentBlockStart, Delta, MessageDelta, MessageResponse, StreamEvent, Usage, + }; + use serde_json::Value; + + /// `MessageStart` event with a synthetic message envelope. + #[must_use] + pub fn message_start(id: &str) -> StreamEvent { + StreamEvent::MessageStart { + message: MessageResponse { + id: id.to_string(), + r#type: "message".to_string(), + role: "assistant".to_string(), + content: vec![], + model: "mock-model".to_string(), + stop_reason: None, + stop_sequence: None, + container: None, + usage: Usage::default(), + }, + } + } + + /// Open a text content block at `index`. + #[must_use] + pub fn text_block_start(index: u32) -> StreamEvent { + StreamEvent::ContentBlockStart { + index, + content_block: ContentBlockStart::Text { + text: String::new(), + }, + } + } + + /// Append `text` to the content block at `index`. + #[must_use] + pub fn text_delta(index: u32, text: &str) -> StreamEvent { + StreamEvent::ContentBlockDelta { + index, + delta: Delta::TextDelta { + text: text.to_string(), + }, + } + } + + /// Open a tool_use content block at `index`. + #[must_use] + pub fn tool_use_block_start(index: u32, id: &str, name: &str) -> StreamEvent { + StreamEvent::ContentBlockStart { + index, + content_block: ContentBlockStart::ToolUse { + id: id.to_string(), + name: name.to_string(), + input: Value::Null, + caller: None, + thought_signature: None, + }, + } + } + + /// Stream partial JSON for a tool's input arguments. + #[must_use] + pub fn tool_input_delta(index: u32, partial_json: &str) -> StreamEvent { + StreamEvent::ContentBlockDelta { + index, + delta: Delta::InputJsonDelta { + partial_json: partial_json.to_string(), + }, + } + } + + /// Close the content block at `index`. + #[must_use] + pub fn block_stop(index: u32) -> StreamEvent { + StreamEvent::ContentBlockStop { index } + } + + /// Emit a `message_delta` carrying `stop_reason` and optional `usage`. + #[must_use] + pub fn message_delta(stop_reason: &str, usage: Option) -> StreamEvent { + StreamEvent::MessageDelta { + delta: MessageDelta { + stop_reason: Some(stop_reason.to_string()), + stop_sequence: None, + }, + usage, + } + } + + /// Final `message_stop` sentinel. + #[must_use] + pub fn message_stop() -> StreamEvent { + StreamEvent::MessageStop + } + + /// Convenience: a complete "assistant emits this text" turn ending with + /// `stop_reason = "end_turn"`. + #[must_use] + pub fn simple_text_turn(text: &str) -> Vec { + vec![ + message_start("mock_msg_1"), + text_block_start(0), + text_delta(0, text), + block_stop(0), + message_delta("end_turn", None), + message_stop(), + ] + } + + /// Convenience: a turn that emits one assistant tool_call and stops. + #[must_use] + pub fn tool_call_turn(call_id: &str, tool_name: &str, args_json: &str) -> Vec { + vec![ + message_start("mock_msg_tool"), + tool_use_block_start(0, call_id, tool_name), + tool_input_delta(0, args_json), + block_stop(0), + message_delta("tool_use", None), + message_stop(), + ] + } +} + +// === Tests === + +#[tokio::test] +async fn partial_input_engine_network_error_before_any_content_can_retry() { + let workspace = tempdir().expect("workspace"); + let model = Arc::new(MockLlmClient::new(vec![canned::simple_text_turn( + "recovered", + )])); + model.initial_network_error.store(true, Ordering::SeqCst); + let config = Config::default(); + let (engine, handle) = Engine::new_with_model_client( + deterministic_engine_config(workspace.path()), + &config, + model.clone(), + ); + let task = tokio::spawn(engine.run()); + handle.send(auto_write_op(&config)).await.expect("send"); + let events = terminal_events(&handle).await; + assert!(events.iter().any(|event| matches!( + event, + Event::TurnComplete { + status: TurnOutcomeStatus::Completed, + .. + } + ))); + assert_eq!( + model.call_count(), + 2, + "a transient stream failure before content may retry" + ); + assert!( + !events + .iter() + .any(|event| matches!(event, Event::ToolCallStarted { .. })) + ); + handle.send(Op::Shutdown).await.expect("shutdown"); + task.await.expect("Engine task"); +} + +/// Headless recovery discards the uncommitted tool fragment. The partial-input +/// history clone in InteractiveNetworkDrop is unreachable while a tool is open: +/// that resume policy requires tool_uses_empty. Pin the actual public contract. +#[tokio::test] +async fn partial_input_engine_headless_resume_discards_partial_history_and_resets_tool_id() { + let workspace = tempdir().expect("workspace"); + fs::write(workspace.path().join("fixture.txt"), "fixture").expect("fixture"); + let model = Arc::new(PartialToolModel { + requests: Mutex::new(Vec::new()), + calls: AtomicUsize::new(0), + entered: Arc::new(tokio::sync::Notify::new()), + dropped: Arc::new(AtomicBool::new(false)), + network_error: true, + }); + let config = Config::default(); + let mut engine_config = deterministic_engine_config(workspace.path()); + engine_config.terminal_chrome_enabled = false; + let (engine, handle) = Engine::new_with_model_client(engine_config, &config, model.clone()); + let task = tokio::spawn(engine.run()); + handle.send(auto_write_op(&config)).await.expect("send"); + let events = terminal_events(&handle).await; + assert!(events.iter().any(|event| matches!( + event, + Event::TurnComplete { + status: TurnOutcomeStatus::Completed, + error: None, + .. + } + ))); + assert_eq!( + model.calls.load(Ordering::SeqCst), + 3, + "one resume followed by tool-result feedback" + ); + let starts: Vec<_> = events + .iter() + .filter_map(|event| match event { + Event::ToolCallStarted { id, input, .. } => Some((id, input)), + _ => None, + }) + .collect(); + assert_eq!(starts.len(), 1, "unfinished tool must never dispatch"); + assert_eq!(starts[0].0, "reused-call"); + let expected = json!({"action":"read","path":"fixture.txt"}); + assert_eq!(starts[0].1, &expected); + assert!(events.iter().any(|event| matches!(event, + Event::ToolCallComplete { id, result: Ok(result), .. } if id == "reused-call" && result.success))); + assert!(!workspace.path().join("never.txt").exists()); + assert!(model.dropped.load(Ordering::SeqCst)); + { + let requests = model.requests.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert_eq!( + serde_json::to_value(&requests[0].messages).unwrap(), + serde_json::to_value(&requests[1].messages).unwrap(), + "headless resume reissues persisted history without partial tool arguments" + ); + assert!( + !requests[1] + .messages + .iter() + .flat_map(|message| &message.content) + .any( + |block| matches!(block, ContentBlock::ToolUse { id, .. } if id == "reused-call") + ) + ); + let history_input = requests[2] + .messages + .iter() + .flat_map(|message| &message.content) + .find_map(|block| match block { + ContentBlock::ToolUse { id, input, .. } if id == "reused-call" => Some(input), + _ => None, + }) + .expect("fresh complete tool arguments in follow-up history"); + assert_eq!(history_input["action"], expected["action"]); + assert_eq!(history_input["path"], expected["path"]); + assert!( + history_input.get("content").is_none(), + "unfinished writer must not contaminate fresh reader history" + ); + } + handle.send(Op::Shutdown).await.expect("shutdown"); + task.await.expect("Engine task"); +}