Skip to content

Commit 751b2b8

Browse files
Merge branch 'gpeal/patch-accumulator' into dh--testing-20250801
2 parents e389a5d + feeb266 commit 751b2b8

File tree

8 files changed

+562
-17
lines changed

8 files changed

+562
-17
lines changed

codex-rs/Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

codex-rs/core/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ serde = { version = "1", features = ["derive"] }
3333
serde_json = "1"
3434
sha1 = "0.10.6"
3535
shlex = "1.3.0"
36+
similar = "2"
3637
strum_macros = "0.27.2"
3738
thiserror = "2.0.12"
3839
time = { version = "0.3", features = ["formatting", "local-offset", "macros"] }

codex-rs/core/src/codex.rs

Lines changed: 87 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -84,11 +84,13 @@ use crate::protocol::SandboxPolicy;
8484
use crate::protocol::SessionConfiguredEvent;
8585
use crate::protocol::Submission;
8686
use crate::protocol::TaskCompleteEvent;
87+
use crate::protocol::TurnDiffEvent;
8788
use crate::rollout::RolloutRecorder;
8889
use crate::safety::SafetyCheck;
8990
use crate::safety::assess_command_safety;
9091
use crate::safety::assess_safety_for_untrusted_command;
9192
use crate::shell;
93+
use crate::turn_diff_tracker::TurnDiffTracker;
9294
use crate::user_notification::UserNotification;
9395
use crate::util::backoff;
9496

@@ -361,7 +363,11 @@ impl Session {
361363
}
362364
}
363365

