fix: preempt root loop on context injection

This commit is contained in:
leokun
2026-08-27 17:01:03 +08:00
parent aa68205735
commit ee915ee760
12 changed files with 1027 additions and 79 deletions
+657 -4
View File
@@ -5,11 +5,15 @@ mod fixtures;
use std::sync::Arc;
use bytes::Bytes;
use cursor_server::{
cursor::prompting::{PromptAssets, PromptCompiler},
cursor::{connect, proto::agent::v1 as pb},
cursor::{CursorCommand, CursorSessionRegistry},
model::{ConversationId, ModelSpec, PreparedRun, PromptSpec, RunAction, RunId, RunKind},
model::{
ConversationId, ModelConfigInput, ModelSpec, ModelType, PreparedRun, PromptSpec, RunAction,
RunId, RunKind, Usage, OPENAI_CHAT_ENDPOINT,
},
provider::{FinishReason, ModelEvent},
run::RunRegistry,
store::RunStatus,
@@ -428,6 +432,472 @@ async fn injected_user_context_restarts_only_the_active_model_cycle() {
);
}
#[tokio::test]
async fn injected_user_context_aborts_pending_tools_and_ignores_late_results() {
let (_directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
provider.push(tool_response("call-1", "Read", "{\"path\":\"/tmp/a\"}"));
let release = provider.push_gated(text_response("continued after tool interruption"));
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("prompt/cursor")
.as_path(),
)
.unwrap();
let registry = CursorSessionRegistry::new(
store,
Arc::new(provider.clone()),
PromptCompiler::new(assets),
Default::default(),
);
let handle = registry
.get_or_create("interrupt-tool-request")
.await
.unwrap();
let mut output = handle.subscribe();
handle
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(client_run_for(
"interrupt-tool-request",
"interrupt-tool-conversation",
)),
})
.await
.unwrap();
let mut append_seqno = 1;
let exec_id = wait_for_exec(&handle, &mut output, &mut append_seqno, "Read").await;
handle
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(runtime_injection_for(
"tool-injection",
"interrupt-tool-request",
)),
})
.await
.unwrap();
append_seqno += 1;
let mut saw_abort = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
while provider.requests().len() < 2 || !saw_abort {
assert!(
tokio::time::Instant::now() < deadline,
"root model did not restart after tool interruption"
);
if let Ok(Some(frame)) =
tokio::time::timeout(std::time::Duration::from_millis(20), output.recv()).await
{
let (_, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
let server = pb::AgentServerMessage::decode(payload).unwrap();
if let Some(pb::agent_server_message::Message::ExecServerControlMessage(control)) =
server.message
{
if let Some(pb::exec_server_control_message::Message::Abort(abort)) =
control.message
{
assert_eq!(abort.id, exec_id);
saw_abort = true;
}
}
acknowledge_kv(&handle, &mut append_seqno, &frame).await;
}
}
handle
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(read_success(exec_id)),
})
.await
.unwrap();
append_seqno += 1;
release.notify_one();
drain_successfully(&handle, &mut output, &mut append_seqno).await;
let requests = provider.requests();
assert_eq!(
requests[0].history,
requests[1].history[..requests[0].history.len()]
);
let history = serde_json::to_string(&requests[1].history).unwrap();
let interrupted = history
.find("Tool execution was interrupted by a newer user message.")
.expect("interrupted tool result missing from provider history");
let injected = history
.find("injected follow-up")
.expect("injected message missing from provider history");
assert!(interrupted < injected);
}
#[tokio::test]
async fn injected_user_context_detaches_subagents_without_cancelling_them() {
let (_directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
provider.push(tool_response(
"task-call",
"Task",
&serde_json::json!({
"description": "Inspect protocol",
"prompt": "Inspect the protocol",
"subagent_type": "generalPurpose",
"run_in_background": false
})
.to_string(),
));
let release = provider.push_gated(text_response("continued while subagent runs"));
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("prompt/cursor")
.as_path(),
)
.unwrap();
let registry = CursorSessionRegistry::new(
store,
Arc::new(provider.clone()),
PromptCompiler::new(assets),
Default::default(),
);
let handle = registry
.get_or_create("detach-subagent-request")
.await
.unwrap();
let mut output = handle.subscribe();
handle
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(client_run_for(
"detach-subagent-request",
"detach-subagent-conversation",
)),
})
.await
.unwrap();
let mut append_seqno = 1;
let exec_id = wait_for_exec(&handle, &mut output, &mut append_seqno, "Task").await;
handle
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(runtime_injection_for(
"subagent-injection",
"detach-subagent-request",
)),
})
.await
.unwrap();
append_seqno += 1;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
while provider.requests().len() < 2 {
assert!(
tokio::time::Instant::now() < deadline,
"root model did not restart while subagent remained active"
);
if let Ok(Some(frame)) =
tokio::time::timeout(std::time::Duration::from_millis(20), output.recv()).await
{
let (_, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
let server = pb::AgentServerMessage::decode(payload).unwrap();
if let Some(pb::agent_server_message::Message::ExecServerControlMessage(control)) =
server.message
{
if let Some(pb::exec_server_control_message::Message::Abort(abort)) =
control.message
{
assert_ne!(abort.id, exec_id, "Task must not be aborted by injection");
}
}
acknowledge_kv(&handle, &mut append_seqno, &frame).await;
}
}
handle
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(subagent_success(exec_id)),
})
.await
.unwrap();
append_seqno += 1;
release.notify_one();
drain_successfully(&handle, &mut output, &mut append_seqno).await;
let history = serde_json::to_string(&provider.requests()[1].history).unwrap();
assert!(history.contains("Tool execution was interrupted by a newer user message."));
assert!(history.contains("injected follow-up"));
}
#[tokio::test]
async fn injected_user_context_interrupts_automatic_compaction() {
let (_directory, store) = fixtures::temp_store().await;
let model = store
.create_model(&ModelConfigInput {
sort_order: 0,
display_name: "Test Model".into(),
model_type: ModelType::OpenAi,
base_url: "https://example.com/v1/chat/completions".into(),
use_full_url: true,
api_key: "test-key".into(),
tooltip_data: "Test Model".into(),
model_id: "test-model".into(),
reasoning_effort: None,
openai_endpoint: OPENAI_CHAT_ENDPOINT.into(),
openai_extra_params_enabled: false,
openai_extra_params: serde_json::json!({}),
custom_headers_enabled: false,
custom_headers: serde_json::json!({}),
anthropic_extra_params_enabled: false,
anthropic_extra_params: serde_json::json!({}),
context_window_tokens: Some(10_001),
max_completion_tokens: None,
anthropic_max_tokens: None,
anthropic_thinking_effort: None,
thinking_budget_tokens: None,
})
.await
.unwrap();
let provider = fake_provider::FakeProvider::default();
provider.push(text_response("seed answer"));
provider.push_pending();
provider.push(text_response("continued after compacting injection"));
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("prompt/cursor")
.as_path(),
)
.unwrap();
let registry = CursorSessionRegistry::new(
store,
Arc::new(provider.clone()),
PromptCompiler::new(assets),
Default::default(),
);
let seed_state = run_to_end(
&registry,
"seed-request",
client_run_for_model(
"seed-request",
"compaction-injection-conversation",
&model.model_hash,
),
)
.await;
let handle = registry
.get_or_create("inject-during-compaction")
.await
.unwrap();
let mut output = handle.subscribe();
let mut compacting_request = client_run_for_model_with_state(
"inject-during-compaction",
"compaction-injection-conversation",
&model.model_hash,
Some(seed_state),
);
let Some(pb::agent_client_message::Message::RunRequest(request)) =
compacting_request.message.as_mut()
else {
panic!("expected RunRequest")
};
request.requested_model.as_mut().unwrap().parameters.push(
pb::requested_model::ModelParameterValue {
id: "context".into(),
value: "10001".into(),
},
);
handle
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(compacting_request),
})
.await
.unwrap();
let mut append_seqno = 1;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
while provider.requests().len() < 2 {
assert!(
tokio::time::Instant::now() < deadline,
"automatic compaction did not start"
);
if let Ok(Some(frame)) =
tokio::time::timeout(std::time::Duration::from_millis(20), output.recv()).await
{
acknowledge_kv(&handle, &mut append_seqno, &frame).await;
}
}
handle
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(runtime_injection_for(
"compaction-injection",
"inject-during-compaction",
)),
})
.await
.unwrap();
append_seqno += 1;
let mut saw_continued = false;
loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
.unwrap()
.expect("RunSSE closed before successful EndStream");
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
if flags & connect::END_STREAM_FLAG != 0 {
assert_eq!(payload.as_ref(), b"{}");
break;
}
let server = pb::AgentServerMessage::decode(payload).unwrap();
if let Some(pb::agent_server_message::Message::InteractionUpdate(update)) = server.message {
if let Some(pb::interaction_update::Message::TextDelta(delta)) = update.message {
saw_continued |= delta.text.contains("continued after compacting injection");
}
}
acknowledge_kv(&handle, &mut append_seqno, &frame).await;
}
let requests = provider.requests();
assert_eq!(requests.len(), 3);
assert!(requests[1]
.prompt
.instructions
.starts_with("Summarize the conversation for the next model turn."));
assert!(!serde_json::to_string(&requests[1].history)
.unwrap()
.contains("injected follow-up"));
assert!(serde_json::to_string(&requests[2].history)
.unwrap()
.contains("injected follow-up"));
assert!(saw_continued);
}
#[tokio::test]
async fn stale_context_injection_is_rejected_without_failing_the_active_run() {
let (_directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
let release = provider.push_gated(vec![
ModelEvent::Start {
model_call_id: "active-cycle".into(),
},
ModelEvent::TextStart,
ModelEvent::TextDelta("active run completed".into()),
ModelEvent::TextEnd,
ModelEvent::Done(FinishReason::Stop),
]);
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("prompt/cursor")
.as_path(),
)
.unwrap();
let registry = CursorSessionRegistry::new(
store,
Arc::new(provider.clone()),
PromptCompiler::new(assets),
Default::default(),
);
let handle = registry.get_or_create("active-request").await.unwrap();
let mut output = handle.subscribe();
handle
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(client_run_for(
"active-request",
"stale-injection-conversation",
)),
})
.await
.unwrap();
let mut append_seqno = 1;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
while provider.requests().is_empty() {
assert!(
tokio::time::Instant::now() < deadline,
"provider did not start"
);
if let Ok(Some(frame)) =
tokio::time::timeout(std::time::Duration::from_millis(20), output.recv()).await
{
acknowledge_kv(&handle, &mut append_seqno, &frame).await;
}
}
handle
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(runtime_injection_for("stale-injection", "replaced-request")),
})
.await
.unwrap();
append_seqno += 1;
handle
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(runtime_injection_for("stale-injection", "replaced-request")),
})
.await
.unwrap();
append_seqno += 1;
let mut rejection_count = 0;
let mut released = false;
loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
.unwrap()
.expect("RunSSE closed before successful EndStream");
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
if flags & connect::END_STREAM_FLAG != 0 {
assert_eq!(payload.as_ref(), b"{}");
break;
}
let server = pb::AgentServerMessage::decode(payload).unwrap();
let rejected = match server.message {
Some(pb::agent_server_message::Message::InteractionUpdate(pb::InteractionUpdate {
message:
Some(pb::interaction_update::Message::ContextInjectionState(
pb::ContextInjectionStateUpdate {
injection_id,
state:
Some(pb::ContextInjectionState {
state:
Some(pb::context_injection_state::State::Rejected(rejected)),
}),
},
)),
..
})) if injection_id == "stale-injection" => {
assert_eq!(
rejected.reason,
"InjectContextAction expected run replaced-request, active run is active-request"
);
true
}
_ => false,
};
acknowledge_kv(&handle, &mut append_seqno, &frame).await;
if rejected {
rejection_count += 1;
if !released {
released = true;
release.notify_one();
}
}
}
assert!(released, "stale injection was not rejected");
assert_eq!(rejection_count, 1);
assert_eq!(provider.requests().len(), 1);
}
#[tokio::test]
async fn cancel_subagent_action_aborts_the_target_task_and_keeps_the_parent_running() {
let (_directory, store) = fixtures::temp_store().await;
@@ -614,6 +1084,23 @@ fn client_run() -> pb::AgentClientMessage {
}
fn client_run_for(request_id: &str, conversation_id: &str) -> pb::AgentClientMessage {
client_run_for_model(request_id, conversation_id, "test-model")
}
fn client_run_for_model(
request_id: &str,
conversation_id: &str,
model_id: &str,
) -> pb::AgentClientMessage {
client_run_for_model_with_state(request_id, conversation_id, model_id, None)
}
fn client_run_for_model_with_state(
request_id: &str,
conversation_id: &str,
model_id: &str,
state: Option<pb::ConversationStateStructure>,
) -> pb::AgentClientMessage {
pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::RunRequest(
pb::AgentRunRequest {
@@ -634,15 +1121,177 @@ fn client_run_for(request_id: &str, conversation_id: &str) -> pb::AgentClientMes
conversation_id: Some(conversation_id.into()),
run_id: Some(request_id.into()),
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
model_id: model_id.into(),
..Default::default()
}),
conversation_state: state,
..Default::default()
},
)),
}
}
fn text_response(text: &str) -> Vec<ModelEvent> {
vec![
ModelEvent::Start {
model_call_id: format!("call-{text}"),
},
ModelEvent::TextStart,
ModelEvent::TextDelta(text.into()),
ModelEvent::TextEnd,
ModelEvent::Usage(Usage {
input_tokens: Some(1),
output_tokens: Some(1),
total_tokens: Some(2),
..Default::default()
}),
ModelEvent::Done(FinishReason::Stop),
]
}
fn tool_response(call_id: &str, name: &str, arguments: &str) -> Vec<ModelEvent> {
vec![
ModelEvent::Start {
model_call_id: format!("call-{call_id}"),
},
ModelEvent::ToolCallStart {
index: 0,
call_id: call_id.into(),
name: name.into(),
},
ModelEvent::ToolCallArgumentsDelta {
index: 0,
delta: arguments.into(),
},
ModelEvent::ToolCallEnd { index: 0 },
ModelEvent::Done(FinishReason::ToolUse),
]
}
async fn wait_for_exec(
handle: &cursor_server::cursor::CursorSessionHandle,
output: &mut tokio::sync::mpsc::UnboundedReceiver<Bytes>,
append_seqno: &mut i64,
tool: &str,
) -> u32 {
loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
.unwrap()
.expect("RunSSE closed before Exec");
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
assert_eq!(flags & connect::END_STREAM_FLAG, 0);
let server = pb::AgentServerMessage::decode(payload).unwrap();
if let Some(pb::agent_server_message::Message::ExecServerMessage(exec)) = server.message {
let matches = match exec.message.as_ref() {
Some(pb::exec_server_message::Message::ReadArgs(_)) => tool == "Read",
Some(pb::exec_server_message::Message::SubagentArgs(_)) => tool == "Task",
_ => false,
};
if matches {
return exec.id;
}
}
acknowledge_kv(handle, append_seqno, &frame).await;
}
}
async fn drain_successfully(
handle: &cursor_server::cursor::CursorSessionHandle,
output: &mut tokio::sync::mpsc::UnboundedReceiver<Bytes>,
append_seqno: &mut i64,
) {
loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
.unwrap()
.expect("RunSSE closed before successful EndStream");
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
if flags & connect::END_STREAM_FLAG != 0 {
assert_eq!(payload.as_ref(), b"{}");
return;
}
acknowledge_kv(handle, append_seqno, &frame).await;
}
}
fn read_success(id: u32) -> pb::AgentClientMessage {
pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::ExecClientMessage(
pb::ExecClientMessage {
id,
message: Some(pb::exec_client_message::Message::ReadResult(
pb::ReadResult {
result: Some(pb::read_result::Result::Success(pb::ReadSuccess {
path: "/tmp/a".into(),
total_lines: 1,
file_size: 1,
output: Some(pb::read_success::Output::Content("late".into())),
..Default::default()
})),
},
)),
..Default::default()
},
)),
}
}
fn subagent_success(id: u32) -> pb::AgentClientMessage {
pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::ExecClientMessage(
pb::ExecClientMessage {
id,
message: Some(pb::exec_client_message::Message::SubagentResult(
pb::SubagentResult {
result: Some(pb::subagent_result::Result::Success(pb::SubagentSuccess {
agent_id: "detached-child".into(),
..Default::default()
})),
},
)),
..Default::default()
},
)),
}
}
async fn run_to_end(
registry: &CursorSessionRegistry,
request_id: &str,
request: pb::AgentClientMessage,
) -> pb::ConversationStateStructure {
let handle = registry.get_or_create(request_id).await.unwrap();
let mut output = handle.subscribe();
handle
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(request),
})
.await
.unwrap();
let mut append_seqno = 1;
let mut state = None;
loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
.unwrap()
.expect("RunSSE closed before EndStream");
let (flags, _) = connect::decode_frames(&frame).unwrap().pop().unwrap();
if flags & connect::END_STREAM_FLAG != 0 {
return state.expect("Run ended without a checkpoint");
}
let (_, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
let server = pb::AgentServerMessage::decode(payload).unwrap();
if let Some(pb::agent_server_message::Message::ConversationCheckpointUpdate(update)) =
server.message
{
state = Some(update);
}
acknowledge_kv(&handle, &mut append_seqno, &frame).await;
}
}
async fn acknowledge_kv(
handle: &cursor_server::cursor::CursorSessionHandle,
append_seqno: &mut i64,
@@ -700,13 +1349,17 @@ fn runtime_user_message() -> pb::AgentClientMessage {
}
fn runtime_injection() -> pb::AgentClientMessage {
runtime_injection_for("injection-1", "inject-request")
}
fn runtime_injection_for(injection_id: &str, expected_run_id: &str) -> pb::AgentClientMessage {
pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::ConversationAction(
pb::ConversationAction {
action: Some(pb::conversation_action::Action::InjectContextAction(
pb::InjectContextAction {
injection_id: "injection-1".into(),
expected_run_id: "inject-request".into(),
injection_id: injection_id.into(),
expected_run_id: expected_run_id.into(),
payload: Some(pb::inject_context_action::Payload::UserContext(
pb::UserContextInjection {
user_message: Some(pb::UserMessage {