From 187d3271446edb918fb64f566c0312285b325eb0 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 11:26:53 +0200 Subject: [PATCH] Write a partial report instead of erroring out when research hits max turns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the researcher agent exhausts its turn budget mid-investigation, rig now surfaces PromptError::MaxTurnsError with the chat history intact rather than nothing at all. Catch it, reconstruct a findings dump from whatever assistant text and tool results the run produced, and still write a report from that — clearly flagged as incomplete — instead of propagating the error and losing all the work. --- deep_research/src/core.rs | 134 +++++++++++++++++++++++++++++++++----- 1 file changed, 119 insertions(+), 15 deletions(-) diff --git a/deep_research/src/core.rs b/deep_research/src/core.rs index 0a55425..dd177a7 100644 --- a/deep_research/src/core.rs +++ b/deep_research/src/core.rs @@ -4,6 +4,8 @@ use crate::stream::write_text_stream; use crate::tools::{FetchPage, SearchWeb}; use clap::Parser; use rig::client::{AgentClientExt, Nothing}; +use rig::completion::message::{ToolResultContent, UserContent}; +use rig::completion::{AssistantContent, Message, PromptError}; use rig::providers::ollama; use rig::streaming::StreamingPrompt; use std::io::Write; @@ -60,9 +62,21 @@ pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result let mut findings = String::new(); let mut feedback: Option = None; + let mut incomplete = false; for round in 1..=MAX_RESEARCH_ROUNDS { - findings = gather_findings(&client, topic, feedback.as_ref(), round, show_progress).await?; + let gathered = + gather_findings(&client, topic, feedback.as_ref(), round, show_progress).await?; + findings = gathered.findings; + incomplete = gathered.incomplete; + + // The researcher ran out of turns mid-investigation rather than + // concluding on its own — another round would just repeat the same + // dead end, so stop and write up whatever was gathered. + if incomplete { + tracing::info!(round, "researcher exhausted its turn budget"); + break; + } let review = review::review_findings(&client, topic, &findings, show_progress).await?; let approved = review.approved; @@ -80,7 +94,61 @@ pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result feedback = Some(review); } - write_report(&client, topic, &findings, show_progress).await + write_report(&client, topic, &findings, incomplete, show_progress).await +} + +/// Findings gathered by a research pass, and whether the researcher was cut +/// off by the turn budget before it could conclude on its own. +struct GatheredFindings { + findings: String, + incomplete: bool, +} + +/// Best-effort reconstruction of research material from a chat history left +/// behind when the researcher hit its turn budget mid-investigation: whatever +/// prose the model wrote between tool calls, plus the raw text of every tool +/// result (search snippets, fetched pages), truncated so one huge page can't +/// crowd out everything else that was found. +fn partial_findings_from_history(chat_history: &[Message]) -> String { + const MAX_TOOL_RESULT_CHARS: usize = 2000; + + let mut sections = Vec::new(); + + for message in chat_history { + match message { + Message::Assistant { content, .. } => { + for item in content.iter() { + if let AssistantContent::Text(text) = item { + sections.push(text.text().to_string()); + } + } + } + Message::User { content } => { + for item in content.iter() { + if let UserContent::ToolResult(result) = item { + for part in result.content.iter() { + if let ToolResultContent::Text(text) = part { + let mut snippet = text.text().to_string(); + if snippet.len() > MAX_TOOL_RESULT_CHARS { + snippet.truncate(MAX_TOOL_RESULT_CHARS); + snippet.push_str(" [...truncated]"); + } + sections.push(snippet); + } + } + } + } + } + Message::System { .. } => {} + } + } + + if sections.is_empty() { + "The researcher exhausted its turn budget before gathering any usable evidence." + .to_string() + } else { + sections.join("\n\n") + } } /// Wraps the tool-calling research loop in its own span so it's visible as a @@ -93,7 +161,7 @@ async fn gather_findings( feedback: Option<&Review>, round: usize, show_progress: bool, -) -> anyhow::Result { +) -> anyhow::Result { let current_date = chrono::offset::Local::now().to_string(); let researcher = client @@ -136,17 +204,24 @@ async fn gather_findings( }; let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching...")); - let findings = researcher - .runner(task) - .max_turns(MAX_RESEARCH_TURNS) - .run() - .await? - .output; + let run_result = researcher.runner(task).max_turns(MAX_RESEARCH_TURNS).run().await; drop(spinner); - tracing::info!(round, findings = %findings, "research phase complete"); + let gathered = match run_result { + Ok(response) => GatheredFindings { findings: response.output, incomplete: false }, + Err(PromptError::MaxTurnsError { chat_history, .. }) => { + tracing::warn!(round, "researcher hit its turn budget before concluding"); + GatheredFindings { + findings: partial_findings_from_history(&chat_history), + incomplete: true, + } + } + Err(err) => return Err(err.into()), + }; - Ok(findings) + tracing::info!(round, findings = %gathered.findings, incomplete = gathered.incomplete, "research phase complete"); + + Ok(gathered) } #[tracing::instrument(skip(client, findings), fields(gen_ai.agent.name = "writer"))] @@ -154,6 +229,7 @@ async fn write_report( client: &ollama::Client, topic: &str, findings: &str, + incomplete: bool, show_progress: bool, ) -> anyhow::Result { let writer = client @@ -171,10 +247,22 @@ async fn write_report( ) .build(); + let task = if incomplete { + format!( + "Topic: {topic}\n\n\ + The researcher ran out of turns before it could finish investigating, so these notes \ + are incomplete and may be a raw, unsorted mix of tool output rather than clean \ + footnoted findings. Write up what the notes actually support, note explicitly where \ + evidence is thin or missing, and open the report with a short callout that this \ + research is inconclusive and further manual searching is needed.\n\n\ + Research notes:\n{findings}" + ) + } else { + format!("Topic: {topic}\n\nResearch notes:\n{findings}") + }; + let spinner = Spinner::start(show_progress, format!("{REPORT_EMOJI} Writing report...")); - let response_stream = writer - .stream_prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}")) - .await; + let response_stream = writer.stream_prompt(task).await; // Locked once for the whole stream rather than per chunk (as print! // would do internally) — chunks arrive in a tight loop, so re-acquiring @@ -183,8 +271,24 @@ async fn write_report( // the gap between sending the prompt and generation actually starting. let stdout = std::io::stdout(); let mut handle = stdout.lock(); + + // Guaranteed regardless of whether the writer model actually heeds the + // instruction above — the reader should never mistake a budget-exhausted + // run for a complete one just because the model forgot to say so. + let disclaimer = incomplete.then_some( + "> **Disclaimer:** this research did not complete — the researcher exhausted its turn \ + budget before it could finish investigating. The findings below are partial and may be \ + incomplete or unbalanced; treat them as a starting point and verify further manually.\n\n", + ); + if let Some(disclaimer) = disclaimer { + write!(handle, "{disclaimer}")?; + } + let report = write_text_stream(response_stream, &mut handle, spinner).await?; writeln!(handle)?; - Ok(report) + Ok(match disclaimer { + Some(disclaimer) => format!("{disclaimer}{report}"), + None => report, + }) }