gitdataai/lib/service/agent/git_tools/branch.rs

use ai::error::{AiError, AiResult};
use ai::tool::tools::FunctionCall;
use async_trait::async_trait;
use git::rpc::proto as p;
use git::rpc::proto::branch_service_client::BranchServiceClient;
use serde_json::{Value, json};

use super::helpers::{arg_str, git_ctx, require_repo_member, rpc_err};
use crate::agent::run::AppAgentContext;

pub struct GitBranchListTool;

impl GitBranchListTool {
    pub fn new() -> Self {
        Self
    }
}

impl Default for GitBranchListTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl FunctionCall for GitBranchListTool {
    type Context = AppAgentContext;

    fn name(&self) -> &'static str {
        "git_branch_list"
    }

    fn description(&self) -> &'static str {
        "List all branches in a repository with their HEAD commit OID."
    }

    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "workspace": { "type": "string", "description": "Workspace name" },
                "repo": { "type": "string", "description": "Repository name" }
            },
            "required": ["workspace", "repo"]
        })
    }

    async fn call(
        &self,
        ctx: &mut AppAgentContext,
        args: Value,
    ) -> AiResult<Value> {
        let git = git_ctx(ctx)?;
        let workspace = arg_str(&args, "workspace")?;
        let repo_name = arg_str(&args, "repo")?;

        let repo =
            require_repo_member(git, ctx.user_id, workspace, repo_name).await?;

        let mut client = BranchServiceClient::new(git.channel.clone());
        let resp = client
            .branch_list(p::BranchListRequest {
                repo_id: repo.id.to_string(),
            })
            .await
            .map_err(rpc_err)?
            .into_inner();

        let branches: Vec<Value> = resp.branches.iter().map(|b| json!({
            "name": b.name,
            "oid": b.oid.as_ref().map(|o| &o.value).unwrap_or(&String::new()),
            "is_head": b.is_head,
            "is_current": b.is_current,
        })).collect();

        Ok(json!({ "branches": branches, "count": branches.len() }))
    }
}

pub struct GitBranchInfoTool;

impl GitBranchInfoTool {
    pub fn new() -> Self {
        Self
    }
}

impl Default for GitBranchInfoTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl FunctionCall for GitBranchInfoTool {
    type Context = AppAgentContext;

    fn name(&self) -> &'static str {
        "git_branch_info"
    }

    fn description(&self) -> &'static str {
        "Get detailed information about a single branch, including its HEAD OID and upstream."
    }

    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "workspace": { "type": "string", "description": "Workspace name" },
                "repo": { "type": "string", "description": "Repository name" },
                "branch": { "type": "string", "description": "Branch name" }
            },
            "required": ["workspace", "repo", "branch"]
        })
    }

    async fn call(
        &self,
        ctx: &mut AppAgentContext,
        args: Value,
    ) -> AiResult<Value> {
        let git = git_ctx(ctx)?;
        let workspace = arg_str(&args, "workspace")?;
        let repo_name = arg_str(&args, "repo")?;
        let branch = arg_str(&args, "branch")?;

        let repo =
            require_repo_member(git, ctx.user_id, workspace, repo_name).await?;

        let mut client = BranchServiceClient::new(git.channel.clone());
        let resp = client
            .branch_info(p::BranchInfoRequest {
                repo_id: repo.id.to_string(),
                branch: branch.to_string(),
            })
            .await
            .map_err(rpc_err)?
            .into_inner();

        let b = resp.branch.ok_or_else(|| {
            AiError::Config(format!("branch '{branch}' not found"))
        })?;
        Ok(json!({
            "name": b.name,
            "oid": b.oid.as_ref().map(|o| &o.value),
            "is_head": b.is_head,
            "is_current": b.is_current,
            "upstream": b.upstream,
        }))
    }
}

pub struct GitBranchAheadBehindTool;

impl GitBranchAheadBehindTool {
    pub fn new() -> Self {
        Self
    }
}

impl Default for GitBranchAheadBehindTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl FunctionCall for GitBranchAheadBehindTool {
    type Context = AppAgentContext;

