goose_agent/

inference.rs

//! Provider inference operation for the unrolled agent loop.

use std::collections::HashSet;
use std::sync::{Arc, Mutex, MutexGuard};

use anyhow::Result;
use async_trait::async_trait;
use futures::StreamExt;
use goose_provider_types::base::Provider;
use goose_provider_types::conversation::message::{InferenceMetadata, Message, MessageContent};
use goose_provider_types::conversation::token_usage::ProviderUsage;
use goose_provider_types::conversation::{
    effective_role, fix_conversation, merge_consecutive_messages_for_request, Conversation,
    EffectiveRole,
};
use goose_provider_types::errors::ProviderError;
use goose_provider_types::model::ModelConfig;
use tracing_futures::Instrument;

use crate::machine::MachineSession;
use crate::operation::{
    applied, messages_since_kickoff, not_applicable, trailing_error, Emitter, Inference,
    InferenceInput, Operation, OperationResult, CLIENT_LOG,
};
use goose_provider_types::maybe_send::{MaybeSend, MaybeSync};

pub struct PreparedInferenceRequest {
    pub system_prompt: String,
    pub tools: Vec<rmcp::model::Tool>,
    pub additional_messages: Vec<Message>,
}

#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
pub trait InferenceRequestPreparer<S>: MaybeSend + MaybeSync {
    async fn prepare_session(&self, _session: &S) -> Result<Option<S>> {
        Ok(None)
    }

    async fn prepare(
        &self,
        session: &S,
        conversation: &Conversation,
        input: InferenceInput,
    ) -> Result<PreparedInferenceRequest>;
}

pub struct IdentityInferenceRequestPreparer;

#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
impl<S: MaybeSync> InferenceRequestPreparer<S> for IdentityInferenceRequestPreparer {
    async fn prepare(
        &self,
        _session: &S,
        _conversation: &Conversation,
        input: InferenceInput,
    ) -> Result<PreparedInferenceRequest> {
        Ok(PreparedInferenceRequest {
            system_prompt: input
                .prompt_parts
                .into_iter()
                .map(|(_, part)| part)
                .collect::<Vec<_>>()
                .join("\n\n"),
            tools: input.tools,
            additional_messages: Vec::new(),
        })
    }
}

pub trait InferenceEffect: From<Message> + MaybeSend + 'static {
    fn record_usage(usage: ProviderUsage) -> Self;
}

pub const EMPTY_RESPONSE_MESSAGE: &str =
    "The model returned an empty response. Please resend your message to continue.";
const MAX_EMPTY_RESPONSE_RETRIES: usize = 3;
const EMPTY_RESPONSE_NOTE_SCOPE: &str = "inference";
const EMPTY_RESPONSE_NOTE: &str = "empty_response";

/// The model's response stayed empty after retries. The fallback message is stored
/// hidden so that operations owning the end of a turn (recipe retries, final output)
/// can take over; when none does, it is revealed to the user.
pub fn is_empty_response_marker(message: &Message) -> bool {
    message
        .metadata
        .operation_note(EMPTY_RESPONSE_NOTE_SCOPE, EMPTY_RESPONSE_NOTE)
        .is_some()
}

fn is_thinking(content: &MessageContent) -> bool {
    matches!(
        content,
        MessageContent::Thinking(_) | MessageContent::RedactedThinking(_)
    )
}

fn drop_repeated_tool_call_thinking(accumulator: &Conversation, chunk: &mut Message) {
    if !chunk
        .content
        .iter()
        .any(|content| matches!(content, MessageContent::ToolRequest(_)))
    {
        return;
    }
    let prior: Vec<&MessageContent> = accumulator
        .iter()
        .filter(|message| message.role == chunk.role)
        .flat_map(|message| message.content.iter())
        .filter(|content| is_thinking(content))
        .collect();
    chunk
        .content
        .retain(|content| !(is_thinking(content) && prior.contains(&content)));
}

