Synced from monorepo
Changes: - Classify clipboard delivery confidence - Add durable session update append - Scope the xAI session bearer to first-party memory embedding endpoints - Persist subagent outputs to disk and bound long-lived agent state - Add MiniSweAgent:bash for mini-swe-agent parity - Revert taking local sessions off the persistent shell - Contextual tip recommending grok wrap on SSH sessions - Voice STT bearer from model BYOK env_key/api_key - Define exact website policies for sandbox - Gate unsafe shell environments - Shared pin hoist; single require_sha gate for marketplace plugins - Server-signed is-managed claim (closes sidecar-removal downgrade) - Optional require_sha pin for remote plugin installs - Show session title and last exchange in the exit resume hint - Gate shell output redirects - Warn when fail_closed is present but not a boolean - Add canonical text editing core (ratatui-textarea) - Keep execution state out of goal scratch - Add acknowledged persistence primitives - Inherit child network restrictions in sandbox - Fail closed when hook matchers fail to recompile - Add MCP setup preferences for plugin MCPs - Gate sourced shell scripts - Gate file-typed project hooks - grok wrap: restore terminal modes on child death - Harden owner-only permissions on auth and MCP credentials - Create crash dump files with owner-only permissions - Write the agent_id cache owner-only (0600) - SessionMetrics mode skips Mixpanel profile sync - Dashboard: slim live-tail peek - Yank full queued prompt text, not (+N lines) - Defeat clock-rollback on the signed managed-config cache - Stop early session/cancel from overtaking the prompt and wedging the turn slot - Self-heal a diverged agent entrypoint on startup - Add matched inference expectations in test-support - Add AuthSingleFlight cancel/successor gap tests - Remove consumer from external OTEL allowlist and pin scrub coverage - Enable /copy in minimal mode - Surface capacity and API-key detail on 429 errors - Single-flight interactive auth - Fix PageUp/PageDown skipping lines behind sticky prompt header
This commit is contained in:
parent
8adf9013a0
commit
98c3b2438a
225 changed files with 18836 additions and 7156 deletions
|
|
@ -18,6 +18,77 @@ use super::embedding::EmbeddingProvider as _;
|
|||
use super::storage::MemoryStorage;
|
||||
use super::watcher::MemoryFileWatcher;
|
||||
|
||||
/// Embedding-client credentials scoped to a trusted endpoint. Only
|
||||
/// [`Self::for_endpoint`] retains a live credential; the empty default fails closed.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct EndpointScopedCredentials {
|
||||
endpoint: Option<reqwest::Url>,
|
||||
auth_credentials: Option<Arc<dyn xai_grok_auth::AuthCredentialProvider>>,
|
||||
api_key_provider: Option<xai_grok_tools::types::SharedApiKeyProvider>,
|
||||
}
|
||||
|
||||
// Manual Debug that redacts the credential handles; only their presence shows.
|
||||
impl std::fmt::Debug for EndpointScopedCredentials {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("EndpointScopedCredentials")
|
||||
.field("endpoint", &self.endpoint)
|
||||
.field("has_auth_credentials", &self.auth_credentials.is_some())
|
||||
.field("has_api_key_provider", &self.api_key_provider.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl EndpointScopedCredentials {
|
||||
pub fn none() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.auth_credentials.is_none() && self.api_key_provider.is_none()
|
||||
}
|
||||
|
||||
/// Retains the credentials only for a trusted, parsable `endpoint`; otherwise drops them.
|
||||
pub fn for_endpoint(
|
||||
endpoint: &str,
|
||||
is_trusted: impl FnOnce(&str) -> bool,
|
||||
auth_credentials: Option<Arc<dyn xai_grok_auth::AuthCredentialProvider>>,
|
||||
api_key_provider: Option<xai_grok_tools::types::SharedApiKeyProvider>,
|
||||
) -> Self {
|
||||
if is_trusted(endpoint)
|
||||
&& let Ok(url) = reqwest::Url::parse(endpoint)
|
||||
{
|
||||
return Self {
|
||||
endpoint: Some(url),
|
||||
auth_credentials,
|
||||
api_key_provider,
|
||||
};
|
||||
}
|
||||
if auth_credentials.is_some() || api_key_provider.is_some() {
|
||||
tracing::info!(
|
||||
target: xai_grok_telemetry::memory_log::TARGET,
|
||||
endpoint,
|
||||
"memory embeddings: session credentials withheld for non-first-party endpoint; its own key, if any, still applies"
|
||||
);
|
||||
}
|
||||
Self::none()
|
||||
}
|
||||
|
||||
fn auth_credentials(&self) -> Option<&Arc<dyn xai_grok_auth::AuthCredentialProvider>> {
|
||||
self.auth_credentials.as_ref()
|
||||
}
|
||||
|
||||
fn api_key_provider(&self) -> Option<&xai_grok_tools::types::SharedApiKeyProvider> {
|
||||
self.api_key_provider.as_ref()
|
||||
}
|
||||
|
||||
fn approved_for(&self, base_url: &str) -> bool {
|
||||
match &self.endpoint {
|
||||
None => self.is_empty(),
|
||||
Some(endpoint) => reqwest::Url::parse(base_url).is_ok_and(|url| &url == endpoint),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// All configuration needed to build a fully-wired [`MemoryBackendImpl`] for a live session.
|
||||
///
|
||||
/// Grouping these in one struct ensures every call site — ToolBridge, first-turn
|
||||
|
|
@ -30,7 +101,8 @@ pub struct MemoryBackendParams {
|
|||
pub session_id: String,
|
||||
/// Embedding provider config — `None` forces FTS-only fallback everywhere.
|
||||
pub embed_config: Option<xai_grok_config_types::MemoryEmbeddingConfig>,
|
||||
/// Base URL for embedding API calls (CLI proxy).
|
||||
/// Base URL for embedding API calls (CLI proxy). Must match the endpoint
|
||||
/// `embedding_credentials` was scoped to; mismatch fails closed.
|
||||
pub embed_base_url: String,
|
||||
/// API key for embedding API calls.
|
||||
pub embed_api_key: Option<String>,
|
||||
|
|
@ -47,10 +119,7 @@ pub struct MemoryBackendParams {
|
|||
/// - `"injection"` — first-turn memory context injection
|
||||
/// - `"compaction_recovery"` — post-compaction context re-injection
|
||||
pub search_source: &'static str,
|
||||
/// Dynamic API key provider — when set, `make_embedding_provider()` resolves
|
||||
/// the key per-call instead of using the static `embed_api_key`.
|
||||
pub api_key_provider: Option<xai_grok_tools::types::SharedApiKeyProvider>,
|
||||
pub auth_credentials: Option<Arc<dyn xai_grok_auth::AuthCredentialProvider>>,
|
||||
pub embedding_credentials: EndpointScopedCredentials,
|
||||
}
|
||||
|
||||
impl MemoryBackendParams {
|
||||
|
|
@ -59,8 +128,7 @@ impl MemoryBackendParams {
|
|||
pub async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
|
||||
build_embedding_provider(
|
||||
self.embed_config.as_ref(),
|
||||
self.auth_credentials.as_ref(),
|
||||
self.api_key_provider.as_ref(),
|
||||
&self.embedding_credentials,
|
||||
self.embed_api_key.as_deref(),
|
||||
&self.embed_base_url,
|
||||
)
|
||||
|
|
@ -70,8 +138,7 @@ impl MemoryBackendParams {
|
|||
|
||||
async fn build_embedding_provider(
|
||||
config: Option<&xai_grok_config_types::MemoryEmbeddingConfig>,
|
||||
auth_credentials: Option<&Arc<dyn xai_grok_auth::AuthCredentialProvider>>,
|
||||
api_key_provider: Option<&xai_grok_tools::types::SharedApiKeyProvider>,
|
||||
credentials: &EndpointScopedCredentials,
|
||||
static_api_key: Option<&str>,
|
||||
base_url: &str,
|
||||
) -> Option<super::embedding::ApiEmbeddingProvider> {
|
||||
|
|
@ -80,9 +147,19 @@ async fn build_embedding_provider(
|
|||
return None;
|
||||
}
|
||||
|
||||
// Prefer the refresh-capable credential provider — the middleware gives
|
||||
// 401 retry for free without any per-call key resolution.
|
||||
if let Some(creds) = auth_credentials {
|
||||
// Enforce at runtime, in release too: a `debug_assert` would compile out of
|
||||
// shipped binaries and let a scoped credential reach an unapproved URL.
|
||||
let credentials_approved = credentials.approved_for(base_url);
|
||||
if !credentials_approved {
|
||||
tracing::error!(
|
||||
target: xai_grok_telemetry::memory_log::TARGET,
|
||||
base_url,
|
||||
approved = ?credentials.endpoint,
|
||||
"memory embeddings: scoped credentials do not match the request URL; dropping them"
|
||||
);
|
||||
}
|
||||
|
||||
if credentials_approved && let Some(creds) = credentials.auth_credentials() {
|
||||
let client = super::embedding::build_middleware_client(creds.clone());
|
||||
return super::embedding::ApiEmbeddingProvider::from_config(
|
||||
config,
|
||||
|
|
@ -91,13 +168,12 @@ async fn build_embedding_provider(
|
|||
);
|
||||
}
|
||||
|
||||
// Fallback: resolve API key per-call, wrap in a static middleware client
|
||||
// (no 401 refresh, but auth header is still stamped by middleware).
|
||||
let api_key = match api_key_provider {
|
||||
Some(p) => p.current_api_key_async().await,
|
||||
None => None,
|
||||
}
|
||||
.or_else(|| static_api_key.map(|s| s.to_owned()))?;
|
||||
let per_call_key = if credentials_approved && let Some(p) = credentials.api_key_provider() {
|
||||
p.current_api_key_async().await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let api_key = per_call_key.or_else(|| static_api_key.map(|s| s.to_owned()))?;
|
||||
super::embedding::ApiEmbeddingProvider::from_session(config, base_url.to_owned(), api_key)
|
||||
}
|
||||
|
||||
|
|
@ -129,10 +205,7 @@ pub struct MemoryBackendImpl {
|
|||
/// Only the ToolBridge backend's counter is shared back to the session actor;
|
||||
/// injection and compaction-recovery backends use their own local counters.
|
||||
pub search_counter: std::sync::Arc<std::sync::atomic::AtomicU64>,
|
||||
/// Dynamic API key provider for embedding requests.
|
||||
api_key_provider: Option<xai_grok_tools::types::SharedApiKeyProvider>,
|
||||
/// Refresh-capable credential provider for embedding HTTP middleware.
|
||||
auth_credentials: Option<Arc<dyn xai_grok_auth::AuthCredentialProvider>>,
|
||||
embedding_credentials: EndpointScopedCredentials,
|
||||
}
|
||||
|
||||
impl MemoryBackendImpl {
|
||||
|
|
@ -150,8 +223,7 @@ impl MemoryBackendImpl {
|
|||
stale_claim_secs: 60,
|
||||
session_id: String::new(),
|
||||
search_source: "tool",
|
||||
api_key_provider: None,
|
||||
auth_credentials: None,
|
||||
embedding_credentials: EndpointScopedCredentials::none(),
|
||||
search_counter: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)),
|
||||
}
|
||||
}
|
||||
|
|
@ -200,8 +272,7 @@ impl MemoryBackendImpl {
|
|||
async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
|
||||
build_embedding_provider(
|
||||
self.embed_config.as_ref(),
|
||||
self.auth_credentials.as_ref(),
|
||||
self.api_key_provider.as_ref(),
|
||||
&self.embedding_credentials,
|
||||
self.embed_api_key.as_deref(),
|
||||
&self.embed_base_url,
|
||||
)
|
||||
|
|
@ -232,8 +303,7 @@ impl MemoryBackendImpl {
|
|||
if let Some(w) = ¶ms.watcher {
|
||||
backend = backend.with_watcher(w.clone(), params.stale_claim_secs);
|
||||
}
|
||||
backend.api_key_provider = params.api_key_provider.clone();
|
||||
backend.auth_credentials = params.auth_credentials.clone();
|
||||
backend.embedding_credentials = params.embedding_credentials.clone();
|
||||
backend
|
||||
}
|
||||
}
|
||||
|
|
@ -543,8 +613,7 @@ mod factory_tests {
|
|||
watcher: None,
|
||||
stale_claim_secs: 60,
|
||||
search_source: "tool",
|
||||
api_key_provider: None,
|
||||
auth_credentials: None,
|
||||
embedding_credentials: EndpointScopedCredentials::none(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1147,9 +1216,13 @@ mod factory_tests {
|
|||
watcher: None,
|
||||
stale_claim_secs: 60,
|
||||
search_source: "tool",
|
||||
api_key_provider: Some(probe),
|
||||
// No auth_credentials — forces the api_key_provider fallback path.
|
||||
auth_credentials: None,
|
||||
// Trusted endpoint + no auth_credentials exercises the api_key_provider path.
|
||||
embedding_credentials: EndpointScopedCredentials::for_endpoint(
|
||||
"http://example/v1",
|
||||
|_| true,
|
||||
None,
|
||||
Some(probe),
|
||||
),
|
||||
};
|
||||
|
||||
let provider = params.make_embedding_provider().await;
|
||||
|
|
@ -1177,6 +1250,15 @@ mod tests {
|
|||
use tempfile::TempDir;
|
||||
use xai_grok_config_types::MemoryIndexConfig;
|
||||
|
||||
/// An api-key provider that fails the test if its key is ever resolved,
|
||||
/// proving a scoped-away credential is never consulted.
|
||||
struct PanicKey;
|
||||
impl xai_grok_tools::types::ApiKeyProvider for PanicKey {
|
||||
fn current_api_key(&self) -> Option<String> {
|
||||
panic!("scoped-away credential must not be resolved");
|
||||
}
|
||||
}
|
||||
|
||||
fn setup_index(tmp: &TempDir) -> (PathBuf, MemoryStorage) {
|
||||
init_sqlite_vec();
|
||||
let global = tmp.path().join("memory");
|
||||
|
|
@ -1222,6 +1304,125 @@ mod tests {
|
|||
assert_send_sync::<MemoryBackendImpl>();
|
||||
}
|
||||
|
||||
/// If credentials approved for one endpoint are used to build against a
|
||||
/// different URL (a wiring bug), they are dropped at build time rather than
|
||||
/// sent to the wrong endpoint. The session provider would panic if resolved.
|
||||
#[tokio::test]
|
||||
async fn test_build_drops_credentials_when_request_url_differs() {
|
||||
let session: xai_grok_tools::types::SharedApiKeyProvider = Arc::new(PanicKey);
|
||||
|
||||
let scoped = EndpointScopedCredentials::for_endpoint(
|
||||
"https://api.x.ai/v1",
|
||||
|_| true,
|
||||
None,
|
||||
Some(session),
|
||||
);
|
||||
assert!(!scoped.is_empty(), "trusted endpoint keeps the credential");
|
||||
|
||||
let config = xai_grok_config_types::MemoryEmbeddingConfig {
|
||||
model: Some("test-embedding-model".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
let provider = build_embedding_provider(
|
||||
Some(&config),
|
||||
&scoped,
|
||||
Some("byok-static-key"),
|
||||
"https://other.example/v1",
|
||||
)
|
||||
.await;
|
||||
assert!(
|
||||
provider.is_some(),
|
||||
"mismatched request URL must fall back to the static key, not the scoped credential"
|
||||
);
|
||||
}
|
||||
|
||||
/// A trusted, URL-matching endpoint builds the provider from the
|
||||
/// refresh-capable session credential and never consults the per-call
|
||||
/// api-key provider. The api-key provider panics if resolved.
|
||||
#[tokio::test]
|
||||
async fn test_trusted_endpoint_prefers_session_credential() {
|
||||
struct StubAuth;
|
||||
impl xai_grok_auth::HttpAuth for StubAuth {
|
||||
fn apply(
|
||||
&self,
|
||||
builder: reqwest::RequestBuilder,
|
||||
_base_url: &str,
|
||||
) -> reqwest::RequestBuilder {
|
||||
builder
|
||||
}
|
||||
}
|
||||
#[async_trait::async_trait]
|
||||
impl xai_grok_auth::AuthCredentialProvider for StubAuth {
|
||||
fn snapshot(&self) -> xai_grok_auth::CredentialSnapshot {
|
||||
xai_grok_auth::CredentialSnapshot::default()
|
||||
}
|
||||
async fn refresh_after_unauthorized(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
let auth: Arc<dyn xai_grok_auth::AuthCredentialProvider> = Arc::new(StubAuth);
|
||||
let api_key: xai_grok_tools::types::SharedApiKeyProvider = Arc::new(PanicKey);
|
||||
let scoped = EndpointScopedCredentials::for_endpoint(
|
||||
"https://api.x.ai/v1",
|
||||
|_| true,
|
||||
Some(auth),
|
||||
Some(api_key),
|
||||
);
|
||||
assert!(!scoped.is_empty(), "trusted endpoint keeps the credential");
|
||||
|
||||
let config = xai_grok_config_types::MemoryEmbeddingConfig {
|
||||
model: Some("test-embedding-model".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
let provider =
|
||||
build_embedding_provider(Some(&config), &scoped, None, "https://api.x.ai/v1").await;
|
||||
assert!(
|
||||
provider.is_some(),
|
||||
"trusted endpoint must build a provider from the session credential"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn endpoint_scoped_credentials_trust_gate_and_url_match() {
|
||||
struct AnyKey;
|
||||
impl xai_grok_tools::types::ApiKeyProvider for AnyKey {
|
||||
fn current_api_key(&self) -> Option<String> {
|
||||
None
|
||||
}
|
||||
}
|
||||
let key = || Arc::new(AnyKey) as xai_grok_tools::types::SharedApiKeyProvider;
|
||||
|
||||
let denied = EndpointScopedCredentials::for_endpoint(
|
||||
"https://byok.example/v1",
|
||||
|_| false,
|
||||
None,
|
||||
Some(key()),
|
||||
);
|
||||
assert!(denied.is_empty(), "untrusted endpoint drops the credential");
|
||||
|
||||
let scoped = EndpointScopedCredentials::for_endpoint(
|
||||
"https://api.x.ai/v1",
|
||||
|_| true,
|
||||
None,
|
||||
Some(key()),
|
||||
);
|
||||
assert!(!scoped.is_empty(), "trusted endpoint keeps the credential");
|
||||
assert!(
|
||||
scoped.approved_for("https://API.x.ai/v1"),
|
||||
"host casing normalizes"
|
||||
);
|
||||
assert!(
|
||||
!scoped.approved_for("https://api.x.ai/v2"),
|
||||
"different path rejected"
|
||||
);
|
||||
assert!(
|
||||
!scoped.approved_for("https://other.example/v1"),
|
||||
"different host rejected"
|
||||
);
|
||||
assert!(!scoped.approved_for("not-a-url"), "unparsable fails closed");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_search_with_punctuation_in_query() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ pub mod storage;
|
|||
pub mod text_utils;
|
||||
pub mod watcher;
|
||||
|
||||
pub use backend::{MemoryBackendImpl, MemoryBackendParams};
|
||||
pub use backend::{EndpointScopedCredentials, MemoryBackendImpl, MemoryBackendParams};
|
||||
pub use index::{MemoryIndex, init_sqlite_vec};
|
||||
pub use storage::{MemoryScope, MemoryStorage};
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue