mirror of
https://wget.la/https://github.com/leookun/cursor-byok
synced 2026-10-06 05:12:03 +08:00
feat(api): enhance statistics storage management with scope options
- Introduced a new type `StatisticsStorageScope` to specify the scope for clearing statistics. - Updated the `clearStatisticsStorage` API method to accept a scope parameter, allowing for selective clearing of detailed records or all statistics. - Modified the demo API to handle the new scope parameter appropriately. - Updated the SettingsPage component to include a selection for clearing scope, enhancing user control over statistics management. - Added new translations for the updated messages related to statistics clearing in both English and Chinese.
This commit is contained in:
@@ -565,6 +565,10 @@ impl ControlService {
|
||||
self.store.clear_statistics_storage().await
|
||||
}
|
||||
|
||||
pub async fn clear_all_statistics_storage(&self) -> Result<StatisticsStorage> {
|
||||
self.store.clear_all_statistics_storage().await
|
||||
}
|
||||
|
||||
pub async fn proxy_settings(&self) -> Result<ProxySettings> {
|
||||
self.store.proxy_settings().await
|
||||
}
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
use crate::Result;
|
||||
use axum::{extract::State, Json};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::store::{
|
||||
DesktopSettings, PortSettings, ProxySettings, ProxySettingsInput, StatisticsStorage,
|
||||
TabSettings,
|
||||
StatisticsStorageScope, TabSettings,
|
||||
};
|
||||
|
||||
use super::{ControlService, ObservabilitySettings};
|
||||
@@ -36,8 +37,20 @@ pub async fn get_storage(State(service): State<ControlService>) -> Result<Json<S
|
||||
|
||||
pub async fn clear_storage(
|
||||
State(service): State<ControlService>,
|
||||
input: Option<Json<ClearStorageInput>>,
|
||||
) -> Result<Json<StatisticsStorage>> {
|
||||
Ok(Json(service.clear_statistics_storage().await?))
|
||||
let scope = input.map(|Json(input)| input.scope).unwrap_or_default();
|
||||
let storage = match scope {
|
||||
StatisticsStorageScope::Details => service.clear_statistics_storage().await?,
|
||||
StatisticsStorageScope::All => service.clear_all_statistics_storage().await?,
|
||||
};
|
||||
Ok(Json(storage))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct ClearStorageInput {
|
||||
#[serde(default)]
|
||||
pub scope: StatisticsStorageScope,
|
||||
}
|
||||
|
||||
pub async fn get_proxy(State(service): State<ControlService>) -> Result<Json<ProxySettings>> {
|
||||
|
||||
@@ -64,6 +64,7 @@ impl CursorActor {
|
||||
};
|
||||
match command {
|
||||
CursorCommand::Abort => {
|
||||
handle.mark_conversation_cancelled();
|
||||
handle.cancel();
|
||||
}
|
||||
CursorCommand::Finished => {
|
||||
@@ -76,6 +77,24 @@ impl CursorActor {
|
||||
Some(pb::agent_client_message::Message::RunRequest(
|
||||
request,
|
||||
)) => {
|
||||
if let Some(conversation_id) =
|
||||
request.conversation_id.as_deref()
|
||||
{
|
||||
if let Err(error) =
|
||||
handle.set_conversation_id(conversation_id)
|
||||
{
|
||||
tracing::error!(
|
||||
request_id = handle.request_id(),
|
||||
%error,
|
||||
"invalid Cursor conversation id"
|
||||
);
|
||||
let _ =
|
||||
crate::cursor::lifecycle::fail(&handle, &error);
|
||||
let _ =
|
||||
handle.command(CursorCommand::Finished).await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
if let Some((results, runtime_actions, dependencies)) =
|
||||
run_resources.take()
|
||||
{
|
||||
@@ -357,8 +376,12 @@ impl CursorActor {
|
||||
) => match action.action {
|
||||
Some(
|
||||
pb::conversation_action::Action::UserMessageAction(_),
|
||||
)
|
||||
| Some(pb::conversation_action::Action::CancelAction(_)) => {
|
||||
) => {
|
||||
handle.mark_conversation_cancelled();
|
||||
handle.cancel();
|
||||
}
|
||||
Some(pb::conversation_action::Action::CancelAction(_)) => {
|
||||
handle.mark_conversation_cancelled();
|
||||
handle.cancel();
|
||||
}
|
||||
Some(
|
||||
|
||||
@@ -43,6 +43,33 @@ impl DecodedAppend {
|
||||
request.conversation_id.as_deref()
|
||||
}
|
||||
|
||||
pub fn is_background_task_completion(&self) -> bool {
|
||||
let Some(agent::agent_client_message::Message::RunRequest(request)) =
|
||||
self.message.message.as_ref()
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
matches!(
|
||||
request
|
||||
.action
|
||||
.as_ref()
|
||||
.and_then(|action| action.action.as_ref()),
|
||||
Some(agent::conversation_action::Action::BackgroundTaskCompletionAction(_))
|
||||
)
|
||||
}
|
||||
|
||||
fn is_runtime_cancellation(&self) -> bool {
|
||||
matches!(
|
||||
self.message.message.as_ref(),
|
||||
Some(agent::agent_client_message::Message::ConversationAction(action))
|
||||
if matches!(
|
||||
action.action.as_ref(),
|
||||
Some(agent::conversation_action::Action::CancelAction(_))
|
||||
| Some(agent::conversation_action::Action::UserMessageAction(_))
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
pub fn trace_metadata(&self) -> serde_json::Value {
|
||||
let Some(message) = self.message.message.as_ref() else {
|
||||
return serde_json::json!({
|
||||
@@ -163,7 +190,29 @@ pub async fn append(
|
||||
request: DecodedAppend,
|
||||
parent: Option<CursorParent>,
|
||||
) -> Result<ai::BidiAppendResponse> {
|
||||
if let Some(conversation_id) = request.conversation_id() {
|
||||
if request.is_background_task_completion()
|
||||
&& registry.conversation_cancelled(conversation_id)
|
||||
{
|
||||
tracing::info!(
|
||||
request_id = %request.request_id,
|
||||
%conversation_id,
|
||||
"dropping background task completion for cancelled conversation"
|
||||
);
|
||||
return Ok(ai::BidiAppendResponse {});
|
||||
}
|
||||
if !request.is_background_task_completion() {
|
||||
registry.clear_conversation_cancelled(conversation_id);
|
||||
}
|
||||
}
|
||||
let handle = registry.get_or_create(&request.request_id).await?;
|
||||
if let Some(conversation_id) = request.conversation_id() {
|
||||
handle.set_conversation_id(conversation_id)?;
|
||||
}
|
||||
if request.is_runtime_cancellation() {
|
||||
handle.mark_conversation_cancelled();
|
||||
handle.cancel();
|
||||
}
|
||||
if let Some(parent) = parent {
|
||||
handle.set_parent(parent)?;
|
||||
}
|
||||
@@ -225,4 +274,21 @@ mod tests {
|
||||
.unwrap();
|
||||
assert_eq!(decoded.model_id(), Some("grok-4.6"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_background_task_completion_actions() {
|
||||
let decoded = decode(&encoded(agent::AgentRunRequest {
|
||||
action: Some(agent::ConversationAction {
|
||||
action: Some(
|
||||
agent::conversation_action::Action::BackgroundTaskCompletionAction(
|
||||
agent::BackgroundTaskCompletionAction::default(),
|
||||
),
|
||||
),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
}))
|
||||
.unwrap();
|
||||
assert!(decoded.is_background_task_completion());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
collections::{HashMap, HashSet},
|
||||
sync::{Arc, OnceLock},
|
||||
};
|
||||
|
||||
@@ -29,6 +29,8 @@ pub struct CursorSessionHandle {
|
||||
commands: mpsc::Sender<CursorCommand>,
|
||||
output: Arc<OutputHub>,
|
||||
cancellation: CancellationToken,
|
||||
conversation_id: Arc<OnceLock<String>>,
|
||||
cancelled_conversations: Arc<parking_lot::Mutex<HashSet<String>>>,
|
||||
parent: Arc<OnceLock<CursorParent>>,
|
||||
trace: Option<CursorTraceRecorder>,
|
||||
}
|
||||
@@ -43,6 +45,35 @@ impl CursorSessionHandle {
|
||||
pub fn request_id(&self) -> &str {
|
||||
&self.request_id
|
||||
}
|
||||
pub fn set_conversation_id(&self, conversation_id: &str) -> Result<()> {
|
||||
if conversation_id.is_empty() {
|
||||
return Err(crate::Error::Protocol(
|
||||
"Cursor conversation id is required".into(),
|
||||
));
|
||||
}
|
||||
if self
|
||||
.conversation_id
|
||||
.get()
|
||||
.is_some_and(|current| current != conversation_id)
|
||||
{
|
||||
return Err(crate::Error::Protocol(format!(
|
||||
"conflicting conversation ids for request {}",
|
||||
self.request_id
|
||||
)));
|
||||
}
|
||||
let _ = self.conversation_id.set(conversation_id.into());
|
||||
Ok(())
|
||||
}
|
||||
pub fn conversation_id(&self) -> Option<&str> {
|
||||
self.conversation_id.get().map(String::as_str)
|
||||
}
|
||||
pub fn mark_conversation_cancelled(&self) {
|
||||
if let Some(conversation_id) = self.conversation_id() {
|
||||
self.cancelled_conversations
|
||||
.lock()
|
||||
.insert(conversation_id.to_owned());
|
||||
}
|
||||
}
|
||||
pub fn subscribe(&self) -> mpsc::UnboundedReceiver<Bytes> {
|
||||
self.output.subscribe()
|
||||
}
|
||||
@@ -160,6 +191,7 @@ struct RegistryInner {
|
||||
store: Store,
|
||||
provider: Arc<dyn Provider>,
|
||||
compiler: PromptCompiler,
|
||||
cancelled_conversations: Arc<parking_lot::Mutex<HashSet<String>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
@@ -188,6 +220,7 @@ impl CursorSessionRegistry {
|
||||
store,
|
||||
provider,
|
||||
compiler,
|
||||
cancelled_conversations: Arc::new(parking_lot::Mutex::new(HashSet::new())),
|
||||
}),
|
||||
}
|
||||
}
|
||||
@@ -205,6 +238,8 @@ impl CursorSessionRegistry {
|
||||
commands,
|
||||
output,
|
||||
cancellation,
|
||||
conversation_id: Arc::new(OnceLock::new()),
|
||||
cancelled_conversations: self.inner.cancelled_conversations.clone(),
|
||||
parent: Arc::new(OnceLock::new()),
|
||||
trace,
|
||||
};
|
||||
@@ -262,6 +297,20 @@ impl CursorSessionRegistry {
|
||||
.contains_key(request_id)
|
||||
}
|
||||
|
||||
pub(crate) fn conversation_cancelled(&self, conversation_id: &str) -> bool {
|
||||
self.inner
|
||||
.cancelled_conversations
|
||||
.lock()
|
||||
.contains(conversation_id)
|
||||
}
|
||||
|
||||
pub(crate) fn clear_conversation_cancelled(&self, conversation_id: &str) {
|
||||
self.inner
|
||||
.cancelled_conversations
|
||||
.lock()
|
||||
.remove(conversation_id);
|
||||
}
|
||||
|
||||
pub(crate) async fn wait_route(&self, request_id: &str) -> CursorRoute {
|
||||
loop {
|
||||
let changed = self.inner.route_changed.notified();
|
||||
|
||||
@@ -49,7 +49,8 @@ impl Provider for AnthropicProvider {
|
||||
let recorder = self.recorder.clone();
|
||||
Box::pin(try_stream! {
|
||||
let ModelInvocation { call_id, request, .. } = invocation;
|
||||
let messages = anthropic_messages(&request.history)?;
|
||||
let mut messages = anthropic_messages(&request.history)?;
|
||||
mark_cache_breakpoint(&mut messages);
|
||||
let max_tokens = request.model.max_output_tokens.or(config.max_output_tokens)
|
||||
.unwrap_or(DEFAULT_MAX_OUTPUT_TOKENS);
|
||||
let mut body = json!({
|
||||
@@ -372,6 +373,24 @@ fn anthropic_messages(messages: &[ProjectedMessage]) -> Result<Vec<Value>> {
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn mark_cache_breakpoint(messages: &mut [Value]) {
|
||||
for message in messages {
|
||||
let Some(content) = message.get_mut("content").and_then(Value::as_array_mut) else {
|
||||
continue;
|
||||
};
|
||||
for block in content {
|
||||
let kind = block.get("type").and_then(Value::as_str);
|
||||
if matches!(kind, Some("thinking" | "redacted_thinking")) {
|
||||
continue;
|
||||
}
|
||||
if let Some(block) = block.as_object_mut() {
|
||||
block.insert("cache_control".into(), json!({"type": "ephemeral"}));
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn anthropic_parts(role: &Role, parts: &[ContentPart]) -> Result<Vec<Value>> {
|
||||
parts
|
||||
.iter()
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! Storage accounting and cleanup for disposable observability data.
|
||||
|
||||
use serde::Serialize;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::Result;
|
||||
|
||||
@@ -13,6 +13,14 @@ pub struct StatisticsStorage {
|
||||
pub trace_count: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum StatisticsStorageScope {
|
||||
#[default]
|
||||
Details,
|
||||
All,
|
||||
}
|
||||
|
||||
impl Store {
|
||||
pub async fn statistics_storage(&self) -> Result<StatisticsStorage> {
|
||||
let (bytes, call_count, trace_count) = sqlx::query_as::<_, (i64, i64, i64)>(
|
||||
@@ -56,6 +64,15 @@ impl Store {
|
||||
pub async fn clear_statistics_storage(&self) -> Result<StatisticsStorage> {
|
||||
let _write = self.writes.lock().await;
|
||||
let mut transaction = self.pool.begin().await?;
|
||||
Self::clear_detail_storage_tx(&mut transaction).await?;
|
||||
transaction.commit().await?;
|
||||
self.statistics_storage().await
|
||||
}
|
||||
|
||||
pub async fn clear_all_statistics_storage(&self) -> Result<StatisticsStorage> {
|
||||
let _write = self.writes.lock().await;
|
||||
let mut transaction = self.pool.begin().await?;
|
||||
Self::clear_trace_artifacts_tx(&mut transaction).await?;
|
||||
sqlx::query("DELETE FROM llm_calls")
|
||||
.execute(&mut *transaction)
|
||||
.await?;
|
||||
@@ -65,6 +82,59 @@ impl Store {
|
||||
transaction.commit().await?;
|
||||
self.statistics_storage().await
|
||||
}
|
||||
|
||||
async fn clear_detail_storage_tx(
|
||||
transaction: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
) -> Result<()> {
|
||||
sqlx::query("DELETE FROM llm_call_requests")
|
||||
.execute(&mut **transaction)
|
||||
.await?;
|
||||
sqlx::query("DELETE FROM llm_call_response_chunks")
|
||||
.execute(&mut **transaction)
|
||||
.await?;
|
||||
Self::clear_trace_artifacts_tx(transaction).await
|
||||
}
|
||||
|
||||
async fn clear_trace_artifacts_tx(
|
||||
transaction: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
) -> Result<()> {
|
||||
sqlx::query(
|
||||
"CREATE TEMP TABLE IF NOT EXISTS clear_statistics_blob_ids(
|
||||
blob_id BLOB PRIMARY KEY
|
||||
)",
|
||||
)
|
||||
.execute(&mut **transaction)
|
||||
.await?;
|
||||
sqlx::query("DELETE FROM clear_statistics_blob_ids")
|
||||
.execute(&mut **transaction)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"INSERT OR IGNORE INTO clear_statistics_blob_ids(blob_id)
|
||||
SELECT blob_id FROM cursor_run_trace_artifacts",
|
||||
)
|
||||
.execute(&mut **transaction)
|
||||
.await?;
|
||||
sqlx::query("DELETE FROM cursor_run_trace_artifacts")
|
||||
.execute(&mut **transaction)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"DELETE FROM blobs
|
||||
WHERE blob_id IN (SELECT blob_id FROM clear_statistics_blob_ids)
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM cursor_run_trace_artifacts a WHERE a.blob_id = blobs.blob_id
|
||||
)
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM blob_edges e
|
||||
WHERE e.parent_blob_id = blobs.blob_id OR e.child_blob_id = blobs.blob_id
|
||||
)",
|
||||
)
|
||||
.execute(&mut **transaction)
|
||||
.await?;
|
||||
sqlx::query("DROP TABLE clear_statistics_blob_ids")
|
||||
.execute(&mut **transaction)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -73,7 +143,7 @@ mod tests {
|
||||
use crate::model::{ModelConfigInput, ModelType, OPENAI_CHAT_ENDPOINT};
|
||||
|
||||
#[tokio::test]
|
||||
async fn clears_observability_without_removing_configuration() {
|
||||
async fn clears_detail_storage_without_removing_configuration() {
|
||||
let store = Store::connect("sqlite::memory:").await.unwrap();
|
||||
store
|
||||
.create_model(&ModelConfigInput {
|
||||
@@ -106,12 +176,16 @@ mod tests {
|
||||
|
||||
assert!(store.statistics_storage().await.unwrap().bytes > 0);
|
||||
let cleared = store.clear_statistics_storage().await.unwrap();
|
||||
assert_eq!(cleared.bytes, 0);
|
||||
assert_eq!(cleared.call_count, 0);
|
||||
assert_eq!(cleared.call_count, 1);
|
||||
assert_eq!(cleared.trace_count, 0);
|
||||
assert!(cleared.bytes > 0);
|
||||
assert!(store.llm_call_request("call-1").await.unwrap().is_none());
|
||||
assert!(store.llm_call_chunks("call-1").await.unwrap().is_empty());
|
||||
let model_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM model_configs")
|
||||
.fetch_one(store.pool())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(model_count, 1);
|
||||
|
||||
store
|
||||
@@ -127,5 +201,11 @@ mod tests {
|
||||
.record_llm_chunk("call-1", 0, 1, b"data", true)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(store.llm_call_request("call-1").await.unwrap().is_some());
|
||||
assert_eq!(store.llm_call_chunks("call-1").await.unwrap().len(), 1);
|
||||
let cleared = store.clear_all_statistics_storage().await.unwrap();
|
||||
assert_eq!(cleared.bytes, 0);
|
||||
assert_eq!(cleared.call_count, 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use cursor_server::{
|
||||
cursor::{
|
||||
bidi_append::{self, DecodedAppend},
|
||||
connect,
|
||||
prompting::{PromptAssets, PromptCompiler},
|
||||
proto::agent::v1 as pb,
|
||||
@@ -20,6 +21,45 @@ use prost::Message;
|
||||
const FOLLOW_UP: &str = "Perform any necessary follow-up actions in response to the subagent completion above. If no follow-up work is needed, no further action is required. If you mention an agent or subagent in your response, link it with the `[Name](id)` Don't use generic label such as `[agent]`, `[worker]`, or `[subagent]`.";
|
||||
const SHELL_FOLLOW_UP: &str = "Briefly inform the user about the task result and perform any follow-up actions (if needed). If there's no follow-ups needed, don't explicitly say that.";
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancelled_conversation_drops_late_background_completion() {
|
||||
let (_directory, store) = fixtures::temp_store().await;
|
||||
let provider = fake_provider::FakeProvider::default();
|
||||
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 active = registry.get_or_create("cancelled-request").await.unwrap();
|
||||
active.set_conversation_id("parent-conversation").unwrap();
|
||||
active.mark_conversation_cancelled();
|
||||
|
||||
bidi_append::append(
|
||||
®istry,
|
||||
DecodedAppend {
|
||||
request_id: "late-completion".into(),
|
||||
seqno: 1,
|
||||
message: completion_run(
|
||||
"child-id",
|
||||
"cancelled-parent-run",
|
||||
pb::ConversationStateStructure::default(),
|
||||
),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(provider.requests().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn background_subagent_completion_starts_a_simulated_parent_turn() {
|
||||
let (_directory, store) = fixtures::temp_store().await;
|
||||
|
||||
@@ -672,7 +672,17 @@ async fn anthropic_raw_stream_uses_explicit_and_default_token_limits() {
|
||||
server.abort();
|
||||
|
||||
assert_eq!(body["max_tokens"], 1234);
|
||||
assert!(body.get("cache_control").is_none());
|
||||
assert_eq!(
|
||||
body["messages"][0]["content"][0]["cache_control"]["type"],
|
||||
"ephemeral"
|
||||
);
|
||||
assert_eq!(default_body["max_tokens"], 65_000);
|
||||
assert!(default_body.get("cache_control").is_none());
|
||||
assert_eq!(
|
||||
default_body["messages"][0]["content"][0]["cache_control"]["type"],
|
||||
"ephemeral"
|
||||
);
|
||||
assert!(body.get("service_tier").is_none());
|
||||
assert_eq!(body["messages"][0]["content"][1]["type"], "image");
|
||||
assert_eq!(
|
||||
|
||||
Reference in New Issue
Block a user