objectiveai-api 2.1.1

//! Mock upstream client for agent completions.

use std::collections::HashMap;
use std::pin::Pin;
use std::time::Duration;
use futures::Stream;
use objectiveai_sdk::agent::completions::response::{Logprob, Logprobs};
use rand::{Rng, SeedableRng};
use super::super::{ContinuationItem, StreamItem, UpstreamClient, ResolvedTool};

/// Mock upstream client that generates random responses with configurable delay.
#[derive(Debug, Clone)]
pub struct Client {
    /// Delay before yielding each chunk.
    pub delay: Duration,
    /// Maximum number of tool calls before returning an error.
    /// Defaults to `MOCK_MAX_TOOL_CALLS` env var, or 1000.
    pub max_tool_calls: u32,
}

/// Resolves the response format for this agent from the request params.
fn resolve_response_format(
    agent_instance_hierarchy: &str,
    params: &objectiveai_sdk::agent::completions::request::AgentCompletionCreateParams,
) -> Option<objectiveai_sdk::agent::completions::request::ResponseFormat> {
    use objectiveai_sdk::agent::completions::request::ResponseFormatParam;
    match params.response_format.as_ref()? {
        ResponseFormatParam::Single(rf) => Some(rf.clone()),
        ResponseFormatParam::PerAgent(map) => map.get(agent_instance_hierarchy).cloned(),
    }
}

/// A single tool call generated by the mock.
pub(super) struct MockToolCall {
    pub(super) tool_name: String,
    pub(super) call_id: String,
    pub(super) arguments: String,
    pub(super) n_deltas: usize,
}

/// The outcome of the tool-vs-content dice roll.
enum MockResponse {
    /// Respond with text content, chunked by logprob token boundaries.
    Content {
        text: String,
        logprobs: Option<Vec<Logprob>>,
    },
    /// Respond with one or more parallel tool calls.
    ToolCalls(Vec<MockToolCall>),
    /// Deterministic-script override response: a single assistant
    /// turn that emits every entry in `tool_calls` (as tool-call
    /// deltas) and then `content` (as content deltas) before
    /// finishing. Sourced from
    /// [`objectiveai_sdk::agent::mock::AgentBase::calls`]. Bypasses
    /// the per-turn RNG selection.
    Override {
        tool_calls: Vec<MockToolCall>,
        content: String,
    },
}

/// Count how many entries from `calls` have already been satisfied,
/// in order, by assistant states in `continuation`. A `Call` is
/// satisfied by a [`ContinuationItem::State`] whose `tool_calls`,
/// compared as `(name, arguments)` pairs (call ids excluded — they're
/// random per-emit), equal the `Call`'s `tool_calls`, AND whose
/// `content` text equals the `Call`'s `content` string (an empty
/// spec string matches `content: None`; multimodal `Parts` never
/// matches since the mock only emits `Text`).
///
/// Returns the index of the next un-matched `Call`. When the return
/// value equals `calls.len()`, every `Call` has been satisfied and
/// the mock should fall through to its normal dispatcher.
///
/// Order-sensitive but interleave-tolerant: an assistant state that
/// doesn't match the current `Call` is skipped (not disqualifying)
/// — the matcher keeps scanning for a later state that does.
fn next_unmatched_call_index(
    calls: &[objectiveai_sdk::agent::mock::Call],
    continuation: Option<&[ContinuationItem<objectiveai_sdk::agent::completions::message::AssistantMessage>]>,
) -> usize {
    use objectiveai_sdk::agent::completions::message::{
        AssistantToolCall, RichContent,
    };

    let Some(items) = continuation else {
        return 0;
    };

    let mut idx = 0usize;
    for item in items {
        if idx >= calls.len() {
            return idx;
        }
        let ContinuationItem::State(state) = item else {
            continue;
        };
        let call = &calls[idx];

        // Compare content. The call spec carries a plain `String`;
        // the wire shape on the state side is `Option<RichContent>`.
        // `None` matches the empty spec string; `Text(s)` matches
        // when `s == call.content`. Multimodal `Parts` never matches
        // — the mock emits text only, so any `Parts` in the
        // continuation is from a different source.
        let state_text = match &state.content {
            None => "",
            Some(RichContent::Text(t)) => t.as_str(),
            Some(RichContent::Parts(_)) => continue,
        };
        if state_text != call.content {
            continue;
        }

        // Compare tool_calls by (name, arguments) only — call ids on
        // the wire are random per emit, so we strip them off both
        // sides before comparing.
        let state_pairs: Vec<(&str, &str)> = state
            .tool_calls
            .as_deref()
            .unwrap_or(&[])
            .iter()
            .map(|tc| match tc {
                AssistantToolCall::Function { function, .. } => {
                    (function.name.as_str(), function.arguments.as_str())
                }
            })
            .collect();
        let call_pairs: Vec<(&str, &str)> = call
            .tool_calls
            .iter()
            .map(|c| (c.name.as_str(), c.arguments.as_str()))
            .collect();
        if state_pairs != call_pairs {
            continue;
        }

        idx += 1;
    }
    idx
}

impl UpstreamClient<objectiveai_sdk::agent::mock::Agent, objectiveai_sdk::agent::mock::Continuation> for Client {
    type State = objectiveai_sdk::agent::completions::message::AssistantMessage;
    type Stream = Pin<
        Box<dyn Stream<Item = StreamItem<Self::State>> + Send + 'static>,
    >;
    type Error = super::Error;

