open-source/cursor-byok
1
0
Fork 0
mirror of https://wget.la/https://github.com/leookun/cursor-byok synced 2026-10-06 21:52:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: enhance bidi request handling and observability tracing

Browse Source 
- 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.
This commit is contained in:
leookun
2026-09-02 01:04:54 +08:00
parent 2dad593263
commit 5cdf642dd1
25 changed files with 1522 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/src/api/cursor/bidi.rs
+5 -1
View File
@@ -184,7 +184,11 @@ pub async fn append(
request: DecodedAppend,
parent: Option<TransportParent>,
) -> Result<ai::BidiAppendResponse> {
let handle = registry.get_or_create(&request.request_id).await?;
let replace_closing = request.model_id().is_some();
let handle = registry
.get_or_create_for_append(&request.request_id, replace_closing)
.await?;
let _admission = handle.admit()?;
if let Some(conversation_id) = request.conversation_id() {
handle.set_conversation_id(conversation_id)?;
}
server/src/api/cursor/handlers.rs
+84 -26
View File
@@ -19,9 +19,7 @@ use crate::{
connect,
proto::{agent::v1 as agent, aiserver::v1 as ai},
},
services::{
account, analytics, knowledge, model_catalog, observability::CursorTraceRecorder, tab,
},
services::{account, analytics, knowledge, model_catalog, tab},
transport::{TransportParent, TransportRegistry},
},
Result,
@@ -123,16 +121,13 @@ async fn run_sse_handler(
let (parts, body) = buffered(request).await?;
let request: agent::BidiRequestId = connect::decode_unary(&body)?;
let route = registry.wait_route(&request.request_id).await;
let trace = CursorTraceRecorder::resume(registry.store().clone(), &request.request_id).await;
if let Some(trace) = &trace {
trace
.request(
"run_sse_request",
&body,
serde_json::json!({"request_id": request.request_id}),
)
.await;
}
let trace = registry.trace(&request.request_id);
trace.resume();
trace.request(
"run_sse_request",
body.clone(),
serde_json::json!({"request_id": request.request_id}),
);
match route {
crate::cursor::transport::TransportRoute::Local => {
run_sse::stream(&registry, &request.request_id).await
@@ -143,7 +138,14 @@ async fn run_sse_handler(
Request::from_parts(parts, Body::from(body)),
)
.await?;
Ok(run_sse::upstream(registry, request.request_id, generation, response, trace).await)
Ok(run_sse::upstream(
registry,
request.request_id,
generation,
response,
Some(trace),
)
.await)
}
}
}
@@ -159,6 +161,7 @@ async fn bidi_handler(
let first_model = decoded.model_id().map(str::to_owned);
let conversation_id = decoded.conversation_id().map(str::to_owned);
let trace_metadata = decoded.trace_metadata();
let trace = registry.trace(&decoded.request_id);
let local = if let Some(model_id) = decoded.model_id() {
// 插件模型 ID 只在本地有意义,永远不转发到 Cursor 官方上游。
if model_id.starts_with(crate::plugin::ADAPTER_ID_PREFIX)
@@ -183,14 +186,18 @@ async fn bidi_handler(
} else if registry.upstream(&decoded.request_id).await {
false
} else {
trace.resume();
trace.request(
"bidi_request",
body.clone(),
trace_outcome(trace_metadata, false, "missing_transport", None),
);
return Err(crate::Error::Protocol(
"first BidiAppend message must select a model".into(),
));
};
let trace = if first_model.is_some() {
CursorTraceRecorder::begin(
registry.store().clone(),
&decoded.request_id,
if first_model.is_some() {
trace.begin(
conversation_id.as_deref(),
if local {
"local_byok"
@@ -198,26 +205,61 @@ async fn bidi_handler(
"cursor_official"
},
first_model.as_deref(),
)
.await
);
} else {
CursorTraceRecorder::resume(registry.store().clone(), &decoded.request_id).await
};
if let Some(trace) = &trace {
trace.request("bidi_request", &body, trace_metadata).await;
trace.resume();
}
if !local {
if first_model.is_some() {
registry.mark_upstream(&decoded.request_id).await;
}
trace.request(
"bidi_request",
body.clone(),
trace_outcome(trace_metadata, true, "upstream", None),
);
return proxy::forward(
Extension(proxy),
Request::from_parts(parts, Body::from(body)),
)
.await;
}
let parent = parent_headers(&parts.headers)?;
bidi::append(&registry, decoded, parent).await?;
let parent = match parent_headers(&parts.headers) {
Ok(parent) => parent,
Err(error) => {
trace.request(
"bidi_request",
body,
trace_outcome(
trace_metadata,
false,
"invalid_parent",
Some(error.to_string()),
),
);
return Err(error);
}
};
match bidi::append(&registry, decoded, parent).await {
Ok(_) => trace.request(
"bidi_request",
body,
trace_outcome(trace_metadata, true, "local", None),
),
Err(error) => {
trace.request(
"bidi_request",
body,
trace_outcome(
trace_metadata,
false,
"command_rejected",
Some(error.to_string()),
),
);
return Err(error);
}
}
let mut response = Response::new(axum::body::Body::empty());
*response.status_mut() = StatusCode::OK;
response.headers_mut().insert(
@@ -227,6 +269,22 @@ async fn bidi_handler(
Ok(response)
}
fn trace_outcome(
mut metadata: serde_json::Value,
accepted: bool,
route_outcome: &str,
error: Option<String>,
) -> serde_json::Value {
if let Some(metadata) = metadata.as_object_mut() {
metadata.insert("accepted".into(), accepted.into());
metadata.insert("route_outcome".into(), route_outcome.into());
if let Some(error) = error {
metadata.insert("error".into(), error.into());
}
}
metadata
}
async fn buffered(request: Request<Body>) -> Result<(axum::http::request::Parts, Bytes)> {
let (parts, body) = request.into_parts();
let body = to_bytes(body, usize::MAX)
server/src/api/cursor/run_sse.rs
+5 -5
View File
@@ -22,7 +22,7 @@ pub async fn stream(registry: &TransportRegistry, request_id: &str) -> Result<Re
let receiver = handle.subscribe();
let trace = handle.trace().cloned();
if let Some(trace) = &trace {
trace.response_started(StatusCode::OK.as_u16()).await;
trace.response_started(StatusCode::OK.as_u16());
}
let body_stream = local_body_stream(receiver, handle, trace);
let mut response = Response::new(Body::from_stream(body_stream));
@@ -133,7 +133,7 @@ pub async fn upstream(
) -> Response<Body> {
let (parts, body) = response.into_parts();
if let Some(trace) = &trace {
trace.response_started(parts.status.as_u16()).await;
trace.response_started(parts.status.as_u16());
}
let stream = async_stream::stream! {
let _guard = UpstreamRunGuard {
@@ -180,15 +180,15 @@ impl TraceStreamSink {
while let Some(event) = receiver.recv().await {
match event {
TraceStreamEvent::Chunk(chunk) => {
trace.response_chunk(source, &chunk).await;
trace.response_chunk(source, chunk);
}
TraceStreamEvent::Finish(error) => {
trace.finish(error.as_deref()).await;
trace.finish(error.as_deref());
return;
}
}
}
trace.finish(None).await;
trace.finish(None);
});
Self {
sender: Some(sender),