gitdataai/libs/service/git_tools/branch.rs

//! Git branch tools.

use super::ctx::GitToolCtx;
use agent::{ToolDefinition, ToolHandler, ToolParam, ToolRegistry, ToolSchema};
use std::collections::HashMap;

async fn git_branch_list_exec(ctx: GitToolCtx, args: serde_json::Value) -> Result<serde_json::Value, String> {
    let p: serde_json::Map<String, serde_json::Value> = serde_json::from_value(args).map_err(|e| e.to_string())?;
    let project_name = p.get("project_name").and_then(|v| v.as_str()).ok_or("missing project_name")?;
    let repo_name = p.get("repo_name").and_then(|v| v.as_str()).ok_or("missing repo_name")?;
    let remote_only = p.get("remote_only").and_then(|v| v.as_bool()).unwrap_or(false);

    let domain = ctx.open_repo(project_name, repo_name).await?;
    let branches = domain.branch_list(remote_only).map_err(|e| e.to_string())?;

    let result: Vec<_> = branches.iter().map(|b| {
        let oid = b.oid.to_string();
        serde_json::json!({
            "name": b.name, "oid": oid.clone(), "short_oid": oid.get(..7).unwrap_or(&oid).to_string(),
            "is_head": b.is_head, "is_remote": b.is_remote, "is_current": b.is_current,
            "upstream": b.upstream.clone()
        })
    }).collect();

    Ok(serde_json::to_value(result).map_err(|e| e.to_string())?)
}

async fn git_branch_info_exec(ctx: GitToolCtx, args: serde_json::Value) -> Result<serde_json::Value, String> {
    let p: serde_json::Map<String, serde_json::Value> = serde_json::from_value(args).map_err(|e| e.to_string())?;
    let project_name = p.get("project_name").and_then(|v| v.as_str()).ok_or("missing project_name")?;
    let repo_name = p.get("repo_name").and_then(|v| v.as_str()).ok_or("missing repo_name")?;
    let name = p.get("name").and_then(|v| v.as_str()).ok_or("missing name")?;

    let domain = ctx.open_repo(project_name, repo_name).await?;
    let info = domain.branch_get(name).map_err(|e| e.to_string())?;

    let ahead_behind = if let Some(ref upstream) = info.upstream {
        let (ahead, behind) = domain.branch_ahead_behind(name, upstream).unwrap_or((0, 0));
        Some(serde_json::json!({ "ahead": ahead, "behind": behind }))
    } else { None };

    Ok(serde_json::json!({
        "branch": { "name": info.name, "oid": info.oid.to_string(), "is_head": info.is_head,
            "is_remote": info.is_remote, "is_current": info.is_current, "upstream": info.upstream },
        "ahead_behind": ahead_behind
    }))
}

async fn git_branches_merged_exec(ctx: GitToolCtx, args: serde_json::Value) -> Result<serde_json::Value, String> {
    let p: serde_json::Map<String, serde_json::Value> = serde_json::from_value(args).map_err(|e| e.to_string())?;
    let project_name = p.get("project_name").and_then(|v| v.as_str()).ok_or("missing project_name")?;
    let repo_name = p.get("repo_name").and_then(|v| v.as_str()).ok_or("missing repo_name")?;
    let branch = p.get("branch").and_then(|v| v.as_str()).ok_or("missing branch")?;
    let into = p.get("into").and_then(|v| v.as_str()).map(|s| s.to_string()).unwrap_or_else(|| "main".to_string());

    let domain = ctx.open_repo(project_name, repo_name).await?;
    let is_merged = domain.branch_is_merged(branch, &into).map_err(|e| e.to_string())?;
    let merge_base = domain.merge_base(&git::commit::types::CommitOid::new(branch), &git::commit::types::CommitOid::new(&into))
        .map(|oid| oid.to_string()).ok();

    Ok(serde_json::json!({ "branch": branch, "into": into, "is_merged": is_merged, "merge_base": merge_base }))
}

async fn git_branch_diff_exec(ctx: GitToolCtx, args: serde_json::Value) -> Result<serde_json::Value, String> {
    let p: serde_json::Map<String, serde_json::Value> = serde_json::from_value(args).map_err(|e| e.to_string())?;
    let project_name = p.get("project_name").and_then(|v| v.as_str()).ok_or("missing project_name")?;
    let repo_name = p.get("repo_name").and_then(|v| v.as_str()).ok_or("missing repo_name")?;
    let local = p.get("local").and_then(|v| v.as_str()).ok_or("missing local")?;
    let remote = p.get("remote").and_then(|v| v.as_str()).unwrap_or(local).to_string();

    let domain = ctx.open_repo(project_name, repo_name).await?;
    let diff = domain.branch_diff(local, &remote).map_err(|e| e.to_string())?;

    Ok(serde_json::json!({ "ahead": diff.ahead, "behind": diff.behind, "diverged": diff.diverged }))
}