    fn create(
        &self,
        id: &str,
        created: u64,
        agent: &objectiveai_sdk::agent::mock::Agent,
        request_continuation: Option<&objectiveai_sdk::agent::mock::Continuation>,
        params: &objectiveai_sdk::agent::completions::request::AgentCompletionCreateParams,
        messages: &[objectiveai_sdk::agent::completions::message::Message],
        mcp_connection: Option<objectiveai_sdk::mcp::Connection>,
        continuation: Option<&[ContinuationItem<Self::State>]>,
        byok: Option<&str>,
        _cost_multiplier: rust_decimal::Decimal,
        tools_enabled: bool,
        invention_type: Option<objectiveai_sdk::functions::inventions::prompts::StepPromptType>,
        invention_step: Option<usize>,
        invention_tasks_min: Option<u64>,
        invention_input_schema: Option<String>,
        agent_instance_hierarchy: &str,
        agent_id_arg: &str,
        agent_full_id: &str,
        agent_remote: Option<&objectiveai_sdk::RemotePath>,
    ) -> impl Future<
        Output = Result<
            Self::Stream,
            Self::Error,
        >,
    > + Send
    + 'static {
        let tools_enabled = tools_enabled;
        let mode = agent.base.mode.unwrap_or_default();
        let id = id.to_string();
        let agent_instance_hierarchy = agent_instance_hierarchy.to_string();
        let agent_id_for_chunks = agent_id_arg.to_string();
        let agent_full_id = agent_full_id.to_string();
        let agent_remote = agent_remote.cloned();
        let error = agent.base.error == Some(true);
        let error_probability = agent.base.error_probability;
        let top_logprobs = agent.base.top_logprobs;
        let response_format = resolve_response_format(&agent.id, params);
        let params_seed = params.seed;
        let delay = self.delay;
        // Build the full message list: request_continuation -> messages -> continuation.
        let rc_len = request_continuation.map_or(0, |rc| rc.messages.len());
        let cont_len = continuation.map_or(0, |c| c.len());
        let mut all_messages: Vec<objectiveai_sdk::agent::completions::message::Message> =
            Vec::with_capacity(rc_len + messages.len() + cont_len);
        if let Some(rc) = request_continuation {
            all_messages.extend_from_slice(&rc.messages);
        }
        all_messages.extend_from_slice(messages);
        if let Some(cont) = continuation {
            all_messages.extend(cont.iter().map(|item| match item {
                ContinuationItem::State(assistant) => objectiveai_sdk::agent::completions::message::Message::Assistant(assistant.clone()),
                ContinuationItem::ToolMessage(t) => objectiveai_sdk::agent::completions::message::Message::Tool(t.clone()),
                ContinuationItem::UserMessage(u) => objectiveai_sdk::agent::completions::message::Message::User(u.clone()),
            }));
        }
        // Compute assistant index from internal continuation (State + ToolMessage items),
        // identical to openrouter logic.
        let assistant_index = continuation
            .map(|c| {
                c.iter()
                    .filter(|item| {
                        matches!(
                            item,
                            ContinuationItem::State(_)
                                | ContinuationItem::ToolMessage(_)
                        )
                    })
                    .count() as u64
            })
            .unwrap_or(0);
        let prior_tool_call_count: u32 = continuation
            .map(|items| {
                items.iter().filter_map(|item| match item {
                    ContinuationItem::State(s) => Some(super::state::tool_call_count(s) as u32),
                    _ => None,
                }).sum()
            })
            .unwrap_or(0);
        // Validate that AppendTask / WriteInputSchema results from the
        // previous turn did not fail. Find the last State in the
        // continuation and check all ToolMessages that follow it.
        let continuation_validation = continuation.and_then(|items| {
            let last_state_idx = items.iter().rposition(|item| {
                matches!(item, ContinuationItem::State(_))
            })?;
            let state = match &items[last_state_idx] {
                ContinuationItem::State(s) => s,
                _ => unreachable!(),
            };
            Some(super::state::validate_continuation(state, items[last_state_idx + 1..].iter()))
        });
        let is_byok = byok.is_some();
        let max_tool_calls = self.max_tool_calls;

        // Clone agent's `calls` script and the continuation slice
        // into owned values so the `'static` async future doesn't
        // borrow from the caller-supplied `&Agent` / `&[Continuation]`.
        let override_calls_owned: Option<
            Vec<objectiveai_sdk::agent::mock::Call>,
        > = agent.base.calls.clone();
        let continuation_owned: Option<Vec<ContinuationItem<Self::State>>> =
            continuation.map(|c| c.to_vec());

