Emit TurnDiff

This commit is contained in:
Gabriel Peal
2025-07-30 21:31:11 -07:00
parent 38e0461af5
commit 756708a671
7 changed files with 64 additions and 8 deletions

View File

@@ -1,12 +1,14 @@
use crate::codex::Session;
use crate::models::FunctionCallOutputPayload;
use crate::models::ResponseInputItem;
use crate::patch_accumulator::PatchAccumulator;
use crate::protocol::Event;
use crate::protocol::EventMsg;
use crate::protocol::FileChange;
use crate::protocol::PatchApplyBeginEvent;
use crate::protocol::PatchApplyEndEvent;
use crate::protocol::ReviewDecision;
use crate::protocol::TurnDiffEvent;
use crate::safety::SafetyCheck;
use crate::safety::assess_patch_safety;
use anyhow::Context;
@@ -42,6 +44,7 @@ impl From<ResponseInputItem> for InternalApplyPatchInvocation {
pub(crate) async fn apply_patch(
sess: &Session,
patch_accumulator: &mut PatchAccumulator,
sub_id: &str,
call_id: &str,
action: ApplyPatchAction,
@@ -148,6 +151,8 @@ pub(crate) async fn apply_patch(
})
.await;
let _ = patch_accumulator.on_patch_begin(&convert_apply_patch_to_protocol(&action));
let mut stdout = Vec::new();
let mut stderr = Vec::new();
// Enforce writable roots. If a write is blocked, collect offending root
@@ -242,6 +247,20 @@ pub(crate) async fn apply_patch(
})
.await;
let _ = patch_accumulator.on_patch_end(convert_apply_patch_to_protocol(&action));
if let Some(unified_diff) = &patch_accumulator.unified_diff {
let _ = sess
.tx_event
.send(Event {
id: sub_id.to_owned(),
msg: EventMsg::TurnDiff(TurnDiffEvent {
unified_diff: unified_diff.clone(),
}),
})
.await;
}
let item = match result {
Ok(_) => ResponseInputItem::FunctionCallOutput {
call_id: call_id.to_owned(),

View File

@@ -59,6 +59,7 @@ use crate::models::ReasoningItemReasoningSummary;
use crate::models::ResponseInputItem;
use crate::models::ResponseItem;
use crate::models::ShellToolCallParams;
use crate::patch_accumulator::PatchAccumulator;
use crate::plan_tool::handle_update_plan;
use crate::project_doc::get_user_instructions;
use crate::protocol::AgentMessageDeltaEvent;
@@ -1074,10 +1075,11 @@ async fn run_turn(
extra_tools,
base_instructions_override: sess.base_instructions.clone(),
};
let mut patch_accumulator = PatchAccumulator::new();
let mut retries = 0;
loop {
match try_run_turn(sess, &sub_id, &prompt).await {
match try_run_turn(sess, &sub_id, &prompt, &mut patch_accumulator).await {
Ok(output) => return Ok(output),
Err(CodexErr::Interrupted) => return Err(CodexErr::Interrupted),
Err(CodexErr::EnvVar(var)) => return Err(CodexErr::EnvVar(var)),
@@ -1125,6 +1127,7 @@ async fn try_run_turn(
sess: &Session,
sub_id: &str,
prompt: &Prompt,
patch_accumulator: &mut PatchAccumulator,
) -> CodexResult<Vec<ProcessedResponseItem>> {
// call_ids that are part of this response.
let completed_call_ids = prompt
@@ -1210,7 +1213,8 @@ async fn try_run_turn(
match event {
ResponseEvent::Created => {}
ResponseEvent::OutputItemDone(item) => {
let response = handle_response_item(sess, sub_id, item.clone()).await?;
let response =
handle_response_item(sess, patch_accumulator, sub_id, item.clone()).await?;
output.push(ProcessedResponseItem { item, response });
}
@@ -1250,6 +1254,7 @@ async fn try_run_turn(
async fn handle_response_item(
sess: &Session,
patch_accumulator: &mut PatchAccumulator,
sub_id: &str,
item: ResponseItem,
) -> CodexResult<Option<ResponseInputItem>> {
@@ -1287,7 +1292,17 @@ async fn handle_response_item(
..
} => {
info!("FunctionCall: {arguments}");
Some(handle_function_call(sess, sub_id.to_string(), name, arguments, call_id).await)
Some(
handle_function_call(
sess,
patch_accumulator,
sub_id.to_string(),
name,
arguments,
call_id,
)
.await,
)
}
ResponseItem::LocalShellCall {
id,
@@ -1322,6 +1337,7 @@ async fn handle_response_item(
handle_container_exec_with_params(
exec_params,
sess,
patch_accumulator,
sub_id.to_string(),
effective_call_id,
)
@@ -1339,6 +1355,7 @@ async fn handle_response_item(
async fn handle_function_call(
sess: &Session,
patch_accumulator: &mut PatchAccumulator,
sub_id: String,
name: String,
arguments: String,
@@ -1352,7 +1369,8 @@ async fn handle_function_call(
return *output;
}
};
handle_container_exec_with_params(params, sess, sub_id, call_id).await
handle_container_exec_with_params(params, sess, patch_accumulator, sub_id, call_id)
.await
}
"update_plan" => handle_update_plan(sess, arguments, sub_id, call_id).await,
_ => {
@@ -1426,6 +1444,7 @@ fn maybe_run_with_user_profile(params: ExecParams, sess: &Session) -> ExecParams
async fn handle_container_exec_with_params(
params: ExecParams,
sess: &Session,
patch_accumulator: &mut PatchAccumulator,
sub_id: String,
call_id: String,
) -> ResponseInputItem {
@@ -1433,7 +1452,9 @@ async fn handle_container_exec_with_params(
let apply_patch_action_for_exec =
match maybe_parse_apply_patch_verified(&params.command, &params.cwd) {
MaybeApplyPatchVerified::Body(changes) => {
match apply_patch::apply_patch(sess, &sub_id, &call_id, changes).await {
match apply_patch::apply_patch(sess, patch_accumulator, &sub_id, &call_id, changes)
.await
{
InternalApplyPatchInvocation::Output(item) => return item,
InternalApplyPatchInvocation::DelegateToExec(action) => Some(action),
}

View File

@@ -6,7 +6,6 @@ use std::process::Command;
use anyhow::Context;
use anyhow::Result;
use anyhow::bail;
use tempfile::TempDir;
use uuid::Uuid;
@@ -24,8 +23,8 @@ pub struct PatchAccumulator {
/// Internal filename -> external path as of baseline commit.
internal_to_baseline_external: HashMap<String, PathBuf>,
/// Internal filename -> external path as of current accumulated state (after applying all changes).
/// This is where renames are tracked.
internal_to_current_external: HashMap<String, PathBuf>,
/// All change sets in the order they occurred.
changes: Vec<HashMap<PathBuf, FileChange>>,
/// Aggregated unified diff for all accumulated changes across files.
pub unified_diff: Option<String>,
@@ -36,7 +35,7 @@ impl PatchAccumulator {
Self::default()
}
/// Ensure we have an initialized repository and a baseline snapshot of any new files.
/// Front-run apply patch calls to track the starting contents of any modified files.
pub fn on_patch_begin(&mut self, changes: &HashMap<PathBuf, FileChange>) -> Result<()> {
self.ensure_repo_init()?;
let repo_dir = self.repo_dir()?.to_path_buf();
@@ -60,6 +59,7 @@ impl PatchAccumulator {
fs::write(&internal_path, contents).with_context(|| {
format!("failed to write baseline file {}", internal_path.display())
})?;
// Stage the starting contents of the file to be included in the baseline snapshot.
run_git(&repo_dir, &["add", &internal])?;
staged_new_baseline = true;
}
@@ -67,6 +67,7 @@ impl PatchAccumulator {
}
// If new baseline files were staged, commit them and update the baseline commit id.
// Only the original file contents are staged so the baseline will always contain only the starting contents of each file.
if staged_new_baseline {
run_git(&repo_dir, &["commit", "-m", "Baseline snapshot"])?;
let id = run_git(&repo_dir, &["rev-parse", "HEAD"])?;

View File

@@ -334,6 +334,8 @@ pub enum EventMsg {
/// Notification that a patch application has finished.
PatchApplyEnd(PatchApplyEndEvent),
TurnDiff(TurnDiffEvent),
/// Response to GetHistoryEntryRequest.
GetHistoryEntryResponse(GetHistoryEntryResponseEvent),
@@ -527,6 +529,11 @@ pub struct PatchApplyEndEvent {
pub changes: HashMap<PathBuf, FileChange>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct TurnDiffEvent {
pub unified_diff: String,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct GetHistoryEntryResponseEvent {
pub offset: usize,

View File

@@ -20,6 +20,7 @@ use codex_core::protocol::PatchApplyEndEvent;
use codex_core::protocol::SessionConfiguredEvent;
use codex_core::protocol::TaskCompleteEvent;
use codex_core::protocol::TokenUsage;
use codex_core::protocol::TurnDiffEvent;
use owo_colors::OwoColorize;
use owo_colors::Style;
use shlex::try_join;
@@ -432,6 +433,10 @@ impl EventProcessor for EventProcessorWithHumanOutput {
println!("{}", line.style(self.dimmed));
}
}
EventMsg::TurnDiff(TurnDiffEvent { unified_diff }) => {
ts_println!(self, "{}", "turn diff:".style(self.magenta));
println!("{unified_diff}");
}
EventMsg::ExecApprovalRequest(_) => {
// Should we exit?
}

View File

@@ -262,6 +262,7 @@ async fn run_codex_tool_session_inner(
| EventMsg::BackgroundEvent(_)
| EventMsg::PatchApplyBegin(_)
| EventMsg::PatchApplyEnd(_)
| EventMsg::TurnDiff(_)
| EventMsg::GetHistoryEntryResponse(_)
| EventMsg::PlanUpdate(_)
| EventMsg::ShutdownComplete => {

View File

@@ -247,6 +247,7 @@ pub enum ClientNotification {
#[allow(clippy::expect_used)]
#[allow(clippy::unwrap_used)]
mod tests {
use std::collections::HashMap;
use std::path::PathBuf;
use super::*;
@@ -903,6 +904,7 @@ mod tests {
stdout: "ok".into(),
stderr: "".into(),
success: true,
changes: HashMap::new(),
}),
};