mirror of
https://github.com/openai/codex.git
synced 2026-09-16 12:13:30 +00:00
More baase
This commit is contained in:
@@ -32,6 +32,8 @@ use async_channel::Sender;
|
||||
use codex_protocol::ConversationId;
|
||||
use codex_protocol::approvals::ExecPolicyAmendment;
|
||||
use codex_protocol::items::TurnItem;
|
||||
use codex_protocol::protocol::AgentId;
|
||||
use codex_protocol::protocol::DEFAULT_AGENT_ID;
|
||||
use codex_protocol::protocol::FileChange;
|
||||
use codex_protocol::protocol::HasLegacyEvent;
|
||||
use codex_protocol::protocol::ItemCompletedEvent;
|
||||
@@ -199,6 +201,7 @@ fn maybe_push_chat_wire_api_deprecation(
|
||||
|
||||
post_session_configured_events.push(Event {
|
||||
id: INITIAL_SUBMIT_ID.to_owned(),
|
||||
agent_id: Some(DEFAULT_AGENT_ID.to_string()),
|
||||
msg: EventMsg::DeprecationNotice(DeprecationNoticeEvent {
|
||||
summary: CHAT_WIRE_API_DEPRECATION_SUMMARY.to_string(),
|
||||
details: None,
|
||||
@@ -354,6 +357,7 @@ pub(crate) struct Session {
|
||||
/// The context needed for a single turn of the conversation.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct TurnContext {
|
||||
pub(crate) agent_id: AgentId,
|
||||
pub(crate) sub_id: String,
|
||||
pub(crate) client: ModelClient,
|
||||
/// The session's current working directory. All relative paths provided by
|
||||
@@ -493,6 +497,7 @@ impl Session {
|
||||
per_turn_config: Config,
|
||||
model_family: ModelFamily,
|
||||
conversation_id: ConversationId,
|
||||
agent_id: AgentId,
|
||||
sub_id: String,
|
||||
) -> TurnContext {
|
||||
let otel_manager = otel_manager.clone().with_model(
|
||||
@@ -519,6 +524,7 @@ impl Session {
|
||||
});
|
||||
|
||||
TurnContext {
|
||||
agent_id,
|
||||
sub_id,
|
||||
client,
|
||||
cwd: session_configuration.cwd.clone(),
|
||||
@@ -619,6 +625,7 @@ impl Session {
|
||||
};
|
||||
post_session_configured_events.push(Event {
|
||||
id: INITIAL_SUBMIT_ID.to_owned(),
|
||||
agent_id: Some(DEFAULT_AGENT_ID.to_string()),
|
||||
msg: EventMsg::DeprecationNotice(DeprecationNoticeEvent { summary, details }),
|
||||
});
|
||||
}
|
||||
@@ -879,6 +886,7 @@ impl Session {
|
||||
|
||||
Ok(self
|
||||
.new_turn_from_configuration(
|
||||
agent_id,
|
||||
sub_id,
|
||||
session_configuration,
|
||||
updates.final_output_json_schema,
|
||||
@@ -889,6 +897,7 @@ impl Session {
|
||||
|
||||
async fn new_turn_from_configuration(
|
||||
&self,
|
||||
agent_id: &str,
|
||||
sub_id: String,
|
||||
session_configuration: SessionConfiguration,
|
||||
final_output_json_schema: Option<Option<Value>>,
|
||||
@@ -927,6 +936,7 @@ impl Session {
|
||||
per_turn_config,
|
||||
model_family,
|
||||
self.conversation_id,
|
||||
agent_id.to_string(),
|
||||
sub_id,
|
||||
);
|
||||
if let Some(final_schema) = final_output_json_schema {
|
||||
@@ -974,6 +984,7 @@ impl Session {
|
||||
let legacy_source = msg.clone();
|
||||
let event = Event {
|
||||
id: turn_context.sub_id.clone(),
|
||||
agent_id: Some(turn_context.agent_id.clone()),
|
||||
msg,
|
||||
};
|
||||
self.send_event_raw(event).await;
|
||||
@@ -982,6 +993,7 @@ impl Session {
|
||||
for legacy in legacy_source.as_legacy_events(show_raw_agent_reasoning) {
|
||||
let legacy_event = Event {
|
||||
id: turn_context.sub_id.clone(),
|
||||
agent_id: Some(turn_context.agent_id.clone()),
|
||||
msg: legacy,
|
||||
};
|
||||
self.send_event_raw(legacy_event).await;
|
||||
@@ -1186,7 +1198,8 @@ impl Session {
|
||||
items: &[ResponseItem],
|
||||
) {
|
||||
self.record_into_history(items, turn_context).await;
|
||||
self.persist_rollout_response_items(items).await;
|
||||
self.persist_rollout_response_items(&turn_context.agent_id, items)
|
||||
.await;
|
||||
self.send_raw_response_items(turn_context, items).await;
|
||||
}
|
||||
|
||||
@@ -1313,7 +1326,7 @@ impl Session {
|
||||
guard.clone()
|
||||
};
|
||||
if let Some(rec) = recorder
|
||||
&& let Err(e) = rec.record_items(items).await
|
||||
&& let Err(e) = rec.record_items(agent_id, items).await
|
||||
{
|
||||
error!("failed to record rollout items: {e:#}");
|
||||
}
|
||||
@@ -1343,7 +1356,7 @@ impl Session {
|
||||
|
||||
pub(crate) async fn recompute_token_usage(&self, turn_context: &TurnContext) {
|
||||
let Some(estimated_total_tokens) = self
|
||||
.clone_history()
|
||||
.clone_history(&turn_context.agent_id)
|
||||
.await
|
||||
.estimate_token_count(turn_context)
|
||||
else {
|
||||
@@ -1592,14 +1605,15 @@ impl Session {
|
||||
|
||||
async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiver<Submission>) {
|
||||
// Seed with context in case there is an OverrideTurnContext first.
|
||||
let mut previous_context: Option<Arc<TurnContext>> = Some(sess.new_default_turn().await);
|
||||
let mut previous_context: Option<Arc<TurnContext>> =
|
||||
Some(sess.new_default_turn_for_agent(&agent_id).await);
|
||||
|
||||
// To break out of this loop, send Op::Shutdown.
|
||||
while let Ok(sub) = rx_sub.recv().await {
|
||||
debug!(?sub, "Submission");
|
||||
match sub.op.clone() {
|
||||
Op::Interrupt => {
|
||||
handlers::interrupt(&sess).await;
|
||||
handlers::interrupt(&sess, &agent_id).await;
|
||||
}
|
||||
Op::OverrideTurnContext {
|
||||
cwd,
|
||||
@@ -1611,6 +1625,7 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
|
||||
} => {
|
||||
handlers::override_turn_context(
|
||||
&sess,
|
||||
&agent_id,
|
||||
sub.id.clone(),
|
||||
SessionSettingsUpdate {
|
||||
cwd,
|
||||
@@ -1625,8 +1640,14 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
|
||||
.await;
|
||||
}
|
||||
Op::UserInput { .. } | Op::UserTurn { .. } => {
|
||||
handlers::user_input_or_turn(&sess, sub.id.clone(), sub.op, &mut previous_context)
|
||||
.await;
|
||||
handlers::user_input_or_turn(
|
||||
&sess,
|
||||
&agent_id,
|
||||
sub.id.clone(),
|
||||
sub.op,
|
||||
&mut previous_context,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Op::ExecApproval { id, decision } => {
|
||||
handlers::exec_approval(&sess, id, decision).await;
|
||||
@@ -1635,30 +1656,38 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
|
||||
handlers::patch_approval(&sess, id, decision).await;
|
||||
}
|
||||
Op::AddToHistory { text } => {
|
||||
handlers::add_to_history(&sess, &config, text).await;
|
||||
handlers::add_to_history(&sess, &agent_id, &config, text).await;
|
||||
}
|
||||
Op::GetHistoryEntryRequest { offset, log_id } => {
|
||||
handlers::get_history_entry_request(&sess, &config, sub.id.clone(), offset, log_id)
|
||||
.await;
|
||||
handlers::get_history_entry_request(
|
||||
&sess,
|
||||
&agent_id,
|
||||
&config,
|
||||
sub.id.clone(),
|
||||
offset,
|
||||
log_id,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Op::ListMcpTools => {
|
||||
handlers::list_mcp_tools(&sess, &config, sub.id.clone()).await;
|
||||
handlers::list_mcp_tools(&sess, &agent_id, &config, sub.id.clone()).await;
|
||||
}
|
||||
Op::ListCustomPrompts => {
|
||||
handlers::list_custom_prompts(&sess, sub.id.clone()).await;
|
||||
handlers::list_custom_prompts(&sess, &agent_id, sub.id.clone()).await;
|
||||
}
|
||||
Op::ListSkills { cwds, force_reload } => {
|
||||
handlers::list_skills(&sess, sub.id.clone(), cwds, force_reload).await;
|
||||
handlers::list_skills(&sess, &agent_id, sub.id.clone(), cwds, force_reload).await;
|
||||
}
|
||||
Op::Undo => {
|
||||
handlers::undo(&sess, sub.id.clone()).await;
|
||||
handlers::undo(&sess, &agent_id, sub.id.clone()).await;
|
||||
}
|
||||
Op::Compact => {
|
||||
handlers::compact(&sess, sub.id.clone()).await;
|
||||
handlers::compact(&sess, &agent_id, sub.id.clone()).await;
|
||||
}
|
||||
Op::RunUserShellCommand { command } => {
|
||||
handlers::run_user_shell_command(
|
||||
&sess,
|
||||
&agent_id,
|
||||
sub.id.clone(),
|
||||
command,
|
||||
&mut previous_context,
|
||||
@@ -1678,7 +1707,7 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
|
||||
}
|
||||
}
|
||||
Op::Review { review_request } => {
|
||||
handlers::review(&sess, &config, sub.id.clone(), review_request).await;
|
||||
handlers::review(&sess, &agent_id, &config, sub.id.clone(), review_request).await;
|
||||
}
|
||||
_ => {} // Ignore unknown ops; enum is non_exhaustive to allow extensions.
|
||||
}
|
||||
@@ -1731,6 +1760,7 @@ mod handlers {
|
||||
|
||||
pub async fn override_turn_context(
|
||||
sess: &Session,
|
||||
agent_id: &str,
|
||||
sub_id: String,
|
||||
updates: SessionSettingsUpdate,
|
||||
) {
|
||||
@@ -1748,6 +1778,7 @@ mod handlers {
|
||||
|
||||
pub async fn user_input_or_turn(
|
||||
sess: &Arc<Session>,
|
||||
agent_id: &str,
|
||||
sub_id: String,
|
||||
op: Op,
|
||||
previous_context: &mut Option<Arc<TurnContext>>,
|
||||
@@ -1804,6 +1835,7 @@ mod handlers {
|
||||
|
||||
pub async fn run_user_shell_command(
|
||||
sess: &Arc<Session>,
|
||||
agent_id: &str,
|
||||
sub_id: String,
|
||||
command: String,
|
||||
previous_context: &mut Option<Arc<TurnContext>>,
|
||||
@@ -1880,7 +1912,12 @@ mod handlers {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn add_to_history(sess: &Arc<Session>, config: &Arc<Config>, text: String) {
|
||||
pub async fn add_to_history(
|
||||
sess: &Arc<Session>,
|
||||
agent_id: &str,
|
||||
config: &Arc<Config>,
|
||||
text: String,
|
||||
) {
|
||||
let id = sess.conversation_id;
|
||||
let config = Arc::clone(config);
|
||||
tokio::spawn(async move {
|
||||
@@ -1892,6 +1929,7 @@ mod handlers {
|
||||
|
||||
pub async fn get_history_entry_request(
|
||||
sess: &Arc<Session>,
|
||||
agent_id: &str,
|
||||
config: &Arc<Config>,
|
||||
sub_id: String,
|
||||
offset: usize,
|
||||
@@ -1927,7 +1965,12 @@ mod handlers {
|
||||
});
|
||||
}
|
||||
|
||||
pub async fn list_mcp_tools(sess: &Session, config: &Arc<Config>, sub_id: String) {
|
||||
pub async fn list_mcp_tools(
|
||||
sess: &Session,
|
||||
agent_id: &str,
|
||||
config: &Arc<Config>,
|
||||
sub_id: String,
|
||||
) {
|
||||
let mcp_connection_manager = sess.services.mcp_connection_manager.read().await;
|
||||
let snapshot = collect_mcp_snapshot_from_manager(
|
||||
&mcp_connection_manager,
|
||||
@@ -1945,7 +1988,7 @@ mod handlers {
|
||||
sess.send_event_raw(event).await;
|
||||
}
|
||||
|
||||
pub async fn list_custom_prompts(sess: &Session, sub_id: String) {
|
||||
pub async fn list_custom_prompts(sess: &Session, agent_id: &str, sub_id: String) {
|
||||
let custom_prompts: Vec<CustomPrompt> =
|
||||
if let Some(dir) = crate::custom_prompts::default_prompts_dir() {
|
||||
crate::custom_prompts::discover_prompts_in(&dir).await
|
||||
@@ -1964,6 +2007,7 @@ mod handlers {
|
||||
|
||||
pub async fn list_skills(
|
||||
sess: &Session,
|
||||
agent_id: &str,
|
||||
sub_id: String,
|
||||
cwds: Vec<PathBuf>,
|
||||
force_reload: bool,
|
||||
@@ -2061,6 +2105,7 @@ mod handlers {
|
||||
|
||||
pub async fn review(
|
||||
sess: &Arc<Session>,
|
||||
agent_id: &str,
|
||||
config: &Arc<Config>,
|
||||
sub_id: String,
|
||||
review_request: ReviewRequest,
|
||||
@@ -2757,12 +2802,15 @@ mod tests {
|
||||
use crate::tools::format_exec_output_str;
|
||||
use codex_protocol::models::FunctionCallOutputPayload;
|
||||
|
||||
use super::AgentState;
|
||||
use crate::protocol::CompactedItem;
|
||||
use crate::protocol::CreditsSnapshot;
|
||||
use crate::protocol::DEFAULT_AGENT_ID;
|
||||
use crate::protocol::InitialHistory;
|
||||
use crate::protocol::RateLimitSnapshot;
|
||||
use crate::protocol::RateLimitWindow;
|
||||
use crate::protocol::ResumedHistory;
|
||||
use crate::protocol::RolloutLine;
|
||||
use crate::state::TaskKind;
|
||||
use crate::tasks::SessionTask;
|
||||
use crate::tasks::SessionTaskContext;
|
||||
@@ -2788,6 +2836,7 @@ mod tests {
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration as StdDuration;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
#[test]
|
||||
fn reconstruct_history_matches_live_compactions() {
|
||||
@@ -3263,6 +3312,7 @@ mod tests {
|
||||
per_turn_config,
|
||||
model_family,
|
||||
conversation_id,
|
||||
agent_id,
|
||||
"turn_id".to_string(),
|
||||
));
|
||||
|
||||
@@ -3502,7 +3552,7 @@ mod tests {
|
||||
|
||||
let initial_context = session.build_initial_context(turn_context);
|
||||
for item in &initial_context {
|
||||
rollout_items.push(RolloutItem::ResponseItem(item.clone()));
|
||||
push_line(RolloutItem::ResponseItem(item.clone()));
|
||||
}
|
||||
live_history.record_items(initial_context.iter(), turn_context.truncation_policy);
|
||||
|
||||
@@ -3514,7 +3564,7 @@ mod tests {
|
||||
}],
|
||||
};
|
||||
live_history.record_items(std::iter::once(&user1), turn_context.truncation_policy);
|
||||
rollout_items.push(RolloutItem::ResponseItem(user1.clone()));
|
||||
push_line(RolloutItem::ResponseItem(user1.clone()));
|
||||
|
||||
let assistant1 = ResponseItem::Message {
|
||||
id: None,
|
||||
@@ -3524,7 +3574,7 @@ mod tests {
|
||||
}],
|
||||
};
|
||||
live_history.record_items(std::iter::once(&assistant1), turn_context.truncation_policy);
|
||||
rollout_items.push(RolloutItem::ResponseItem(assistant1.clone()));
|
||||
push_line(RolloutItem::ResponseItem(assistant1.clone()));
|
||||
|
||||
let summary1 = "summary one";
|
||||
let snapshot1 = live_history.get_history();
|
||||
@@ -3535,7 +3585,7 @@ mod tests {
|
||||
summary1,
|
||||
);
|
||||
live_history.replace(rebuilt1);
|
||||
rollout_items.push(RolloutItem::Compacted(CompactedItem {
|
||||
push_line(RolloutItem::Compacted(CompactedItem {
|
||||
message: summary1.to_string(),
|
||||
replacement_history: None,
|
||||
}));
|
||||
@@ -3548,7 +3598,7 @@ mod tests {
|
||||
}],
|
||||
};
|
||||
live_history.record_items(std::iter::once(&user2), turn_context.truncation_policy);
|
||||
rollout_items.push(RolloutItem::ResponseItem(user2.clone()));
|
||||
push_line(RolloutItem::ResponseItem(user2.clone()));
|
||||
|
||||
let assistant2 = ResponseItem::Message {
|
||||
id: None,
|
||||
@@ -3558,7 +3608,7 @@ mod tests {
|
||||
}],
|
||||
};
|
||||
live_history.record_items(std::iter::once(&assistant2), turn_context.truncation_policy);
|
||||
rollout_items.push(RolloutItem::ResponseItem(assistant2.clone()));
|
||||
push_line(RolloutItem::ResponseItem(assistant2.clone()));
|
||||
|
||||
let summary2 = "summary two";
|
||||
let snapshot2 = live_history.get_history();
|
||||
@@ -3569,7 +3619,7 @@ mod tests {
|
||||
summary2,
|
||||
);
|
||||
live_history.replace(rebuilt2);
|
||||
rollout_items.push(RolloutItem::Compacted(CompactedItem {
|
||||
push_line(RolloutItem::Compacted(CompactedItem {
|
||||
message: summary2.to_string(),
|
||||
replacement_history: None,
|
||||
}));
|
||||
@@ -3582,7 +3632,7 @@ mod tests {
|
||||
}],
|
||||
};
|
||||
live_history.record_items(std::iter::once(&user3), turn_context.truncation_policy);
|
||||
rollout_items.push(RolloutItem::ResponseItem(user3.clone()));
|
||||
push_line(RolloutItem::ResponseItem(user3.clone()));
|
||||
|
||||
let assistant3 = ResponseItem::Message {
|
||||
id: None,
|
||||
@@ -3592,7 +3642,7 @@ mod tests {
|
||||
}],
|
||||
};
|
||||
live_history.record_items(std::iter::once(&assistant3), turn_context.truncation_policy);
|
||||
rollout_items.push(RolloutItem::ResponseItem(assistant3.clone()));
|
||||
push_line(RolloutItem::ResponseItem(assistant3.clone()));
|
||||
|
||||
(rollout_items, live_history.get_history())
|
||||
}
|
||||
|
||||
@@ -20,8 +20,10 @@ use codex_protocol::ConversationId;
|
||||
use codex_protocol::items::TurnItem;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_protocol::openai_models::ModelPreset;
|
||||
use codex_protocol::protocol::DEFAULT_AGENT_ID;
|
||||
use codex_protocol::protocol::InitialHistory;
|
||||
use codex_protocol::protocol::RolloutItem;
|
||||
use codex_protocol::protocol::RolloutLine;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
Reference in New Issue
Block a user