mirror of
https://github.com/Mesh-LLM/mesh-llm.git
synced 2026-08-08 22:23:19 -04:00
fix(moa): route tool-result follow-ups to reducer, not fan-out
PR #566 review feedback (Apr 2026): > The tool-result path isn't ready for agent loops: > - A tool-result follow-up was treated like another fanout turn. > - It wasn't handled like a controlled reducer/synthesis turn. > - Tool results should be handled carefully and predictably, not > sprayed back through the whole fanout path. `Session::classify_turn` only routed to `TurnType::ToolResult` when the very last message had `role: "tool"`. Many agent harnesses send the tool result followed by a short `user` nudge ("continue", "what did you find?"). That landed at the very-last-message check as `user`, so the gateway classified the turn as Continuation, fanned out to all workers, and invited a worker to re-propose the same tool call whose result was already in context. The session-state fallback at `last_was_tool_call && has_unprocessed_tool_results` was dead code in production: the gateway never invokes `record_assistant_response` between turns, so `last_was_tool_call` is always false. ## Test (added first, observed failing) `crates/mesh-mixture-of-agents/tests/sim_tool_result_routes_to_reducer.rs` \u2014 three scenarios, all using mock backends that count calls: * OpenAI canonical shape (last msg role=tool) \u2014 must classify as `ToolResult`, exactly one backend call (reducer only). Already passed pre-fix; pinned to prevent regression. * Trailing-user-after-unsynthesised-tool-result \u2014 must also classify as `ToolResult`, exactly one backend call. **Failed pre-fix with `TurnKind::EarlyExit`** (fanned out, multiple worker calls). * Plain fresh user question \u2014 must still fan out. Pins that we don't over-trigger the tool-result path. ## Fix Scan messages from the end in `Session::classify_turn`: * First message we hit with `role: "tool"` \u2192 classify as `ToolResult`. The tool result has not yet been synthesised by an assistant message after it. * First message we hit with `role: "assistant"` \u2192 stop. The assistant has already spoken since the last tool result; the next turn is a normal continuation. * Other roles (`user`, `system`) \u2192 keep scanning. A user nudge after an unsynthesised tool result still belongs in the reducer-only path. If the scan reaches the start without hitting either, fall through to the existing `Fresh`/`Continuation` classification. ## Validation `cargo test -p mesh-mixture-of-agents` \u2014 70 unit + 5 integration tests pass (this PR\u2019s 3 new tests + the two earlier sim files). `cargo test -p mesh-llm-host-runtime --lib` \u2014 1434/1434 pass. `cargo clippy -p mesh-mixture-of-agents --all-targets -- -D warnings` \u2014 clean. `cargo fmt --all -- --check` \u2014 clean.
This commit is contained in:
Michael Neale
parent
a396ab1ed9
commit
000ae50cf6
2 changed files with 306 additions and 5 deletions
|
|
@ -120,15 +120,40 @@ impl Session {
|
|||
}
|
||||
|
||||
/// Classify what kind of turn this is.
|
||||
///
|
||||
/// Anything that ends with an unsynthesised tool result is a
|
||||
/// `ToolResult` turn: the gateway must skip fan-out and go
|
||||
/// straight to the reducer, so the workers don't re-broadcast
|
||||
/// the same tool call whose result we already have in context.
|
||||
///
|
||||
/// We scan from the end of the conversation backwards. The first
|
||||
/// message we hit decides:
|
||||
///
|
||||
/// * `role: "tool"` first — OpenAI canonical: classify as
|
||||
/// `ToolResult`.
|
||||
/// * `role: "assistant"` first — the assistant has already
|
||||
/// spoken since the last tool result. Hand the next turn
|
||||
/// to fan-out normally.
|
||||
/// * `role: "user"` first — keep scanning past it. A user
|
||||
/// nudge after an unsynthesised tool result is still a
|
||||
/// tool-result turn; the model needs to consume the tool
|
||||
/// output and answer the nudge in one synthesis pass. A
|
||||
/// user message that *predates* any tool result reaches the
|
||||
/// start of the history and we fall through to the normal
|
||||
/// Fresh/Continuation classification.
|
||||
pub fn classify_turn(&self) -> TurnType {
|
||||
// If the last message is a tool result, this is a tool-result turn
|
||||
if let Some(last) = self.messages.last() {
|
||||
if last.get("role").and_then(|r| r.as_str()) == Some("tool") {
|
||||
return TurnType::ToolResult;
|
||||
for msg in self.messages.iter().rev() {
|
||||
let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
|
||||
match role {
|
||||
"tool" => return TurnType::ToolResult,
|
||||
"assistant" => break,
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
|
||||
// Also check if we just got tool results back after our tool_call
|
||||
// Also check the session-state fallback (set by
|
||||
// record_assistant_response, which the gateway doesn't invoke
|
||||
// yet but tests do).
|
||||
if self.last_was_tool_call && self.has_unprocessed_tool_results() {
|
||||
return TurnType::ToolResult;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,276 @@
|
|||
//! Pin the tool-result turn classification.
|
||||
//!
|
||||
//! Background — PR #566 review feedback (Apr 2026):
|
||||
//!
|
||||
//! > The tool-result path isn't ready for agent loops:
|
||||
//! > - A tool-result follow-up was treated like another fanout turn.
|
||||
//! > - It wasn't handled like a controlled reducer/synthesis turn.
|
||||
//! > - Tool results should be handled carefully and predictably, not
|
||||
//! > sprayed back through the whole fanout path.
|
||||
//!
|
||||
//! When the conversation has a recent tool result that has not yet
|
||||
//! been synthesized into an assistant answer, the gateway must take
|
||||
//! the reducer-only path (`TurnKind::ToolResult`), not fan-out to all
|
||||
//! workers. Fanning out wastes a round-trip per worker, drowns the
|
||||
//! reducer in worker outputs that ignore the tool result, and \u2014 most
|
||||
//! dangerously \u2014 invites a worker to re-propose the same tool call
|
||||
//! whose result we already have in-context.
|
||||
//!
|
||||
//! Two shapes of agent conversation must classify as ToolResult:
|
||||
//!
|
||||
//! 1. **OpenAI canonical shape.** The last message has role `tool`
|
||||
//! (the harness sent the tool result and expects the next assistant
|
||||
//! turn to interpret it). This is the simplest and most explicit
|
||||
//! shape. Classifying this is straightforward.
|
||||
//!
|
||||
//! 2. **Trailing-user shape.** The conversation ends with assistant
|
||||
//! tool_calls + tool result + a user message that just nudges
|
||||
//! ("continue", "what did you find?"). The harness has left the
|
||||
//! tool result in-context for the model to consume. There is no
|
||||
//! new tool result in *this* turn, but the previous one was never
|
||||
//! synthesized into an assistant message. Today this classifies
|
||||
//! as `Continuation` and fans out \u2014 wrong per the review.
|
||||
//!
|
||||
//! This file pins both shapes.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use mesh_mixture_of_agents as moa;
|
||||
use serde_json::{json, Value};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Backend that records every call it receives so the test can assert
|
||||
/// fan-out did or did not happen.
|
||||
struct RecordingBackend {
|
||||
text: String,
|
||||
delay: Duration,
|
||||
calls: AtomicUsize,
|
||||
}
|
||||
|
||||
impl RecordingBackend {
|
||||
fn new(text: impl Into<String>) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
text: text.into(),
|
||||
delay: Duration::from_millis(10),
|
||||
calls: AtomicUsize::new(0),
|
||||
})
|
||||
}
|
||||
|
||||
fn calls(&self) -> usize {
|
||||
self.calls.load(Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl moa::ModelBackend for RecordingBackend {
|
||||
async fn chat_completion(
|
||||
&self,
|
||||
_model: &str,
|
||||
_messages: &[Value],
|
||||
_tools: Option<&Value>,
|
||||
_max_tokens: u32,
|
||||
_timeout: Duration,
|
||||
_sampling: moa::SamplingParams,
|
||||
) -> Result<Value, String> {
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
tokio::time::sleep(self.delay).await;
|
||||
Ok(json!({
|
||||
"choices": [{"message": {"content": self.text}}],
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
fn config_with_three_recording_workers() -> (
|
||||
moa::GatewayConfig,
|
||||
Arc<RecordingBackend>,
|
||||
Arc<RecordingBackend>,
|
||||
Arc<RecordingBackend>,
|
||||
) {
|
||||
let fast = RecordingBackend::new("synthesised: README says 'Hello World'");
|
||||
let mid = RecordingBackend::new("synthesised: README says 'Hello World'");
|
||||
let strong = RecordingBackend::new("synthesised: README says 'Hello World'");
|
||||
|
||||
let backends: Vec<Arc<dyn moa::ModelBackend>> = vec![fast.clone(), mid.clone(), strong.clone()];
|
||||
let models = vec![
|
||||
moa::ModelEntry {
|
||||
name: "fast-3b".into(),
|
||||
backend_index: 0,
|
||||
},
|
||||
moa::ModelEntry {
|
||||
name: "mid-13b".into(),
|
||||
backend_index: 1,
|
||||
},
|
||||
moa::ModelEntry {
|
||||
name: "strong-32b".into(),
|
||||
backend_index: 2,
|
||||
},
|
||||
];
|
||||
|
||||
let config = moa::GatewayConfig {
|
||||
backends,
|
||||
models,
|
||||
worker_timeout: Duration::from_secs(2),
|
||||
hedge_delay: Duration::from_millis(50),
|
||||
reducer_timeout: Duration::from_secs(2),
|
||||
};
|
||||
(config, fast, mid, strong)
|
||||
}
|
||||
|
||||
fn read_file_tool() -> Value {
|
||||
json!([{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"path": {"type": "string"}},
|
||||
"required": ["path"],
|
||||
}
|
||||
}
|
||||
}])
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn last_message_role_tool_classifies_as_tool_result() {
|
||||
// OpenAI canonical shape. This is already supposed to work today
|
||||
// and pins the existing behavior so we don't regress when we
|
||||
// tighten the "trailing user" shape below.
|
||||
let (config, fast, mid, strong) = config_with_three_recording_workers();
|
||||
|
||||
let body = json!({
|
||||
"model": "mesh",
|
||||
"tools": read_file_tool(),
|
||||
"messages": [
|
||||
{"role": "user", "content": "Read README.md"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": null,
|
||||
"tool_calls": [{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "read_file", "arguments": "{\"path\":\"README.md\"}"}
|
||||
}]
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "# Hello World\n"},
|
||||
],
|
||||
"max_tokens": 64,
|
||||
});
|
||||
|
||||
let result = moa::handle_turn(&config, &body).await;
|
||||
|
||||
assert_eq!(
|
||||
result.turn_kind,
|
||||
moa::TurnKind::ToolResult,
|
||||
"last-msg-role=tool must classify as TurnKind::ToolResult; got {:?}",
|
||||
result.turn_kind
|
||||
);
|
||||
assert!(
|
||||
result.reducer_used,
|
||||
"tool-result turn must invoke the reducer"
|
||||
);
|
||||
// No fanout: only the reducer (a single backend in the candidate
|
||||
// ladder) should have been called.
|
||||
let total_calls = fast.calls() + mid.calls() + strong.calls();
|
||||
assert_eq!(
|
||||
total_calls,
|
||||
1,
|
||||
"tool-result turn must not fan out — expected 1 reducer call, got {total_calls} \
|
||||
(fast={}, mid={}, strong={})",
|
||||
fast.calls(),
|
||||
mid.calls(),
|
||||
strong.calls()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn trailing_user_after_unsynthesised_tool_result_classifies_as_tool_result() {
|
||||
// The bug from the PR review. The harness has appended a `user`
|
||||
// message AFTER an unsynthesised tool result. The last message is
|
||||
// now `user`, not `tool`. The gateway today classifies this as
|
||||
// Continuation and fans out to every worker — wasting a round-trip
|
||||
// per worker and risking duplicate tool calls. It must instead
|
||||
// take the reducer-only path: synthesize the tool result, address
|
||||
// the user nudge, return one coherent response.
|
||||
let (config, fast, mid, strong) = config_with_three_recording_workers();
|
||||
|
||||
let body = json!({
|
||||
"model": "mesh",
|
||||
"tools": read_file_tool(),
|
||||
"messages": [
|
||||
{"role": "user", "content": "Read README.md and tell me what it says."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": null,
|
||||
"tool_calls": [{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "read_file", "arguments": "{\"path\":\"README.md\"}"}
|
||||
}]
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "# Hello World\n"},
|
||||
// Harness leaves the tool result in-context and asks the
|
||||
// model to continue. There is NO new assistant message
|
||||
// synthesizing the tool result yet.
|
||||
{"role": "user", "content": "Go on."},
|
||||
],
|
||||
"max_tokens": 64,
|
||||
});
|
||||
|
||||
let result = moa::handle_turn(&config, &body).await;
|
||||
|
||||
assert_eq!(
|
||||
result.turn_kind,
|
||||
moa::TurnKind::ToolResult,
|
||||
"trailing-user after unsynthesised tool result must classify as \
|
||||
TurnKind::ToolResult to avoid spraying through fanout; got {:?}",
|
||||
result.turn_kind
|
||||
);
|
||||
assert!(
|
||||
result.reducer_used,
|
||||
"unsynthesised-tool-result turn must invoke the reducer"
|
||||
);
|
||||
let total_calls = fast.calls() + mid.calls() + strong.calls();
|
||||
assert_eq!(
|
||||
total_calls,
|
||||
1,
|
||||
"tool-result follow-up must not fan out — expected 1 reducer call, got \
|
||||
{total_calls} (fast={}, mid={}, strong={})",
|
||||
fast.calls(),
|
||||
mid.calls(),
|
||||
strong.calls()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fresh_user_question_still_classifies_as_fan_out() {
|
||||
// Counterpart: a fresh conversation with just a user message must
|
||||
// continue to fan out. We must not over-trigger the tool-result
|
||||
// path and drop fan-out for normal questions.
|
||||
let (config, fast, mid, strong) = config_with_three_recording_workers();
|
||||
|
||||
let body = json!({
|
||||
"model": "mesh",
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is the capital of Japan? One word only."},
|
||||
],
|
||||
"max_tokens": 32,
|
||||
});
|
||||
|
||||
let result = moa::handle_turn(&config, &body).await;
|
||||
|
||||
assert!(
|
||||
matches!(
|
||||
result.turn_kind,
|
||||
moa::TurnKind::Fanout | moa::TurnKind::EarlyExit
|
||||
),
|
||||
"fresh user question must fan out (Fanout or EarlyExit); got {:?}",
|
||||
result.turn_kind
|
||||
);
|
||||
let total_calls = fast.calls() + mid.calls() + strong.calls();
|
||||
assert!(
|
||||
total_calls >= 1,
|
||||
"fresh user question must reach at least one worker; got 0 calls"
|
||||
);
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue