Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions .repository-projection.json
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,11 @@
"projection": "deixic-code",
"projectionSchemaVersion": 1,
"sourceRepository": "dx-corp/mono",
"sourceSha": "083eacdbac0fef20630986791529393aa4e2683c",
"sourceSha": "be83027d2b5902447ca1ba418d2b4f85df867bbc",
"destinationRepository": "dx-corp/code",
"priorProjectedBase": "402b1d7e517279b33d3f6d6aef9c77a0633a8266",
"priorProjectedBase": "d5e2b3dd85095b0190afc25bb7492414ad27a29b",
"definitionDigest": "cb9d429542ebb0a2de9b42a7aad60d9d8696a648ceba47c30f05c0b285ca0db7",
"toolDigest": "f8cb071b0f27267120ccf45a00d0982f45113bd23535bef6a1555b4933f99f13",
"contentDigest": "b19cd9d7b8b22856b6565af0c70d61f3b0125ade7f5dd8e4494610dce7cd538d",
"contentDigest": "4676d676f45ea745c21a6ebb2d20aedcf3d83bcf978ca61d2f303c5ca4a09c82",
"publicationEligible": true
}
59 changes: 52 additions & 7 deletions packages/runtime-gateway-rs/src/chat.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use super::*;
use crate::turn_diffs::{begin_chat_snapshot, finish_chat_snapshot};
use maestro_local_host::embedding::{EmbeddedAgent, EmbeddedAgentBuilder};

