mirror of
https://wget.la/https://github.com/leookun/cursor-byok
synced 2026-10-04 02:52:55 +08:00
feat: webfetch
This commit is contained in:
@@ -31,6 +31,7 @@ pub fn router(registry: TransportRegistry) -> Result<Router> {
|
||||
}
|
||||
|
||||
fn router_with_proxy(registry: TransportRegistry, proxy: CursorProxy) -> Router {
|
||||
let web_cache = registry.web_cache().router();
|
||||
Router::new()
|
||||
.route("/__byok-api__/healthz", get(health))
|
||||
.route("/agent.v1.AgentService/RunSSE", post(run_sse_handler))
|
||||
@@ -80,6 +81,7 @@ fn router_with_proxy(registry: TransportRegistry, proxy: CursorProxy) -> Router
|
||||
.method_not_allowed_fallback(proxy::forward)
|
||||
.layer(Extension(proxy))
|
||||
.with_state(registry)
|
||||
.merge(web_cache)
|
||||
}
|
||||
|
||||
async fn health() -> StatusCode {
|
||||
|
||||
+8
-1
@@ -14,6 +14,7 @@ use crate::{
|
||||
},
|
||||
local_app::CursorHarness,
|
||||
provider::ProviderRouter,
|
||||
search::WebCache,
|
||||
store::Store,
|
||||
Result,
|
||||
};
|
||||
@@ -40,7 +41,12 @@ impl App {
|
||||
store.clone(),
|
||||
config.provider_request_timeout,
|
||||
));
|
||||
let registry = TransportRegistry::new(store.clone(), provider.clone(), compiler);
|
||||
let registry = TransportRegistry::with_web_cache(
|
||||
store.clone(),
|
||||
provider.clone(),
|
||||
compiler,
|
||||
WebCache::managed()?,
|
||||
);
|
||||
let control = control::ControlService::new(store.clone(), provider)?;
|
||||
let harness = control.cursor_harness().clone();
|
||||
let mut router = api::router(registry.clone())?;
|
||||
@@ -104,6 +110,7 @@ impl App {
|
||||
|
||||
pub async fn serve_on(self, listener: TcpListener, shutdown: CancellationToken) -> Result<()> {
|
||||
let address = listener.local_addr()?;
|
||||
self.registry.web_cache().set_service_addr(address);
|
||||
self.harness.set_backend_addr(address);
|
||||
tracing::info!(%address, "cursor server listening");
|
||||
let registry = self.registry;
|
||||
|
||||
@@ -9,6 +9,7 @@ use crate::{
|
||||
model::{ConversationId, RunId},
|
||||
provider::Provider,
|
||||
run::{CommandResult, RunHandle},
|
||||
search::WebCache,
|
||||
store::Store,
|
||||
};
|
||||
|
||||
@@ -24,6 +25,7 @@ pub(crate) struct ConversationDependencies {
|
||||
pub store: Store,
|
||||
pub provider: Arc<dyn Provider>,
|
||||
pub compiler: PromptCompiler,
|
||||
pub web_cache: WebCache,
|
||||
}
|
||||
|
||||
struct RegistryInner {
|
||||
@@ -40,7 +42,12 @@ struct ActiveRun {
|
||||
}
|
||||
|
||||
impl ConversationRegistry {
|
||||
pub fn new(store: Store, provider: Arc<dyn Provider>, compiler: PromptCompiler) -> Self {
|
||||
pub fn new(
|
||||
store: Store,
|
||||
provider: Arc<dyn Provider>,
|
||||
compiler: PromptCompiler,
|
||||
web_cache: WebCache,
|
||||
) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(RegistryInner {
|
||||
current: Mutex::new(HashMap::new()),
|
||||
@@ -50,6 +57,7 @@ impl ConversationRegistry {
|
||||
store,
|
||||
provider,
|
||||
compiler,
|
||||
web_cache,
|
||||
},
|
||||
}),
|
||||
}
|
||||
|
||||
@@ -147,6 +147,7 @@ impl ConversationRuntime {
|
||||
tool_runtime.clone(),
|
||||
results.clone(),
|
||||
dependencies.store.clone(),
|
||||
dependencies.web_cache.clone(),
|
||||
);
|
||||
let generation = RunGeneration {
|
||||
superseded: CancellationToken::new(),
|
||||
|
||||
@@ -18,7 +18,7 @@ pub(crate) mod tool_call_result;
|
||||
|
||||
use crate::{
|
||||
model::{CanonicalMessage, MessageContent, Role, ToolCall},
|
||||
search::{WebFetch, WebSearch},
|
||||
search::{WebCache, WebFetch, WebSearch},
|
||||
store::Store,
|
||||
Error, Result,
|
||||
};
|
||||
@@ -72,12 +72,13 @@ impl ToolDispatcher {
|
||||
runtime: CursorToolRuntime,
|
||||
results: ToolResultSender,
|
||||
store: Store,
|
||||
web_cache: WebCache,
|
||||
) -> Self {
|
||||
Self {
|
||||
runtime,
|
||||
results,
|
||||
search: WebSearch::managed(store.clone()),
|
||||
fetch: WebFetch::managed(store.clone()),
|
||||
fetch: WebFetch::managed(store.clone(), web_cache),
|
||||
store: Some(store),
|
||||
edit_schedule: Arc::new(Mutex::new(EditSchedule::default())),
|
||||
}
|
||||
|
||||
@@ -235,12 +235,23 @@ pub(crate) fn complete_web_fetch(
|
||||
};
|
||||
let (output, is_error) = match outcome {
|
||||
Ok(page) => {
|
||||
let output = page.markdown.clone();
|
||||
let output = match page.cache.as_ref() {
|
||||
Some(cache) => format!(
|
||||
"<system_reminder>Web content has been downloaded to: {}. If the content is omitted and you need it, use Shell to download it to a temporary directory, then use an appropriate tool to read it in pages.</system_reminder>\n{}",
|
||||
cache.url, page.markdown
|
||||
),
|
||||
None => page.markdown.clone(),
|
||||
};
|
||||
let output_location = page.cache.map(|cache| pb::OutputLocation {
|
||||
file_path: cache.file_path,
|
||||
size_bytes: cache.size_bytes,
|
||||
line_count: cache.line_count,
|
||||
});
|
||||
tool.result = Some(pb::WebFetchResult {
|
||||
result: Some(pb::web_fetch_result::Result::Success(pb::WebFetchSuccess {
|
||||
url: page.url,
|
||||
markdown: page.markdown,
|
||||
output_location: None,
|
||||
output_location,
|
||||
})),
|
||||
});
|
||||
(output, false)
|
||||
@@ -328,3 +339,75 @@ fn switch_mode_result(
|
||||
fn missing(name: &str) -> Error {
|
||||
Error::Protocol(format!("{name} returned no result"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{complete_web_fetch, PendingInteraction};
|
||||
use crate::{
|
||||
cursor::protocol::proto::agent::v1 as pb,
|
||||
model::ToolCall,
|
||||
search::{FetchedPage, WebCacheEntry},
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn web_fetch_result_leads_with_cache_reminder_and_bounds_cursor_payload() {
|
||||
let markdown = "x".repeat(40 * 1024);
|
||||
let location = "http://127.0.0.1:4312/web-cache/550e8400-e29b-41d4-a716-446655440000.txt";
|
||||
let completion = complete_web_fetch(
|
||||
pending_fetch(),
|
||||
Ok(FetchedPage {
|
||||
url: "https://example.com/final".into(),
|
||||
markdown: markdown.clone(),
|
||||
cache: Some(WebCacheEntry {
|
||||
url: location.into(),
|
||||
file_path: "C:/Users/test/.cursor-byok-v3/cache/web/page.txt".into(),
|
||||
size_bytes: markdown.len() as i64,
|
||||
line_count: 1,
|
||||
}),
|
||||
}),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(completion.result().content.starts_with(&format!(
|
||||
"<system_reminder>Web content has been downloaded to: {location}."
|
||||
)));
|
||||
assert!(completion.result().content.contains("[truncated: WebFetch"));
|
||||
let Some(pb::tool_call::Tool::WebFetchToolCall(tool)) =
|
||||
completion.tool_call().tool.as_ref()
|
||||
else {
|
||||
panic!("expected WebFetchToolCall")
|
||||
};
|
||||
let Some(pb::web_fetch_result::Result::Success(success)) = tool
|
||||
.result
|
||||
.as_ref()
|
||||
.and_then(|result| result.result.as_ref())
|
||||
else {
|
||||
panic!("expected WebFetchSuccess")
|
||||
};
|
||||
assert!(success.markdown.len() <= 32 * 1024);
|
||||
assert!(success.markdown.contains("[truncated: WebFetch"));
|
||||
assert_eq!(
|
||||
success
|
||||
.output_location
|
||||
.as_ref()
|
||||
.map(|location| location.file_path.as_str()),
|
||||
Some("C:/Users/test/.cursor-byok-v3/cache/web/page.txt")
|
||||
);
|
||||
}
|
||||
|
||||
fn pending_fetch() -> PendingInteraction {
|
||||
PendingInteraction {
|
||||
call: ToolCall {
|
||||
index: 0,
|
||||
call_id: "fetch-call".into(),
|
||||
model_call_id: "model-call".into(),
|
||||
name: "WebFetch".into(),
|
||||
arguments_text: r#"{"url":"https://example.com"}"#.into(),
|
||||
arguments: json!({"url": "https://example.com"}),
|
||||
},
|
||||
started_at_ms: 1,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ use crate::{
|
||||
services::observability::CursorTraceRecorder,
|
||||
},
|
||||
provider::Provider,
|
||||
search::WebCache,
|
||||
store::Store,
|
||||
Result,
|
||||
};
|
||||
@@ -26,6 +27,7 @@ struct RegistryInner {
|
||||
upstream: Mutex<HashMap<String, u64>>,
|
||||
route_changed: Notify,
|
||||
store: Store,
|
||||
web_cache: WebCache,
|
||||
conversations: ConversationRegistry,
|
||||
}
|
||||
|
||||
@@ -37,13 +39,28 @@ pub enum TransportRoute {
|
||||
|
||||
impl TransportRegistry {
|
||||
pub fn new(store: Store, provider: Arc<dyn Provider>, compiler: PromptCompiler) -> Self {
|
||||
Self::with_web_cache(store, provider, compiler, WebCache::default())
|
||||
}
|
||||
|
||||
pub fn with_web_cache(
|
||||
store: Store,
|
||||
provider: Arc<dyn Provider>,
|
||||
compiler: PromptCompiler,
|
||||
web_cache: WebCache,
|
||||
) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(RegistryInner {
|
||||
local: Mutex::new(HashMap::new()),
|
||||
upstream: Mutex::new(HashMap::new()),
|
||||
route_changed: Notify::new(),
|
||||
conversations: ConversationRegistry::new(store.clone(), provider, compiler),
|
||||
conversations: ConversationRegistry::new(
|
||||
store.clone(),
|
||||
provider,
|
||||
compiler,
|
||||
web_cache.clone(),
|
||||
),
|
||||
store,
|
||||
web_cache,
|
||||
}),
|
||||
}
|
||||
}
|
||||
@@ -52,6 +69,10 @@ impl TransportRegistry {
|
||||
&self.inner.store
|
||||
}
|
||||
|
||||
pub fn web_cache(&self) -> &WebCache {
|
||||
&self.inner.web_cache
|
||||
}
|
||||
|
||||
pub fn conversations(&self) -> &ConversationRegistry {
|
||||
&self.inner.conversations
|
||||
}
|
||||
|
||||
+32
-14
@@ -2,7 +2,7 @@
|
||||
use std::{
|
||||
fs,
|
||||
net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
|
||||
path::{Path, PathBuf},
|
||||
path::PathBuf,
|
||||
sync::Arc,
|
||||
};
|
||||
|
||||
@@ -15,6 +15,14 @@ use crate::{config::managed_data_dir, Error, Result};
|
||||
|
||||
const CACHE_ROUTE: &str = "/web-cache";
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct WebCacheEntry {
|
||||
pub url: String,
|
||||
pub file_path: String,
|
||||
pub size_bytes: i64,
|
||||
pub line_count: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct WebCache {
|
||||
inner: Option<Arc<WebCacheInner>>,
|
||||
@@ -27,7 +35,7 @@ struct WebCacheInner {
|
||||
|
||||
impl WebCache {
|
||||
pub fn managed() -> Result<Self> {
|
||||
Self::at(managed_data_dir()?.join("cache/web"))
|
||||
Self::at(managed_data_dir()?.join("cache").join("web"))
|
||||
}
|
||||
|
||||
pub fn at(directory: PathBuf) -> Result<Self> {
|
||||
@@ -46,7 +54,7 @@ impl WebCache {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn store(&self, content: &str) -> Result<Option<String>> {
|
||||
pub async fn store(&self, content: &str) -> Result<Option<WebCacheEntry>> {
|
||||
let Some(inner) = &self.inner else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -55,11 +63,19 @@ impl WebCache {
|
||||
})?;
|
||||
let file_name = format!("{}.txt", Uuid::new_v4());
|
||||
let path = inner.directory.join(&file_name);
|
||||
let file_path = path.to_string_lossy().into_owned();
|
||||
let size_bytes = content.len() as i64;
|
||||
let line_count = content.lines().count() as i64;
|
||||
let bytes = content.as_bytes().to_vec();
|
||||
tokio::task::spawn_blocking(move || fs::write(path, bytes))
|
||||
.await
|
||||
.map_err(|error| Error::Store(format!("web cache write task failed: {error}")))??;
|
||||
Ok(Some(format!("http://{address}{CACHE_ROUTE}/{file_name}")))
|
||||
Ok(Some(WebCacheEntry {
|
||||
url: format!("http://{address}{CACHE_ROUTE}/{file_name}"),
|
||||
file_path,
|
||||
size_bytes,
|
||||
line_count,
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn router(&self) -> Router {
|
||||
@@ -70,7 +86,7 @@ impl WebCache {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn directory(&self) -> &Path {
|
||||
fn directory(&self) -> &std::path::Path {
|
||||
&self.inner.as_ref().expect("enabled web cache").directory
|
||||
}
|
||||
}
|
||||
@@ -102,22 +118,24 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn stores_uuid_named_content_and_serves_it_from_existing_router() {
|
||||
let directory = tempdir().unwrap();
|
||||
let cache = WebCache::at(directory.path().join("cache/web")).unwrap();
|
||||
cache
|
||||
.set_service_addr("0.0.0.0:4312".parse().unwrap());
|
||||
let cache = WebCache::at(directory.path().join("cache").join("web")).unwrap();
|
||||
cache.set_service_addr("0.0.0.0:4312".parse().unwrap());
|
||||
|
||||
let location = cache.store("complete fetched content").await.unwrap().unwrap();
|
||||
let name = location.rsplit('/').next().unwrap();
|
||||
let entry = cache
|
||||
.store("complete fetched content")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let name = entry.url.rsplit('/').next().unwrap();
|
||||
let id = name.strip_suffix(".txt").unwrap();
|
||||
assert!(Uuid::parse_str(id).is_ok());
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(cache.directory().join(name)).unwrap(),
|
||||
"complete fetched content"
|
||||
);
|
||||
assert_eq!(
|
||||
location,
|
||||
format!("http://127.0.0.1:4312/web-cache/{name}")
|
||||
);
|
||||
assert_eq!(entry.url, format!("http://127.0.0.1:4312/web-cache/{name}"));
|
||||
assert_eq!(entry.size_bytes, 24);
|
||||
assert_eq!(entry.line_count, 1);
|
||||
|
||||
let response = cache
|
||||
.router()
|
||||
|
||||
@@ -8,7 +8,7 @@ use crate::store::Store;
|
||||
use super::{catalog, SearchEngine, SearchHit};
|
||||
|
||||
const RRF_K: f64 = 60.0;
|
||||
const MAX_RESULTS: usize = 30;
|
||||
const MAX_RESULTS: usize = 15;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct WebSearch {
|
||||
|
||||
@@ -15,7 +15,9 @@ use reqwest::{
|
||||
use tokio::{net::lookup_host, time::timeout};
|
||||
use url::{Host, Url};
|
||||
|
||||
use crate::{search::WebCache, store::Store};
|
||||
use crate::store::Store;
|
||||
|
||||
use super::{WebCache, WebCacheEntry};
|
||||
|
||||
const MAX_RESPONSE_SIZE: usize = 5 * 1024 * 1024;
|
||||
const MAX_REDIRECTS: usize = 5;
|
||||
@@ -25,6 +27,7 @@ const FETCH_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
pub struct FetchedPage {
|
||||
pub url: String,
|
||||
pub markdown: String,
|
||||
pub cache: Option<WebCacheEntry>,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
@@ -34,6 +37,7 @@ pub struct FetchError(String);
|
||||
#[derive(Clone)]
|
||||
pub struct WebFetch {
|
||||
client: FetchClient,
|
||||
cache: WebCache,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -46,19 +50,27 @@ impl WebFetch {
|
||||
pub fn built_in() -> Self {
|
||||
Self {
|
||||
client: FetchClient::Direct,
|
||||
cache: WebCache::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn managed(store: Store) -> Self {
|
||||
pub(crate) fn managed(store: Store, cache: WebCache) -> Self {
|
||||
Self {
|
||||
client: FetchClient::Managed(store),
|
||||
cache,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn fetch(&self, value: &str) -> Result<FetchedPage, FetchError> {
|
||||
timeout(FETCH_TIMEOUT, self.fetch_inner(value))
|
||||
let mut page = timeout(FETCH_TIMEOUT, self.fetch_inner(value))
|
||||
.await
|
||||
.map_err(|_| failure("request timed out"))?
|
||||
.map_err(|_| failure("request timed out"))??;
|
||||
page.cache = self
|
||||
.cache
|
||||
.store(&page.markdown)
|
||||
.await
|
||||
.map_err(|error| failure(format!("cannot cache fetched content: {error}")))?;
|
||||
Ok(page)
|
||||
}
|
||||
|
||||
async fn fetch_inner(&self, value: &str) -> Result<FetchedPage, FetchError> {
|
||||
@@ -190,7 +202,11 @@ async fn page(response: Response) -> Result<FetchedPage, FetchError> {
|
||||
if markdown.trim().is_empty() {
|
||||
return Err(failure("response contains no readable content"));
|
||||
}
|
||||
Ok(FetchedPage { url, markdown })
|
||||
Ok(FetchedPage {
|
||||
url,
|
||||
markdown,
|
||||
cache: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn limited_body(response: Response) -> Result<BytesMut, FetchError> {
|
||||
|
||||
@@ -6,7 +6,7 @@ mod federation;
|
||||
mod fetch;
|
||||
mod search_provider;
|
||||
|
||||
pub use cache::WebCache;
|
||||
pub use cache::{WebCache, WebCacheEntry};
|
||||
pub use engine::{HtmlEngine, JsonEngine, SearchEngine, SearchHit};
|
||||
pub use federation::{SearchError, WebSearch};
|
||||
pub use fetch::{FetchError, FetchedPage, WebFetch};
|
||||
|
||||
Reference in New Issue
Block a user