fn pending_operation_logs(messages: &[Message]) -> Vec<String> {
    let mut seen = messages
        .iter()
        .flat_map(|message| message.metadata.operation_logs.iter().cloned())
        .collect::<HashSet<_>>();
    let mut logs = Vec::new();

    for (operation, line) in messages
        .iter()
        .filter_map(|message| message.metadata.operations.as_deref())
        .flat_map(|operations| operations.iter())
        .filter_map(|(operation, notes)| {
            notes
                .get(CLIENT_LOG)
                .and_then(serde_json::Value::as_str)
                .map(|line| (operation, line))
        })
    {
        let log = format!("ops_{operation}: {line}");
        if seen.insert(log.clone()) {
            logs.push(log);
        }
    }

    logs
}

fn attach_operation_logs(message: &mut Message, logs: &mut Vec<String>) {
    if message.role == rmcp::model::Role::Assistant && !logs.is_empty() {
        message.metadata.operation_logs = std::mem::take(logs);
    }
}

pub fn chat_span(
    provider: &dyn Provider,
    model_config: &ModelConfig,
    session_id: &str,
    purpose: &'static str,
) -> tracing::Span {
    let span = tracing::info_span!(
        target: "goose::state_machine",
        "chat",
        "gen_ai.operation.name" = "chat",
        "gen_ai.provider.name" = %provider.get_name(),
        "gen_ai.request.model" = %model_config.model_name,
        "gen_ai.request.temperature" = tracing::field::Empty,
        "gen_ai.request.max_tokens" = tracing::field::Empty,
        "gen_ai.response.model" = tracing::field::Empty,
        "gen_ai.response.finish_reasons" = tracing::field::Empty,
        "gen_ai.response.id" = tracing::field::Empty,
        "gen_ai.request.id" = tracing::field::Empty,
        "gen_ai.usage.input_tokens" = tracing::field::Empty,
        "gen_ai.usage.output_tokens" = tracing::field::Empty,
        "goose.chat.purpose" = purpose,
        "error.type" = tracing::field::Empty,
        session.id = %session_id,
    );
    record_request_params(&span, model_config);
    span
}

fn is_empty_response(message: &Message) -> bool {
    message.content.iter().all(|content| match content {
        MessageContent::Text(text) => text.text.trim().is_empty(),
        MessageContent::Thinking(thinking) => {
            thinking.thinking.trim().is_empty() && thinking.signature.is_empty()
        }
        _ => false,
    })
}

pub fn ends_with_successful_tool_response(messages: &[Message]) -> bool {
    let Some(message) = messages.last() else {
        return false;
    };
    let mut responses = message
        .content
        .iter()
        .filter_map(MessageContent::as_tool_response)
        .peekable();
    responses.peek().is_some()
        && responses.all(|response| {
            response
                .tool_result
                .as_ref()
                .is_ok_and(|result| !result.is_error.unwrap_or(false))
        })
}

fn record_request_params(span: &tracing::Span, model_config: &ModelConfig) {
    if let Some(temperature) = model_config.temperature {
        span.record("gen_ai.request.temperature", temperature as f64);
    }
    if let Some(max_tokens) = model_config.max_tokens {
        span.record("gen_ai.request.max_tokens", max_tokens as i64);
    }
}

pub fn record_chat_usage(span: &tracing::Span, usage: &ProviderUsage) {
    span.record("gen_ai.response.model", usage.model.as_str());
    if let Some(tokens) = usage.usage.input_tokens {
        span.record("gen_ai.usage.input_tokens", tokens);
    }
    if let Some(tokens) = usage.usage.output_tokens {
        span.record("gen_ai.usage.output_tokens", tokens);
    }
    if let Some(tokens) = usage.usage.cache_read_input_tokens {
        span.record("gen_ai.usage.cache_read.input_tokens", tokens);
    }
    if let Some(tokens) = usage.usage.cache_write_input_tokens {
        span.record("gen_ai.usage.cache_creation.input_tokens", tokens);
    }
    if let Some(reasons) = &usage.finish_reasons {
        let reasons_json = serde_json::to_string(reasons).unwrap_or_default();
        span.record("gen_ai.response.finish_reasons", reasons_json.as_str());
    }
    if let Some(id) = &usage.response_id {
        span.record("gen_ai.response.id", id.as_str());
    }
    if let Some(id) = &usage.request_id {
        span.record("gen_ai.request.id", id.as_str());
    }
}

#[derive(Default)]
struct InferenceOutput {
    accumulator: Conversation,
    additional_messages: Vec<Message>,
    usage: Vec<ProviderUsage>,
}

