huskies/server/src/agents/runtime/gemini/api.rs

//! Gemini API types, request builders, and response parsers.
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};

use super::super::super::TokenUsage;
use super::super::RuntimeContext;

// ── Gemini API types ─────────────────────────────────────────────────

#[derive(Debug, Serialize, Deserialize)]
pub(super) struct GeminiFunctionDeclaration {
    pub name: String,
    pub description: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parameters: Option<Value>,
}

// ── Request builders ─────────────────────────────────────────────────

/// Build the system instruction content from the RuntimeContext.
pub(super) fn build_system_instruction(ctx: &RuntimeContext) -> Value {
    // Use system_prompt from args if provided via --append-system-prompt,
    // otherwise use a sensible default.
    let system_text = ctx
        .args
        .iter()
        .position(|a| a == "--append-system-prompt")
        .and_then(|i| ctx.args.get(i + 1))
        .cloned()
        .unwrap_or_else(|| {
            format!(
                "You are an AI coding agent working on story {}. \
                 You have access to tools via function calling. \
                 Use them to complete the task. \
                 Work in the directory: {}",
                ctx.story_id, ctx.cwd
            )
        });

    json!({
        "parts": [{ "text": system_text }]
    })
}

/// Build the full `generateContent` request body.
pub(super) fn build_generate_content_request(
    system_instruction: &Value,
    contents: &[Value],
    gemini_tools: &[GeminiFunctionDeclaration],
) -> Value {
    let mut body = json!({
        "system_instruction": system_instruction,
        "contents": contents,
        "generationConfig": {
            "temperature": 0.2,
            "maxOutputTokens": 65536,
        }
    });

    if !gemini_tools.is_empty() {
        body["tools"] = json!([{
            "functionDeclarations": gemini_tools
        }]);
    }

    body
}

// ── Response parsing ─────────────────────────────────────────────────

/// Parse token usage metadata from a Gemini API response.
pub(super) fn parse_usage_metadata(response: &Value) -> Option<TokenUsage> {
    let metadata = response.get("usageMetadata")?;
    Some(TokenUsage {
        input_tokens: metadata
            .get("promptTokenCount")
            .and_then(|v| v.as_u64())
            .unwrap_or(0),
        output_tokens: metadata
            .get("candidatesTokenCount")
            .and_then(|v| v.as_u64())
            .unwrap_or(0),
        // Gemini doesn't have cache token fields, but we keep the struct uniform.
        cache_creation_input_tokens: 0,
        cache_read_input_tokens: 0,
        // Google AI API doesn't report cost; leave at 0.
        total_cost_usd: 0.0,
    })
}

// ── Tests ────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::http::context::AppContext;
    use std::sync::Arc;

    fn test_app_ctx() -> Arc<AppContext> {
        let tmp = tempfile::tempdir().unwrap();
        Arc::new(AppContext::new_test(tmp.path().to_path_buf()))
    }

    #[test]
    fn build_system_instruction_uses_args() {
        let ctx = RuntimeContext {
            story_id: "42_story_test".to_string(),
            agent_name: "coder-1".to_string(),
            command: "gemini-2.5-pro".to_string(),
            args: vec![
                "--append-system-prompt".to_string(),
                "Custom system prompt".to_string(),
            ],
            prompt: "Do the thing".to_string(),
            cwd: "/tmp/wt".to_string(),
            inactivity_timeout_secs: 300,
            app_ctx: Some(test_app_ctx()),
            session_id_to_resume: None,
            fresh_prompt: None,
        };

        let instruction = build_system_instruction(&ctx);
        assert_eq!(instruction["parts"][0]["text"], "Custom system prompt");
    }

    #[test]
    fn build_system_instruction_default() {
        let ctx = RuntimeContext {
            story_id: "42_story_test".to_string(),
            agent_name: "coder-1".to_string(),
            command: "gemini-2.5-pro".to_string(),
            args: vec![],
            prompt: "Do the thing".to_string(),
            cwd: "/tmp/wt".to_string(),
            inactivity_timeout_secs: 300,
            app_ctx: Some(test_app_ctx()),
            session_id_to_resume: None,
            fresh_prompt: None,
        };

        let instruction = build_system_instruction(&ctx);
        let text = instruction["parts"][0]["text"].as_str().unwrap();
        assert!(text.contains("42_story_test"));
        assert!(text.contains("/tmp/wt"));
    }

    #[test]
    fn build_generate_content_request_includes_tools() {
        let system = json!({"parts": [{"text": "system"}]});
        let contents = vec![json!({"role": "user", "parts": [{"text": "hello"}]})];
        let tools = vec![GeminiFunctionDeclaration {
            name: "my_tool".to_string(),
            description: "A tool".to_string(),
            parameters: Some(json!({"type": "object", "properties": {"x": {"type": "string"}}})),
        }];

        let body = build_generate_content_request(&system, &contents, &tools);
        assert!(body["tools"][0]["functionDeclarations"].is_array());
        assert_eq!(
            body["tools"][0]["functionDeclarations"][0]["name"],
            "my_tool"
        );
    }

    #[test]
    fn build_generate_content_request_no_tools() {
        let system = json!({"parts": [{"text": "system"}]});
        let contents = vec![json!({"role": "user", "parts": [{"text": "hello"}]})];
        let tools: Vec<GeminiFunctionDeclaration> = vec![];

        let body = build_generate_content_request(&system, &contents, &tools);
        assert!(body.get("tools").is_none());
    }

    #[test]
    fn parse_usage_metadata_valid() {
        let response = json!({
            "usageMetadata": {
                "promptTokenCount": 100,
                "candidatesTokenCount": 50,
                "totalTokenCount": 150
            }
        });

        let usage = parse_usage_metadata(&response).unwrap();
        assert_eq!(usage.input_tokens, 100);
        assert_eq!(usage.output_tokens, 50);
        assert_eq!(usage.cache_creation_input_tokens, 0);
        assert_eq!(usage.total_cost_usd, 0.0);
    }

    #[test]
    fn parse_usage_metadata_missing() {
        let response = json!({"candidates": []});
        assert!(parse_usage_metadata(&response).is_none());
    }
}