pub(crate) fn is_chat_endpoint(head: &RequestHead) -> bool {
Expand Down Expand Up @@ -179,19 +180,21 @@ async fn handle_codex_app_server_chat(
model: &str,
prompt: &str,
attachment_paths: &[String],
scope: Option<&str>,
) -> Result<bool, String> {
handle_codex_app_server_chat_transport(
handle_codex_app_server_chat_transport_scoped(
stream,
state,
session_id,
model,
prompt,
attachment_paths,
CodexBridgeTransport::Sse,
(CodexBridgeTransport::Sse, scope),
)
.await
}

#[cfg(test)]
pub(crate) async fn handle_codex_app_server_chat_transport(
stream: &mut TcpStream,
state: &AppState,
Expand All @@ -201,6 +204,28 @@ pub(crate) async fn handle_codex_app_server_chat_transport(
attachment_paths: &[String],
transport: CodexBridgeTransport,
) -> Result<bool, String> {
handle_codex_app_server_chat_transport_scoped(
stream,
state,
session_id,
model,
prompt,
attachment_paths,
(transport, None),
)
.await
}

async fn handle_codex_app_server_chat_transport_scoped(
stream: &mut TcpStream,
state: &AppState,
session_id: Option<&str>,
model: &str,
prompt: &str,
attachment_paths: &[String],
(transport, scope): (CodexBridgeTransport, Option<&str>),
) -> Result<bool, String> {
let mut turn_snapshot = begin_chat_snapshot(state, session_id, prompt, scope).await;
let session_approval_mode = approval_mode_for_session(state, session_id).await;
let approval_mode = codex_app_server_approval_mode(&session_approval_mode);
send_codex_bridge_event(
Expand Down Expand Up @@ -263,7 +288,7 @@ pub(crate) async fn handle_codex_app_server_chat_transport(
send_codex_bridge_tool_event(stream, transport, tool_event).await?;
}

let message =
let mut message =
composer_assistant_message(&assistant_output.text, "", assistant_output.usage.clone());
if !assistant_output.text.is_empty() {
send_codex_bridge_event(
Expand All @@ -281,6 +306,7 @@ pub(crate) async fn handle_codex_app_server_chat_transport(
)
.await?;
}
finish_chat_snapshot(&mut turn_snapshot, &mut message).await;
record_chat_assistant_message(state, session_id, message.clone()).await;
record_usage_entry(
state,
Expand Down Expand Up @@ -327,15 +353,16 @@ async fn handle_codex_app_server_chat_ws(
model: &str,
prompt: &str,
attachment_paths: &[String],
scope: Option<&str>,
) -> Result<bool, String> {
handle_codex_app_server_chat_transport(
handle_codex_app_server_chat_transport_scoped(
stream,
state,
session_id,
model,
prompt,
attachment_paths,
CodexBridgeTransport::WebSocket,
(CodexBridgeTransport::WebSocket, scope),
)
.await
}
Expand Down Expand Up @@ -620,6 +647,7 @@ pub(crate) async fn handle_chat_endpoint(
&codex_model,
&prompt,
&prepared_attachments.paths,
turn_scope.as_deref(),
)
.await?;
if prompt_succeeded {
Expand All @@ -635,6 +663,13 @@ pub(crate) async fn handle_chat_endpoint(
cleanup_prepared_attachments(prepared_attachments).await;
return Ok(());
}
let mut turn_snapshot = begin_chat_snapshot(
&state,
session_id.as_deref(),
&prompt,
turn_scope.as_deref(),
)
.await;
let (usage_provider, usage_model) = usage_provider_model(&chat, &state, &model).await;
let (mut client_tools, client_tool_names) = client_tool_definitions(&chat);
// Gateway-handled session-messaging tools give an agent turn the same reach
Expand Down Expand Up @@ -1304,12 +1339,13 @@ pub(crate) async fn handle_chat_endpoint(
take_client_tool_results(&state, &client_tool_call_ids).await;
finish_client_tool_metadata(&mut assistant_tools, &client_tool_results);
let usage = last_usage.take();
let message = composer_assistant_message_with_tools(
let mut message = composer_assistant_message_with_tools(
&assistant_text,
&thinking_text,
usage,
&assistant_tools,
);
finish_chat_snapshot(&mut turn_snapshot, &mut message).await;
record_chat_assistant_message(&state, session_id.as_deref(), message.clone()).await;
send_sse(
&mut stream,
Expand Down Expand Up @@ -1588,6 +1624,7 @@ pub(crate) async fn handle_chat_websocket_endpoint(
&codex_model,
&prompt,
&prepared_attachments.paths,
turn_scope.as_deref(),
)
.await?;
if prompt_succeeded {
Expand All @@ -1604,6 +1641,13 @@ pub(crate) async fn handle_chat_websocket_endpoint(
cleanup_prepared_attachments(prepared_attachments).await;
return Ok(());
}
let mut turn_snapshot = begin_chat_snapshot(
&state,
session_id.as_deref(),
&prompt,
turn_scope.as_deref(),
)
.await;
let (usage_provider, usage_model) = usage_provider_model(&chat, &state, &model).await;
let (mut client_tools, client_tool_names) = client_tool_definitions(&chat);
// Gateway-handled session-messaging tools give an agent turn the same reach
Expand Down Expand Up @@ -2251,12 +2295,13 @@ pub(crate) async fn handle_chat_websocket_endpoint(
take_client_tool_results(&state, &client_tool_call_ids).await;
finish_client_tool_metadata(&mut assistant_tools, &client_tool_results);
let usage = last_usage.take();
let message = composer_assistant_message_with_tools(
let mut message = composer_assistant_message_with_tools(
&assistant_text,
&thinking_text,
usage,
&assistant_tools,
);
finish_chat_snapshot(&mut turn_snapshot, &mut message).await;
record_chat_assistant_message(&state, session_id.as_deref(), message.clone()).await;
send_ws_json(
&mut stream,
Expand Down
13 changes: 9 additions & 4 deletions packages/runtime-gateway-rs/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ mod migrations;
mod model_catalog;
mod session_messaging;
mod sessions;
mod turn_diffs;

#[allow(unused_imports)]
pub(crate) use a2a::{
Expand Down Expand Up @@ -89,16 +90,18 @@ use a2a_platform_registration::{
#[cfg(test)]
use a2a_skill_catalog::A2A_SUBAGENT_REQUEST_METADATA_PATH;
use auth::*;
#[cfg(test)]
use chat::handle_codex_app_server_chat_transport;
#[allow(unused_imports)]
pub(crate) use chat::{
ChatAttachment, ChatMessage, ChatRequest, ExtractAttachmentRequest, ExtractDocumentOutput,
PreparedAttachments, approval_blocked_tool_event, build_prompt_from_chat,
composer_assistant_message, composer_assistant_message_with_tools, composer_text_content,
finish_client_tool_metadata, finish_tool_metadata, handle_chat_endpoint,
handle_chat_websocket_endpoint, handle_codex_app_server_chat_transport, is_chat_endpoint,
is_chat_websocket_endpoint, prepare_chat_attachments, record_chat_user_message,
record_tool_call_metadata, send_sse, send_ws_json, sse_headers, strip_data_url_prefix,
try_parse_websocket_text_message, update_tool_metadata_status, websocket_accept_key,
handle_chat_websocket_endpoint, is_chat_endpoint, is_chat_websocket_endpoint,
prepare_chat_attachments, record_chat_user_message, record_tool_call_metadata, send_sse,
send_ws_json, sse_headers, strip_data_url_prefix, try_parse_websocket_text_message,
update_tool_metadata_status, websocket_accept_key,
};
pub(crate) use codex_bridge::*;
use extended::{ExtendedApiState, handle_extended_endpoint, is_extended_endpoint};
Expand Down Expand Up @@ -461,6 +464,7 @@ struct AppState {
shared_sessions: Arc<Mutex<HashMap<String, SharedSessionGrant>>>,
approval_modes: Arc<Mutex<HashMap<String, String>>>,
pending_tool_responses: Arc<Mutex<HashMap<String, PendingToolResponseSender>>>,
native_snapshot_registry: Arc<turn_diffs::NativeSnapshotRegistry>,
// Maps a pending-request id to the session or authenticated principal that
// owns the blocked agent turn. Sessionless client-tool turns retain a
// resumable owner without weakening the cross-tenant resume check.
Expand Down Expand Up @@ -796,6 +800,7 @@ pub async fn serve_listener(
shared_sessions: Arc::new(Mutex::new(shared_sessions)),
approval_modes: Arc::new(Mutex::new(HashMap::new())),
pending_tool_responses: Arc::new(Mutex::new(HashMap::new())),
native_snapshot_registry: Arc::new(turn_diffs::NativeSnapshotRegistry::default()),
pending_tool_response_sessions: Arc::new(Mutex::new(HashMap::new())),
completed_client_tool_results: Arc::new(Mutex::new(HashMap::new())),
a2a_tasks: Arc::new(Mutex::new(a2a_tasks)),
Expand Down
3 changes: 3 additions & 0 deletions packages/runtime-gateway-rs/src/sessions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -692,6 +692,9 @@ pub(super) async fn handle_session_get(

match session_path.tail {
None => json_response(200, &session_full_value(&session)),
Some("turn-diff") => {
crate::turn_diffs::session_turn_diff_response(head, state, &session).await
}
Some("timeline") => json_response(200, &session_timeline_value(&session)),
Some("share") => json_response(
200,
Expand Down
81 changes: 6 additions & 75 deletions packages/runtime-gateway-rs/src/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2739,6 +2739,7 @@ fn test_app_state_with_sessions(sessions: HashMap<String, SessionRecord>) -> App
shared_sessions: Arc::new(Mutex::new(HashMap::new())),
approval_modes: Arc::new(Mutex::new(HashMap::new())),
pending_tool_responses: Arc::new(Mutex::new(HashMap::new())),
native_snapshot_registry: Arc::new(turn_diffs::NativeSnapshotRegistry::default()),
pending_tool_response_sessions: Arc::new(Mutex::new(HashMap::new())),
completed_client_tool_results: Arc::new(Mutex::new(HashMap::new())),
extended_api: Arc::new(Mutex::new(ExtendedApiState::default())),
Expand Down Expand Up @@ -13272,6 +13273,7 @@ async fn delete_session_subpath_returns_404_without_removing_session() {
shared_sessions: Arc::new(Mutex::new(HashMap::new())),
approval_modes: Arc::new(Mutex::new(HashMap::new())),
pending_tool_responses: Arc::new(Mutex::new(HashMap::new())),
native_snapshot_registry: Arc::new(turn_diffs::NativeSnapshotRegistry::default()),
pending_tool_response_sessions: Arc::new(Mutex::new(HashMap::new())),
completed_client_tool_results: Arc::new(Mutex::new(HashMap::new())),
extended_api: Arc::new(Mutex::new(ExtendedApiState::default())),
Expand Down Expand Up @@ -13363,6 +13365,7 @@ async fn invalid_session_store_is_left_untouched_and_future_writes_are_blocked()
shared_sessions: Arc::new(Mutex::new(HashMap::new())),
approval_modes: Arc::new(Mutex::new(HashMap::new())),
pending_tool_responses: Arc::new(Mutex::new(HashMap::new())),
native_snapshot_registry: Arc::new(turn_diffs::NativeSnapshotRegistry::default()),
pending_tool_response_sessions: Arc::new(Mutex::new(HashMap::new())),
completed_client_tool_results: Arc::new(Mutex::new(HashMap::new())),
extended_api: Arc::new(Mutex::new(ExtendedApiState::default())),
Expand Down Expand Up @@ -13502,81 +13505,6 @@ fn enterprise_policy_admin_routes_are_implemented() {
assert!(is_extended_endpoint(&head));
}
}
#[test]
fn undo_endpoint_reads_and_consumes_tui_checkpoint_store() {
use maestro_local_host::checkpoints::{Checkpoint, CheckpointStore, EntryKind, FileEntry};
use sha2::{Digest, Sha256};

let temp = unique_test_dir("maestro-undo-checkpoint");
std::fs::create_dir_all(&temp).unwrap();
// Match checkpoint capture: persist the resolved root, not a platform alias.
let temp = dunce::canonicalize(temp).unwrap();
let file = temp.join("src.txt");
let before = b"before";
let after = b"after";
std::fs::write(&file, after).unwrap();
let before_hash = Sha256::digest(before)
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
let after_hash = Sha256::digest(after)
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();

let store = CheckpointStore::new(&temp.join("sessions"), "session-1");
let checkpoint_dir = store.root().join("checkpoint-1");
std::fs::create_dir_all(checkpoint_dir.join("blobs")).unwrap();
std::fs::write(checkpoint_dir.join("blobs").join(&before_hash), before).unwrap();
let checkpoint = Checkpoint {
id: "checkpoint-1".to_string(),
created_at: "2026-08-05T00:00:00Z".to_string(),
prompt: "edit src.txt".to_string(),
repo_root: temp.clone(),
head: None,
user_turn_index: None,
entries: vec![FileEntry {
path: "src.txt".to_string(),
kind: EntryKind::Modified,
pre_blob: Some(before_hash),
post_hash: Some(after_hash),
}],
};
std::fs::write(
checkpoint_dir.join("checkpoint.json"),
serde_json::to_vec(&checkpoint).unwrap(),
)
.unwrap();

let head = RequestHead {
method: "GET".to_string(),
path: "/api/undo".to_string(),
query: HashMap::from([(String::from("sessionId"), String::from("session-1"))]),
headers: HashMap::new(),
};
let summary = undo_response_for_store(&head, &store);
assert_eq!(summary["totalChanges"], 1);
assert_eq!(summary["canUndo"], true);

let restored = restore_undo_response_for_store(&store);
assert_eq!(restored["success"], true, "{restored}");
assert_eq!(std::fs::read(&file).unwrap(), before);
assert_eq!(store.list().len(), 0);
// A missing restore blob must not become success after other file work.
std::fs::create_dir_all(&checkpoint_dir).unwrap();
std::fs::write(
checkpoint_dir.join("checkpoint.json"),
serde_json::to_vec(&checkpoint).unwrap(),
)
.unwrap();
std::fs::write(&file, after).unwrap();
let failed = restore_undo_response_for_store(&store);
assert_eq!(failed["success"], false);
assert_eq!(failed["failedFiles"].as_array().unwrap().len(), 1);
assert_eq!(std::fs::read(&file).unwrap(), after);
assert_eq!(store.list().len(), 1);
let _ = std::fs::remove_dir_all(temp);
}

// ---------------------------------------------------------------------------
// Session-to-session messaging
Expand Down Expand Up @@ -14599,3 +14527,6 @@ async fn platform_a2a_push_evicts_terminal_payloads_and_replay_history() {
assert!(!tasks.contains_key("callback-0000"));
assert!(!histories.contains_key("callback-0000"));
}

#[path = "tests/turn_diffs.rs"]
mod turn_diff_tests;
Loading
Loading