        async move {
            use objectiveai_sdk::agent::completions::request::ResponseFormat;

            if error && error_probability.is_none() {
                return Err(super::Error::ExpectedError);
            }

            if matches!(mode, objectiveai_sdk::agent::mock::Mode::Invention) && invention_type.is_none() {
                return Err(super::Error::InventionAgentWithoutInventionTools);
            }

            if let Some(result) = continuation_validation {
                result?;
            }

            // Reject Grammar and Python response formats.
            if let Some(ref rf) = response_format {
                match rf {
                    ResponseFormat::Grammar { .. } => {
                        return Err(super::Error::UnsupportedResponseFormat(
                            "grammar".into(),
                        ));
                    }
                    ResponseFormat::Python => {
                        return Err(super::Error::UnsupportedResponseFormat(
                            "python".into(),
                        ));
                    }
                    _ => {}
                }
            }

            // Reject required tool call when tools are not allowed.
            if !tools_enabled {
                if let Some(ResponseFormat::ToolCall { required: Some(true), .. }) = &response_format {
                    return Err(super::Error::ToolsNotAllowedWithRequiredToolCall);
                }
            }

            // Source MCP tools from the per-agent proxy connection (if any)
            // and merge with the response-format ToolCall name. The proxy
            // fans out to the agent's declared upstream MCP servers and
            // the invention server, so a single list_tools() (inside
            // resolve_tools) returns the union — no separate invention_tools
            // path.
            let (tool_names, tool_map) = super::super::resolved_tool::resolve_tools(
                mcp_connection.as_ref(),
                response_format.as_ref(),
            )
            .await
            .map_err(|e| super::Error::McpListTools {
                url: e.url,
                error: e.error,
            })?;

            // Hash full prompt + tool names so that adding/removing tools
            // produces a different deterministic output.
            let seed = params_seed.map(|s| {
                use std::hash::Hasher;
                let mut hasher = twox_hash::XxHash3_64::with_seed(s as u64);
                {
                    let mut prompt = all_messages.clone();
                    objectiveai_sdk::agent::completions::message::prompt::prepare(&mut prompt);
                    let pid = objectiveai_sdk::agent::completions::message::prompt::id(&prompt);
                    hasher.write(pid.as_bytes());
                }
                for tn in &tool_names {
                    hasher.write(tn.as_bytes());
                }
                hasher.finish()
            });

            let mut rng = match seed {
                Some(s) => rand::rngs::StdRng::seed_from_u64(s),
                None => rand::rngs::StdRng::from_os_rng(),
            };

            // --- Probabilistic error: roll against error_probability ---
            if error_probability.is_some_and(|p| rng.random_range(0u8..100) < p) {
                return Err(super::Error::ExpectedError);
            }

            // --- Reasoning: roll 0-5 chunks ---
            let n_reasoning = rng.random_range(0u32..=5);
            let reasoning_chunks: Vec<String> = (0..n_reasoning)
                .map(|_| random_string(&mut rng, 20, 200))
                .collect();

            // --- Tool call vs content ---
            //
            // `agent.base.calls` is a deterministic-script preceding
            // override. When set and the continuation hasn't yet
            // satisfied every `Call`, the next un-matched `Call` is
            // what this turn emits — regardless of mode. Once the
            // override is exhausted (all `Call`s matched), we fall
            // through to the normal mode-driven dispatcher below.
            let override_calls = override_calls_owned.as_ref();
            let override_response = override_calls.and_then(|calls| {
                let idx = next_unmatched_call_index(
                    calls,
                    continuation_owned.as_deref(),
                );
                if idx < calls.len() {
                    let call = &calls[idx];
                    let tool_calls: Vec<MockToolCall> = call
                        .tool_calls
                        .iter()
                        .map(|c| MockToolCall {
                            tool_name: c.name.clone(),
                            // Same shape as the existing random-mint
                            // path. The matcher ignores id, so the
                            // value here only matters for downstream
                            // tool-response wiring.
                            call_id: format!(
                                "call_mock_{}",
                                random_string(&mut rng, 16, 16),
                            ),
                            arguments: c.arguments.clone(),
                            n_deltas: 1,
                        })
                        .collect();
                    Some(MockResponse::Override {
                        tool_calls,
                        content: call.content.clone(),
                    })
                } else {
                    None
                }
            });
            let mock_response = if let Some(r) = override_response {
                r
            } else if matches!(mode, objectiveai_sdk::agent::mock::Mode::LaboratoryEvaluation) {
                // Extract schema from last user message: "## evaluation schema\n\n{json}"
                let schema_json = {
                    use objectiveai_sdk::agent::completions::message::RichContent;
                    let extract = |c: &RichContent| match c {
                        RichContent::Text(t) => t.clone(),
                        RichContent::Parts(parts) => parts.iter().filter_map(|p| match p {
                            objectiveai_sdk::agent::completions::message::RichContentPart::Text { text } => Some(text.as_str()),
                            _ => None,
                        }).collect::<Vec<_>>().join(""),
                    };
                    all_messages.iter().rev().find_map(|m| match m {
                        objectiveai_sdk::agent::completions::message::Message::User(u) => {
                            let text = extract(&u.content);
                            text.find("## evaluation schema\n\n").map(|pos| {
                                text[pos + "## evaluation schema\n\n".len()..].to_string()
                            })
                        }
                        _ => None,
                    }).unwrap_or_default()
                };
                let schema: objectiveai_sdk::functions::expression::InputSchema = serde_json::from_str(&schema_json)
                    .unwrap_or(objectiveai_sdk::functions::expression::InputSchema::String(
                        objectiveai_sdk::functions::expression::StringInputSchema {
                            r#type: objectiveai_sdk::functions::expression::StringInputSchemaType::String,
                            description: None,
                            r#enum: None,
                        },
                    ));
                let input_value = objectiveai_sdk::functions::check::example_inputs::generate_seeded(
                    &schema,
                    rng.clone(),
                ).next().unwrap_or(objectiveai_sdk::functions::expression::InputValue::String("mock".to_string()));
                let text = serde_json::to_string(&input_value).unwrap();
                MockResponse::Content { text, logprobs: None }
            } else if matches!(mode, objectiveai_sdk::agent::mock::Mode::LaboratoryBuilder) && tools_enabled && prior_tool_call_count == 0 {
                MockResponse::ToolCalls(vec![super::builder::write_tool_call(&mut rng)])
            } else if matches!(mode, objectiveai_sdk::agent::mock::Mode::Invention) && tools_enabled {
                resolve_invention_response(
                    invention_type.unwrap(),
                    invention_step.unwrap(),
                    invention_tasks_min.unwrap_or(3),
                    invention_input_schema.as_deref(),
                    &tool_names,
                    &tool_map,
                    &mut rng,
                )
            } else {
                let effective_tool_names = if tools_enabled { &tool_names[..] } else { &[] };
                resolve_mock_response(
                    &response_format,
                    effective_tool_names,
                    &tool_map,
                    top_logprobs,
                    &mut rng,
                )
            };
            let current_tool_calls: Vec<(String, String)> = match &mock_response {
                MockResponse::ToolCalls(calls)
                | MockResponse::Override { tool_calls: calls, .. } => calls.iter()
                    .map(|c| (c.tool_name.clone(), c.call_id.clone()))
                    .collect(),
                _ => Vec::new(),
            };
            let total_tool_call_count = prior_tool_call_count + current_tool_calls.len() as u32;

            if total_tool_call_count > max_tool_calls {
                return Err(super::Error::MaxToolCallsExceeded(max_tool_calls));
            }

            let state = {
                use objectiveai_sdk::agent::completions::message::{
                    AssistantMessage, AssistantToolCall, AssistantToolCallFunction, RichContent,
                };
                let reasoning = if reasoning_chunks.is_empty() {
                    None
                } else {
                    Some(reasoning_chunks.join(""))
                };
                match &mock_response {
                    MockResponse::Content { text, .. } => AssistantMessage {
                        content: Some(RichContent::Text(text.clone())),
                        name: None,
                        refusal: None,
                        tool_calls: None,
                        reasoning,
                    },
                    MockResponse::ToolCalls(calls) => AssistantMessage {
                        content: None,
                        name: None,
                        refusal: None,
                        tool_calls: Some(calls.iter().map(|c| {
                            AssistantToolCall::Function {
                                id: c.call_id.clone(),
                                function: AssistantToolCallFunction {
                                    name: c.tool_name.clone(),
                                    arguments: c.arguments.clone(),
                                },
                            }
                        }).collect()),
                        reasoning,
                    },
                    MockResponse::Override { tool_calls, content } => AssistantMessage {
                        // Empty content string round-trips through the
                        // matcher as `None` — see
                        // `next_unmatched_call_index` for the
                        // "empty string ↔ None" convention.
                        content: if content.is_empty() {
                            None
                        } else {
                            Some(RichContent::Text(content.clone()))
                        },
                        name: None,
                        refusal: None,
                        tool_calls: Some(tool_calls.iter().map(|c| {
                            AssistantToolCall::Function {
                                id: c.call_id.clone(),
                                function: AssistantToolCallFunction {
                                    name: c.tool_name.clone(),
                                    arguments: c.arguments.clone(),
                                },
                            }
                        }).collect()),
                        reasoning,
                    },
                }
            };

            let stream = async_stream::stream! {
                use objectiveai_sdk::agent::completions::message::{
                    AssistantToolCallDelta, AssistantToolCallFunctionDelta,
                    AssistantToolCallType, RichContent,
                };
                use objectiveai_sdk::agent::completions::response::streaming::{
                    AgentCompletionChunk, AssistantResponseChunk, MessageChunk,
                };
                use objectiveai_sdk::agent::completions::response::FinishReason;

                // --- Yield reasoning chunks ---
                for reasoning_text in &reasoning_chunks {
                    if !delay.is_zero() {
                        tokio::time::sleep(delay).await;
                    }
                    yield StreamItem::Chunk(AgentCompletionChunk {
                        id: id.clone(),
                        agent_instance_hierarchy: agent_instance_hierarchy.clone(),
                        agent_id: agent_id_for_chunks.clone(),
                        agent_full_id: agent_full_id.clone(),
                        agent_remote: agent_remote.clone(),
                        created,
                        messages: vec![MessageChunk::Assistant(AssistantResponseChunk {
                            index: assistant_index,
                            created,
                            model: "mock".into(),
                            upstream_id: id.clone(),
                            reasoning: Some(reasoning_text.clone()),
                            ..Default::default()
                        })],
                        object: Default::default(),
                        usage: None,
                        upstream: objectiveai_sdk::agent::Upstream::Mock,
                        error: None,
                        continuation: None,
                        messages_queued: None,
                    });
                }

                // --- Yield content or tool call chunks ---
                match &mock_response {
                    MockResponse::Content { text, logprobs } => {
                        let chunks = chunk_by_logprobs(text, logprobs.as_deref(), &mut rng);

                        for (i, (chunk_text, chunk_logprobs)) in chunks.iter().enumerate() {
                            let is_last = i == chunks.len() - 1;
                            if !delay.is_zero() {
                                tokio::time::sleep(delay).await;
                            }
                            yield StreamItem::Chunk(AgentCompletionChunk {
                                id: id.clone(),
                                agent_instance_hierarchy: agent_instance_hierarchy.clone(),
                                agent_id: agent_id_for_chunks.clone(),
                                agent_full_id: agent_full_id.clone(),
                                agent_remote: agent_remote.clone(),
                                created,
                                messages: vec![MessageChunk::Assistant(AssistantResponseChunk {
                                    index: assistant_index,
                                    created,
                                    model: "mock".into(),
                                    upstream_id: id.clone(),
                                    content: Some(RichContent::Text(chunk_text.clone())),
                                    logprobs: chunk_logprobs.as_ref().map(|lps| Logprobs {
                                        content: Some(lps.clone()),
                                        refusal: None,
                                    }),
                                    finish_reason: if is_last {
                                        Some(FinishReason::Stop)
                                    } else {
                                        None
                                    },
                                    ..Default::default()
                                })],
                                object: Default::default(),
                                usage: None,
                                upstream: objectiveai_sdk::agent::Upstream::Mock,
                                error: None,
                                continuation: None,
                                messages_queued: None,
                            });
                        }
                    }
                    MockResponse::ToolCalls(calls) => {
                        for (tc_idx, tc) in calls.iter().enumerate() {
                            let chunk_size = (tc.arguments.len() + tc.n_deltas - 1) / tc.n_deltas;
                            let parts: Vec<&str> = if tc.arguments.is_empty() {
                                vec![""]
                            } else {
                                tc.arguments.as_bytes()
                                    .chunks(chunk_size.max(1))
                                    .map(|b| std::str::from_utf8(b).unwrap_or(""))
                                    .collect()
                            };

                            for (i, part) in parts.iter().enumerate() {
                                let is_first = i == 0;
                                let is_last = i == parts.len() - 1;
                                if !delay.is_zero() {
                                    tokio::time::sleep(delay).await;
                                }
                                yield StreamItem::Chunk(AgentCompletionChunk {
                                    id: id.clone(),
                                    agent_instance_hierarchy: agent_instance_hierarchy.clone(),
                                    agent_id: agent_id_for_chunks.clone(),
                                    agent_full_id: agent_full_id.clone(),
                                    agent_remote: agent_remote.clone(),
                                    created,
                                    messages: vec![MessageChunk::Assistant(AssistantResponseChunk {
                                        index: assistant_index,
                                        created,
                                        model: "mock".into(),
                                        upstream_id: id.clone(),
                                        tool_calls: Some(vec![AssistantToolCallDelta {
                                            index: tc_idx as u64,
                                            r#type: if is_first {
                                                Some(AssistantToolCallType::Function)
                                            } else {
                                                None
                                            },
                                            id: if is_first {
                                                Some(tc.call_id.clone())
                                            } else {
                                                None
                                            },
                                            function: Some(AssistantToolCallFunctionDelta {
                                                name: if is_first {
                                                    Some(tc.tool_name.clone())
                                                } else {
                                                    None
                                                },
                                                arguments: Some(part.to_string()),
                                            }),
                                        }]),
                                        finish_reason: if is_last && tc_idx == calls.len() - 1 {
                                            Some(FinishReason::ToolCalls)
                                        } else {
                                            None
                                        },
                                        ..Default::default()
                                    })],
                                    object: Default::default(),
                                    usage: None,
                                    upstream: objectiveai_sdk::agent::Upstream::Mock,
                                    error: None,
                                    continuation: None,
                                    messages_queued: None,
                                });
                            }
                        }
                    }
                    MockResponse::Override { tool_calls, content } => {
                        // Two phases: emit every tool-call delta first,
                        // then emit content text chunks. The `finish_reason`
                        // lands on the last chunk overall — if `content`
                        // is non-empty, that's the last text chunk
                        // (`Stop`); if `content` is empty, that's the
                        // last tool-call chunk (`ToolCalls`).
                        let has_content = !content.is_empty();

                        // --- Phase 1: tool-call deltas ---
                        for (tc_idx, tc) in tool_calls.iter().enumerate() {
                            let chunk_size = (tc.arguments.len() + tc.n_deltas - 1) / tc.n_deltas;
                            let parts: Vec<&str> = if tc.arguments.is_empty() {
                                vec![""]
                            } else {
                                tc.arguments.as_bytes()
                                    .chunks(chunk_size.max(1))
                                    .map(|b| std::str::from_utf8(b).unwrap_or(""))
                                    .collect()
                            };

                            for (i, part) in parts.iter().enumerate() {
                                let is_first = i == 0;
                                let is_last_in_call = i == parts.len() - 1;
                                let is_last_tool_call =
                                    tc_idx == tool_calls.len() - 1;
                                if !delay.is_zero() {
                                    tokio::time::sleep(delay).await;
                                }
                                yield StreamItem::Chunk(AgentCompletionChunk {
                                    id: id.clone(),
                                    agent_instance_hierarchy: agent_instance_hierarchy.clone(),
                                    agent_id: agent_id_for_chunks.clone(),
                                    agent_full_id: agent_full_id.clone(),
                                    agent_remote: agent_remote.clone(),
                                    created,
                                    messages: vec![MessageChunk::Assistant(AssistantResponseChunk {
                                        index: assistant_index,
                                        created,
                                        model: "mock".into(),
                                        upstream_id: id.clone(),
                                        tool_calls: Some(vec![AssistantToolCallDelta {
                                            index: tc_idx as u64,
                                            r#type: if is_first {
                                                Some(AssistantToolCallType::Function)
                                            } else {
                                                None
                                            },
                                            id: if is_first {
                                                Some(tc.call_id.clone())
                                            } else {
                                                None
                                            },
                                            function: Some(AssistantToolCallFunctionDelta {
                                                name: if is_first {
                                                    Some(tc.tool_name.clone())
                                                } else {
                                                    None
                                                },
                                                arguments: Some(part.to_string()),
                                            }),
                                        }]),
                                        // No content following → terminate
                                        // on the last tool-call chunk with
                                        // `ToolCalls`, same as the
                                        // `MockResponse::ToolCalls` arm.
                                        finish_reason: if !has_content
                                            && is_last_in_call
                                            && is_last_tool_call
                                        {
                                            Some(FinishReason::ToolCalls)
                                        } else {
                                            None
                                        },
                                        ..Default::default()
                                    })],
                                    object: Default::default(),
                                    usage: None,
                                    upstream: objectiveai_sdk::agent::Upstream::Mock,
                                    error: None,
                                    continuation: None,
                                    messages_queued: None,
                                });
                            }
                        }

                        // --- Phase 2: content text chunks (skipped
                        // entirely when `content` is empty) ---
                        if has_content {
                            // Override doesn't carry logprobs — text-only
                            // chunking via the same helper as the
                            // `Content` arm (no logprobs alignment).
                            let chunks = chunk_by_logprobs(content, None, &mut rng);
                            for (i, (chunk_text, _)) in chunks.iter().enumerate() {
                                let is_last = i == chunks.len() - 1;
                                if !delay.is_zero() {
                                    tokio::time::sleep(delay).await;
                                }
                                yield StreamItem::Chunk(AgentCompletionChunk {
                                    id: id.clone(),
                                    agent_instance_hierarchy: agent_instance_hierarchy.clone(),
                                    agent_id: agent_id_for_chunks.clone(),
                                    agent_full_id: agent_full_id.clone(),
                                    agent_remote: agent_remote.clone(),
                                    created,
                                    messages: vec![MessageChunk::Assistant(AssistantResponseChunk {
                                        index: assistant_index,
                                        created,
                                        model: "mock".into(),
                                        upstream_id: id.clone(),
                                        content: Some(RichContent::Text(chunk_text.clone())),
                                        finish_reason: if is_last {
                                            Some(FinishReason::Stop)
                                        } else {
                                            None
                                        },
                                        ..Default::default()
                                    })],
                                    object: Default::default(),
                                    usage: None,
                                    upstream: objectiveai_sdk::agent::Upstream::Mock,
                                    error: None,
                                    continuation: None,
                                    messages_queued: None,
                                });
                            }
                        }
                    }
                }

