llm_pipeline/

tool_loop.rs

use crate::{
    client::LlmConfig, exec_ctx::ExecCtx, streaming::StreamingDecoder, PipelineError, Result,
};
use futures::StreamExt;
use llm_tool_runtime::{
    render_ollama_tool, render_openai_tool, ApprovalGrant, ToolBudgetContext, ToolCall, ToolCtx,
    ToolDescriptor, ToolError, ToolExecutionPermit, ToolExposureDecision, ToolExposureRequest,
    ToolOriginKind, ToolPlannerStage, ToolReceiptSink, ToolRegistry, ToolResult, ToolRetryOwner,
    ToolRuntime,
};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use stack_ids::{AttemptId, ScopeKey, TraceCtx, TrialId};
use std::collections::BTreeMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

type JsonFetchFuture<'a> = Pin<Box<dyn Future<Output = Result<Value>> + 'a>>;
type JsonFetcher<'a> = dyn FnMut(Value) -> JsonFetchFuture<'a> + 'a;

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolLoopChoice {
    Auto,
    None,
    Required { tool_name: String },
}

#[derive(Debug, Clone)]
pub struct ToolLoopRequest {
    pub model: String,
    pub user_input: String,
    pub instructions: Option<String>,
    pub config: LlmConfig,
    pub attempt_id: Option<AttemptId>,
    pub planner_stage: ToolPlannerStage,
    pub retry_owner: ToolRetryOwner,
    pub allowed_tools: Option<Vec<String>>,
    pub tool_choice: ToolLoopChoice,
    pub max_round_trips: u32,
    pub strict: bool,
    pub parallel_tool_calls: bool,
    pub dry_run: bool,
    pub scope: Option<ScopeKey>,
    pub approval_grant: Option<ApprovalGrant>,
    pub execution_permit: Option<std::sync::Arc<ToolExecutionPermit>>,
    pub stream: bool,
    pub api_key: Option<String>,
    pub organization: Option<String>,
}