364-
async fn notify_exec_command_begin(&self, exec_command_context: ExecCommandContext) {
366+
async fn on_exec_command_begin(
367+
&self,
368+
turn_diff_tracker: &mut TurnDiffTracker,
369+
exec_command_context: ExecCommandContext,
370+
) {
365371
let ExecCommandContext {
366372
sub_id,
367373
call_id,
@@ -373,11 +379,15 @@ impl Session {
373379
Some(ApplyPatchCommandContext {
374380
user_explicitly_approved_this_action,
375381
changes,
376-
}) => EventMsg::PatchApplyBegin(PatchApplyBeginEvent {
377-
call_id,
378-
auto_approved: !user_explicitly_approved_this_action,
379-
changes,
380-
}),
382+
}) => {
383+
let _ = turn_diff_tracker.on_patch_begin(&changes);
384+
385+
EventMsg::PatchApplyBegin(PatchApplyBeginEvent {
386+
call_id,
387+
auto_approved: !user_explicitly_approved_this_action,
388+
changes,
389+
})
390+
}
381391
None => EventMsg::ExecCommandBegin(ExecCommandBeginEvent {
382392
call_id,
383393
command: command_for_display.clone(),
@@ -391,8 +401,10 @@ impl Session {
391401
let _ = self.tx_event.send(event).await;
392402
}
393403

394-
async fn notify_exec_command_end(
404+
#[allow(clippy::too_many_arguments)]
405+
async fn on_exec_command_end(
395406
&self,
407+
turn_diff_tracker: &mut TurnDiffTracker,
396408
sub_id: &str,
397409
call_id: &str,
398410
stdout: &str,
@@ -427,6 +439,20 @@ impl Session {
427439
msg,
428440
};
429441
let _ = self.tx_event.send(event).await;
442+
443+
// If this is an apply_patch, after we emit the end patch, emit a second event
444+
// with the full turn diff if there is one.
445+
if is_apply_patch {
446+
let unified_diff = turn_diff_tracker.get_unified_diff();
447+
if let Ok(Some(unified_diff)) = unified_diff {
448+
let msg = EventMsg::TurnDiff(TurnDiffEvent { unified_diff });
449+
let event = Event {
450+
id: sub_id.into(),
451+
msg,
452+
};
453+
let _ = self.tx_event.send(event).await;
454+
}
455+
}
430456
}
431457

432458
/// Helper that emits a BackgroundEvent with the given message. This keeps
@@ -1000,6 +1026,10 @@ async fn run_task(sess: Arc<Session>, sub_id: String, input: Vec<InputItem>) {
10001026
.await;
10011027

10021028
let last_agent_message: Option<String>;
1029+
// Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains
1030+
// many turns, from the perspective of the user, it is a single turn.
1031+
let mut turn_diff_tracker = TurnDiffTracker::new();
1032+
10031033
loop {
10041034
// Note that pending_input would be something like a message the user
10051035
// submitted through the UI while the model was running. Though the UI
@@ -1031,7 +1061,7 @@ async fn run_task(sess: Arc<Session>, sub_id: String, input: Vec<InputItem>) {
10311061
})
10321062
})
10331063
.collect();
1034-
match run_turn(&sess, sub_id.clone(), turn_input).await {
1064+
match run_turn(&sess, &mut turn_diff_tracker, sub_id.clone(), turn_input).await {
10351065
Ok(turn_output) => {
10361066
let mut items_to_record_in_conversation_history = Vec::<ResponseItem>::new();
10371067
let mut responses = Vec::<ResponseInputItem>::new();
@@ -1157,6 +1187,7 @@ async fn run_task(sess: Arc<Session>, sub_id: String, input: Vec<InputItem>) {
11571187

11581188
async fn run_turn(
11591189
sess: &Session,
1190+
turn_diff_tracker: &mut TurnDiffTracker,
11601191
sub_id: String,
11611192
input: Vec<ResponseItem>,
11621193
) -> CodexResult<Vec<ProcessedResponseItem>> {
@@ -1171,7 +1202,7 @@ async fn run_turn(
11711202

11721203
let mut retries = 0;
11731204
loop {
1174-
match try_run_turn(sess, &sub_id, &prompt).await {
1205+
match try_run_turn(sess, turn_diff_tracker, &sub_id, &prompt).await {
11751206
Ok(output) => return Ok(output),
11761207
Err(CodexErr::Interrupted) => return Err(CodexErr::Interrupted),
11771208
Err(CodexErr::EnvVar(var)) => return Err(CodexErr::EnvVar(var)),
@@ -1217,6 +1248,7 @@ struct ProcessedResponseItem {
12171248

12181249
async fn try_run_turn(
12191250
sess: &Session,
1251+
turn_diff_tracker: &mut TurnDiffTracker,
12201252
sub_id: &str,
12211253
prompt: &Prompt,
12221254
) -> CodexResult<Vec<ProcessedResponseItem>> {
@@ -1310,7 +1342,8 @@ async fn try_run_turn(
13101342
match event {
13111343
ResponseEvent::Created => {}
13121344
ResponseEvent::OutputItemDone(item) => {
1313-
let response = handle_response_item(sess, sub_id, item.clone()).await?;
1345+
let response =
1346+
handle_response_item(sess, turn_diff_tracker, sub_id, item.clone()).await?;
13141347

13151348
output.push(ProcessedResponseItem { item, response });
13161349
}
@@ -1328,6 +1361,16 @@ async fn try_run_turn(
13281361
.ok();
13291362
}
13301363

1364+
let unified_diff = turn_diff_tracker.get_unified_diff();
1365+
if let Ok(Some(unified_diff)) = unified_diff {
1366+
let msg = EventMsg::TurnDiff(TurnDiffEvent { unified_diff });
1367+
let event = Event {
1368+
id: sub_id.to_string(),
1369+
msg,
1370+
};
1371+
let _ = sess.tx_event.send(event).await;
1372+
}
1373+
13311374
return Ok(output);
13321375
}
13331376
ResponseEvent::OutputTextDelta(delta) => {
@@ -1432,6 +1475,7 @@ async fn run_compact_task(
14321475

14331476
async fn handle_response_item(
14341477
sess: &Session,
1478+
turn_diff_tracker: &mut TurnDiffTracker,
14351479
sub_id: &str,
14361480
item: ResponseItem,
14371481
) -> CodexResult<Option<ResponseInputItem>> {
@@ -1469,7 +1513,17 @@ async fn handle_response_item(
14691513
..
14701514
} => {
14711515
info!("FunctionCall: {arguments}");
1472-
Some(handle_function_call(sess, sub_id.to_string(), name, arguments, call_id).await)
1516+
Some(
1517+
handle_function_call(
1518+
sess,
1519+
turn_diff_tracker,
1520+
sub_id.to_string(),
1521+
name,
1522+
arguments,
1523+
call_id,
1524+
)
1525+
.await,
1526+
)
14731527
}
14741528
ResponseItem::LocalShellCall {
14751529
id,
@@ -1506,6 +1560,7 @@ async fn handle_response_item(
15061560
handle_container_exec_with_params(
15071561
exec_params,
15081562
sess,
1563+
turn_diff_tracker,
15091564
sub_id.to_string(),
15101565
effective_call_id,
15111566
)
@@ -1523,6 +1578,7 @@ async fn handle_response_item(
15231578

15241579
async fn handle_function_call(
15251580
sess: &Session,
1581+
turn_diff_tracker: &mut TurnDiffTracker,
15261582
sub_id: String,
15271583
name: String,
15281584
arguments: String,
@@ -1536,7 +1592,8 @@ async fn handle_function_call(
15361592
return *output;
15371593
}
15381594
};
1539-
handle_container_exec_with_params(params, sess, sub_id, call_id).await
1595+
handle_container_exec_with_params(params, sess, turn_diff_tracker, sub_id, call_id)
1596+
.await
15401597
}
15411598
"update_plan" => handle_update_plan(sess, arguments, sub_id, call_id).await,
15421599
_ => {
@@ -1612,6 +1669,7 @@ fn maybe_run_with_user_profile(params: ExecParams, sess: &Session) -> ExecParams
16121669
async fn handle_container_exec_with_params(
16131670
params: ExecParams,
16141671
sess: &Session,
1672+
turn_diff_tracker: &mut TurnDiffTracker,
16151673
sub_id: String,
16161674
call_id: String,
16171675
) -> ResponseInputItem {
@@ -1766,7 +1824,7 @@ async fn handle_container_exec_with_params(
17661824
},
17671825
),
17681826
};
1769-
sess.notify_exec_command_begin(exec_command_context.clone())
1827+
sess.on_exec_command_begin(turn_diff_tracker, exec_command_context.clone())
17701828
.await;
17711829

17721830
let params = maybe_run_with_user_profile(params, sess);
@@ -1788,7 +1846,8 @@ async fn handle_container_exec_with_params(
17881846
duration,
17891847
} = output;
17901848

1791-
sess.notify_exec_command_end(
1849+
sess.on_exec_command_end(
1850+
turn_diff_tracker,
17921851
&sub_id,
17931852
&call_id,
17941853
&stdout,
@@ -1814,7 +1873,15 @@ async fn handle_container_exec_with_params(
18141873
}
18151874
}
18161875
Err(CodexErr::Sandbox(error)) => {
1817-
handle_sandbox_error(params, exec_command_context, error, sandbox_type, sess).await
1876+
handle_sandbox_error(
1877+
turn_diff_tracker,
1878+
params,
1879+
exec_command_context,
1880+
error,
1881+
sandbox_type,
1882+
sess,
1883+
)
1884+
.await
18181885
}
18191886
Err(e) => {
18201887
// Handle non-sandbox errors
@@ -1830,6 +1897,7 @@ async fn handle_container_exec_with_params(
18301897
}
18311898

18321899
async fn handle_sandbox_error(
1900+
turn_diff_tracker: &mut TurnDiffTracker,
18331901
params: ExecParams,
18341902
exec_command_context: ExecCommandContext,
18351903
error: SandboxErr,
@@ -1889,7 +1957,8 @@ async fn handle_sandbox_error(
18891957
sess.notify_background_event(&sub_id, "retrying command without sandbox")
18901958
.await;
18911959

1892-
sess.notify_exec_command_begin(exec_command_context).await;
1960+
sess.on_exec_command_begin(turn_diff_tracker, exec_command_context)
1961+
.await;
18931962

18941963
// This is an escalated retry; the policy will not be
18951964
// examined and the sandbox has been set to `None`.
@@ -1911,7 +1980,8 @@ async fn handle_sandbox_error(
19111980
duration,
19121981
} = retry_output;
19131982

1914-
sess.notify_exec_command_end(
1983+
sess.on_exec_command_end(
1984+
turn_diff_tracker,
19151985
&sub_id,
19161986
&call_id,
19171987
&stdout,

codex-rs/core/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ mod safety;
4242
pub mod seatbelt;
4343
pub mod shell;
4444
pub mod spawn;
45+
pub mod turn_diff_tracker;
4546
mod user_notification;
4647
pub mod util;
4748

codex-rs/core/src/protocol.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -343,6 +343,8 @@ pub enum EventMsg {
343343
/// Notification that a patch application has finished.
344344
PatchApplyEnd(PatchApplyEndEvent),
345345

346+
TurnDiff(TurnDiffEvent),
347+
346348
/// Response to GetHistoryEntryRequest.
347349
GetHistoryEntryResponse(GetHistoryEntryResponseEvent),
348350

@@ -534,6 +536,11 @@ pub struct PatchApplyEndEvent {
534536
pub success: bool,
535537
}
536538

539+
#[derive(Debug, Clone, Deserialize, Serialize)]
540+
pub struct TurnDiffEvent {
541+
pub unified_diff: String,
542+
}
543+
537544
#[derive(Debug, Clone, Deserialize, Serialize)]
538545
pub struct GetHistoryEntryResponseEvent {
539546
pub offset: usize,

0 commit comments

Comments
 (0)