                // --- Yield final state ---
                yield StreamItem::State(state);
            };

            let boxed: Pin<Box<dyn Stream<Item = StreamItem<Self::State>> + Send>> =
                Box::pin(stream);
            Ok(boxed)
        }
    }

    fn response_continuation(
        &self,
        mcp_sessions: indexmap::IndexMap<String, String>,
        request_continuation: Option<&objectiveai_sdk::agent::mock::Continuation>,
        messages: &[objectiveai_sdk::agent::completions::message::Message],
        continuation: Option<&[ContinuationItem<Self::State>]>,
        agent_instance_hierarchy: &str,
    ) -> objectiveai_sdk::agent::mock::Continuation {
        use objectiveai_sdk::agent::completions::message::Message;
        let rc_len = request_continuation.map_or(0, |rc| rc.messages.len());
        let cont_len = continuation.map_or(0, |c| c.len());
        let mut all_messages = Vec::with_capacity(rc_len + messages.len() + cont_len);
        if let Some(rc) = request_continuation {
            all_messages.extend_from_slice(&rc.messages);
        }
        all_messages.extend_from_slice(messages);
        if let Some(cont) = continuation {
            all_messages.extend(cont.iter().map(|item| match item {
                ContinuationItem::State(assistant) => Message::Assistant(assistant.clone()),
                ContinuationItem::ToolMessage(t) => Message::Tool(t.clone()),
                ContinuationItem::UserMessage(u) => Message::User(u.clone()),
            }));
        }
        objectiveai_sdk::agent::mock::Continuation {
            upstream: objectiveai_sdk::agent::mock::Upstream::default(),
            agent_instance_hierarchy: agent_instance_hierarchy.to_string(),
            messages: all_messages,
            mcp_sessions,
        }
    }
}

