codewhale-hq/Codewhale
1
0 Fork 0
mirror of https://github.com/codewhale-hq/Codewhale.git synced 2026-11-09 09:59:45 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

fix(core): bound streamed tool batches and migrate cancellable delivery

Browse Source 
C02-11: migrate every16 process_stream observation plus the finalizer warning to the engine cancellation-aware sender. Eight actual decoder observation shapes are exercised against a full queue; live drain preserves output ordering and usage. Terminal settlement still requires consumer progress on the existing lossless channel and remains a separately owned boundary.

C02-14: bound each response to256 tool calls before retaining either native/server ToolStart257. The text fallback enforces the same ceiling before conversion, observations or planning. A response rejected by the stream guard uses the existing failed-response admission path; configured lower per-turn limits and protocol ID validation remain in force.

Source checkpoint: rustfmt --check and git diff --check pass. New actual decoder tests cover exactly256 versus257 empty native/server calls. Focused governed compilation and copied-production negative controls pending; no passing runtime or full C02-11 closure claimed. Refs #6561.

Signed-off-by: CodeWhale Bot <[email protected]>
This commit is contained in:
CodeWhale Bot committed 2026-09-30 02:40:02 -07:00
1 parent d23d5590e9
commit e44cf31229
2 files changed
+427 -27
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/tui/src/core/engine/streaming.rs
+17
View File
@@ -55,6 +55,23 @@ impl ToolUseState {
/// Maximum total bytes of text, reasoning and tool-argument content before aborting the stream.
pub(super) const STREAM_MAX_CONTENT_BYTES: usize = 10 * 1024 * 1024; // 10 MB
/// A response can contain many empty tool starts without spending the byte
/// budget. Bound that batch before any call is retained or admitted. A lower
/// configured per-turn tool budget remains authoritative at execution.
pub(super) const MAX_TOOL_CALLS_PER_RESPONSE: usize = 256;
pub(super) fn tool_call_limit_error() -> crate::error_taxonomy::ErrorEnvelope {
crate::error_taxonomy::ErrorEnvelope::new(
crate::error_taxonomy::ErrorCategory::InvalidInput,
crate::error_taxonomy::ErrorSeverity::Error,
false,
"response_tool_call_limit",
format!(
"Model response exceeded the maximum of {MAX_TOOL_CALLS_PER_RESPONSE} tool calls; no call from this response was executed"
),
)
}
/// Sanity backstop for total stream wall-clock duration. **Not** a routine
/// kill switch — the stream chunk idle timeout is the primary stall
/// detector. The wall-clock cap is here only to bound pathological cases
crates/tui/src/core/engine/turn_loop.rs
+410 -27
View File
@@ -2414,6 +2414,16 @@ impl Engine {
let mut final_text = current_text_visible.clone();
if tool_uses.is_empty() && tool_parser::has_tool_call_markers(&current_text_raw) {
let parsed = tool_parser::parse_tool_calls(&current_text_raw);
if parsed.tool_calls.len() > super::streaming::MAX_TOOL_CALLS_PER_RESPONSE {
let envelope = super::streaming::tool_call_limit_error();
let error = envelope.message.clone();
turn.stop_diagnostics.last_response_tool_calls_suppressed =
Some(parsed.tool_calls.len());
let _ = self.send_stream_event(Event::error(envelope)).await;
self.add_interrupted_assistant_text(&current_text_visible)
.await;
return (TurnOutcomeStatus::Failed, Some(error));
}
final_text = parsed.clean_text;
for call in parsed.tool_calls {
tool_uses.push(ToolUseState {
@@ -5940,7 +5950,7 @@ impl Engine {
// of ending "Completed" over a frozen block.
stream_errors = stream_errors.saturating_add(1);
stream_error.get_or_insert(envelope.message.clone());
let _ = self.tx_event.send(Event::error(envelope)).await;
let _ = self.send_stream_event(Event::error(envelope)).await;
None
}
}
@@ -5957,8 +5967,7 @@ impl Engine {
let preview = summarize_text(pending.content.trim(), 120);
pending_steers.push(pending);
let _ = self
.tx_event
.send(Event::status(format!("Steer input queued: {preview}")))
.send_stream_event(Event::status(format!("Steer input queued: {preview}")))
.await;
}
@@ -5974,7 +5983,7 @@ impl Engine {
.into_envelope();
crate::logging::warn(&envelope.message);
stream_error.get_or_insert(envelope.message.clone());
let _ = self.tx_event.send(Event::error(envelope)).await;
let _ = self.send_stream_event(Event::error(envelope)).await;
break;
}
@@ -6078,8 +6087,7 @@ impl Engine {
));
stream_error.get_or_insert(retry_msg.clone());
let _ = self
.tx_event
.send(Event::error(
.send_stream_event(Event::error(
crate::error_taxonomy::envelope_for_llm_error(
retry_err, retry_msg,
),
@@ -6167,7 +6175,7 @@ impl Engine {
// card. Recoverable classes (rate limit, network) keep
// the bounded retry tail.
let terminal = !envelope.recoverable;
let _ = self.tx_event.send(Event::error(envelope)).await;
let _ = self.send_stream_event(Event::error(envelope)).await;
if terminal || stream_errors >= retry_limits.max_errors {
break;
}
@@ -6189,7 +6197,22 @@ impl Engine {
.into_envelope();
crate::logging::warn(&envelope.message);
stream_error.get_or_insert(envelope.message.clone());
let _ = self.tx_event.send(Event::error(envelope)).await;
let _ = self.send_stream_event(Event::error(envelope)).await;
break;
}
if matches!(
&event,
StreamEvent::ContentBlockStart {
content_block: ContentBlockStart::ToolUse { .. }
| ContentBlockStart::ServerToolUse { .. },
..
}
) && tool_uses.len() >= super::streaming::MAX_TOOL_CALLS_PER_RESPONSE
{
let envelope = super::streaming::tool_call_limit_error();
stream_error.get_or_insert(envelope.message.clone());
let _ = self.send_stream_event(Event::error(envelope)).await;
break;
}
@@ -6200,8 +6223,7 @@ impl Engine {
omitted_tool_count,
} => {
let _ = self
.tx_event
.send(Event::ToolProjectionWarning {
.send_stream_event(Event::ToolProjectionWarning {
provider,
omitted_tool_names,
omitted_tool_count,
@@ -6231,15 +6253,16 @@ impl Engine {
&& filtered.len() < current_text_raw.len()
&& contains_fake_tool_wrapper(&current_text_raw)
{
let _ = self.tx_event.send(Event::status(FAKE_WRAPPER_NOTICE)).await;
let _ = self
.send_stream_event(Event::status(FAKE_WRAPPER_NOTICE))
.await;
fake_wrapper_notice_emitted = true;
}
current_text_visible.push_str(&filtered);
current_block_kind = Some(ContentBlockKind::Text);
last_text_index = Some(index as usize);
let _ = self
.tx_event
.send(Event::MessageStarted {
.send_stream_event(Event::MessageStarted {
index: index as usize,
})
.await;
@@ -6250,8 +6273,7 @@ impl Engine {
current_thinking_state = None;
current_block_kind = Some(ContentBlockKind::Thinking);
let _ = self
.tx_event
.send(Event::ThinkingStarted {
.send_stream_event(Event::ThinkingStarted {
index: index as usize,
})
.await;
@@ -6310,14 +6332,15 @@ impl Engine {
&& filtered.len() < text.len()
&& contains_fake_tool_wrapper(&current_text_raw)
{
let _ = self.tx_event.send(Event::status(FAKE_WRAPPER_NOTICE)).await;
let _ = self
.send_stream_event(Event::status(FAKE_WRAPPER_NOTICE))
.await;
fake_wrapper_notice_emitted = true;
}
if !filtered.is_empty() {
current_text_visible.push_str(&filtered);
let _ = self
.tx_event
.send(Event::MessageDelta {
.send_stream_event(Event::MessageDelta {
index: index as usize,
content: filtered,
})
@@ -6328,8 +6351,7 @@ impl Engine {
current_thinking.push_str(&thinking);
if !thinking.is_empty() {
let _ = self
.tx_event
.send(Event::ThinkingDelta {
.send_stream_event(Event::ThinkingDelta {
index: index as usize,
content: thinking,
})
@@ -6380,8 +6402,7 @@ impl Engine {
if !flushed.is_empty() {
current_text_visible.push_str(&flushed);
let _ = self
.tx_event
.send(Event::MessageDelta {
.send_stream_event(Event::MessageDelta {
index: index as usize,
content: flushed,
})
@@ -6392,8 +6413,7 @@ impl Engine {
}
Some(ContentBlockKind::Thinking) => {
let _ = self
.tx_event
.send(Event::ThinkingComplete {
.send_stream_event(Event::ThinkingComplete {
index: index as usize,
})
.await;
@@ -6448,7 +6468,7 @@ impl Engine {
.unwrap_or("provider stream error");
let envelope = ErrorEnvelope::classify(message.to_string(), true);
crate::logging::warn(format!("Provider stream error event: {message}"));
let _ = self.tx_event.send(Event::error(envelope)).await;
let _ = self.send_stream_event(Event::error(envelope)).await;
stream_error.get_or_insert(message.to_string());
break;
}
@@ -6552,8 +6572,7 @@ impl Engine {
tool_state.input_parse_error = Some(error);
tool_state.input = malformed_tool_arguments_input(&tool_state.input_buffer);
let _ = self
.tx_event
.send(Event::status(format!(
.send_stream_event(Event::status(format!(
"⚠ Tool '{}' received malformed arguments from model",
tool_state.name
)))
@@ -7709,6 +7728,370 @@ mod tests {
use std::time::Duration;
use tempfile::tempdir;
fn stream_backpressure_fixture(
workspace: &std::path::Path,
capacity: usize,
) -> (
Engine,
Arc<crate::llm_client::mock::MockLlmClient>,
mpsc::Receiver<Event>,
) {
let model = Arc::new(crate::llm_client::mock::MockLlmClient::new(Vec::new()));
let (mut engine, _handle) = Engine::new_with_model_client(
EngineConfig {
workspace: workspace.into(),
snapshots_enabled: false,
subagents_enabled: false,
terminal_chrome_enabled: false,
..Default::default()
},
&Config::default(),
model.clone(),
);
let (tx, rx) = mpsc::channel(capacity);
engine.tx_event = tx;
(engine, model, rx)
}
fn stream_backpressure_request() -> codewhale_models::MessageRequest {
prepare_primary_turn_request(PrimaryTurnRequest {
model: "mock-model".into(),
messages: Vec::new(),
max_tokens: 128,
system: None,
tools: None,
tool_choice: None,
reasoning_effort: None,
})
}
/// Hold the actual stream decoder in a full host queue, then cancel
/// without draining that queue. The provider suffix must never be polled.
#[tokio::test]
async fn stream_backpressure_cancellation_releases_every_observation_kind() {
use crate::llm_client::mock::{MockLlmClient, canned};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
struct StreamDrop(Arc<AtomicBool>);
impl Drop for StreamDrop {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
let thinking_start = StreamEvent::ContentBlockStart {
index: 0,
content_block: ContentBlockStart::Thinking {
thinking: String::new(),
},
};
let cases = [
("text start", vec![canned::text_block_start(0)], 0),
("text delta", vec![canned::text_delta(0, "partial")], 0),
("thinking start", vec![thinking_start.clone()], 0),
(
"thinking delta",
vec![canned::thinking_delta(0, "partial")],
0,
),
(
"thinking stop",
vec![thinking_start, canned::block_stop(0)],
1,
),
(
"projection warning",
vec![StreamEvent::ToolProjectionWarning {
provider: "mock".into(),
omitted_tool_names: vec!["omitted".into()],
omitted_tool_count: 1,
}],
0,
),
(
"provider error",
vec![StreamEvent::Error {
error: json!({"message": "invalid provider request"}),
}],
0,
),
(
"malformed tool arguments",
vec![
canned::tool_use_block_start(0, "call_1", "read_file"),
canned::tool_input_delta(0, "not-json"),
canned::block_stop(0),
],
0,
),
];
for (label, events, observations_before_block) in cases {
let tmp = tempdir().expect("tempdir");
let (mut engine, model, mut rx) =
stream_backpressure_fixture(tmp.path(), observations_before_block + 1);
engine
.tx_event
.send(Event::status("queue already occupied"))
.await
.unwrap();
let cancel = engine.cancel_token.clone();
let mut start = canned::message_start("backpressure");
if let StreamEvent::MessageStart { message } = &mut start {
message.usage.input_tokens = 17;
}
let expected_polls = events.len() + 1;
let polls = Arc::new(AtomicUsize::new(0));
let counted = Arc::clone(&polls);
let dropped = Arc::new(AtomicBool::new(false));
let drop_probe = StreamDrop(Arc::clone(&dropped));
let stream = futures_util::stream::iter(
std::iter::once(start)
.chain(events)
.chain(std::iter::once(canned::text_delta(0, "UNREAD-SUFFIX")))
.map(Ok),
)
.inspect(move |_| {
let _ = &drop_probe;
counted.fetch_add(1, Ordering::SeqCst);
});
let request = stream_backpressure_request();
let mut diagnostics = crate::tool_inspection::TurnStopDiagnostics::default();
let mut process = Box::pin(engine.process_stream(
model.as_ref(),
Box::pin(stream),
&request,
Instant::now(),
0,
&mut diagnostics,
));
assert!(
tokio::time::timeout(Duration::from_millis(25), &mut process)
.await
.is_err(),
"{label}: a live turn must wait for capacity"
);
assert_eq!(polls.load(Ordering::SeqCst), expected_polls, "{label}");
assert!(
!dropped.load(Ordering::SeqCst),
"{label}: stream is in flight"
);
cancel.cancel();
let outcome = tokio::time::timeout(Duration::from_secs(1), &mut process)
.await
.unwrap_or_else(|_| {
panic!("{label}: cancellation must release a full event queue")
});
drop(process);
assert!(
dropped.load(Ordering::SeqCst),
"{label}: provider stream released"
);
assert_eq!(
polls.load(Ordering::SeqCst),
expected_polls,
"{label}: suffix unread"
);
assert_eq!(
outcome.usage.input_tokens, 17,
"{label}: billed usage retained"
);
assert!(
outcome.pending_resume.is_none(),
"{label}: cancellation cannot retry"
);
assert_eq!(model.call_count(), 0, "{label}: no provider retry");
assert_eq!(
rx.len(),
observations_before_block + 1,
"{label}: no drain was needed"
);
while let Ok(event) = rx.try_recv() {
assert!(
!matches!(event, Event::MessageDelta { content, .. } if content == "UNREAD-SUFFIX")
);
}
if label == "malformed tool arguments" {
assert!(outcome.tool_uses[0].input_parse_error.is_some());
}
}
}
#[tokio::test]
async fn stream_backpressure_live_delivery_preserves_order_and_usage() {
use crate::llm_client::mock::{MockLlmClient, canned};
let tmp = tempdir().expect("tempdir");
let model = Arc::new(MockLlmClient::new(Vec::new()));
let (mut engine, _handle) = Engine::new_with_model_client(
EngineConfig {
workspace: tmp.path().into(),
snapshots_enabled: false,
subagents_enabled: false,
..Default::default()
},
&Config::default(),
model.clone(),
);
let (tx, mut rx) = mpsc::channel(1);
tx.send(Event::status("occupied")).await.unwrap();
engine.tx_event = tx;
let stream = futures_util::stream::iter([
Ok(canned::text_delta(0, "first")),
Ok(canned::text_delta(0, "second")),
Ok(canned::message_delta(
"end_turn",
Some(Usage {
output_tokens: 9,
..Default::default()
}),
)),
]);
let request = stream_backpressure_request();
let mut diagnostics = crate::tool_inspection::TurnStopDiagnostics::default();
let mut process = Box::pin(engine.process_stream(
model.as_ref(),
Box::pin(stream),
&request,
Instant::now(),
0,
&mut diagnostics,
));
assert!(
tokio::time::timeout(Duration::from_millis(25), &mut process)
.await
.is_err()
);
assert!(matches!(rx.recv().await, Some(Event::Status { .. })));
let drain = async {
let first = rx.recv().await.expect("first delta");
let second = rx.recv().await.expect("second delta");
[first, second]
};
let (outcome, events) = tokio::time::timeout(Duration::from_secs(1), async {
tokio::join!(&mut process, drain)
})
.await
.expect("draining the queue must resume lossless delivery");
assert!(matches!(&events[0], Event::MessageDelta { content, .. } if content == "first"));
assert!(matches!(&events[1], Event::MessageDelta { content, .. } if content == "second"));
assert_eq!(outcome.current_text_visible, "firstsecond");
assert_eq!(outcome.usage.output_tokens, 9);
assert_eq!(outcome.stop_reason.as_deref(), Some("end_turn"));
}
#[tokio::test]
async fn stream_response_tool_limit_bounds_empty_native_and_server_calls() {
use crate::llm_client::mock::{MockLlmClient, canned};
use std::sync::atomic::{AtomicUsize, Ordering};
for server_tool in [false, true] {
for count in [
super::super::streaming::MAX_TOOL_CALLS_PER_RESPONSE,
super::super::streaming::MAX_TOOL_CALLS_PER_RESPONSE + 1,
] {
let tmp = tempdir().expect("tempdir");
let model = Arc::new(MockLlmClient::new(Vec::new()));
let (mut engine, _handle) = Engine::new_with_model_client(
EngineConfig {
workspace: tmp.path().into(),
snapshots_enabled: false,
subagents_enabled: false,
..Default::default()
},
&Config::default(),
model.clone(),
);
let (tx, mut rx) = mpsc::channel(4);
engine.tx_event = tx;
let events = (0..count).map(|index| StreamEvent::ContentBlockStart {
index: u32::try_from(index).unwrap(),
content_block: if server_tool {
ContentBlockStart::ServerToolUse {
id: format!("call_{index}"),
name: "web_search".into(),
input: json!({}),
}
} else {
ContentBlockStart::ToolUse {
id: format!("call_{index}"),
name: "read_file".into(),
input: json!({}),
caller: None,
thought_signature: None,
}
},
});
let polls = Arc::new(AtomicUsize::new(0));
let counted = Arc::clone(&polls);
let stream = futures_util::stream::iter(
events
.chain(std::iter::once(canned::text_delta(
0,
"SUFFIX-AFTER-TOOL-BATCH",
)))
.map(Ok),
)
.inspect(move |_| {
counted.fetch_add(1, Ordering::SeqCst);
});
let request = stream_backpressure_request();
let mut diagnostics = crate::tool_inspection::TurnStopDiagnostics::default();
let outcome = tokio::time::timeout(
Duration::from_secs(1),
engine.process_stream(
model.as_ref(),
Box::pin(stream),
&request,
Instant::now(),
0,
&mut diagnostics,
),
)
.await
.expect("empty tool starts must be bounded");
assert_eq!(
outcome.tool_uses.len(),
super::super::streaming::MAX_TOOL_CALLS_PER_RESPONSE
);
if count > super::super::streaming::MAX_TOOL_CALLS_PER_RESPONSE {
assert!(
outcome
.stream_error
.as_deref()
.is_some_and(|error| error.contains("256 tool calls"))
);
assert_eq!(
polls.load(Ordering::SeqCst),
count,
"overflow stops before the suffix"
);
assert!(
matches!(rx.try_recv(), Ok(Event::Error { envelope, .. }) if envelope.code == "response_tool_call_limit" && !envelope.recoverable)
);
assert!(outcome.current_text_raw.is_empty());
} else {
assert!(
outcome.stream_error.is_none(),
"exactly256 calls remain valid"
);
assert_eq!(polls.load(Ordering::SeqCst), count + 1);
}
while let Ok(event) = rx.try_recv() {
assert!(
!matches!(event, Event::ToolCallStarted { .. }),
"decoding cannot observe/execute a rejected batch"
);
}
assert_eq!(
model.call_count(),
0,
"tool cardinality overflow cannot retry"
);
}
}
}
#[tokio::test]
async fn rlm_tool_context_inherits_the_spent_parent_clock() {
let tmp = tempdir().unwrap();