Merge pull request #144 from nearai/feat/agent-loop-web-context-search

feat(agent-loop): server-side web_context_search loop for /v1/chat/completions

Author: Evrard-Nil

benches/e2e.rs

@@ -87,6 +87,10 @@ fn build_test_app(mock_url: &str) -> axum::Router {
         dstack_socket_path: "/var/run/dstack.sock".to_string(),
         gpu_evidence_delegate_url: None,
         gpu_evidence_delegate_timeout_secs: 30,
+        web_context_search_url: None,
+        web_context_search_api_key: None,
+        agent_loop_max_iterations: 5,
+        web_context_search_timeout_secs: 30,
     };
 
     let ecdsa = signing::EcdsaContext::from_key_bytes(&ECDSA_KEY).unwrap();

src/agent_loop.rs

@@ -160,6 +160,21 @@ pub struct Config {
     /// unreachable (otherwise `/v1/attestation/report` silently 500s while
     /// `/v1/models` still passes). Default: `/var/run/dstack.sock`.
     pub dstack_socket_path: String,
+
+    // Agent loop (server-side web_context_search tool)
+    /// Brave LLM Context API endpoint. When unset, requests advertising the
+    /// `{"type":"web_context_search"}` tool are rejected with 400. All
+    /// tool execution happens inside the CVM; the query is the only thing
+    /// that egresses, going directly to Brave under TLS.
+    pub web_context_search_url: Option<String>,
+    /// Brave subscription token for the LLM Context endpoint. Sent as
+    /// `X-Subscription-Token`. Required when `web_context_search_url` is set.
+    pub web_context_search_api_key: Option<String>,
+    /// Hard cap on tool-call iterations within a single chat completion.
+    /// Once hit, the loop emits a synthetic terminator chunk and stops.
+    pub agent_loop_max_iterations: u32,
+    /// Per-tool-call timeout for the Brave HTTP request.
+    pub web_context_search_timeout_secs: u64,
 }
 
 impl Config {
@@ -316,6 +331,18 @@ impl Config {
             rerank_url_override,
             score_url_override,
             dstack_socket_path: env_or("DSTACK_SOCKET_PATH", "/var/run/dstack.sock"),
+            web_context_search_url: env::var("WEB_CONTEXT_SEARCH_URL")
+                .ok()
+                .filter(|s| !s.is_empty()),
+            web_context_search_api_key: env::var("WEB_CONTEXT_SEARCH_API_KEY")
+                .ok()
+                .filter(|s| !s.is_empty()),
+            // `env_int` returns `usize`; on 64-bit hosts a user-supplied value
+            // > u32::MAX would silently wrap. `try_from` surfaces it as a
+            // config error instead so a typo can't become a tiny iteration cap.
+            agent_loop_max_iterations: u32::try_from(env_int("AGENT_LOOP_MAX_ITERATIONS", 5))
+                .map_err(|_| anyhow::anyhow!("AGENT_LOOP_MAX_ITERATIONS exceeds the u32 range"))?,
+            web_context_search_timeout_secs: env_int("WEB_CONTEXT_SEARCH_TIMEOUT_SECS", 30) as u64,
         };
 
         // Validate attestation cache TTL (TTL/2 is used as refresh interval, so TTL < 2 would cause a busy loop)
@@ -334,6 +361,19 @@ impl Config {
             anyhow::bail!("STARTUP_CHECK_TIMEOUT_SECS must be greater than 0");
         }
 
+        // Agent loop: URL and key must both be set or both unset; iteration cap must be positive.
+        if config.web_context_search_url.is_some() != config.web_context_search_api_key.is_some() {
+            anyhow::bail!(
+                "WEB_CONTEXT_SEARCH_URL and WEB_CONTEXT_SEARCH_API_KEY must both be set or both unset"
+            );
+        }
+        if config.agent_loop_max_iterations == 0 {
+            anyhow::bail!("AGENT_LOOP_MAX_ITERATIONS must be at least 1");
+        }
+        if config.web_context_search_timeout_secs == 0 {
+            anyhow::bail!("WEB_CONTEXT_SEARCH_TIMEOUT_SECS must be greater than 0");
+        }
+
         Ok(config)
     }
 }

src/config.rs

@@ -718,6 +718,19 @@ fn encrypt_chat_response_choices(
                     encrypt_field(audio, "data", ctx, signing)?;
                 }
 
