Files
cursor-byok/server/src/cursor/compile/run.rs
T
leookun 5cdf642dd1 feat: enhance bidi request handling and observability tracing
- Updated the `append` function to include a flag for replacing closing requests, improving request handling.
- Refactored the `run_sse_handler` and `bidi_handler` functions to utilize a new tracing mechanism, enhancing observability.
- Introduced a new `trace_outcome` function to standardize tracing outcomes for requests.
- Removed the `CursorTraceRecorder` in favor of a new `CursorTraceService` for better performance and non-blocking behavior.
- Added tests to validate the new tracing functionality and ensure correct behavior during request processing.
2026-09-02 01:04:54 +08:00

615 lines
21 KiB
Rust

//! Compiles an AgentRunRequest into a PreparedRun.
use std::collections::BTreeMap;
use uuid::Uuid;
use crate::{
cursor::prompting::{Mode, PromptCompiler},
cursor::{
checkpoint::messages,
checkpoint::CheckpointBuilder,
protocol::proto::agent::v1 as pb,
services::blob_sync::BlobSynchronizer,
services::context_sync::RequestContextSynchronizer,
tools::runtime::{ExecContext, SubagentModel},
},
model::{
CanonicalMessage, ContentPart, ConversationId, MessageContent, Origin, PreparedRun,
PromptSpec, Role, RunAction, RunId, RunKind,
},
store::{BlobId, Store},
Error, Result,
};
use super::{break_messages, context, insert_messages, model};
struct ActionProjection {
mode: i32,
turn_user: Option<pb::UserMessage>,
action_context: String,
event_id: Option<String>,
input_id: Option<String>,
starts_turn: bool,
compacting: bool,
background_completion: bool,
}
pub struct CursorRunContext {
pub request_id: String,
pub mode: i32,
pub turn_user: Option<pb::UserMessage>,
pub exec: ExecContext,
pub dynamic_tools: BTreeMap<String, pb::McpToolDefinition>,
pub checkpoint_prompt: PromptSpec,
pub compacting: bool,
pub background_completion: bool,
}
pub(crate) struct PrepareDependencies<'a> {
pub compiler: &'a PromptCompiler,
pub store: &'a Store,
pub checkpoint: &'a CheckpointBuilder,
pub blob_sync: &'a BlobSynchronizer,
pub context_sync: &'a RequestContextSynchronizer,
pub local_rules_dir: Option<&'a std::path::Path>,
}
pub(crate) async fn prepare(
request_id: &str,
request: &pb::AgentRunRequest,
dependencies: PrepareDependencies<'_>,
) -> Result<(PreparedRun, CursorRunContext)> {
let PrepareDependencies {
compiler,
store,
checkpoint,
blob_sync,
context_sync,
local_rules_dir,
} = dependencies;
checkpoint
.import_prefetched(&request.pre_fetched_blobs)
.await?;
let conversation_id = ConversationId::new(
request
.conversation_id
.clone()
.unwrap_or_else(|| request_id.into()),
);
let run_id = execution_run_id(request_id);
let mut base_messages = if request.conversation_state.is_some() {
Some(
checkpoint
.hydrate_messages(request.conversation_state.as_ref())
.await?,
)
} else {
None
};
if let Some(trace) = blob_sync.trace() {
let hydrated_messages = base_messages.as_deref().unwrap_or_default();
let hydrated_images = hydrated_messages
.iter()
.map(|message| match &message.content {
MessageContent::Parts { parts } => parts
.iter()
.filter(|part| matches!(part, ContentPart::Image { .. }))
.count(),
_ => 0,
})
.sum::<usize>();
let history = request
.action
.as_ref()
.and_then(|action| action.action.as_ref())
.and_then(|action| match action {
pb::conversation_action::Action::UserMessageAction(action) => {
action.conversation_history.as_ref()
}
_ => None,
});
let summary = serde_json::json!({
"checkpoint_root_count": request.conversation_state.as_ref().map_or(0, |state| state.root_prompt_messages_json.len()),
"checkpoint_turn_count": request.conversation_state.as_ref().map_or(0, |state| state.turns.len()),
"conversation_history_message_count": history.map_or(0, |history| history.messages.len()),
"hydrated_message_count": hydrated_messages.len(),
"hydrated_image_count": hydrated_images,
"selected_source": "root_prompt_messages_json",
});
let encoded = serde_json::to_vec(&summary)?;
trace.artifact("history_projection", "byok_server", &encoded, summary);
}
let mut request_context = context::hydrate(request, context_sync).await?;
if let Some(rules_dir) = local_rules_dir {
context::merge_local_rules(&mut request_context, rules_dir);
}
let request_context = request_context;
let ActionProjection {
mode: mode_number,
mut turn_user,
action_context,
mut event_id,
input_id,
starts_turn,
compacting,
background_completion,
} = action(request)?;
let checkpoint_mode = if request.subagent_type_name.is_some() {
Mode::Subagent
} else {
mode_from_proto(mode_number)?
};
let mut model = model::requested_model(request)?;
if let Some(configured_model) = store.model(&model.model_id).await? {
configured_model.configure(&mut model);
}
let dynamic = context::dynamic_mcp(request, &request_context)?;
let subagent_model_overrides = model::overrides(request)?;
let subagents_disabled = subagent_model_overrides
.first()
.is_some_and(|(_, selection)| {
matches!(selection, crate::model::SubagentModelOverride::Disabled)
});
let mut checkpoint_prompt = compiler.prompt_spec(
checkpoint_mode,
&model,
&dynamic
.values()
.map(|(_, definition)| definition.clone())
.collect::<Vec<_>>(),
request.suppress_subagent_progress_update_tool == Some(true),
)?;
if subagents_disabled {
checkpoint_prompt.tools.retain(|tool| tool.name != "Task");
}
let prompt = if compacting {
compiler.prompt_spec(Mode::Compaction, &model, &[], false)?
} else {
checkpoint_prompt.clone()
};
let proposed_base_checkpoint_id = match base_messages.as_mut() {
Some(messages) if !messages.is_empty() => {
validate_prompt_root(messages)?;
messages.retain(|message| {
!(message.role == Role::System && message.origin == Origin::Prompt)
});
store.import_checkpoint(&conversation_id, messages).await?
}
Some(_) | None => store.ensure_conversation(&conversation_id).await?,
};
let base_checkpoint_id = match input_id.as_deref() {
Some(input_id) => {
store
.anchor_input(&conversation_id, input_id, proposed_base_checkpoint_id)
.await?
}
None => proposed_base_checkpoint_id,
};
let mut projected_user_context = if input_id.is_some() && !compacting && !background_completion
{
break_messages::compile_request_context(
"identity",
&request_context,
base_messages.as_deref().unwrap_or_default(),
)?
} else {
None
};
if event_id.is_none() {
if let (Some(input_id), Some(user)) = (input_id.as_deref(), turn_user.as_ref()) {
event_id = Some(
break_messages::user_event_id(
input_id,
checkpoint_mode,
user,
&request_context,
&action_context,
projected_user_context
.as_ref()
.map(|message| &message.content),
compiler,
blob_sync,
)
.await?,
);
}
}
let existing_runtime = match event_id.as_deref() {
Some(event_id) => {
store
.message(&conversation_id, &format!("runtime:{event_id}"))
.await?
}
_ => None,
};
let request_context_message = match event_id.as_deref() {
Some(event_id) if !compacting && !background_completion => {
let message_id = format!("request-context:{event_id}");
match store.message(&conversation_id, &message_id).await? {
Some(message) => Some(message),
None if input_id.is_some() => projected_user_context.take().map(|mut message| {
message.message_id = message_id;
message
}),
None => break_messages::compile_request_context(
event_id,
&request_context,
base_messages.as_deref().unwrap_or_default(),
)?,
}
}
_ => None,
};
let mut initial_messages = if compacting {
Vec::new()
} else {
match (turn_user.clone(), event_id) {
(Some(mut user), Some(event_id)) if background_completion => {
let (message, text) = match existing_runtime {
Some(message) => {
let text = runtime_message_text(&message)?;
(message, text)
}
None => {
break_messages::compile_background(
event_id,
&user,
&request_context,
&action_context,
blob_sync,
)
.await?
}
};
user.text = text;
turn_user = Some(user);
vec![message]
}
(Some(user), Some(event_id)) => {
let runtime = match existing_runtime {
Some(message) => message,
None => {
break_messages::compile(
event_id,
checkpoint_mode,
&user,
&request_context,
&action_context,
compiler,
blob_sync,
)
.await?
}
};
request_context_message
.into_iter()
.chain(std::iter::once(runtime))
.collect()
}
(None, None) => Vec::new(),
_ => {
return Err(Error::Protocol(
"Cursor action has an incomplete runtime event".into(),
))
}
}
};
let (base_checkpoint_id, reused) = store
.match_checkpoint_prefix(&conversation_id, base_checkpoint_id, &initial_messages)
.await?;
initial_messages.drain(..reused);
let action = if compacting {
RunAction::Compact
} else if starts_turn {
RunAction::Start
} else {
let pending_tool_round = match request
.conversation_state
.as_ref()
.map(|state| state.pending_tool_calls.as_slice())
.unwrap_or_default()
{
[] => None,
[pending] => Some(messages::decode_pending(pending)?),
pending => {
return Err(Error::Protocol(format!(
"Cursor resume contains {} pending assistant messages",
pending.len()
)))
}
};
RunAction::Resume { pending_tool_round }
};
let exec = exec_context(
request,
&request_context,
&conversation_id,
&model.model_id,
subagents_disabled,
&subagent_model_overrides,
);
Ok((
PreparedRun {
run_id,
cursor_request_id: Some(request_id.into()),
conversation_id,
kind: RunKind::Root,
model,
prompt,
initial_messages,
action,
base_checkpoint_id,
},
CursorRunContext {
request_id: request_id.into(),
mode: mode_number,
turn_user,
exec,
dynamic_tools: dynamic
.into_iter()
.map(|(name, (wire, _))| (name, wire))
.collect(),
checkpoint_prompt,
compacting,
background_completion,
},
))
}
fn runtime_message_text(message: &CanonicalMessage) -> Result<String> {
let MessageContent::Parts { parts } = &message.content else {
return Err(Error::Protocol(
"stored runtime message does not contain parts".into(),
));
};
let Some(ContentPart::Text { text }) = parts.first() else {
return Err(Error::Protocol(
"stored runtime message does not start with text".into(),
));
};
Ok(text.clone())
}
fn validate_prompt_root(messages: &[CanonicalMessage]) -> Result<()> {
let prompts = messages
.iter()
.filter(|message| message.role == Role::System && message.origin == Origin::Prompt)
.collect::<Vec<_>>();
let [prompt] = prompts.as_slice() else {
return Err(Error::Protocol(format!(
"Cursor history contains {} system prompt roots",
prompts.len()
)));
};
let MessageContent::Parts { parts } = &prompt.content else {
return Err(Error::Protocol(
"Cursor system prompt root is not textual content".into(),
));
};
let [ContentPart::Text { .. }] = parts.as_slice() else {
return Err(Error::Protocol(
"Cursor system prompt root is not one text part".into(),
));
};
Ok(())
}
fn execution_run_id(request_id: &str) -> RunId {
let execution_id = Uuid::new_v4().simple().to_string();
RunId::new(format!("{request_id}:{}", &execution_id[..8]))
}
fn action(request: &pb::AgentRunRequest) -> Result<ActionProjection> {
let conversation_mode = request
.conversation_state
.as_ref()
.and_then(|state| state.mode);
let mode = conversation_mode.unwrap_or(pb::AgentMode::Agent as i32);
let Some(action) = request
.action
.as_ref()
.and_then(|action| action.action.as_ref())
else {
return Ok(ActionProjection {
mode,
turn_user: None,
action_context: String::new(),
event_id: None,
input_id: None,
starts_turn: false,
compacting: false,
background_completion: false,
});
};
match action {
pb::conversation_action::Action::UserMessageAction(action) => {
let user = action.user_message.as_ref().ok_or_else(|| {
Error::Protocol("Cursor user message action has no UserMessage".into())
})?;
let mode = if user.mode == pb::AgentMode::Unspecified as i32 {
conversation_mode.unwrap_or(user.mode)
} else {
user.mode
};
if user.message_id.is_empty() {
return Err(Error::Protocol(
"Cursor user message action has no message_id".into(),
));
}
if user.text.trim() == "/summarize" {
return Ok(ActionProjection {
mode,
turn_user: Some(user.clone()),
action_context: String::new(),
event_id: None,
input_id: None,
starts_turn: false,
compacting: true,
background_completion: false,
});
}
let mut context = action
.prepend_user_messages
.iter()
.map(|message| message.text.trim())
.filter(|text| !text.is_empty())
.map(str::to_string)
.collect::<Vec<_>>();
context.extend(
user.subagent_system_reminder
.iter()
.filter(|text| !text.is_empty())
.cloned(),
);
let input_id = format!("cursor:user:{}", user.message_id);
Ok(ActionProjection {
mode,
turn_user: Some(user.clone()),
action_context: context.join("\n\n"),
event_id: None,
input_id: Some(input_id),
starts_turn: true,
compacting: false,
background_completion: false,
})
}
pb::conversation_action::Action::BackgroundTaskCompletionAction(action) => {
let projection = insert_messages::project(action, mode)?;
let event_id = projection.turn_user.message_id.clone();
Ok(ActionProjection {
mode,
action_context: projection.context,
event_id: Some(event_id),
input_id: None,
turn_user: Some(projection.turn_user),
starts_turn: true,
compacting: false,
background_completion: true,
})
}
pb::conversation_action::Action::ExecutePlanAction(action) => execute_plan(action),
pb::conversation_action::Action::SummarizeAction(_) => Ok(ActionProjection {
mode,
turn_user: None,
action_context: String::new(),
event_id: None,
input_id: None,
starts_turn: false,
compacting: true,
background_completion: false,
}),
_ => Ok(ActionProjection {
mode,
turn_user: None,
action_context: String::new(),
event_id: None,
input_id: None,
starts_turn: false,
compacting: false,
background_completion: false,
}),
}
}
fn execute_plan(action: &pb::ExecutePlanAction) -> Result<ActionProjection> {
let plan = action
.plan_file_content
.as_deref()
.or_else(|| action.plan.as_ref().map(|plan| plan.plan.as_str()))
.filter(|plan| !plan.trim().is_empty())
.ok_or_else(|| Error::Protocol("ExecutePlan is missing plan content".into()))?;
let source = action
.plan_file_uri
.as_deref()
.or(action.plan_file_path.as_deref())
.filter(|source| !source.is_empty());
let action_context = match source {
Some(source) => {
format!("<approved_plan>\n<plan_file>{source}</plan_file>\n{plan}\n</approved_plan>")
}
None => format!("<approved_plan>\n{plan}\n</approved_plan>"),
};
let identity = BlobId::digest(
format!(
"{}\0{}\0{}\0{}\0{}",
action.execution_mode,
action.plan_id.as_deref().unwrap_or_default(),
action.kickoff_message_id.as_deref().unwrap_or_default(),
source.unwrap_or_default(),
plan,
)
.as_bytes(),
)
.to_base64();
let event_id = format!("execute-plan:{identity}");
Ok(ActionProjection {
mode: action.execution_mode,
turn_user: Some(pb::UserMessage {
text: "Execute the approved plan.".into(),
message_id: event_id.clone(),
mode: action.execution_mode,
..Default::default()
}),
action_context,
event_id: Some(event_id),
input_id: None,
starts_turn: true,
compacting: false,
background_completion: false,
})
}
pub(super) fn mode_from_proto(mode: i32) -> Result<Mode> {
let mode = pb::AgentMode::try_from(mode)
.map_err(|_| Error::Protocol(format!("unknown Cursor agent mode: {mode}")))?;
match mode {
pb::AgentMode::Agent => Ok(Mode::Agent),
pb::AgentMode::Ask => Ok(Mode::Ask),
pb::AgentMode::Plan => Ok(Mode::Plan),
pb::AgentMode::Debug => Ok(Mode::Debug),
pb::AgentMode::Multitask => Ok(Mode::Multitask),
mode => Err(Error::Protocol(format!(
"unsupported Cursor agent mode: {}",
mode.as_str_name()
))),
}
}
fn exec_context(
request: &pb::AgentRunRequest,
request_context: &pb::RequestContext,
conversation_id: &ConversationId,
model_id: &str,
subagents_disabled: bool,
overrides: &[(
crate::model::SubagentKind,
crate::model::SubagentModelOverride,
)],
) -> ExecContext {
let subagent_model = overrides.first().map(|(_, value)| match value {
crate::model::SubagentModelOverride::Explicit(model) => {
SubagentModel::Model(model.model_id.clone())
}
crate::model::SubagentModelOverride::Inherit => SubagentModel::Model(model_id.into()),
crate::model::SubagentModelOverride::Disabled => SubagentModel::Disabled,
});
ExecContext {
conversation_id: conversation_id.to_string(),
root_conversation_id: request
.conversation_group_id
.clone()
.unwrap_or_else(|| conversation_id.to_string()),
default_subagent_model: model_id.into(),
subagent_model,
allow_subagents: request.subagent_type_name.is_none() && !subagents_disabled,
subagents_disabled,
terminals_folder: request_context
.env
.as_ref()
.map(|env| env.terminals_folder.clone())
.unwrap_or_default(),
admin_command_denylist: request_context.admin_command_denylist.clone(),
mcp_routes: context::meta_mcp_routes(request_context),
}
}