/// Splits content into chunks aligned to logprob token boundaries.
/// Each chunk consumes 1-5 logprob tokens. When logprobs is None,
/// falls back to simple byte-based chunking.
fn chunk_by_logprobs(
    text: &str,
    logprobs: Option<&[Logprob]>,
    rng: &mut impl Rng,
) -> Vec<(String, Option<Vec<Logprob>>)> {
    let logprobs = match logprobs {
        Some(lps) if !lps.is_empty() => lps,
        _ => {
            // No logprobs: simple chunking like before.
            let n_chunks = rng.random_range(1u32..=5) as usize;
            let chunk_size = (text.len() + n_chunks - 1) / n_chunks.max(1);
            let parts: Vec<&str> = if text.is_empty() {
                vec![""]
            } else {
                text.as_bytes()
                    .chunks(chunk_size.max(1))
                    .map(|b| std::str::from_utf8(b).unwrap_or(""))
                    .collect()
            };
            return parts.into_iter().map(|p| (p.to_string(), None)).collect();
        }
    };

    let mut chunks = Vec::new();
    let mut pos = 0;

    while pos < logprobs.len() {
        let n_tokens = rng.random_range(1u32..=5) as usize;
        let end = (pos + n_tokens).min(logprobs.len());
        let chunk_lps = &logprobs[pos..end];
        let chunk_text: String = chunk_lps.iter().map(|lp| lp.token.as_str()).collect();
        chunks.push((chunk_text, Some(chunk_lps.to_vec())));
        pos = end;
    }

    if chunks.is_empty() {
        chunks.push((String::new(), Some(Vec::new())));
    }

    chunks
}

