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<Review> = 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<String> {
+) -> anyhow::Result<GatheredFindings> {
     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
+    let run_result = researcher.runner(task).max_turns(MAX_RESEARCH_TURNS).run().await;
-        .runner(task)
-        .max_turns(MAX_RESEARCH_TURNS)
-        .run()
-        .await?
-        .output;
     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<String> {
     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
+    let response_stream = writer.stream_prompt(task).await;
-        .stream_prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}"))
-        .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,
+    })
 }