    fn name(&self) -> &'static str {
        "git_branch_ahead_behind"
    }

    fn description(&self) -> &'static str {
        "Compare a local branch with its remote tracking branch. Returns commits ahead and behind."
    }

    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "workspace": { "type": "string", "description": "Workspace name" },
                "repo": { "type": "string", "description": "Repository name" },
                "local_branch": { "type": "string", "description": "Local branch name" },
                "remote_branch": { "type": "string", "description": "Remote tracking branch name" }
            },
            "required": ["workspace", "repo", "local_branch", "remote_branch"]
        })
    }

    async fn call(
        &self,
        ctx: &mut AppAgentContext,
        args: Value,
    ) -> AiResult<Value> {
        let git = git_ctx(ctx)?;
        let workspace = arg_str(&args, "workspace")?;
        let repo_name = arg_str(&args, "repo")?;
        let local_branch = arg_str(&args, "local_branch")?;
        let remote_branch = arg_str(&args, "remote_branch")?;

        let repo =
            require_repo_member(git, ctx.user_id, workspace, repo_name).await?;

        let mut client = BranchServiceClient::new(git.channel.clone());
        let resp = client
            .branch_ahead_behind(p::BranchAheadBehindRequest {
                repo_id: repo.id.to_string(),
                local_branch: local_branch.to_string(),
                remote_branch: remote_branch.to_string(),
            })
            .await
            .map_err(rpc_err)?
            .into_inner();

        Ok(json!({ "ahead": resp.ahead, "behind": resp.behind }))
    }
}

pub struct GitBranchDeleteTool;

impl GitBranchDeleteTool {
    pub fn new() -> Self {
        Self
    }
}

impl Default for GitBranchDeleteTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl FunctionCall for GitBranchDeleteTool {
    type Context = AppAgentContext;

    fn name(&self) -> &'static str {
        "git_branch_delete"
    }

    fn description(&self) -> &'static str {
        "Delete a branch from the repository. Requires write access."
    }

    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "workspace": { "type": "string", "description": "Workspace name" },
                "repo": { "type": "string", "description": "Repository name" },
                "name": { "type": "string", "description": "Branch name to delete" },
                "force": { "type": "boolean", "description": "Force delete (even if not merged)" }
            },
            "required": ["workspace", "repo", "name"]
        })
    }

    async fn call(
        &self,
        ctx: &mut AppAgentContext,
        args: Value,
    ) -> AiResult<Value> {
        let git = git_ctx(ctx)?;
        let workspace = arg_str(&args, "workspace")?;
        let repo_name = arg_str(&args, "repo")?;
        let name = arg_str(&args, "name")?;
        let force =
            args.get("force").and_then(|v| v.as_bool()).unwrap_or(false);

        let repo =
            require_repo_member(git, ctx.user_id, workspace, repo_name).await?;

        let mut client = BranchServiceClient::new(git.channel.clone());
        client
            .branch_delete(p::BranchDeleteRequest {
                repo_id: repo.id.to_string(),
                params: Some(p::BranchDeleteParams {
                    name: name.to_string(),
                    force,
                }),
            })
            .await
            .map_err(rpc_err)?;

        Ok(json!({ "success": true, "branch": name }))
    }
}

pub struct GitCreateBranchTool;

impl GitCreateBranchTool {
    pub fn new() -> Self {
        Self
    }
}

impl Default for GitCreateBranchTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl FunctionCall for GitCreateBranchTool {
    type Context = AppAgentContext;

    fn name(&self) -> &'static str {
        "git_create_branch"
    }

    fn description(&self) -> &'static str {
        "Create a new branch in a repository. Requires write access."
    }

    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "workspace": { "type": "string", "description": "Workspace name" },
                "repo": { "type": "string", "description": "Repository name" },
                "name": { "type": "string", "description": "New branch name" },
                "oid": { "type": "string", "description": "Commit OID to branch from" },
                "force": { "type": "boolean", "description": "Force create (overwrite existing)" }
            },
            "required": ["workspace", "repo", "name", "oid"]
        })
    }

    async fn call(
        &self,
        ctx: &mut AppAgentContext,
        args: Value,
    ) -> AiResult<Value> {
        let git = git_ctx(ctx)?;
        let workspace = arg_str(&args, "workspace")?;
        let repo_name = arg_str(&args, "repo")?;
        let name = arg_str(&args, "name")?;
        let oid = arg_str(&args, "oid")?;
        let force =
            args.get("force").and_then(|v| v.as_bool()).unwrap_or(false);

        let repo =
            require_repo_member(git, ctx.user_id, workspace, repo_name).await?;

        let mut client = BranchServiceClient::new(git.channel.clone());
        client
            .branch_fork(p::BranchForkRequest {
                repo_id: repo.id.to_string(),
                params: Some(p::BranchForkParams {
                    name: name.to_string(),
                    oid: Some(p::ObjectId {
                        value: oid.to_string(),
                    }),
                    force,
                }),
            })
            .await
            .map_err(rpc_err)?;

        Ok(json!({ "success": true, "branch": name, "oid": oid }))
    }
}