/// Decides whether to call a tool or respond with content, and generates the data.
fn resolve_mock_response(
    response_format: &Option<objectiveai_sdk::agent::completions::request::ResponseFormat>,
    tool_names: &[String],
    tool_map: &HashMap<String, ResolvedTool>,
    top_logprobs: Option<u64>,
    rng: &mut impl Rng,
) -> MockResponse {
    use objectiveai_sdk::agent::completions::request::ResponseFormat;

    // Check for required tool call from response format.
    if let Some(ResponseFormat::ToolCall {
        name, required: Some(true), ..
    }) = response_format
    {
        if tool_map.contains_key(name) {
            let arguments = generate_tool_arguments(tool_map, name, rng);
            return MockResponse::ToolCalls(vec![MockToolCall {
                tool_name: name.clone(),
                call_id: format!("call_mock_{}", rng.random_range(0u64..u64::MAX)),
                arguments,
                n_deltas: rng.random_range(1u32..=5) as usize,
            }]);
        }
    }

    if !tool_names.is_empty() {
        // Roll the dice: equal probability for each tool or respond as-is,
        // with respond-as-is having a minimum 25% chance.
        let n_tools = tool_names.len();
        let respond_as_is_weight = if n_tools >= 3 {
            25u32
        } else {
            (100 / (n_tools as u32 + 1)).max(25)
        };
        let tool_weight = (100 - respond_as_is_weight) / n_tools as u32;
        let roll = rng.random_range(0u32..100);

        if roll >= respond_as_is_weight {
            let tool_index = ((roll - respond_as_is_weight) / tool_weight.max(1))
                .min(n_tools as u32 - 1) as usize;
            let tool_name = &tool_names[tool_index];
            let arguments = generate_tool_arguments(tool_map, tool_name, rng);
            let mut calls = vec![MockToolCall {
                tool_name: tool_name.clone(),
                call_id: format!("call_mock_{}", rng.random_range(0u64..u64::MAX)),
                arguments,
                n_deltas: rng.random_range(1u32..=5) as usize,
            }];
            // 50/50 coin flip for each additional parallel tool call.
            while rng.random_range(0u32..2) == 1 {
                let extra_index = rng.random_range(0..n_tools);
                let extra_name = &tool_names[extra_index];
                let extra_args = generate_tool_arguments(tool_map, extra_name, rng);
                calls.push(MockToolCall {
                    tool_name: extra_name.clone(),
                    call_id: format!("call_mock_{}", rng.random_range(0u64..u64::MAX)),
                    arguments: extra_args,
                    n_deltas: rng.random_range(1u32..=5) as usize,
                });
            }
            return MockResponse::ToolCalls(calls);
        }
    }

    // Respond as-is with content.
    let (text, logprobs) = generate_content(response_format, top_logprobs, rng);
    MockResponse::Content { text, logprobs }
}