+                // Synthetic agent-loop tool result chunk (streaming only).
+                // Emitted by `agent_loop::run_chat_completion` between
+                // iterations as `delta.nearai_tool_result.output`. Encrypted
+                // unconditionally whenever an encryption context is active —
+                // the agent loop is the privacy-critical path, so the search
+                // output must travel encrypted whether or not the client set
+                // `X-Encrypt-All-Fields`. The field only appears when the
+                // client also opted into server-side tool execution by
+                // sending `tools: [{"type":"web_context_search"}]`.
+                if let Some(tool_result) = msg.get_mut("nearai_tool_result") {
+                    encrypt_field(tool_result, "output", ctx, signing)?;
+                }
+
                 // Extended fields — only when client opts in
                 if ctx.encrypt_all_fields {
                     encrypt_field(msg, "refusal", ctx, signing)?;

src/encryption.rs

@@ -5,6 +5,7 @@ use axum::middleware::Next;
 use axum::response::Response;
 use tracing::Instrument;
 
+pub mod agent_loop;
 pub mod attestation;
 pub mod attestation_sdk;
 pub mod auth;

src/lib.rs

@@ -310,7 +310,7 @@ fn try_report_usage(response_data: &serde_json::Value, id: &str, opts: &ProxyOpt
 /// service-token path (`/v1/internal/usage`) when [`UsageReporter::can_use_service_token_path`]
 /// is true; otherwise falls back to the legacy sk-bearer path
 /// (`/v1/usage`) verbatim.
-fn spawn_usage_report(reporter: &UsageReporter, mut body: serde_json::Value) {
+pub(crate) fn spawn_usage_report(reporter: &UsageReporter, mut body: serde_json::Value) {
     let client = reporter.http_client.clone();
     let (url, auth, mode) = if reporter.can_use_service_token_path() {
         // Inject subject identity into the body. Cloud-api's
@@ -1169,7 +1169,7 @@ pub async fn proxy_streaming_request(
 /// (vLLM `qwen3` parser as of v0.10) so downstream clients see the standard
 /// `delta.reasoning_content` field consistently across reasoning models.
 /// If both fields are present the existing `reasoning_content` is kept.
-fn normalize_chat_chunk(val: &mut serde_json::Value) {
+pub(crate) fn normalize_chat_chunk(val: &mut serde_json::Value) {
     let Some(choices) = val.get_mut("choices").and_then(|c| c.as_array_mut()) else {
         return;
     };
@@ -1652,10 +1652,10 @@ pub async fn proxy_streaming_response(
 
 /// Drop guard that tracks the streaming_connections gauge.
 /// Increments on creation, decrements on drop — guarantees they stay paired.
-struct StreamingGuard;
+pub(crate) struct StreamingGuard;
 
 impl StreamingGuard {
-    fn new() -> Self {
+    pub(crate) fn new() -> Self {
         metrics::gauge!("streaming_connections").increment(1);
         Self
     }

src/proxy.rs

@@ -6,6 +6,7 @@ use axum::Extension;
 
 use sha2::Digest;
 
+use crate::agent_loop;
 use crate::auth::RequireAuth;
 use crate::encryption::{self, Endpoint};
 use crate::error::AppError;
@@ -35,11 +36,8 @@ pub async fn chat_completions(
     // honored when authenticated with config.token (trusted gateway); sk- clients
     // always bind signatures to the wire body so they cannot forge a hash for a
     // different payload.
-    let original_request_hash = Some(resolve_request_hash_for_signing(
-        &headers,
-        &request_body,
-        auth.cloud_api_key.is_none(),
-    ));
+    let request_hash =
+        resolve_request_hash_for_signing(&headers, &request_body, auth.cloud_api_key.is_none());
 
     // Decrypt request fields if encryption is active
     if let Some(ref ctx) = enc_ctx {
@@ -56,6 +54,50 @@ pub async fn chat_completions(
         .and_then(|v| v.as_bool())
         .unwrap_or(false);
 
+    // Server-side agent loop opt-in: the request advertises exactly
+    // `{"type":"web_context_search"}` and nothing else. Requires streaming
+    // (we splice tool-result chunks between iterations) and Brave creds
+    // configured on this CVM. Anything that doesn't match falls through to
+    // the existing pass-through path below, byte-for-byte identical.
+    if agent_loop::is_web_context_search_request(&request_json) {
+        if !is_stream {
+            return Err(AppError::BadRequest(
+                "web_context_search requires stream:true".to_string(),
+            ));
+        }
+        if state.config.web_context_search_url.is_none()
+            || state.config.web_context_search_api_key.is_none()
+        {
+            return Err(AppError::BadRequest(
+                "web_context_search is not configured on this deployment".to_string(),
+            ));
+        }
+
+        // Build chunk transform if E2EE is active. The agent loop is the
+        // privacy-critical path — both our synthetic `nearai_tool_result`
+        // chunks AND the model's own `tool_calls[].function.{name,arguments}`
+        // (which contain the search query the model just generated from the
+        // user's E2EE-decrypted prompt) must travel encrypted. Force
+        // `encrypt_all_fields: true` on the context used to build this
+        // transform so clients don't need to remember to send
+        // `X-Encrypt-All-Fields: true` to get the full privacy guarantee.
+        // This only affects the agent-loop path; the regular chat path
+        // below still honors the client's `X-Encrypt-All-Fields` choice.
+        let chunk_transform = enc_ctx.map(|mut ctx| {
+            ctx.encrypt_all_fields = true;
+            encryption::make_chunk_transform(Endpoint::ChatCompletions, ctx, state.signing.clone())
+        });
+        return agent_loop::run_chat_completion(
+            state,
+            auth,
+            tracing_ids,
+            request_hash,
+            request_json,
+            chunk_transform,
+        )
+        .await;
+    }
+
     // For cloud API key requests with streaming, force include_usage
     // so the backend always sends token counts for billing.
     // (Non-streaming requests also stream internally via proxy_json_request,
@@ -102,7 +144,7 @@ pub async fn chat_completions(
         model_name: state.config.model_name.clone(),
         usage_reporter: make_usage_reporter(&auth, &state),
         usage_type: UsageType::ChatCompletion,
-        request_hash: original_request_hash,
+        request_hash: Some(request_hash),
         response_transform,
         chunk_transform,
         backend_guard: Some(guard),