pub struct InferenceRunner<'a, S, E> {
    provider: Arc<dyn Provider>,
    model_config: ModelConfig,
    request_preparer: Arc<dyn InferenceRequestPreparer<S> + 'a>,
    output: Mutex<InferenceOutput>,
    effect: std::marker::PhantomData<fn() -> E>,
}

/// The agent-visible conversation as the provider sees it: tool requests left
/// unanswered by an earlier turn are dropped, since nothing will answer them now.
fn messages_for_provider(
    conversation: &Conversation,
    turn: &[Message],
    keep_empty_messages: bool,
) -> Vec<Message> {
    let answered: HashSet<&str> = conversation
        .messages()
        .iter()
        .flat_map(|message| message.get_tool_response_ids())
        .collect();
    let start = conversation.len() - turn.len();
    conversation
        .messages()
        .iter()
        .enumerate()
        .filter(|(_, message)| message.is_agent_visible())
        .map(|(index, message)| {
            let mut message = message.agent_visible_content();
            if index < start {
                message.content.retain(|content| match content {
                    MessageContent::ToolRequest(request) => answered.contains(request.id.as_str()),
                    _ => true,
                });
            }
            message
        })
        .filter(|message| keep_empty_messages || !message.content.is_empty())
        .collect()
}

fn latest_provider_session_id<'a>(
    conversation: &'a Conversation,
    provider: &str,
) -> Option<&'a str> {
    conversation
        .messages()
        .iter()
        .rev()
        .find_map(|message| message.metadata.inference.as_ref())
        .filter(|inference| inference.provider == provider)
        .and_then(|inference| inference.provider_session_id.as_deref())
}

fn ends_with_provider_turn(messages: &[Message]) -> bool {
    messages.last().is_some_and(|message| {
        matches!(
            effective_role(message),
            EffectiveRole::User | EffectiveRole::Tool
        )
    })
}

fn should_infer(conversation: &Conversation, turn: &[Message]) -> bool {
    let projected = messages_for_provider(conversation, turn, true);
    if projected
        .last()
        .is_some_and(|message| message.content.is_empty())
    {
        return false;
    }
    ends_with_provider_turn(&messages_for_provider(conversation, turn, false))
}

fn inference_span(provider: &dyn Provider, model_config: &ModelConfig) -> tracing::Span {
    let span = tracing::info_span!(
        target: "goose::state_machine",
        "chat",
        "gen_ai.operation.name" = "chat",
        "gen_ai.provider.name" = %provider.get_name(),
        "gen_ai.request.model" = %model_config.model_name,
        "gen_ai.request.temperature" = tracing::field::Empty,
        "gen_ai.request.max_tokens" = tracing::field::Empty,
        "gen_ai.response.model" = tracing::field::Empty,
        "gen_ai.response.finish_reasons" = tracing::field::Empty,
        "gen_ai.response.id" = tracing::field::Empty,
        "gen_ai.request.id" = tracing::field::Empty,
        "gen_ai.usage.input_tokens" = tracing::field::Empty,
        "gen_ai.usage.output_tokens" = tracing::field::Empty,
        "error.type" = tracing::field::Empty,
    );
    record_request_params(&span, model_config);
    span
}

impl<'a, S: MaybeSync, E: InferenceEffect> InferenceRunner<'a, S, E> {
    pub fn new(provider: Arc<dyn Provider>, model_config: ModelConfig) -> Self {
        Self {
            provider,
            model_config,
            request_preparer: Arc::new(IdentityInferenceRequestPreparer),
            output: Mutex::new(InferenceOutput::default()),
            effect: std::marker::PhantomData,
        }
    }

    pub fn with_request_preparer(
        mut self,
        request_preparer: Arc<dyn InferenceRequestPreparer<S> + 'a>,
    ) -> Self {
        self.request_preparer = request_preparer;
        self
    }