impl ToolLoopRequest {
    pub fn new(model: impl Into<String>, user_input: impl Into<String>) -> Self {
        Self {
            model: model.into(),
            user_input: user_input.into(),
            instructions: None,
            config: LlmConfig::default(),
            attempt_id: None,
            planner_stage: ToolPlannerStage::Execution,
            retry_owner: ToolRetryOwner::LlmPipeline,
            allowed_tools: None,
            tool_choice: ToolLoopChoice::Auto,
            max_round_trips: 8,
            strict: true,
            parallel_tool_calls: false,
            dry_run: false,
            scope: None,
            approval_grant: None,
            execution_permit: None,
            stream: false,
            api_key: None,
            organization: None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolInvocation {
    pub call: ToolCall,
    pub outcome: std::result::Result<ToolResult, ToolError>,
    pub receipt: llm_tool_runtime::ToolReceipt,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolLoopResponse {
    pub final_text: String,
    pub rounds: u32,
    pub trace_ctx: TraceCtx,
    pub invocations: Vec<ToolInvocation>,
    pub exposure_decisions: Vec<ToolExposureDecision>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub provider_response_id: Option<String>,
}

pub struct ToolLoopRunner {
    runtime: Arc<ToolRuntime>,
}

struct PreparedToolLoop {
    descriptors_by_name: BTreeMap<String, ToolDescriptor>,
    tools: Vec<Value>,
    exposure_decisions: Vec<ToolExposureDecision>,
}

impl ToolLoopRunner {
    pub fn new(runtime: Arc<ToolRuntime>) -> Self {
        Self { runtime }
    }

    pub fn from_registry(registry: ToolRegistry) -> Self {
        Self::new(Arc::new(ToolRuntime::new(registry)))
    }

    pub fn from_registry_with_receipt_sink(
        registry: ToolRegistry,
        receipt_sink: Arc<dyn ToolReceiptSink>,
    ) -> Self {
        Self::new(Arc::new(
            ToolRuntime::new(registry).with_receipt_sink(receipt_sink),
        ))
    }

    pub fn runtime(&self) -> &Arc<ToolRuntime> {
        &self.runtime
    }

    fn prepare_tool_loop<F>(
        &self,
        request: &ToolLoopRequest,
        mut render_tool: F,
    ) -> Result<PreparedToolLoop>
    where
        F: FnMut(&ToolDescriptor) -> std::result::Result<Value, ToolError>,
    {
        let exposure = self.runtime.registry().plan_exposure(&ToolExposureRequest {
            allowed_names: request.allowed_tools.clone(),
            planner_stage: request.planner_stage.clone(),
            include_hidden: false,
            max_tools: None,
        });
        let descriptors = exposure
            .tools
            .iter()
            .map(|tool| tool.descriptor().clone())
            .collect::<Vec<_>>();
        let tools = descriptors
            .iter()
            .map(&mut render_tool)
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(tool_error_to_pipeline)?;
        let descriptors_by_name = descriptors
            .into_iter()
            .map(|descriptor| (descriptor.name.clone(), descriptor))
            .collect::<BTreeMap<_, _>>();

        Ok(PreparedToolLoop {
            descriptors_by_name,
            tools,
            exposure_decisions: exposure.decisions,
        })
    }

    async fn execute_tool_call(
        &self,
        ctx: &ExecCtx,
        request: &ToolLoopRequest,
        attempt_id: &AttemptId,
        call: ToolCall,
        invocations: &mut Vec<ToolInvocation>,
    ) -> String {
        let tool_ctx = build_tool_ctx(ctx, request, attempt_id);
        let execution = self
            .runtime
            .execute(
                &tool_ctx,
                &call,
                request.execution_permit.clone(),
                ctx.cancel_flag(),
            )
            .await;
        let tool_output = execution
            .result
            .as_ref()
            .map(ToolResult::to_model_output)
            .unwrap_or_else(|error| {
                json!({
                    "error_class": error.class,
                    "message": error.message,
                })
                .to_string()
            });
        invocations.push(ToolInvocation {
            call,
            outcome: execution.result,
            receipt: execution.receipt,
        });
        tool_output
    }

    async fn execute_openai_function_calls(
        &self,
        ctx: &ExecCtx,
        request: &ToolLoopRequest,
        attempt_id: &AttemptId,
        descriptors_by_name: &BTreeMap<String, ToolDescriptor>,
        function_calls: Vec<OpenAiFunctionCall>,
        invocations: &mut Vec<ToolInvocation>,
    ) -> Result<Vec<Value>> {
        let mut outputs = Vec::with_capacity(function_calls.len());
        for function_call in function_calls {
            let descriptor = descriptors_by_name
                .get(&function_call.name)
                .ok_or_else(|| {
                    PipelineError::Other(format!(
                        "provider requested unknown or unexposed tool {}",
                        function_call.name
                    ))
                })?;
            let call = ToolCall {
                descriptor_name: descriptor.name.clone(),
                descriptor_version: descriptor.version.clone(),
                arguments: function_call.arguments.clone(),
                origin_kind: ToolOriginKind::OpenAiResponses,
                provider_call_id: Some(function_call.call_id.clone()),
                tool_run_id: uuid::Uuid::new_v4().to_string(),
            };
            let tool_output = self
                .execute_tool_call(ctx, request, attempt_id, call, invocations)
                .await;
            outputs.push(json!({
                "type": "function_call_output",
                "call_id": function_call.call_id,
                "output": tool_output,
            }));
        }
        Ok(outputs)
    }

    async fn execute_ollama_tool_calls(
        &self,
        ctx: &ExecCtx,
        request: &ToolLoopRequest,
        attempt_id: &AttemptId,
        descriptors_by_name: &BTreeMap<String, ToolDescriptor>,
        tool_calls: Vec<OllamaFunctionCall>,
        invocations: &mut Vec<ToolInvocation>,
    ) -> Result<Vec<Value>> {
        let mut outputs = Vec::with_capacity(tool_calls.len());
        for tool_call in tool_calls {
            let descriptor = descriptors_by_name.get(&tool_call.name).ok_or_else(|| {
                PipelineError::Other(format!(
                    "provider requested unknown or unexposed tool {}",
                    tool_call.name
                ))
            })?;
            let call = ToolCall {
                descriptor_name: descriptor.name.clone(),
                descriptor_version: descriptor.version.clone(),
                arguments: tool_call.arguments.clone(),
                origin_kind: ToolOriginKind::OllamaChat,
                provider_call_id: None,
                tool_run_id: uuid::Uuid::new_v4().to_string(),
            };
            let tool_output = self
                .execute_tool_call(ctx, request, attempt_id, call, invocations)
                .await;
            outputs.push(json!({
                "role": "tool",
                "name": tool_call.name,
                "content": tool_output,
            }));
        }
        Ok(outputs)
    }

    pub async fn run_openai_responses(
        &self,
        ctx: &ExecCtx,
        request: ToolLoopRequest,
    ) -> Result<ToolLoopResponse> {
        let api_key = request.api_key.clone();
        let organization = request.organization.clone();
        let mut fetch = |body: Value| -> JsonFetchFuture<'_> {
            let api_key = api_key.clone();
            let organization = organization.clone();
            Box::pin(async move {
                post_json_with_backoff(
                    ctx,
                    "/v1/responses",
                    &body,
                    api_key.as_deref(),
                    organization.as_deref(),
                )
                .await
            })
        };
        self.run_openai_responses_with_fetcher(ctx, request, &mut fetch)
            .await
    }

    async fn run_openai_responses_with_fetcher(
        &self,
        ctx: &ExecCtx,
        request: ToolLoopRequest,
        fetch: &mut JsonFetcher<'_>,
    ) -> Result<ToolLoopResponse> {
        let prepared = self.prepare_tool_loop(&request, |descriptor| {
            render_openai_tool(descriptor, request.strict)
        })?;

        let mut input_items = vec![json!({
            "role": "user",
            "content": [{"type": "input_text", "text": request.user_input}],
        })];
        let mut invocations = Vec::new();
        let attempt_id = request
            .attempt_id
            .clone()
            .unwrap_or_else(AttemptId::generate);

        for round in 0..request.max_round_trips {
            ctx.check_cancelled()?;
            let response =
                fetch(openai_request_body(&request, &input_items, &prepared.tools)).await?;

            let provider_response_id = response
                .get("id")
                .and_then(|value| value.as_str())
                .map(|value| value.to_string());

            let output_items = response
                .get("output")
                .and_then(|value| value.as_array())
                .cloned()
                .unwrap_or_default();
            let function_calls = parse_openai_function_calls(&output_items)?;

            if function_calls.is_empty() {
                return Ok(ToolLoopResponse {
                    final_text: openai_output_text(&response),
                    rounds: round + 1,
                    trace_ctx: ctx.trace_ctx.clone(),
                    invocations,
                    exposure_decisions: prepared.exposure_decisions.clone(),
                    provider_response_id,
                });
            }

            if function_calls.len() > 1 && !request.parallel_tool_calls {
                return Err(PipelineError::InvalidConfig(
                    "parallel tool calls are disabled by default".into(),
                ));
            }

            input_items.extend(output_items);
            input_items.extend(
                self.execute_openai_function_calls(
                    ctx,
                    &request,
                    &attempt_id,
                    &prepared.descriptors_by_name,
                    function_calls,
                    &mut invocations,
                )
                .await?,
            );
        }

        Err(PipelineError::Other(format!(
            "tool loop exceeded max_round_trips={}",
            request.max_round_trips
        )))
    }

    pub async fn run_ollama(
        &self,
        ctx: &ExecCtx,
        request: ToolLoopRequest,
    ) -> Result<ToolLoopResponse> {
        if request.stream {
            return self.run_ollama_streaming(ctx, request).await;
        }

        let mut fetch = |body: Value| -> JsonFetchFuture<'_> {
            Box::pin(
                async move { post_json_with_backoff(ctx, "/api/chat", &body, None, None).await },
            )
        };
        self.run_ollama_with_fetcher(ctx, request, &mut fetch).await
    }

    async fn run_ollama_with_fetcher(
        &self,
        ctx: &ExecCtx,
        request: ToolLoopRequest,
        fetch: &mut JsonFetcher<'_>,
    ) -> Result<ToolLoopResponse> {
        let prepared = self.prepare_tool_loop(&request, render_ollama_tool)?;

        let mut messages = Vec::new();
        if let Some(instructions) = &request.instructions {
            messages.push(json!({"role": "system", "content": instructions}));
        }
        messages.push(json!({"role": "user", "content": request.user_input}));

        let mut invocations = Vec::new();
        let attempt_id = request
            .attempt_id
            .clone()
            .unwrap_or_else(AttemptId::generate);
        for round in 0..request.max_round_trips {
            ctx.check_cancelled()?;
            let response = fetch(ollama_request_body(
                &request,
                &messages,
                &prepared.tools,
                false,
            ))
            .await?;

            let message = response
                .get("message")
                .cloned()
                .ok_or_else(|| PipelineError::Other("Ollama response missing message".into()))?;
            let tool_calls = parse_ollama_tool_calls(&message)?;
            messages.push(message.clone());

            if tool_calls.is_empty() {
                return Ok(ToolLoopResponse {
                    final_text: message
                        .get("content")
                        .and_then(|value| value.as_str())
                        .unwrap_or_default()
                        .to_string(),
                    rounds: round + 1,
                    trace_ctx: ctx.trace_ctx.clone(),
                    invocations,
                    exposure_decisions: prepared.exposure_decisions.clone(),
                    provider_response_id: None,
                });
            }

            if tool_calls.len() > 1 && !request.parallel_tool_calls {
                return Err(PipelineError::InvalidConfig(
                    "parallel tool calls are disabled by default".into(),
                ));
            }

            messages.extend(
                self.execute_ollama_tool_calls(
                    ctx,
                    &request,
                    &attempt_id,
                    &prepared.descriptors_by_name,
                    tool_calls,
                    &mut invocations,
                )
                .await?,
            );
        }

        Err(PipelineError::Other(format!(
            "tool loop exceeded max_round_trips={}",
            request.max_round_trips
        )))
    }

    async fn run_ollama_streaming(
        &self,
        ctx: &ExecCtx,
        request: ToolLoopRequest,
    ) -> Result<ToolLoopResponse> {
        let prepared = self.prepare_tool_loop(&request, render_ollama_tool)?;

        let mut messages = Vec::new();
        if let Some(instructions) = &request.instructions {
            messages.push(json!({"role": "system", "content": instructions}));
        }
        messages.push(json!({"role": "user", "content": request.user_input}));

        let mut invocations = Vec::new();
        let attempt_id = request
            .attempt_id
            .clone()
            .unwrap_or_else(AttemptId::generate);
        for round in 0..request.max_round_trips {
            ctx.check_cancelled()?;
            let response = post_ollama_stream(ctx, &request, &messages, &prepared.tools).await?;

            let message = response
                .get("message")
                .cloned()
                .ok_or_else(|| PipelineError::Other("Ollama response missing message".into()))?;
            let tool_calls = parse_ollama_tool_calls(&message)?;
            messages.push(message.clone());

            if tool_calls.is_empty() {
                return Ok(ToolLoopResponse {
                    final_text: message
                        .get("content")
                        .and_then(|value| value.as_str())
                        .unwrap_or_default()
                        .to_string(),
                    rounds: round + 1,
                    trace_ctx: ctx.trace_ctx.clone(),
                    invocations,
                    exposure_decisions: prepared.exposure_decisions.clone(),
                    provider_response_id: None,
                });
            }

            if tool_calls.len() > 1 && !request.parallel_tool_calls {
                return Err(PipelineError::InvalidConfig(
                    "parallel tool calls are disabled by default".into(),
                ));
            }

            messages.extend(
                self.execute_ollama_tool_calls(
                    ctx,
                    &request,
                    &attempt_id,
                    &prepared.descriptors_by_name,
                    tool_calls,
                    &mut invocations,
                )
                .await?,
            );
        }

        Err(PipelineError::Other(format!(
            "tool loop exceeded max_round_trips={}",
            request.max_round_trips
        )))
    }
}

#[derive(Debug, Clone)]
struct OpenAiFunctionCall {
    name: String,
    call_id: String,
    arguments: Value,
}

#[derive(Debug, Clone)]
struct OllamaFunctionCall {
    name: String,
    arguments: Value,
}

fn build_tool_ctx(ctx: &ExecCtx, request: &ToolLoopRequest, attempt_id: &AttemptId) -> ToolCtx {
    let deadline = chrono::Utc::now()
        + chrono::Duration::from_std(ctx.limits.request_timeout)
            .unwrap_or_else(|_| chrono::Duration::seconds(60));
    ToolCtx {
        trace_ctx: ctx.trace_ctx.clone(),
        attempt_id: attempt_id.clone(),
        trial_id: TrialId::generate(),
        deadline: Some(deadline.to_rfc3339()),
        workload_class: Some("llm_pipeline_tool_loop".into()),
        budget_context: Some(ToolBudgetContext {
            budget_kind: Some("tool_loop".into()),
            max_steps: Some(request.max_round_trips),
            time_budget_ms: Some(ctx.limits.request_timeout.as_millis() as u64),
            cost_budget_units: None,
        }),
        scope: request.scope.clone(),
        dry_run: request.dry_run,
        approval_grant: request.approval_grant.clone(),
        execution_permit: request.execution_permit.clone(),
        idempotency_key: None,
        caller: format!("llm-pipeline:{}", request.model),
        planner_stage: request.planner_stage.clone(),
        parent_receipt_id: None,
        family_receipt_id: None,
        replay_parent_receipt_id: None,
        remote_oracle_lease_id: None,
        remote_slice_result_id: None,
        attestation_envelope_id: None,
        cross_runtime_replay_ticket_id: None,
        retry_owner: Some(request.retry_owner.clone()),
    }
}

fn openai_request_body(request: &ToolLoopRequest, input_items: &[Value], tools: &[Value]) -> Value {
    let mut body = json!({
        "model": request.model,
        "input": input_items,
        "parallel_tool_calls": request.parallel_tool_calls,
    });

    if let Some(instructions) = &request.instructions {
        body["instructions"] = json!(instructions);
    }

    if !tools.is_empty() {
        body["tools"] = Value::Array(tools.to_vec());
        body["tool_choice"] = match &request.tool_choice {
            ToolLoopChoice::Auto => json!("auto"),
            ToolLoopChoice::None => json!("none"),
            ToolLoopChoice::Required { tool_name } => json!({
                "type": "function",
                "name": tool_name,
            }),
        };
    }

    body
}

fn ollama_request_body(
    request: &ToolLoopRequest,
    messages: &[Value],
    tools: &[Value],
    stream: bool,
) -> Value {
    let mut body = json!({
        "model": request.model,
        "messages": messages,
        "stream": stream,
        "options": {
            "temperature": request.config.temperature,
            "num_predict": request.config.max_tokens,
        }
    });

    if request.config.thinking {
        body["think"] = json!(true);
    }
    if request.config.json_mode {
        body["format"] = json!("json");
    }
    if let Some(options) = request.config.options.clone() {
        body["options"] = merge_json_objects(body["options"].clone(), options);
    }
    if !tools.is_empty() && !matches!(request.tool_choice, ToolLoopChoice::None) {
        body["tools"] = Value::Array(tools.to_vec());
    }

    body
}

async fn post_json_with_backoff(
    ctx: &ExecCtx,
    path: &str,
    body: &Value,
    api_key: Option<&str>,
    organization: Option<&str>,
) -> Result<Value> {
    let url = format!("{}{}", ctx.base_url.trim_end_matches('/'), path);
    let mut attempt = 0u32;

    loop {
        ctx.check_cancelled()?;
        let mut request = ctx.client.post(&url).json(body);
        if let Some(api_key) = api_key {
            request = request.header("Authorization", format!("Bearer {}", api_key));
        }
        if let Some(organization) = organization {
            request = request.header("OpenAI-Organization", organization);
        }
        let response = request.send().await.map_err(PipelineError::Request)?;

        if response.status().is_success() {
            return response.json().await.map_err(PipelineError::Request);
        }

        let retry_after = response
            .headers()
            .get("retry-after")
            .and_then(|value| value.to_str().ok())
            .and_then(|value| value.parse::<u64>().ok())
            .map(Duration::from_secs);
        let error = PipelineError::HttpError {
            status: response.status().as_u16(),
            body: response.text().await.unwrap_or_default(),
            retry_after,
        };

        if attempt >= ctx.backoff.max_retries || !crate::backend::is_retryable(&error, &ctx.backoff)
        {
            return Err(error);
        }

        let delay = retry_after
            .filter(|_| ctx.backoff.respect_retry_after)
            .unwrap_or_else(|| ctx.backoff.delay_for_attempt(attempt));
        tokio::time::sleep(delay).await;
        if let Some(flag) = ctx.cancel_flag() {
            if flag.load(Ordering::Relaxed) {
                return Err(PipelineError::Cancelled);
            }
        }
        attempt += 1;
    }
}

async fn post_ollama_stream(
    ctx: &ExecCtx,
    request: &ToolLoopRequest,
    messages: &[Value],
    tools: &[Value],
) -> Result<Value> {
    let url = format!("{}/api/chat", ctx.base_url.trim_end_matches('/'));
    let response = ctx
        .client
        .post(&url)
        .json(&ollama_request_body(request, messages, tools, true))
        .send()
        .await
        .map_err(PipelineError::Request)?;

    if !response.status().is_success() {
        return Err(PipelineError::HttpError {
            status: response.status().as_u16(),
            body: response.text().await.unwrap_or_default(),
            retry_after: None,
        });
    }

    let stream = response.bytes_stream().map(|chunk| {
        chunk
            .map(|chunk| chunk.to_vec())
            .map_err(PipelineError::Request)
    });

    consume_ollama_stream(ctx, stream).await
}

async fn consume_ollama_stream<S>(ctx: &ExecCtx, mut stream: S) -> Result<Value>
where
    S: futures::Stream<Item = Result<Vec<u8>>> + Unpin,
{
    let mut decoder = StreamingDecoder::new();
    let mut content = String::new();
    let mut thinking = String::new();
    let mut tool_calls = Vec::new();

    loop {
        ctx.check_cancelled()?;
        let Some(chunk) = stream.next().await else {
            break;
        };
        let chunk = chunk?;
        for value in decoder.decode(&chunk) {
            accumulate_ollama_stream_message(&value, &mut content, &mut thinking, &mut tool_calls);
        }
    }
    ctx.check_cancelled()?;
    if let Some(value) = decoder.flush() {
        accumulate_ollama_stream_message(&value, &mut content, &mut thinking, &mut tool_calls);
    }

    Ok(json!({
        "message": {
            "role": "assistant",
            "content": content,
            "thinking": thinking,
            "tool_calls": tool_calls,
        }
    }))
}

fn accumulate_ollama_stream_message(
    value: &Value,
    content: &mut String,
    thinking: &mut String,
    tool_calls: &mut Vec<Value>,
) {
    let Some(message) = value.get("message") else {
        return;
    };
    if let Some(partial_content) = message.get("content").and_then(|value| value.as_str()) {
        content.push_str(partial_content);
    }
    if let Some(partial_thinking) = message.get("thinking").and_then(|value| value.as_str()) {
        thinking.push_str(partial_thinking);
    }
    if let Some(partial_tool_calls) = message.get("tool_calls").and_then(|value| value.as_array()) {
        tool_calls.extend(partial_tool_calls.iter().cloned());
    }
}

fn parse_openai_function_calls(output_items: &[Value]) -> Result<Vec<OpenAiFunctionCall>> {
    let mut calls = Vec::new();
    for item in output_items {
        if item.get("type").and_then(|value| value.as_str()) != Some("function_call") {
            continue;
        }

        let name = item
            .get("name")
            .and_then(|value| value.as_str())
            .ok_or_else(|| PipelineError::Other("OpenAI function call missing name".into()))?;
        let call_id = item
            .get("call_id")
            .or_else(|| item.get("id"))
            .and_then(|value| value.as_str())
            .ok_or_else(|| PipelineError::Other("OpenAI function call missing call_id".into()))?;
        let arguments = match item.get("arguments") {
            Some(Value::String(arguments)) if arguments.trim().is_empty() => {
                Value::Object(Default::default())
            }
            Some(Value::String(arguments)) => serde_json::from_str(arguments)
                .map_err(|error| PipelineError::Other(error.to_string()))?,
            Some(arguments) => arguments.clone(),
            None => Value::Object(Default::default()),
        };

        calls.push(OpenAiFunctionCall {
            name: name.to_string(),
            call_id: call_id.to_string(),
            arguments,
        });
    }

    Ok(calls)
}

fn parse_ollama_tool_calls(message: &Value) -> Result<Vec<OllamaFunctionCall>> {
    let raw_calls = message
        .get("tool_calls")
        .and_then(|value| value.as_array())
        .cloned()
        .unwrap_or_default();
    let mut calls = Vec::new();

    for item in raw_calls {
        let function = item
            .get("function")
            .ok_or_else(|| PipelineError::Other("Ollama tool call missing function".into()))?;
        let name = function
            .get("name")
            .and_then(|value| value.as_str())
            .ok_or_else(|| PipelineError::Other("Ollama tool call missing name".into()))?;
        let arguments = function
            .get("arguments")
            .cloned()
            .unwrap_or_else(|| Value::Object(Default::default()));
        calls.push(OllamaFunctionCall {
            name: name.to_string(),
            arguments,
        });
    }

    Ok(calls)
}

fn openai_output_text(response: &Value) -> String {
    if let Some(output_text) = response.get("output_text").and_then(|value| value.as_str()) {
        if !output_text.is_empty() {
            return output_text.to_string();
        }
    }

    let mut rendered = String::new();
    if let Some(items) = response.get("output").and_then(|value| value.as_array()) {
        for item in items {
            if item.get("type").and_then(|value| value.as_str()) != Some("message") {
                continue;
            }
            if let Some(content) = item.get("content").and_then(|value| value.as_array()) {
                for part in content {
                    if let Some(text) = part.get("text").and_then(|value| value.as_str()) {
                        rendered.push_str(text);
                    }
                }
            }
        }
    }
    rendered
}

fn merge_json_objects(left: Value, right: Value) -> Value {
    let mut merged = left.as_object().cloned().unwrap_or_default();
    for (key, value) in right.as_object().cloned().unwrap_or_default() {
        merged.insert(key, value);
    }
    Value::Object(merged)
}

fn tool_error_to_pipeline(error: ToolError) -> PipelineError {
    PipelineError::Other(error.to_string())
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use forge_engine::{ForgeStore, ForgeToolReceiptSink};
    use llm_tool_runtime::{
        Tool, ToolApprovalKind, ToolBackendKind, ToolDescriptor, ToolExposureMode,
        ToolExposurePolicy, ToolIdempotencyClass, ToolOutputMode, ToolReceiptPersistence,
        ToolSideEffectClass,
    };
    use std::collections::VecDeque;

    #[derive(Clone)]
    struct AddTool {
        descriptor: ToolDescriptor,
    }

    #[async_trait]
    impl Tool for AddTool {
        fn descriptor(&self) -> &ToolDescriptor {
            &self.descriptor
        }

        async fn invoke(
            &self,
            _ctx: &ToolCtx,
            call: &ToolCall,
        ) -> std::result::Result<ToolResult, ToolError> {
            let a = call
                .arguments
                .get("a")
                .and_then(|value| value.as_i64())
                .unwrap_or(0);
            let b = call
                .arguments
                .get("b")
                .and_then(|value| value.as_i64())
                .unwrap_or(0);
            Ok(ToolResult::text((a + b).to_string()))
        }
    }

    fn add_descriptor() -> ToolDescriptor {
        ToolDescriptor {
            name: "add".into(),
            version: "1.0.0".into(),
            description: Some("Add two numbers".into()),
            backend_kind: ToolBackendKind::LocalFunction,
            input_schema: json!({
                "type": "object",
                "required": ["a", "b"],
                "properties": {
                    "a": {"type": "integer"},
                    "b": {"type": "integer"}
                },
                "additionalProperties": false
            }),
            output_mode: ToolOutputMode::Text,
            read_only: true,
            side_effect_class: ToolSideEffectClass::ReadOnly,
            idempotency_class: ToolIdempotencyClass::Idempotent,
            approval_kind: ToolApprovalKind::None,
            timeout_ms: 5_000,
            concurrency_key: None,
            cache_ttl_ms: None,
            exposure_mode: ToolExposureMode::Auto,
            mcp_surface_kind: llm_tool_runtime::McpSurfaceKind::Tool,
            exposure_policy: ToolExposurePolicy::default(),
            receipt_persistence: ToolReceiptPersistence::Ephemeral,
            output_size_limit_bytes: None,
            provider_payload: None,
            effect_target: Default::default(),
            rollback_contract: None,
        }
    }

    fn runner() -> ToolLoopRunner {
        let mut registry = llm_tool_runtime::ToolRegistry::new();
        registry.register(AddTool {
            descriptor: add_descriptor(),
        });
        ToolLoopRunner::from_registry(registry)
    }

    fn durable_runner(store: Arc<ForgeStore>) -> ToolLoopRunner {
        let mut registry = llm_tool_runtime::ToolRegistry::new();
        let mut descriptor = add_descriptor();
        descriptor.receipt_persistence = ToolReceiptPersistence::ForgeRaw;
        registry.register(AddTool { descriptor });
        ToolLoopRunner::from_registry_with_receipt_sink(
            registry,
            Arc::new(ForgeToolReceiptSink::new(store)),
        )
    }

    fn canned_fetcher(responses: Vec<Value>) -> impl FnMut(Value) -> JsonFetchFuture<'static> {
        let mut responses = VecDeque::from(responses);
        move |_body: Value| {
            let response = responses
                .pop_front()
                .expect("canned provider response should be available");
            Box::pin(async move { Ok(response) })
        }
    }

    #[tokio::test]
    async fn openai_tool_loop_executes_local_function_calls() {
        let ctx = ExecCtx::builder("http://provider.invalid").build();
        let mut fetch = canned_fetcher(vec![
            json!({
                "id": "resp_1",
                "output": [
                    {
                        "type": "function_call",
                        "call_id": "call_1",
                        "name": "add",
                        "arguments": "{\"a\":2,\"b\":3}"
                    }
                ]
            }),
            json!({
                "id": "resp_2",
                "output_text": "5"
            }),
        ]);
        let output = runner()
            .run_openai_responses_with_fetcher(
                &ctx,
                ToolLoopRequest::new("gpt-4.1", "add 2 and 3"),
                &mut fetch,
            )
            .await
            .unwrap();

        assert_eq!(output.final_text, "5");
        assert_eq!(output.invocations.len(), 1);
    }

    #[tokio::test]
    async fn ollama_tool_loop_executes_local_function_calls() {
        let ctx = ExecCtx::builder("http://provider.invalid").build();
        let mut fetch = canned_fetcher(vec![
            json!({
                "message": {
                    "role": "assistant",
                    "content": "",
                    "tool_calls": [
                        {
                            "type": "function",
                            "function": {
                                "index": 0,
                                "name": "add",
                                "arguments": {"a": 4, "b": 6}
                            }
                        }
                    ]
                }
            }),
            json!({
                "message": {
                    "role": "assistant",
                    "content": "10"
                }
            }),
        ]);
        let output = runner()
            .run_ollama_with_fetcher(
                &ctx,
                ToolLoopRequest::new("qwen3", "add 4 and 6"),
                &mut fetch,
            )
            .await
            .unwrap();

        assert_eq!(output.final_text, "10");
        assert_eq!(output.invocations.len(), 1);
    }

    #[tokio::test]
    async fn openai_tool_loop_persists_forge_raw_receipts_via_receipt_sink() {
        let dir = tempfile::TempDir::new().unwrap();
        let store = Arc::new(ForgeStore::open(&dir.path().join("forge.db")).unwrap());
        let ctx = ExecCtx::builder("http://provider.invalid").build();
        let mut fetch = canned_fetcher(vec![
            json!({
                "id": "resp_receipt_1",
                "output": [
                    {
                        "type": "function_call",
                        "call_id": "call_receipt_1",
                        "name": "add",
                        "arguments": "{\"a\":1,\"b\":2}"
                    }
                ]
            }),
            json!({
                "id": "resp_receipt_2",
                "output_text": "3"
            }),
        ]);

        let output = durable_runner(store.clone())
            .run_openai_responses_with_fetcher(
                &ctx,
                ToolLoopRequest::new("gpt-4.1", "add 1 and 2"),
                &mut fetch,
            )
            .await
            .unwrap();

        assert_eq!(output.final_text, "3");
        assert_eq!(output.invocations.len(), 1);

        let receipt_id = &output.invocations[0].receipt.receipt_id;
        let stored = store
            .get_tool_receipt(receipt_id)
            .unwrap()
            .expect("receipt should persist to Forge");

        assert_eq!(stored.tool_name, "add");
        assert_eq!(stored.tool_version, "1.0.0");
        assert_eq!(stored.provider_call_id.as_deref(), Some("call_receipt_1"));
        assert_eq!(stored.trace_id, output.trace_ctx.trace_id);
    }

    #[tokio::test]
    async fn ollama_streaming_checks_cancellation_between_chunks() {
        use futures::stream::poll_fn;
        use std::task::Poll;

        let cancel = Arc::new(std::sync::atomic::AtomicBool::new(false));
        let ctx = ExecCtx::builder("http://provider.invalid")
            .cancellation(Some(cancel.clone()))
            .build();
        let first = b"{\"message\":{\"content\":\"hello\"}}\n".to_vec();
        let second = b"{\"message\":{\"content\":\" world\"}}\n".to_vec();
        let cancel_for_stream = cancel.clone();
        let mut step = 0;

        let stream = poll_fn(move |_cx| match step {
            0 => {
                step = 1;
                cancel_for_stream.store(true, Ordering::Relaxed);
                Poll::Ready(Some(Ok(first.clone())))
            }
            1 => {
                step = 2;
                Poll::Ready(Some(Ok(second.clone())))
            }
            _ => Poll::Ready(None),
        });

        let err = consume_ollama_stream(&ctx, stream).await.unwrap_err();
        assert!(matches!(err, PipelineError::Cancelled));
    }
}