pub fn register_git_tools(registry: &mut ToolRegistry) {
    let mut p = HashMap::from([
        ("project_name".into(), ToolParam {
            name: "project_name".into(), param_type: "string".into(),
            description: Some("Project name (slug)".into()),
            required: true, properties: None, items: None,
        }),
        ("repo_name".into(), ToolParam {
            name: "repo_name".into(), param_type: "string".into(),
            description: Some("Repository name".into()),
            required: true, properties: None, items: None,
        }),
    ]);

    // git_branch_list
    p.insert("remote_only".into(), ToolParam {
        name: "remote_only".into(), param_type: "boolean".into(),
        description: Some("If true, list only remote-tracking branches".into()),
        required: false, properties: None, items: None,
    });
    let schema = ToolSchema { schema_type: "object".into(), properties: Some(std::mem::take(&mut p)), required: Some(vec!["project_name".into(), "repo_name".into()]) };
    p = HashMap::from([
        ("project_name".into(), ToolParam { name: "project_name".into(), param_type: "string".into(), description: Some("Project name (slug)".into()), required: true, properties: None, items: None }),
        ("repo_name".into(), ToolParam { name: "repo_name".into(), param_type: "string".into(), description: Some("Repository name".into()), required: true, properties: None, items: None }),
    ]);
    registry.register(
        ToolDefinition::new("git_branch_list").description("List all local or remote branches with their current HEAD commit.").parameters(schema),
        ToolHandler::new(move |ctx, args| {
            let gctx = super::ctx::GitToolCtx::new(ctx);
            Box::pin(async move {
                git_branch_list_exec(gctx, args).await.map_err(agent::ToolError::ExecutionError)
            })
        }),
    );

    // git_branch_info
    p.insert("name".into(), ToolParam {
        name: "name".into(), param_type: "string".into(),
        description: Some("Branch name (e.g. main, feature/my-branch)".into()),
        required: true, properties: None, items: None,
    });
    let schema = ToolSchema { schema_type: "object".into(), properties: Some(std::mem::take(&mut p)), required: Some(vec!["project_name".into(), "repo_name".into(), "name".into()]) };
    registry.register(
        ToolDefinition::new("git_branch_info").description("Get detailed information about a specific branch including ahead/behind status vs upstream.").parameters(schema),
        ToolHandler::new(move |ctx, args| {
            let gctx = super::ctx::GitToolCtx::new(ctx);
            Box::pin(async move {
                git_branch_info_exec(gctx, args).await.map_err(agent::ToolError::ExecutionError)
            })
        }),
    );

    // git_branches_merged
    p = HashMap::from([
        ("project_name".into(), ToolParam { name: "project_name".into(), param_type: "string".into(), description: Some("Project name (slug)".into()), required: true, properties: None, items: None }),
        ("repo_name".into(), ToolParam { name: "repo_name".into(), param_type: "string".into(), description: Some("Repository name".into()), required: true, properties: None, items: None }),
        ("branch".into(), ToolParam { name: "branch".into(), param_type: "string".into(), description: Some("Branch to check (source)".into()), required: true, properties: None, items: None }),
        ("into".into(), ToolParam { name: "into".into(), param_type: "string".into(), description: Some("Branch to check against (target, default: main)".into()), required: false, properties: None, items: None }),
    ]);
    let schema = ToolSchema { schema_type: "object".into(), properties: Some(p), required: Some(vec!["project_name".into(), "repo_name".into(), "branch".into()]) };
    registry.register(
        ToolDefinition::new("git_branches_merged").description("Check whether a branch has been merged into another branch, and find the merge base.").parameters(schema),
        ToolHandler::new(move |ctx, args| {
            let gctx = super::ctx::GitToolCtx::new(ctx);
            Box::pin(async move {
                git_branches_merged_exec(gctx, args).await.map_err(agent::ToolError::ExecutionError)
            })
        }),
    );

    // git_branch_diff
    let p = HashMap::from([
        ("project_name".into(), ToolParam { name: "project_name".into(), param_type: "string".into(), description: Some("Project name (slug)".into()), required: true, properties: None, items: None }),
        ("repo_name".into(), ToolParam { name: "repo_name".into(), param_type: "string".into(), description: Some("Repository name".into()), required: true, properties: None, items: None }),
        ("local".into(), ToolParam { name: "local".into(), param_type: "string".into(), description: Some("Local branch name".into()), required: true, properties: None, items: None }),
        ("remote".into(), ToolParam { name: "remote".into(), param_type: "string".into(), description: Some("Remote branch name (defaults to local)".into()), required: false, properties: None, items: None }),
    ]);
    let schema = ToolSchema { schema_type: "object".into(), properties: Some(p), required: Some(vec!["project_name".into(), "repo_name".into(), "local".into()]) };
    registry.register(
        ToolDefinition::new("git_branch_diff").description("Compare a local branch against its remote counterpart to see how many commits ahead/behind they are.").parameters(schema),
        ToolHandler::new(move |ctx, args| {
            let gctx = super::ctx::GitToolCtx::new(ctx);
            Box::pin(async move {
                git_branch_diff_exec(gctx, args).await.map_err(agent::ToolError::ExecutionError)
            })
        }),
    );
}