mirror of
https://wget.la/https://github.com/leookun/cursor-byok
synced 2026-10-05 03:56:45 +08:00
feat: add update indicators and stabilize Cursor streams
This commit is contained in:
@@ -8,6 +8,7 @@ use crate::{
|
||||
blob_sync::BlobSynchronizer,
|
||||
checkpoint::CheckpointBuilder,
|
||||
context_sync::RequestContextSynchronizer,
|
||||
interaction,
|
||||
proto::agent::v1 as pb,
|
||||
request,
|
||||
session::CursorSession,
|
||||
@@ -272,6 +273,9 @@ impl CursorActor {
|
||||
)) => {
|
||||
let _ = blob_sync.handle_client(message).await;
|
||||
}
|
||||
Some(pb::agent_client_message::Message::ClientHeartbeat(_)) => {
|
||||
let _ = handle.emit(&interaction::heartbeat());
|
||||
}
|
||||
// TODO: ConversationAction has two different delivery paths that
|
||||
// must not be conflated:
|
||||
//
|
||||
|
||||
@@ -66,6 +66,12 @@ pub fn thinking_completed(elapsed: Duration) -> pb::AgentServerMessage {
|
||||
))
|
||||
}
|
||||
|
||||
pub fn heartbeat() -> pb::AgentServerMessage {
|
||||
server_interaction(pb::interaction_update::Message::Heartbeat(
|
||||
pb::HeartbeatUpdate {},
|
||||
))
|
||||
}
|
||||
|
||||
pub fn arguments_delta(call: &ToolCall, delta: &str) -> Result<pb::AgentServerMessage> {
|
||||
Ok(server_interaction(
|
||||
pb::interaction_update::Message::PartialToolCall(pb::PartialToolCallUpdate {
|
||||
|
||||
+102
-11
@@ -3,29 +3,24 @@ use axum::{
|
||||
http::{header, HeaderValue, Response, StatusCode},
|
||||
};
|
||||
use bytes::Bytes;
|
||||
use std::convert::Infallible;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_stream::StreamExt;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::{
|
||||
cursor::{observability::CursorTraceRecorder, CursorSessionRegistry},
|
||||
cursor::{connect::END_STREAM_FLAG, observability::CursorTraceRecorder, CursorSessionRegistry},
|
||||
Result,
|
||||
};
|
||||
|
||||
pub async fn stream(registry: &CursorSessionRegistry, request_id: &str) -> Result<Response<Body>> {
|
||||
let handle = registry.get_or_create(request_id).await?;
|
||||
let mut receiver = handle.subscribe();
|
||||
let receiver = handle.subscribe();
|
||||
let trace = handle.trace().cloned();
|
||||
if let Some(trace) = &trace {
|
||||
trace.response_started(StatusCode::OK.as_u16()).await;
|
||||
}
|
||||
let body_stream = async_stream::stream! {
|
||||
let mut trace = TraceStreamSink::new(trace, "byok_server");
|
||||
while let Some(chunk) = receiver.recv().await {
|
||||
trace.chunk(&chunk);
|
||||
yield Ok::<Bytes, std::convert::Infallible>(chunk);
|
||||
}
|
||||
trace.finish(None);
|
||||
};
|
||||
let body_stream = local_body_stream(receiver, handle.cancellation(), trace);
|
||||
let mut response = Response::new(Body::from_stream(body_stream));
|
||||
*response.status_mut() = StatusCode::OK;
|
||||
response.headers_mut().insert(
|
||||
@@ -41,6 +36,63 @@ pub async fn stream(registry: &CursorSessionRegistry, request_id: &str) -> Resul
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
fn local_body_stream(
|
||||
mut receiver: mpsc::UnboundedReceiver<Bytes>,
|
||||
cancellation: CancellationToken,
|
||||
trace: Option<CursorTraceRecorder>,
|
||||
) -> impl tokio_stream::Stream<Item = std::result::Result<Bytes, Infallible>> {
|
||||
async_stream::stream! {
|
||||
let mut guard = LocalRunGuard::new(cancellation);
|
||||
let mut trace = TraceStreamSink::new(trace, "byok_server");
|
||||
while let Some(chunk) = receiver.recv().await {
|
||||
let terminal = is_end_stream_frame(&chunk);
|
||||
trace.chunk(&chunk);
|
||||
if terminal {
|
||||
guard.complete();
|
||||
trace.finish(None);
|
||||
}
|
||||
yield Ok::<Bytes, Infallible>(chunk);
|
||||
if terminal {
|
||||
return;
|
||||
}
|
||||
}
|
||||
guard.complete();
|
||||
trace.finish(None);
|
||||
}
|
||||
}
|
||||
|
||||
fn is_end_stream_frame(frame: &Bytes) -> bool {
|
||||
frame
|
||||
.first()
|
||||
.is_some_and(|flags| flags & END_STREAM_FLAG != 0)
|
||||
}
|
||||
|
||||
struct LocalRunGuard {
|
||||
cancellation: CancellationToken,
|
||||
completed: bool,
|
||||
}
|
||||
|
||||
impl LocalRunGuard {
|
||||
fn new(cancellation: CancellationToken) -> Self {
|
||||
Self {
|
||||
cancellation,
|
||||
completed: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn complete(&mut self) {
|
||||
self.completed = true;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for LocalRunGuard {
|
||||
fn drop(&mut self) {
|
||||
if !self.completed {
|
||||
self.cancellation.cancel();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn upstream(
|
||||
registry: CursorSessionRegistry,
|
||||
request_id: String,
|
||||
@@ -127,7 +179,11 @@ impl TraceStreamSink {
|
||||
|
||||
impl Drop for TraceStreamSink {
|
||||
fn drop(&mut self) {
|
||||
self.finish(None);
|
||||
if self.sender.is_some() {
|
||||
self.finish(Some(
|
||||
"response stream dropped before completion".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -143,3 +199,38 @@ impl Drop for UpstreamRunGuard {
|
||||
.finish_upstream(self.request_id.clone(), self.generation);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::cursor::{connect, proto::agent::v1 as pb};
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_stream_cancels_when_the_client_disconnects() {
|
||||
let (sender, receiver) = mpsc::unbounded_channel();
|
||||
let cancellation = CancellationToken::new();
|
||||
sender
|
||||
.send(connect::encode_message(&pb::AgentServerMessage::default()).unwrap())
|
||||
.unwrap();
|
||||
let mut stream = Box::pin(local_body_stream(receiver, cancellation.clone(), None));
|
||||
|
||||
stream.next().await.unwrap().unwrap();
|
||||
|
||||
drop(sender);
|
||||
drop(stream);
|
||||
assert!(cancellation.is_cancelled());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_frame_does_not_cancel_a_completed_local_run() {
|
||||
let (sender, receiver) = mpsc::unbounded_channel();
|
||||
let cancellation = CancellationToken::new();
|
||||
sender.send(connect::encode_end_stream()).unwrap();
|
||||
let mut stream = Box::pin(local_body_stream(receiver, cancellation.clone(), None));
|
||||
|
||||
let terminal = stream.next().await.unwrap().unwrap();
|
||||
assert!(is_end_stream_frame(&terminal));
|
||||
drop(stream);
|
||||
assert!(!cancellation.is_cancelled());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -140,11 +140,11 @@ impl Provider for OpenAiResponsesProvider {
|
||||
}
|
||||
if text_open { text_open = false; yield ModelEvent::TextEnd; }
|
||||
}
|
||||
"response.reasoning_summary_text.delta" => {
|
||||
"response.reasoning_summary_text.delta" | "response.reasoning_text.delta" => {
|
||||
if !thinking_open { thinking_open = true; yield ModelEvent::ThinkingStart; }
|
||||
if let Some(delta) = value.get("delta").and_then(Value::as_str) { yield ModelEvent::ThinkingDelta(delta.into()); }
|
||||
}
|
||||
"response.reasoning_summary_text.done" => {
|
||||
"response.reasoning_summary_text.done" | "response.reasoning_text.done" => {
|
||||
if thinking_open { thinking_open = false; yield ModelEvent::ThinkingEnd; }
|
||||
}
|
||||
"response.output_item.added" => {
|
||||
|
||||
Reference in New Issue
Block a user