use anyhow::Result;
use crossterm::event;
use ratatui::{backend::CrosstermBackend, Terminal};
use std::io::Stdout;
use tokio::sync::mpsc;
use crate::ai::adapter::{CommentSeverity, ReviewComment as AiReviewComment};
use crate::ai::orchestrator::{OrchestratorCommand, RallyEvent};
use crate::ai::prompt_loader::{PromptLoader, PromptSource};
use crate::ai::{
Context, Orchestrator, RallyState, ReviewAction as AiReviewAction, ReviewerOutput,
};
use crate::cache::load_local_review_comments;
use crate::keybinding::{event_to_keybinding, SequenceMatch};
use crate::ui;
use super::types::*;
use super::{App, AppState};
impl App {
pub(crate) async fn handle_ai_rally_input(
&mut self,
key: event::KeyEvent,
terminal: &mut Terminal<CrosstermBackend<Stdout>>,
) -> Result<()> {
let kb = self.config.keybindings.clone();
let in_log_detail = self
.ai_rally_state
.as_ref()
.is_some_and(|s| s.showing_log_detail);
if in_log_detail {
if self.matches_single_key(&key, &kb.quit)
|| self.matches_single_key(&key, &kb.open_panel)
{
if let Some(ref mut rally_state) = self.ai_rally_state {
rally_state.showing_log_detail = false;
}
}
return Ok(());
}
if let Some(kb_event) = event_to_keybinding(&key) {
self.check_sequence_timeout();
if !self.pending_keys.is_empty() {
self.push_pending_key(kb_event);
if self.try_match_sequence(&kb.jump_to_first) == SequenceMatch::Full {
self.clear_pending_keys();
if let Some(ref mut rally_state) = self.ai_rally_state {
if !rally_state.logs.is_empty() {
rally_state.selected_log_index = Some(0);
rally_state.log_scroll_offset = 1;
}
}
return Ok(());
}
self.clear_pending_keys();
} else if self.key_could_match_sequence(&key, &kb.jump_to_first) {
self.push_pending_key(kb_event);
return Ok(());
}
}
if self.matches_single_key(&key, &kb.rally_background) {
if self
.ai_rally_state
.as_ref()
.and_then(|s| s.pending_config_warning.as_ref())
.is_some()
{
return Ok(());
}
self.state = AppState::FileList;
} else if self.matches_single_key(&key, &kb.quit) {
if self
.ai_rally_state
.as_ref()
.and_then(|s| s.pending_config_warning.as_ref())
.is_some()
{
self.pending_rally_context = None;
self.cleanup_rally_state();
self.state = AppState::FileList;
return Ok(());
}
if let Some(ref state) = self.ai_rally_state {
if matches!(
state.state,
RallyState::WaitingForClarification
| RallyState::WaitingForPermission
| RallyState::WaitingForPostConfirmation
) {
self.send_rally_command(OrchestratorCommand::Abort);
}
}
if let Some(handle) = self.rally_abort_handle.take() {
handle.abort();
}
self.cleanup_rally_state();
self.state = AppState::FileList;
} else if self.matches_single_key(&key, &kb.confirm_yes) {
if self
.ai_rally_state
.as_ref()
.and_then(|s| s.pending_config_warning.as_ref())
.is_some()
{
if let Some(ref mut rally_state) = self.ai_rally_state {
rally_state.pending_config_warning = None;
}
if let Some(context) = self.pending_rally_context.take() {
if let Some(prompt_loader) = self.pending_rally_prompt_loader.take() {
self.spawn_rally_orchestrator(context, prompt_loader, None);
}
}
return Ok(());
}
let current_state = self
.ai_rally_state
.as_ref()
.map(|s| s.state)
.unwrap_or(RallyState::Error);
match current_state {
RallyState::WaitingForPermission => {
self.send_rally_command(OrchestratorCommand::PermissionResponse(true));
if let Some(ref mut rally_state) = self.ai_rally_state {
rally_state.pending_permission = None;
rally_state.state = RallyState::RevieweeFix;
rally_state.push_log(LogEntry::new(
LogEventType::Info,
"Permission granted, continuing...".to_string(),
));
}
}
RallyState::WaitingForClarification => {
let question = self
.ai_rally_state
.as_ref()
.and_then(|s| s.pending_question.clone())
.unwrap_or_default();
self.open_clarification_editor_sync(&question, terminal)?;
}
RallyState::WaitingForPostConfirmation => {
self.send_rally_command(OrchestratorCommand::PostConfirmResponse(true));
if let Some(ref mut rally_state) = self.ai_rally_state {
rally_state.pending_review_post = None;
rally_state.pending_fix_post = None;
rally_state.state = RallyState::RevieweeFix;
rally_state.push_log(LogEntry::new(
LogEventType::Info,
"Post approved, posting to PR...".to_string(),
));
}
}
_ => {}
}
} else if self.matches_single_key(&key, &kb.confirm_no) {
if self
.ai_rally_state
.as_ref()
.and_then(|s| s.pending_config_warning.as_ref())
.is_some()
{
self.pending_rally_context = None;
self.cleanup_rally_state();
self.state = AppState::FileList;
return Ok(());
}
let current_state = self
.ai_rally_state
.as_ref()
.map(|s| s.state)
.unwrap_or(RallyState::Error);
match current_state {
RallyState::WaitingForPermission => {
self.send_rally_command(OrchestratorCommand::PermissionResponse(false));
if let Some(ref mut rally_state) = self.ai_rally_state {
rally_state.pending_permission = None;
rally_state.push_log(LogEntry::new(
LogEventType::Info,
"Permission denied, continuing without it...".to_string(),
));
}
}
RallyState::WaitingForClarification => {
self.send_rally_command(OrchestratorCommand::SkipClarification);
if let Some(ref mut rally_state) = self.ai_rally_state {
rally_state.pending_question = None;
rally_state.push_log(LogEntry::new(
LogEventType::Info,
"Clarification skipped, continuing with best judgment...".to_string(),
));
}
}
RallyState::WaitingForPostConfirmation => {
self.send_rally_command(OrchestratorCommand::PostConfirmResponse(false));
if let Some(ref mut rally_state) = self.ai_rally_state {
rally_state.pending_review_post = None;
rally_state.pending_fix_post = None;
rally_state.state = RallyState::RevieweeFix;
rally_state.push_log(LogEntry::new(
LogEventType::Info,
"Post skipped, continuing...".to_string(),
));
}
}
_ => {}
}
} else if self.matches_single_key(&key, &kb.retry) {
if let Some(ref state) = self.ai_rally_state {
if state.state == RallyState::Error {
if let Some(handle) = self.rally_abort_handle.take() {
handle.abort();
}
self.ai_rally_state = None;
self.rally_event_receiver = None;
self.state = AppState::FileList;
self.start_ai_rally();
}
}
} else if self.matches_single_key(&key, &kb.move_down) {
if let Some(ref mut rally_state) = self.ai_rally_state {
let total_logs = rally_state.logs.len();
if total_logs == 0 {
return Ok(());
}
let current = rally_state.selected_log_index.unwrap_or(0);
let new_index = (current + 1).min(total_logs.saturating_sub(1));
rally_state.selected_log_index = Some(new_index);
self.adjust_log_scroll_to_selection();
}
} else if self.matches_single_key(&key, &kb.move_up) {
if let Some(ref mut rally_state) = self.ai_rally_state {
let total_logs = rally_state.logs.len();
if total_logs == 0 {
return Ok(());
}
let current = rally_state
.selected_log_index
.unwrap_or(total_logs.saturating_sub(1));
let new_index = current.saturating_sub(1);
rally_state.selected_log_index = Some(new_index);
self.adjust_log_scroll_to_selection();
}
} else if self.matches_single_key(&key, &kb.page_down)
|| Self::is_shift_char_shortcut(&key, 'j')
{
if let Some(ref mut rally_state) = self.ai_rally_state {
let total_logs = rally_state.logs.len();
if total_logs == 0 {
return Ok(());
}
let page_step = rally_state.last_visible_log_height.saturating_sub(1).max(1);
let current = rally_state.selected_log_index.unwrap_or(0);
let new_index = (current + page_step).min(total_logs.saturating_sub(1));
rally_state.selected_log_index = Some(new_index);
self.adjust_log_scroll_to_selection();
}
} else if self.matches_single_key(&key, &kb.page_up)
|| Self::is_shift_char_shortcut(&key, 'k')
{
if let Some(ref mut rally_state) = self.ai_rally_state {
let total_logs = rally_state.logs.len();
if total_logs == 0 {
return Ok(());
}
let page_step = rally_state.last_visible_log_height.saturating_sub(1).max(1);
let current = rally_state
.selected_log_index
.unwrap_or(total_logs.saturating_sub(1));
let new_index = current.saturating_sub(page_step);
rally_state.selected_log_index = Some(new_index);
self.adjust_log_scroll_to_selection();
}
} else if self.matches_single_key(&key, &kb.open_panel) {
if let Some(ref mut rally_state) = self.ai_rally_state {
if rally_state.selected_log_index.is_some() && !rally_state.logs.is_empty() {
rally_state.showing_log_detail = true;
}
}
} else if self.matches_single_key(&key, &kb.jump_to_last) {
if let Some(ref mut rally_state) = self.ai_rally_state {
let total_logs = rally_state.logs.len();
if total_logs > 0 {
rally_state.selected_log_index = Some(total_logs.saturating_sub(1));
rally_state.log_scroll_offset = 0;
}
}
} else if self.matches_single_key(&key, &kb.rally_pause) {
if self
.ai_rally_state
.as_ref()
.and_then(|s| s.pending_config_warning.as_ref())
.is_some()
{
return Ok(());
}
if let Some(ref mut rally_state) = self.ai_rally_state {
if !matches!(
rally_state.state,
RallyState::ReviewerReviewing | RallyState::RevieweeFix
) {
return Ok(());
}
match rally_state.pause_state {
PauseState::Running => {
rally_state.pause_state = PauseState::PauseRequested;
self.send_rally_command(OrchestratorCommand::Pause);
if let Some(ref mut rs) = self.ai_rally_state {
rs.push_log(LogEntry::new(
LogEventType::Info,
"Pausing after current step...".to_string(),
));
}
}
PauseState::PauseRequested | PauseState::Paused => {
rally_state.pause_state = PauseState::Running;
self.send_rally_command(OrchestratorCommand::Resume);
if let Some(ref mut rs) = self.ai_rally_state {
rs.push_log(LogEntry::new(
LogEventType::Info,
"Resuming...".to_string(),
));
}
}
}
}
}
Ok(())
}
pub(crate) fn adjust_log_scroll_to_selection(&mut self) {
if let Some(ref mut rally_state) = self.ai_rally_state {
let Some(selected) = rally_state.selected_log_index else {
return;
};
let visible_height = rally_state.last_visible_log_height;
let total_logs = rally_state.logs.len();
let scroll_offset = if rally_state.log_scroll_offset == 0 {
total_logs.saturating_sub(visible_height)
} else {
rally_state.log_scroll_offset
};
if selected < scroll_offset {
rally_state.log_scroll_offset = selected.max(1);
} else if selected >= scroll_offset + visible_height {
rally_state.log_scroll_offset = selected.saturating_sub(visible_height - 1).max(1);
}
}
}
pub(crate) fn send_rally_command(&mut self, cmd: OrchestratorCommand) {
if let Some(ref sender) = self.rally_command_sender {
if sender.try_send(cmd).is_err() {
self.cleanup_rally_state();
}
}
}
pub(crate) fn cleanup_rally_state(&mut self) {
self.ai_rally_state = None;
self.rally_command_sender = None;
self.rally_event_receiver = None;
self.pending_rally_context = None;
self.pending_rally_prompt_loader = None;
self.pending_rally_seed_review = None;
if let Some(handle) = self.rally_abort_handle.take() {
handle.abort();
}
}
pub(crate) fn open_clarification_editor_sync(
&mut self,
question: &str,
terminal: &mut Terminal<CrosstermBackend<Stdout>>,
) -> Result<()> {
ui::restore_terminal(terminal)?;
let answer =
crate::editor::open_clarification_editor(self.config.editor.as_deref(), question)?;
*terminal = ui::setup_terminal()?;
if let Some(ref mut rally_state) = self.ai_rally_state {
rally_state.pending_question = None;
}
match answer {
Some(text) if !text.trim().is_empty() => {
self.send_rally_command(OrchestratorCommand::ClarificationResponse(text.clone()));
if let Some(ref mut rally_state) = self.ai_rally_state {
rally_state.push_log(LogEntry::new(
LogEventType::Info,
format!("Clarification provided: {}", text),
));
}
}
_ => {
self.send_rally_command(OrchestratorCommand::Abort);
if let Some(ref mut rally_state) = self.ai_rally_state {
rally_state.push_log(LogEntry::new(
LogEventType::Info,
"Clarification cancelled by user".to_string(),
));
}
}
}
Ok(())
}
pub(crate) fn resume_or_start_ai_rally(&mut self) {
if self.ai_rally_state.is_some() {
self.state = AppState::AiRally;
return;
}
self.start_ai_rally();
}
#[allow(dead_code)]
pub fn is_rally_running_in_background(&self) -> bool {
self.state != AppState::AiRally
&& self
.ai_rally_state
.as_ref()
.map(|s| s.state.is_active())
.unwrap_or(false)
}
pub fn has_background_rally(&self) -> bool {
self.state != AppState::AiRally && self.ai_rally_state.is_some()
}
#[allow(dead_code)]
pub fn is_background_rally_finished(&self) -> bool {
self.state != AppState::AiRally
&& self
.ai_rally_state
.as_ref()
.map(|s| s.state.is_finished())
.unwrap_or(false)
}
pub(crate) fn start_ai_rally(&mut self) {
let Some(pr) = self.pr() else {
return;
};
let seed_review = match self.build_seed_review_from_local_comments() {
Ok(seed_review) => seed_review,
Err(e) => {
self.cmt.submission_result = Some((
false,
format!("Failed to load local comments for AI Rally: {}", e),
));
self.cmt.submission_result_time = Some(std::time::Instant::now());
return;
}
};
let file_patches: Vec<(String, String)> = self
.files()
.iter()
.filter_map(|f| f.patch.as_ref().map(|p| (f.filename.clone(), p.clone())))
.collect();
let diff = file_patches
.iter()
.map(|(_, p)| p.as_str())
.collect::<Vec<_>>()
.join("\n");
let base_branch = if self.local_mode {
Self::detect_local_base_branch(self.working_dir.as_deref())
.unwrap_or_else(|| "main".to_string())
} else {
pr.base.ref_name.clone()
};
let context = Context {
repo: self.repo.clone(),
pr_number: self.pr_number(),
pr_title: pr.title.clone(),
pr_body: pr.body.clone(),
diff,
working_dir: self.working_dir.clone(),
head_sha: pr.head.sha.clone(),
base_branch,
external_comments: Vec::new(),
local_mode: self.local_mode,
file_patches,
};
let mut warnings: Vec<(String, String)> = crate::config::SENSITIVE_AI_KEYS
.iter()
.filter(|key| self.config.local_overrides.contains(**key))
.map(|key| {
let value = self.get_config_value_for_key(key);
(key.to_string(), value)
})
.collect();
let prompt_loader = PromptLoader::new(&self.config.ai, &self.config.project_root);
for (filename, source) in prompt_loader.resolve_all_sources() {
if let PromptSource::Local(path) = source {
warnings.push((
format!("local prompt: {}", filename),
path.display().to_string(),
));
}
}
self.ai_rally_state = Some(AiRallyState {
iteration: 0,
max_iterations: self.config.ai.max_iterations,
state: RallyState::Initializing,
history: Vec::new(),
logs: Vec::new(),
log_scroll_offset: 0,
selected_log_index: None,
showing_log_detail: false,
pending_question: None,
pending_permission: None,
pending_review_post: None,
pending_fix_post: None,
last_visible_log_height: 10,
pending_config_warning: if warnings.is_empty() {
None
} else {
Some(warnings)
},
pause_state: PauseState::Running,
});
self.state = AppState::AiRally;
if self
.ai_rally_state
.as_ref()
.and_then(|s| s.pending_config_warning.as_ref())
.is_some()
{
self.pending_rally_context = Some(context);
self.pending_rally_prompt_loader = Some(prompt_loader);
self.pending_rally_seed_review = seed_review;
return;
}
self.spawn_rally_orchestrator(context, prompt_loader, seed_review);
}
pub(crate) fn build_seed_review_from_local_comments(
&self,
) -> anyhow::Result<Option<ReviewerOutput>> {
if !self.local_mode {
return Ok(None);
}
let comments = load_local_review_comments(&self.repo, self.working_dir.as_deref())?;
let comments: Vec<AiReviewComment> = comments
.into_iter()
.filter_map(|entry| {
if entry.meta.is_resolved {
return None;
}
let line = entry.comment.line?;
Some(AiReviewComment {
path: entry.comment.path,
line,
body: entry.comment.body,
severity: CommentSeverity::Major,
})
})
.collect();
if comments.is_empty() {
return Ok(None);
}
let summary = format!(
"Address {} local comment{} from the user before re-reviewing.",
comments.len(),
if comments.len() == 1 { "" } else { "s" }
);
Ok(Some(ReviewerOutput {
action: AiReviewAction::RequestChanges,
summary,
blocking_issues: vec!["Resolve the user-provided local comments.".to_string()],
comments,
}))
}
pub(crate) fn spawn_rally_orchestrator(
&mut self,
context: Context,
prompt_loader: PromptLoader,
seed_review: Option<ReviewerOutput>,
) {
let (event_tx, event_rx) = mpsc::channel(100);
let (cmd_tx, cmd_rx) = mpsc::channel(10);
self.rally_event_receiver = Some(event_rx);
self.rally_command_sender = Some(cmd_tx);
let config = self.config.ai.clone();
let repo = self.repo.clone();
let pr_number = self.pr_number();
let handle = tokio::spawn(async move {
let orchestrator_result = Orchestrator::new(
&repo,
pr_number,
config,
event_tx.clone(),
Some(cmd_rx),
prompt_loader,
);
match orchestrator_result {
Ok(mut orchestrator) => {
orchestrator.set_context(context);
if let Some(seed_review) = seed_review {
orchestrator.set_seed_review(seed_review);
}
let _ = orchestrator.run().await;
}
Err(e) => {
let _ = event_tx
.send(RallyEvent::Error(format!(
"Failed to create orchestrator: {}",
e
)))
.await;
}
}
});
self.rally_abort_handle = Some(handle.abort_handle());
}
fn get_config_value_for_key(&self, key: &str) -> String {
match key {
"ai.reviewer_additional_tools" => {
format!("{:?}", self.config.ai.reviewer_additional_tools)
}
"ai.reviewee_additional_tools" => {
format!("{:?}", self.config.ai.reviewee_additional_tools)
}
"ai.auto_post" => format!("{}", self.config.ai.auto_post),
"ai.reviewer" => self.config.ai.reviewer.clone(),
"ai.reviewee" => self.config.ai.reviewee.clone(),
"ai.prompt_dir" => self
.config
.ai
.prompt_dir
.clone()
.unwrap_or_else(|| "(none)".to_string()),
_ => "(unknown)".to_string(),
}
}
}