    fn output(&self) -> MutexGuard<'_, InferenceOutput> {
        self.output.lock().unwrap()
    }

    fn take_output(&self) -> Vec<E> {
        let output = std::mem::take(&mut *self.output());
        let mut effects: Vec<E> = output
            .additional_messages
            .into_iter()
            .map(E::from)
            .collect();
        effects.extend(output.usage.into_iter().map(E::record_usage));
        effects.extend(output.accumulator.into_iter().map(E::from));
        effects
    }

    fn emit_message(&self, message: Message, emit: &Emitter) {
        let message = message.with_generated_id_if_missing();
        self.output().accumulator.push(message.clone());
        emit.message(message);
    }

    fn error_outcome(
        &self,
        err: &ProviderError,
        operation_logs: &mut Vec<String>,
        emit: &Emitter,
    ) -> Result<OperationResult<E>> {
        tracing::Span::current().record("error.type", err.telemetry_type());
        tracing::error!("LLM provider error: {err}");
        let mut message = Message::from_provider_error(err);
        attach_operation_logs(&mut message, operation_logs);
        self.emit_message(message, emit);
        applied(self.take_output())
    }
}

#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
impl<S: MaybeSync, E: InferenceEffect> Operation<S, E> for InferenceRunner<'_, S, E> {
    fn name(&self) -> &'static str {
        "llm"
    }

    async fn finalize_cancellation(
        &self,
        _session: &S,
        _conversation: &Conversation,
        _emit: &Emitter,
    ) -> Vec<E> {
        self.take_output()
    }
}