/// Generates content and optional logprobs based on response format and top_logprobs setting.
fn generate_content(
    response_format: &Option<objectiveai_sdk::agent::completions::request::ResponseFormat>,
    top_logprobs: Option<u64>,
    rng: &mut impl Rng,
) -> (String, Option<Vec<Logprob>>) {
    use objectiveai_sdk::agent::completions::request::ResponseFormat;

    let permutations = match top_logprobs {
        None | Some(0) => 1,
        Some(n) => n as usize,
    };
    let yield_logprobs = top_logprobs.is_some_and(|n| n >= 1);

    match response_format {
        Some(ResponseFormat::JsonObject) => {
            if yield_logprobs {
                let serialized = vec!["{}".to_string(); permutations];
                let (text, logprobs) = super::json_schema::generate_logprobs_from_serialized(&serialized, rng);
                (text, Some(logprobs))
            } else {
                ("{}".into(), None)
            }
        }
        Some(ResponseFormat::JsonSchema { schema }) | Some(ResponseFormat::ToolCall { schema, .. }) => {
            generate_from_schema(schema, permutations, yield_logprobs, rng)
        }
        _ => {
            // Plain text content.
            let text = random_string(rng, 10, 100);
            if yield_logprobs {
                let serialized: Vec<String> = std::iter::once(text)
                    .chain((1..permutations).map(|_| random_string(rng, 10, 100)))
                    .collect();
                let (text, logprobs) = super::json_schema::generate_logprobs_from_serialized(&serialized, rng);
                (text, Some(logprobs))
            } else {
                (text, None)
            }
        }
    }
}

