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::{json, Value}; 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 { 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 = 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 { 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 { 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 { 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 { 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 })) } }