#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
impl<S: MachineSession, E: InferenceEffect> Inference<S, E> for InferenceRunner<'_, S, E> {
    fn applies(&self, conversation: &Conversation) -> bool {
        let Ok(turn) = messages_since_kickoff(conversation) else {
            return false;
        };
        trailing_error(conversation).is_none() && should_infer(conversation, turn)
    }

    async fn prepare_session(&self, session: &S) -> Result<Option<S>> {
        self.request_preparer.prepare_session(session).await
    }

    async fn infer(
        &self,
        session: &S,
        conversation: &Conversation,
        input: InferenceInput,
        emit: &Emitter,
    ) -> Result<OperationResult<E>> {
        let messages = messages_since_kickoff(conversation)?;
        if trailing_error(conversation).is_some() {
            return not_applicable();
        }

        if !should_infer(conversation, messages) {
            return not_applicable();
        }
        let mut operation_logs = pending_operation_logs(messages);
        let mut messages_for_provider = messages_for_provider(conversation, messages, false);

        let model_config = session
            .thinking_effort()
            .map(|effort| self.model_config.clone().with_thinking_effort(effort))
            .unwrap_or_else(|| self.model_config.clone());
        let span = inference_span(self.provider.as_ref(), &model_config);

        async {
            let PreparedInferenceRequest {
                system_prompt,
                tools,
                additional_messages,
            } = self
                .request_preparer
                .prepare(session, conversation, input)
                .await?;

            for message in &additional_messages {
                messages_for_provider.push(message.clone());
            }
            self.output().additional_messages = additional_messages;

            let provider_name = self.provider.get_name();
            if let Some(session_id) = latest_provider_session_id(conversation, provider_name) {
                if let Err(error) = self.provider.resume(session_id).await {
                    tracing::warn!(
                        provider = provider_name,
                        %error,
                        "Could not resume provider session; continuing with a handoff"
                    );
                }
            }

            let projected =
                Conversation::new_unvalidated(messages_for_provider).agent_visible_messages();
            let (fixed, _) = fix_conversation(Conversation::new_unvalidated(projected));
            let conversation_for_provider = Conversation::new_unvalidated(
                merge_consecutive_messages_for_request(fixed.messages().clone()),
            );
            let successful_tool_response =
                ends_with_successful_tool_response(conversation.messages());
            let mut empty_responses = 0;
            let empty_output = loop {
            let attempt_start = self.output().usage.len();
            let stream = self
                .provider
                .stream(
                    &model_config,
                    &system_prompt,
                    conversation_for_provider.messages(),
                    &tools,
                )
                .await;

            let mut stream = match stream {
                Ok(stream) => stream,
                Err(err) => return self.error_outcome(&err, &mut operation_logs, emit),
            };

            let requested_model = model_config.model_name.clone();
            let resolved_model = self
                .provider
                .fetch_model_info(&requested_model)
                .await
                .ok()
                .and_then(|model_info| model_info.resolved_model);
            let provider_session_id = self.provider.provider_session_id();
            let inference = Some(InferenceMetadata {
                provider: self.provider.get_name().to_string(),
                requested_model,
                resolved_model,
                provider_session_id,
            });

            let mut tool_request_ids = HashSet::new();
            while let Some(result) = stream.next().await {
                let (msg_opt, usage_opt) = match result {
                    Ok(chunk) => chunk,
                    Err(err) => return self.error_outcome(&err, &mut operation_logs, emit),
                };
                if let Some(usage) = usage_opt {
                    record_chat_usage(&tracing::Span::current(), &usage);
                    let mut output = self.output();
                    output.usage.truncate(attempt_start);
                    output.usage.push(usage);
                }
                if let Some(mut chunk) = msg_opt {
                    if let Some(inference) = &inference {
                        chunk = chunk.with_inference_if_assistant(inference.clone());
                    }
                    chunk.content.retain(|content| match content {
                        MessageContent::ToolRequest(request) => {
                            tool_request_ids.insert(request.id.clone())
                        }
                        _ => true,
                    });
                    drop_repeated_tool_call_thinking(&self.output().accumulator, &mut chunk);
                    if chunk.content.is_empty() && !chunk.metadata.output_token_limit_reached {
                        self.output().accumulator.push(chunk);
                    } else {
                        attach_operation_logs(&mut chunk, &mut operation_logs);
                        self.emit_message(chunk, emit);
                    }
                }
            }

            let empty_output = {
                let output = self.output();
                !output
                    .accumulator
                    .iter()
                    .any(|message| message.metadata.output_token_limit_reached)
                    && output.accumulator.iter().all(is_empty_response)
            };
            if !empty_output || successful_tool_response || emit.cancel_token().is_cancelled() {
                break empty_output;
            }
            self.output().accumulator.clear();
            if empty_responses < MAX_EMPTY_RESPONSE_RETRIES {
                empty_responses += 1;
                tracing::warn!(
                    "Provider returned an empty response; retrying ({empty_responses}/{MAX_EMPTY_RESPONSE_RETRIES})"
                );
                continue;
            }
            let mut marker = Message::assistant()
                .with_text(EMPTY_RESPONSE_MESSAGE)
                .with_visibility(false, false);
            attach_operation_logs(&mut marker, &mut operation_logs);
            marker.metadata.set_operation_note(
                EMPTY_RESPONSE_NOTE_SCOPE,
                EMPTY_RESPONSE_NOTE,
                serde_json::Value::Bool(true),
            );
            self.output().accumulator.push(marker);
            return applied(self.take_output());
            };

            if empty_output && successful_tool_response {
                let message = {
                    let mut output = self.output();
                    let mut message = output
                        .accumulator
                        .last()
                        .cloned()
                        .unwrap_or_else(Message::assistant);
                    message.content.clear();
                    message.metadata.user_visible = false;
                    message.metadata.agent_visible = true;
                    attach_operation_logs(&mut message, &mut operation_logs);
                    output.accumulator.clear();
                    message
                };
                self.emit_message(message, emit);
            }
            applied(self.take_output())
        }
        .instrument(span)
        .await
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn provider_session_id_comes_only_from_latest_inference() {
        let conversation = Conversation::new_unvalidated([
            Message::assistant().with_inference(InferenceMetadata {
                provider: "provider-a".to_string(),
                requested_model: "model".to_string(),
                resolved_model: None,
                provider_session_id: Some("session-a".to_string()),
            }),
            Message::assistant().with_inference(InferenceMetadata {
                provider: "provider-b".to_string(),
                requested_model: "model".to_string(),
                resolved_model: None,
                provider_session_id: Some("session-b".to_string()),
            }),
        ]);

        assert_eq!(
            latest_provider_session_id(&conversation, "provider-b"),
            Some("session-b")
        );
        assert_eq!(
            latest_provider_session_id(&conversation, "provider-a"),
            None
        );
    }

    #[test]
    fn signed_thinking_without_text_is_not_an_empty_response() {
        assert!(is_empty_response(
            &Message::assistant().with_content(MessageContent::thinking("", ""))
        ));
        assert!(!is_empty_response(
            &Message::assistant().with_content(MessageContent::thinking("", "sig-omitted"))
        ));
    }

    #[test]
    fn whitespace_only_text_is_an_empty_response() {
        assert!(is_empty_response(&Message::assistant().with_text(" \n\t ")));
    }
}