agent-runtime-infra-adapter 0.1.1

use std::sync::Arc;

use agent_infra_sdk::{CallOptions, InfraClient};
use async_trait::async_trait;
use runtime_ports::{
    ModelContent, ModelFinish, ModelMessage, ModelRequest, ModelResponse, ModelRole, ModelSession,
    ModelToolChoice, OperationContext, PortFailure, PortFailureKind, TokenUsage,
};
use serde_json::{Value, json};

use crate::map_sdk_error;

pub struct AgentInfraModelSession {
    client: Arc<InfraClient>,
}
impl AgentInfraModelSession {
    pub fn new(client: Arc<InfraClient>) -> Self {
        Self { client }
    }
}

#[async_trait]
impl ModelSession for AgentInfraModelSession {
    async fn invoke(
        &self,
        operation: &OperationContext,
        request: ModelRequest,
    ) -> Result<ModelResponse, PortFailure> {
        let remaining = operation.remaining()?;
        // Auto is represented by omission. A forced provider-shaped value
        // would leak OpenAI semantics through Agent Infra's canonical port.
        // `none` has already been enforced by the Kernel's empty tool list.
        let tool_choice = match request.generation.tool_choice {
            ModelToolChoice::Auto | ModelToolChoice::None => Value::Null,
        };
        let body = json!({
            "model": request.model,
            "messages": request.messages.iter().map(encode_message).collect::<Vec<_>>(),
            "tools": request.tools.into_iter().map(|tool| json!({"type":"function","function":{"name":tool.name,"description":tool.description,"parameters":tool.input_schema}})).collect::<Vec<_>>(),
            "tool_choice": tool_choice,
            "response_format": request.generation.response_format,
            "temperature": request.generation.temperature,
            "stop": request.generation.stop_sequences,
            "max_completion_tokens": request.generation.max_output_tokens,
            "stream": false
        });
        let response = self
            .client
            .model()
            .chat_completions(
                &body,
                CallOptions::default()
                    .request_id(request.invocation_id)
                    .deadline(remaining),
            )
            .await
            .map_err(map_sdk_error)?;
        parse_response(response)
    }
}

fn encode_message(message: &ModelMessage) -> Value {
    let role = match message.role {
        ModelRole::System => "system",
        ModelRole::User => "user",
        ModelRole::Assistant => "assistant",
        ModelRole::Tool => "tool",
    };
    let text = message
        .content
        .iter()
        .filter_map(|part| match part {
            ModelContent::Text { text } => Some(text.as_str()),
            _ => None,
        })
        .collect::<Vec<_>>()
        .join("");
    let calls = message.content.iter().filter_map(|part| match part {
        ModelContent::ToolUse { id, name, arguments } => Some(json!({"id":id,"type":"function","function":{"name":name,"arguments":arguments.to_string()}})), _ => None,
    }).collect::<Vec<_>>();
    if let Some(ModelContent::ToolResult {
        tool_use_id,
        content,
        ..
    }) = message.content.first()
    {
        return json!({"role":"tool","tool_call_id":tool_use_id,"content":content});
    }
    let mut value = json!({"role":role,"content":text});
    if !calls.is_empty() {
        value["tool_calls"] = Value::Array(calls);
    }
    value
}

fn parse_response(value: Value) -> Result<ModelResponse, PortFailure> {
    let choice = value
        .pointer("/choices/0")
        .ok_or_else(|| protocol("model response omitted choices[0]"))?;
    let message = choice
        .get("message")
        .ok_or_else(|| protocol("model response omitted message"))?;
    let mut output = Vec::new();
    if let Some(text) = message.get("content").and_then(Value::as_str)
        && !text.is_empty()
    {
        output.push(ModelContent::Text { text: text.into() });
    }
    if let Some(calls) = message.get("tool_calls").and_then(Value::as_array) {
        for call in calls {
            let id = call
                .get("id")
                .and_then(Value::as_str)
                .ok_or_else(|| protocol("tool call omitted id"))?;
            let function = call
                .get("function")
                .ok_or_else(|| protocol("tool call omitted function"))?;
            let name = function
                .get("name")
                .and_then(Value::as_str)
                .ok_or_else(|| protocol("tool call omitted name"))?;
            let raw = function
                .get("arguments")
                .and_then(Value::as_str)
                .ok_or_else(|| protocol("tool call arguments must be a JSON string"))?;
            let arguments = serde_json::from_str(raw)
                .map_err(|_| protocol("tool call arguments contained malformed JSON"))?;
            output.push(ModelContent::ToolUse {
                id: id.into(),
                name: name.into(),
                arguments,
            });
        }
    }
    let finish = match choice
        .get("finish_reason")
        .and_then(Value::as_str)
        .unwrap_or("other")
    {
        "stop" => ModelFinish::Stop,
        "tool_calls" => ModelFinish::ToolCalls,
        "length" => ModelFinish::Length,
        "content_filter" => ModelFinish::ContentFilter,
        _ => ModelFinish::Other,
    };
    let usage = TokenUsage {
        input_tokens: value
            .pointer("/usage/prompt_tokens")
            .and_then(Value::as_u64),
        output_tokens: value
            .pointer("/usage/completion_tokens")
            .and_then(Value::as_u64),
    };
    Ok(ModelResponse {
        output,
        finish,
        usage,
    })
}

fn protocol(message: &str) -> PortFailure {
    PortFailure::new(
        PortFailureKind::Protocol,
        "MODEL_PROTOCOL_VIOLATION",
        message,
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn rejects_malformed_tool_arguments() {
        let value = json!({"choices":[{"message":{"tool_calls":[{"id":"c","function":{"name":"read","arguments":"{"}}]},"finish_reason":"tool_calls"}]});
        assert_eq!(
            parse_response(value).unwrap_err().code,
            "MODEL_PROTOCOL_VIOLATION"
        );
    }

    #[test]
    fn preserves_missing_usage_as_unknown() {
        let response = parse_response(json!({
            "choices":[{"message":{"content":"done"},"finish_reason":"stop"}]
        }))
        .unwrap();
        assert_eq!(response.usage, TokenUsage::default());
    }
}