/// Generates content from an IndexMap schema, with optional logprobs.
fn generate_from_schema(
    schema: &indexmap::IndexMap<String, serde_json::Value>,
    permutations: usize,
    yield_logprobs: bool,
    rng: &mut impl Rng,
) -> (String, Option<Vec<Logprob>>) {
    match serde_json::from_value::<super::json_schema::JsonSchema>(
        serde_json::Value::Object(
            schema.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
        ),
    ) {
        Ok(js) => {
            let (text, logprobs) = js.generate_content_from_rng(rng, permutations);
            (text, if yield_logprobs { Some(logprobs) } else { None })
        }
        Err(_) => ("{}".into(), None),
    }
}

/// Generates a random alphanumeric string (with spaces) of length between `min` and `max`.
pub(super) fn random_string(rng: &mut impl Rng, min: usize, max: usize) -> String {
    let len = rng.random_range(min..=max);
    (0..len)
        .map(|_| {
            const CHARS: &[u8] =
                b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 ";
            CHARS[rng.random_range(0..CHARS.len())] as char
        })
        .collect()
}

/// Generates tool call arguments by parsing the tool's parameter schema as JsonSchema.
/// Tool calls never include logprobs.
pub(super) fn generate_tool_arguments(
    tool_map: &HashMap<String, ResolvedTool>,
    tool_name: &str,
    rng: &mut impl Rng,
) -> String {
    let schema_value = match tool_map.get(tool_name) {
        Some(ResolvedTool::Mcp { tool, .. }) => {
            let mut map = serde_json::Map::new();
            map.insert("type".into(), serde_json::json!("object"));
            if let Some(props) = &tool.input_schema.properties {
                map.insert(
                    "properties".into(),
                    serde_json::Value::Object(
                        props.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
                    ),
                );
            }
            Some(serde_json::Value::Object(map))
        }
        Some(ResolvedTool::ResponseFormat { schema, .. }) => {
            Some(serde_json::Value::Object(
                schema.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
            ))
        }
        None => None,
    };

    match schema_value {
        Some(sv) => {
            match serde_json::from_value::<super::json_schema::JsonSchema>(sv) {
                Ok(js) => js.generate_content_from_rng(rng, 1).0,
                Err(_) => "{}".into(),
            }
        }
        None => "{}".into(),
    }
}

/// Generates tool calls for an invention agent based on explicit type and step.
fn resolve_invention_response(
    invention_type: objectiveai_sdk::functions::inventions::prompts::StepPromptType,
    invention_step: usize,
    tasks_min: u64,
    invention_input_schema: Option<&str>,
    tool_names: &[String],
    tool_map: &HashMap<String, ResolvedTool>,
    rng: &mut impl Rng,
) -> MockResponse {
    use objectiveai_sdk::functions::inventions::prompts::StepPromptType::*;

    let tc = match (invention_type, invention_step) {
        // Step 0: Essay
        (AlphaScalarBranchFunction | AlphaScalarLeafFunction, 0) => {
            super::invention::alpha_scalar::essay_tool_call(tool_names, tool_map, rng)
        }
        (AlphaVectorBranchFunction | AlphaVectorLeafFunction, 0) => {
            super::invention::alpha_vector::essay_tool_call(tool_names, tool_map, rng)
        }
        // Step 1: Input Schema
        (AlphaScalarBranchFunction | AlphaScalarLeafFunction, 1) => {
            super::invention::alpha_scalar::input_schema_tool_call(tool_names, tool_map, rng)
        }
        (AlphaVectorBranchFunction | AlphaVectorLeafFunction, 1) => {
            super::invention::alpha_vector::input_schema_tool_call(tool_names, tool_map, rng)
        }
        // Step 2: Essay Tasks
        (AlphaScalarBranchFunction | AlphaScalarLeafFunction, 2) => {
            super::invention::alpha_scalar::essay_tasks_tool_call(tool_names, tool_map, rng)
        }
        (AlphaVectorBranchFunction | AlphaVectorLeafFunction, 2) => {
            super::invention::alpha_vector::essay_tasks_tool_call(tool_names, tool_map, rng)
        }
        // Step 3: Tasks
        (AlphaScalarLeafFunction, 3) => {
            let fallback = r#"{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}"#;
            let schema = invention_input_schema.unwrap_or(fallback);
            super::invention::alpha_scalar_leaf::tasks_tool_call(schema, tasks_min, tool_names, tool_map, rng)
        }
        (AlphaScalarBranchFunction, 3) => {
            let fallback = r#"{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}"#;
            let schema = invention_input_schema.unwrap_or(fallback);
            super::invention::alpha_scalar_branch::tasks_tool_call(schema, tasks_min, tool_names, tool_map, rng)
        }
        (AlphaVectorLeafFunction, 3) => {
            let fallback = r#"{"items":{"type":"string"}}"#;
            let schema = invention_input_schema.unwrap_or(fallback);
            super::invention::alpha_vector_leaf::tasks_tool_call(schema, tasks_min, tool_names, tool_map, rng)
        }
        (AlphaVectorBranchFunction, 3) => {
            let fallback = r#"{"items":{"type":"string"}}"#;
            let schema = invention_input_schema.unwrap_or(fallback);
            super::invention::alpha_vector_branch::tasks_tool_call(schema, 0, 0, tasks_min, tool_names, tool_map, rng)
        }
        // Step 4: Description
        (_, 4) => {
            super::invention::description_tool_call(tool_names, tool_map, rng)
        }
        // Unknown step — description fallback
        _ => {
            super::invention::description_tool_call(tool_names, tool_map, rng)
        }
    };


    MockResponse::ToolCalls(vec![tc])
}