jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-10-27 12:40:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(buzz-agent): add DatabricksAuthCoordinator single-flight OAuth (#5545)

Browse Source 
## What

Consolidates all Databricks OAuth acquisition behind one coordinator on
`PkceOAuthTokenSource`. Every entry point — the four `TokenSource`
methods (`bearer`, `bearer_no_browser`, `refresh_now`,
`interactive_login`) plus the public `acquire_with_intent` — routes
through a single `acquire()`/`acquire_locked()` core that owns browser
and cooldown policy.

Before this, acquisition logic was scattered across those methods with
no coordination: concurrent callers (Desktop discovery, the saved-agent
model picker, managed-runtime inference) could each pop their own
browser, and a just-denied attempt would immediately re-prompt on the
next passive read.

## How

- **Intent policy.** `AuthIntent::{Auto, UserInitiated, Headless}`
decides whether a caller may open a browser and whether it honors the
cooldown. `Headless` never browses; `Auto` browses but honors an
unexpired cooldown; `UserInitiated` browses and bypasses+clears the
cooldown.
- **Two-layer single-flight per cache key.** An in-process registry
(`INFLIGHT`) coalesces same-key, same-intent callers onto one leader's
attempt before the file lock. The slot key is `(lock_path, AuthIntent)`,
so a `UserInitiated` sign-in never inherits an `Auto` leader's result. A
joined result is revalidated against the waiter's own contract. Across
processes, callers serialize on a `flock`-based advisory lock and share
success through the on-disk cache. RAII `Drop` releases both lock and
leader slot.
- **Joiner credential-state reconciliation.** `SlotPublish` carries the
full `CachedToken` on success. Each joiner reconciles its own
independent `state` cell under `state.lock().await` before returning:
adopt when absent, expired, or matching the joiner's rejected
credential; preserve any distinct newer usable credential. On a matching
shared failure, neutralize the joiner's in-memory rejected entry under
lock via `expire_rejected_memory` — durable disk mutation is reserved
for `acquire_locked` under the cross-process file lock. Without this, a
joining source's state remains stale or empty and subsequent plain
`bearer()` calls resurface the rejected or absent credential.
- **Validate-before-persist boundary.** `finish()` is the
candidate-token persistence boundary for refresh and browser results.
Before a token is written to cache or the cooldown is cleared, a bearer
equal to the caller's rejected bytes yields a typed failure.
- **Token neutralization.** When `acquire_locked` enters with `rejected
= Some(bytes)`, it calls `expire_rejected()` under the state lock before
any cache check.
- **Cross-process failure single-flight.** An `AttemptRecord` sidecar
records a monotonically-increasing generation, intent, result code, and
SHA-256 digest of the completing caller's rejected token. Adoption is
temporal (pre-queue snapshot predates current generation) and
digest-matched.
- **Typed outcomes.** `AuthError` with stable `code()`/`from_code()`
replaces display-text matching.
- **Durable cooldown sidecar.** Every failed browser attempt is recorded
next to the cache key. 5-minute expiry.
- **Windows disk persistence disabled.** On non-Unix platforms,
`persist()` is a no-op. Lock, cooldown, and attempt sidecars are active
on all platforms. Tests that seed or assert on the on-disk token cache
are `#[cfg(unix)]`-gated.
- **Injected browser opener** invoked while the localhost callback
listener is live.

## Tests

- `crates/buzz-agent/tests/databricks_auth_coordinator.rs`:
browser/cooldown/classification acceptance matrix with a scripted
`BrowserOpener` and stub OIDC provider. P1 regressions exercise the full
`finish()` → `acquire_locked()` → `acquire_leader()` →
`LeaderGuard::complete()` → joiner wiring:
- `test_inprocess_joiner_reconciles_stale_state_after_shared_success`
(Unix): two real sources both loaded locally-fresh-but-rejected X; after
shared success Y, subsequent plain `bearer()` on both returns Y, not X.
-
`test_inprocess_joiner_neutralizes_rejected_on_matching_shared_failure`
(Unix): B's matching rejected X is force-expired in memory after shared
`RefreshRejected`; subsequent read cannot return X.
- `test_inprocess_joiner_populates_empty_state_no_second_acquisition`
(non-Unix): empty A/B join a browser success; B's subsequent headless
read returns Y without a second browser (no disk fallback on non-Unix
exposes the regression).
- `test_crossprocess_userinitiated_waiter_adopts_predecessor_denial`
(snapshot-marker barrier replacing an earlier sleep for deterministic
generation ordering).
- `auth.rs` in-crate tests: lock-primitive edges, disk-recheck on shared
failure (`#[cfg(unix)]`), and:
- `test_joiner_reconciliation_blocked_until_state_lock_released`:
deterministic direct-poll proof that awaited reconciliation requires
`state.lock().await`. The test task holds B's state mutex and manually
polls a pinned real `acquire()` future — Poll 2 (slot published, mutex
still held) must return `Pending` because `lock().await` blocks; with
`try_lock` instead, Poll 2 returns `Ready`, failing the assertion.
- `test_joiner_preserve_distinct_newer_credential`: deterministic direct
polling parks B at `slot.wait()`, then writes Z directly into B's state
in the same task, then publishes Y and awaits completion. B must return
Y but leave state == Z. Mutation check: unconditional adoption
overwrites Z with Y, failing the state assertion.
- `test_joiner_shared_failure_recovers_disk_replacement` (Unix): the
matching-failure joiner enters the recovery branch — after
`expire_rejected_memory` (in-memory, empty state no-op) it reads a
sibling-written disk replacement via `usable_from_disk` and returns it.
Mutation check: removing the `usable_from_disk` recovery branch returns
`Err(RefreshRejected)`.
- `test_joiner_failure_does_not_write_disk` (Unix): byte-for-byte
disk-invariance regression — a matching-failure joiner calls
`expire_rejected_memory` and must not touch the on-disk cache. An
independent process C may write a valid replacement between A's failure
and B's reconciliation; this guard ensures B's unfenced in-memory
neutralization cannot overwrite C's concurrent disk write. Mutation
check: reverting to `expire_rejected` rewrites the file (`expires_at =
0`), changing the bytes and failing the assertion.
- `test_lock_timeout_leaves_cooldown_sidecar_byte_for_byte_untouched`: a
waiter past its deadline returns `LockTimeout` before entering
`acquire_locked`; the cooldown sidecar bytes are unchanged. Documents a
pre-lock limitation: `LockTimeout` callers do not neutralize state or
sidecars.

## Scope / follow-ups

- **Runtime 401 handling is deferred.** This PR owns acquisition
single-flight and policy.
- **Desktop wiring is Phase 2** (not in this PR's boundary). Confined to
`crates/buzz-agent/`.
- **Windows DACL** is a follow-up once the `windows-sys` binding is
available.

## Stack

Built on [#5534](https://github.com/block/buzz/pull/5534)
(`hayt/databricks-oauth-cache-hardening`), now merged. Retargeted to
`main`.

---------

Signed-off-by: Will Pfleger <pfleger.will@gmail.com>
Signed-off-by: Duncan <dcfd242e557282d7a1e2cf2e6877522682f1e5c6156dc92ca7d90eaedd3b0f95@buzz.block.builderlab.xyz>
Co-authored-by: Duncan <dcfd242e557282d7a1e2cf2e6877522682f1e5c6156dc92ca7d90eaedd3b0f95@buzz.block.builderlab.xyz>
This commit is contained in:
Will Pfleger
2026-09-01 15:59:21 -04:00
committed by GitHub
parent e5a7e26a10
commit 5aed49b505
10 changed files with 5979 additions and 263 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
.github/workflows/ci.yml vendored
View File

@@ -1170,6 +1170,20 @@ jobs:
# Serial: windows_resolver_tests mutate process-global env
# (BUZZ_SHELL/GIT_BASH/SystemRoot) that SharedState::new reads.
run: cargo test -p buzz-dev-mcp --target $env:TARGET -- --test-threads=1
- name: Test (buzz-agent auth coordinator)
# The auth coordinator single-flights on an OS advisory lock, which is
# LockFileEx on Windows; this integration suite drives real second
# processes on the same lock file, so it only exercises the Windows
# lock runtime if it runs ON Windows. Every other job compiles it but
# never executes it. Tests exercised on Windows: lock serialization
# (two coordinators race for the same key), cooldown sidecar sharing
# across processes, attempt-sidecar adoption (UserInitiated waiter
# adopts a predecessor's denial), and the in-process single-flight for
# same-key coalescing. Tests that are UNIX-ONLY and NOT executed here:
# crash-release (flock drop on SIGKILL, guarded by #[cfg(unix)]) and
# cross-process cache success/race (on-disk token handoff, also
# #[cfg(unix)]).
run: cargo test -p buzz-agent --target $env:TARGET --test databricks_auth_coordinator
# Smoke-test the new host-prereq contract: Git for Windows (which provides
# bash) is available on the runner, a shell command round-trips, and bash
# does NOT resolve from System32 (so WSL's launcher is never picked up).

11
Cargo.lock generated
View File

@@ -894,6 +894,7 @@ dependencies = [
"axum",
"base64 0.22.1",
"dirs",
"fs2",
"getrandom 0.4.3",
"hex",
"nix 0.31.3",
@@ -2971,6 +2972,16 @@ dependencies = [
"percent-encoding",
]
[[package]]
name = "fs2"
version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213"
dependencies = [
"libc",
"winapi",
]
[[package]]
name = "fs_extra"
version = "1.3.0"

23
Justfile
View File

@@ -391,14 +391,21 @@ test-unit:
# because nothing in CI runs `cargo test --workspace` — workspace
# membership alone buys clippy/check, not a single executed test.
cargo nextest run -p buzz-backend-kubernetes
# buzz-agent model-capabilities corpus: the Rust half of the
# cross-language drift guard. `model_capabilities.rs` embeds
# scripts/model-capabilities.json + scripts/normative-corpus.json via
# include_str! and replays the full locked corpus as pure in-process tests (no
# infra). Enumerated explicitly because nothing in CI runs
# `cargo test --workspace`; without this step a manifest edit that
# diverges Rust from the corpus ships green.
cargo nextest run -p buzz-agent --lib
# buzz-agent: two infra-free concerns run together by executing the
# whole crate (lib + integration tests), because nothing in CI runs
# `cargo test --workspace`, so without this stanza neither the crate's
# library tests nor its integration tests execute remotely.
# * model-capabilities corpus (lib): the Rust half of the
# cross-language drift guard. `model_capabilities.rs` embeds
# scripts/model-capabilities.json + scripts/normative-corpus.json via
# include_str! and replays the full locked corpus as pure in-process
# tests; without it a manifest edit that diverges Rust from the
# corpus ships green.
# * OAuth auth coordinator (lib concurrency matrix + databricks
# integration tests): lock single-flight, cooldown, cross-process
# crash recovery — infra-free via a stub OIDC provider and an
# injected browser opener, no network or Postgres.
cargo nextest run -p buzz-agent
# Admin API auth-boundary tests (api::admin in buzz-relay): the NIP-98
# duplicate-tag rejections, the Host/Origin replay-ordering causal pair,
# the admin.localhost origin/advertisement/canonical-URL pins, and the

23
crates/buzz-agent/Cargo.toml
View File

@@ -24,6 +24,24 @@ path = "src/main.rs"
name = "fake-mcp"
path = "tests/bin/fake_mcp.rs"
# Test-only lock holder: a real second process that takes the coordinator's
# cross-process advisory lock, so the auth tests can prove genuine
# inter-process single-flight and crash-release rather than same-process
# handles. Tiny; only used by the databricks auth integration tests.
[[bin]]
name = "lock-holder"
path = "tests/bin/lock_holder.rs"
# Test-only auth worker: a real second process that runs the PUBLIC auth
# coordinator API (`acquire_with_intent`) with a scripted browser opener and a
# shared temp cache, so the auth tests can prove the cross-process single-flight
# contract end-to-end — durable cooldown sharing and one-grant/one-cache races
# across a genuine process boundary, not two in-process handles. Only used by
# the databricks auth integration tests.
[[bin]]
name = "auth-worker"
path = "tests/bin/auth_worker.rs"
[dependencies]
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "io-std", "io-util", "sync", "process", "time", "net"] }
serde = { workspace = true }
@@ -45,6 +63,11 @@ url = { workspace = true }
urlencoding = "2"
webbrowser = "1"
dirs = "6"
# Cross-process advisory file lock (flock on Unix, LockFileEx on Windows) for
# the auth coordinator's single-flight. Kept off std's `File::try_lock` so the
# crate stays buildable on the repo's declared 1.88 MSRV (those std APIs are
# 1.89+).
fs2 = "0.4"
[target.'cfg(unix)'.dependencies]
nix = { version = "0.31", default-features = false, features = ["signal", "process"] }

2357
crates/buzz-agent/src/auth.rs
View File

File diff suppressed because it is too large Load Diff

252
crates/buzz-agent/tests/bin/auth_worker.rs Normal file
View File

@@ -0,0 +1,252 @@
//! Test-only helper: a real second process that runs the PUBLIC auth
//! coordinator (`PkceOAuthTokenSource::acquire_with_intent`) against a shared
//! temp cache, so the auth tests can prove the *cross-process* single-flight
//! contract end-to-end rather than with two in-process handles.
//!
//! The in-process `INFLIGHT` registry coalesces same-key callers within one
//! process before they ever reach the file lock, so two `PkceOAuthTokenSource`
//! instances in one test do NOT exercise the cross-process protocol (the OS
//! advisory lock and the on-disk cache re-read). This binary is a genuine
//! second process: it contends on the same `flock`/`LockFileEx` and reads/writes
//! the same private cache file the parent coordinator does.
//!
//! The browser step is scripted (no real window): the opener drives the
//! loopback callback exactly as a real browser would, and its launch count is
//! reported back so a test can assert "exactly one browser across processes".
//!
//! Env contract (all required unless noted):
//! AUTH_WORKER_DISCOVERY_URL — OIDC discovery URL (the parent stub).
//! AUTH_WORKER_CACHE_DIR — shared cache dir (`cache_dir_override`).
//! AUTH_WORKER_NAMESPACE — cache namespace.
//! AUTH_WORKER_CLIENT_ID — OAuth client id.
//! AUTH_WORKER_SCOPES — comma-separated scopes.
//! AUTH_WORKER_INTENT — auto | userinitiated | headless.
//! AUTH_WORKER_SCRIPT — approve | deny | failopen.
//! AUTH_WORKER_RESULT — path to write the JSON outcome to.
//! AUTH_WORKER_REJECTED — (optional) rejected token bytes passed to
//! `acquire_with_intent`; absent means no rejection.
//! AUTH_WORKER_READY_MARKER — (optional) written once the source is built,
//! before acquisition, so the parent can release
//! several workers into a genuine lock race.
//! AUTH_WORKER_START_MARKER — (optional) acquisition blocks until this file
//! exists, so multiple workers begin together.
//! AUTH_WORKER_LAUNCHED_MARKER — (optional) written when the browser opener
//! fires (i.e. this process holds the lock and is
//! mid-flow), so the parent can queue behind it.
//! AUTH_WORKER_PROCEED_MARKER — (optional) the scripted callback is withheld
//! until this file exists, so the parent can
//! confirm another process is already waiting on
//! the lock before this one resolves.
//! AUTH_WORKER_SNAPSHOT_MARKER — (optional) a file path; when set, a tracing
//! layer intercepts the `acquire_leader_snapshot`
//! event emitted by `auth.rs` after the attempt-
//! generation snapshot is taken (and before the
//! cross-process lock is acquired) and writes this
//! file once. Lets the parent observe that this
//! process has committed its snapshot-gen and is
//! about to queue on the lock.
//!
//! Result JSON: `{ "result": "ok"|"<error_code>", "bearer": <string|null>,
//! "launches": <u64> }`.
use std::fs;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use buzz_agent::auth::{AuthIntent, BrowserOpener, PkceOAuthConfig, PkceOAuthTokenSource};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
/// Tracing layer that writes a file once when it sees the
/// `buzz_agent::auth::acquire_leader_snapshot` event emitted by
/// `acquire_leader` immediately after the attempt-generation snapshot is fixed
/// and before the cross-process lock is acquired. Installed only when
/// `AUTH_WORKER_SNAPSHOT_MARKER` is set, so normal test runs incur no overhead.
struct SnapshotMarkerLayer {
path: PathBuf,
written: AtomicBool,
}
impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for SnapshotMarkerLayer {
fn on_event(
&self,
event: &tracing::Event<'_>,
_ctx: tracing_subscriber::layer::Context<'_, S>,
) {
if event.metadata().target() == "buzz_agent::auth::acquire_leader_snapshot"
&& !self.written.swap(true, Ordering::SeqCst)
{
let _ = fs::write(&self.path, b"snapshotted");
}
}
}
/// What the scripted "user" does when the coordinator opens a browser.
#[derive(Clone, Copy)]
enum Script {
Approve,
Deny,
FailToOpen,
}
/// A [`BrowserOpener`] that counts launches and drives the loopback callback on
/// a background thread — the same technique as the in-crate test opener, but
/// with two optional cross-process barriers so the parent can order events:
/// `launched_marker` announces that this process holds the lock and has opened
/// the browser, and `proceed_marker` withholds the callback until the parent
/// signals it has queued another process behind the lock.
struct WorkerOpener {
script: Script,
calls: Arc<AtomicU64>,
launched_marker: Option<PathBuf>,
proceed_marker: Option<PathBuf>,
}
impl BrowserOpener for WorkerOpener {
fn open(&self, url: &str) -> Result<(), String> {
self.calls.fetch_add(1, Ordering::SeqCst);
if let Some(marker) = &self.launched_marker {
fs::write(marker, b"launched").expect("write launched marker");
}
let query = match self.script {
Script::FailToOpen => return Err("no browser available".into()),
Script::Approve => "code=scripted-code",
Script::Deny => "error=access_denied",
};
let parsed = url::Url::parse(url).expect("authorize URL must parse");
let redirect = parsed
.query_pairs()
.find(|(k, _)| k == "redirect_uri")
.map(|(_, v)| v.into_owned())
.expect("authorize URL carries redirect_uri");
let state = parsed
.query_pairs()
.find(|(k, _)| k == "state")
.map(|(_, v)| v.into_owned())
.expect("authorize URL carries state");
let redirect = url::Url::parse(&redirect).expect("redirect_uri must parse");
let port = redirect.port().expect("loopback redirect carries a port");
let request = format!(
"GET /?{query}&state={state} HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\nConnection: close\r\n\r\n"
);
let proceed = self.proceed_marker.clone();
std::thread::spawn(move || {
// Hold the callback until the parent has confirmed another process
// is already queued behind the lock (bounded so a missing signal
// can't wedge the test past the browser timeout).
if let Some(marker) = proceed {
for _ in 0..6000 {
if marker.exists() {
break;
}
std::thread::sleep(Duration::from_millis(10));
}
}
if let Ok(mut sock) = TcpStream::connect(("127.0.0.1", port)) {
let _ = sock.write_all(request.as_bytes());
let _ = sock.flush();
let mut discard = Vec::new();
let _ = sock.read_to_end(&mut discard);
}
});
Ok(())
}
}
fn env(key: &str) -> String {
std::env::var(key).unwrap_or_else(|_| panic!("{key} set"))
}
#[tokio::main]
async fn main() {
// If the parent test set AUTH_WORKER_SNAPSHOT_MARKER, install a tracing
// subscriber layer that fires when the coordinator emits its pre-lock
// snapshot event and writes the marker file.
if let Ok(marker_path) = std::env::var("AUTH_WORKER_SNAPSHOT_MARKER") {
tracing_subscriber::registry()
.with(SnapshotMarkerLayer {
path: PathBuf::from(marker_path),
written: AtomicBool::new(false),
})
.init();
}
let intent = match env("AUTH_WORKER_INTENT").as_str() {
"auto" => AuthIntent::Auto,
"userinitiated" => AuthIntent::UserInitiated,
"headless" => AuthIntent::Headless,
other => panic!("unknown AUTH_WORKER_INTENT: {other}"),
};
let script = match env("AUTH_WORKER_SCRIPT").as_str() {
"approve" => Script::Approve,
"deny" => Script::Deny,
"failopen" => Script::FailToOpen,
other => panic!("unknown AUTH_WORKER_SCRIPT: {other}"),
};
let result_path = PathBuf::from(env("AUTH_WORKER_RESULT"));
let start_marker = std::env::var("AUTH_WORKER_START_MARKER")
.ok()
.map(PathBuf::from);
let ready_marker = std::env::var("AUTH_WORKER_READY_MARKER")
.ok()
.map(PathBuf::from);
let calls = Arc::new(AtomicU64::new(0));
let opener = WorkerOpener {
script,
calls: calls.clone(),
launched_marker: std::env::var("AUTH_WORKER_LAUNCHED_MARKER")
.ok()
.map(PathBuf::from),
proceed_marker: std::env::var("AUTH_WORKER_PROCEED_MARKER")
.ok()
.map(PathBuf::from),
};
let cfg = PkceOAuthConfig {
discovery_url: env("AUTH_WORKER_DISCOVERY_URL"),
client_id: env("AUTH_WORKER_CLIENT_ID"),
scopes: env("AUTH_WORKER_SCOPES")
.split(',')
.map(str::to_owned)
.collect(),
cache_namespace: env("AUTH_WORKER_NAMESPACE"),
cache_dir_override: Some(PathBuf::from(env("AUTH_WORKER_CACHE_DIR"))),
};
let src = PkceOAuthTokenSource::new_with(cfg, Arc::new(opener)).expect("build token source");
// Announce readiness, then wait for the parent's release so several workers
// hit the lock together — a genuine race rather than staggered spawns.
if let Some(marker) = &ready_marker {
fs::write(marker, b"ready").expect("write ready marker");
}
if let Some(marker) = start_marker {
for _ in 0..6000 {
if marker.exists() {
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
}
let (result, bearer) = match src
.acquire_with_intent(
intent,
std::env::var("AUTH_WORKER_REJECTED").ok().as_deref(),
)
.await
{
Ok(token) => ("ok".to_owned(), Some(token)),
Err(e) => (e.code().to_owned(), None),
};
let body = serde_json::json!({
"result": result,
"bearer": bearer,
"launches": calls.load(Ordering::SeqCst),
});
fs::write(&result_path, serde_json::to_vec(&body).unwrap()).expect("write result file");
}

50
crates/buzz-agent/tests/bin/lock_holder.rs Normal file
View File

@@ -0,0 +1,50 @@
//! Test-only helper: a real second process that takes the coordinator's
//! cross-process advisory lock and holds it until killed.
//!
//! The auth coordinator single-flights per cache key on an `fs2` advisory lock
//! (`flock` on Unix, `LockFileEx` on Windows). To prove the *cross-process*
//! contract — a genuine other process serializes the flow, and its death
//! releases the lock with no PID files or lock-breaking — a test needs an
//! actual separate process on the same lock file, not a second in-process
//! handle. This binary is that process.
//!
//! Driven by two env vars:
//! LOCK_HELPER_PATH — the lock file to acquire (the coordinator's
//! `<hash>.json.lock`).
//! LOCK_HELPER_READY — a marker file created *after* the lock is held, so
//! the parent test can synchronize on ownership before
//! racing the coordinator.
//!
//! After signaling readiness it blocks forever; the parent kills it to model a
//! crash mid-flow.
use std::fs;
use fs2::FileExt;
fn main() {
let lock_path = std::env::var("LOCK_HELPER_PATH").expect("LOCK_HELPER_PATH set");
let ready_path = std::env::var("LOCK_HELPER_READY").expect("LOCK_HELPER_READY set");
if let Some(parent) = std::path::Path::new(&lock_path).parent() {
fs::create_dir_all(parent).expect("create lock parent dir");
}
// Open exactly as the coordinator does so we contend on the same inode.
let file = fs::OpenOptions::new()
.create(true)
.truncate(false)
.write(true)
.open(&lock_path)
.expect("open lock file");
file.lock_exclusive()
.expect("hold the exclusive advisory lock");
// Signal ownership only once the lock is truly held.
fs::write(&ready_path, b"held").expect("write ready marker");
// Hold the lock until the parent kills us (crash stand-in). The kernel
// releases the advisory lock on process death.
loop {
std::thread::sleep(std::time::Duration::from_secs(3600));
}
}

3410
crates/buzz-agent/tests/databricks_auth_coordinator.rs Normal file
View File

File diff suppressed because it is too large Load Diff

91
crates/buzz-agent/tests/fake_llm.rs
View File

@@ -80,19 +80,35 @@ async fn spawn_capturing_fake_llm(responses: Vec<Value>) -> (String, Arc<Mutex<V
async fn spawn_capturing_fake_llm_with_statuses(
responses: Vec<CannedResponse>,
) -> (String, Arc<Mutex<Vec<Value>>>) {
let captures: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let url = spawn_capturing_fake_llm_core(responses, captures.clone(), None).await;
(url, captures)
}
/// Shared connection loop for the capturing fake LLM: reads each request,
/// records its JSON body into `captures`, and replies with the next canned
/// response. When `gate` is `Some`, the FIRST request's response is withheld
/// until the gate fires; when `None`, every response is served immediately.
async fn spawn_capturing_fake_llm_core(
responses: Vec<CannedResponse>,
captures: Arc<Mutex<Vec<Value>>>,
gate: Option<Arc<Mutex<Option<tokio::sync::oneshot::Receiver<()>>>>>,
) -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = format!("http://{}", listener.local_addr().unwrap());
let queue = Arc::new(Mutex::new(VecDeque::from(responses)));
let captures: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let captures_clone = captures.clone();
tokio::spawn(async move {
let mut request_num = 0usize;
loop {
let (mut sock, _) = match listener.accept().await {
Ok(p) => p,
Err(_) => return,
};
let queue = queue.clone();
let captures = captures_clone.clone();
let captures = captures.clone();
let gate = gate.clone();
request_num += 1;
let req_num = request_num;
tokio::spawn(async move {
// Read headers.
let mut buf = Vec::new();
@@ -141,6 +157,15 @@ async fn spawn_capturing_fake_llm_with_statuses(
captures.lock().await.push(parsed);
}
// Hold the first request's response until the gate opens.
if req_num == 1 {
if let Some(gate) = &gate {
if let Some(rx) = gate.lock().await.take() {
let _ = rx.await;
}
}
}
// Send canned response.
let response = queue.lock().await.pop_front().unwrap_or(CannedResponse {
status: 500,
@@ -164,6 +189,20 @@ async fn spawn_capturing_fake_llm_with_statuses(
});
}
});
url
}
/// A capturing fake LLM whose FIRST provider response is withheld until
/// `gate` fires. Later responses are served immediately. Used to make
/// round-boundary races deterministic: hold round 1 open until a client action
/// (e.g. a steer) is confirmed, so the second round observes it. Request bodies
/// are recorded into `captures` exactly as `spawn_capturing_fake_llm` does.
async fn spawn_gated_capturing_fake_llm(
responses: Vec<CannedResponse>,
captures: Arc<Mutex<Vec<Value>>>,
gate: Arc<Mutex<Option<tokio::sync::oneshot::Receiver<()>>>>,
) -> (String, Arc<Mutex<Vec<Value>>>) {
let url = spawn_capturing_fake_llm_core(responses, captures.clone(), Some(gate)).await;
(url, captures)
}
@@ -774,14 +813,37 @@ async fn recv_active_run_id(h: &mut Harness) -> String {
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn steer_folds_into_active_turn_without_cancelling() {
use tokio::sync::oneshot;
// A two-round turn (tool call → text). A steer sent once the run is live
// must (a) be accepted with the matching runId, (b) NOT cancel the turn —
// it still ends with end_turn — and (c) reach the provider as a user turn.
let (url, captures) = spawn_capturing_fake_llm(vec![
openai_tool_call("call_steer", "fake__noop", json!({})),
openai_text("acknowledged the steer"),
])
.await;
//
// The steer is drained only at a round boundary (before the next provider
// request), so it must be enqueued before round 2 begins. Without
// synchronization a fast worker can complete round 1, drain an empty steer
// queue at the round-2 boundary, and dispatch round 2 before the steer is
// even sent — the steer then lands after the turn ends and never reaches
// the provider. To make this deterministic, the FIRST provider response is
// gated: it is withheld until the steer has been sent AND observed
// accepted, so round 1 cannot complete (and round 2 cannot start its drain)
// until the steer is already queued.
let (gate_tx, gate_rx) = oneshot::channel::<()>();
let gate_rx = Arc::new(Mutex::new(Some(gate_rx)));
let responses = vec![
CannedResponse {
status: 200,
body: openai_tool_call("call_steer", "fake__noop", json!({})),
},
CannedResponse {
status: 200,
body: openai_text("acknowledged the steer"),
},
];
let captures: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let (url, _) = spawn_gated_capturing_fake_llm(responses, captures.clone(), gate_rx).await;
let mut h = Harness::spawn(&url).await;
let sid = init_session(&mut h).await;
@@ -795,7 +857,8 @@ async fn steer_folds_into_active_turn_without_cancelling() {
)
.await;
// Learn the run id, then steer into it before the turn finishes.
// Learn the run id (advertised before the gated round-1 request), then steer
// into the live turn while round 1 is still held.
let run_id = recv_active_run_id(&mut h).await;
let steer_text = "STEER-CANARY: also consider the edge case";
let s_id = h
@@ -809,9 +872,12 @@ async fn steer_folds_into_active_turn_without_cancelling() {
)
.await;
// Steer is accepted and echoes the run id it landed in.
// Steer is accepted and echoes the run id it landed in. Only after this
// confirmation do we release the gate, so the steer is guaranteed queued
// before round 2's boundary drains it.
let mut steer_ok = false;
let mut end_turn = false;
let mut gate = Some(gate_tx);
for _ in 0..40 {
let v = h.recv().await;
if v["id"] == json!(s_id) {
@@ -827,6 +893,11 @@ async fn steer_folds_into_active_turn_without_cancelling() {
"steer reply carries a messageId"
);
steer_ok = true;
// Steer accepted — release round 1 so the turn proceeds to round 2,
// whose boundary now drains the queued steer.
if let Some(tx) = gate.take() {
let _ = tx.send(());
}
} else if v["id"] == json!(p_id) {
// The turn was NOT cancelled — it completed normally.
assert_eq!(v["result"]["stopReason"], "end_turn");

11
desktop/src-tauri/Cargo.lock generated
View File

@@ -1041,6 +1041,7 @@ dependencies = [
"axum",
"base64 0.22.1",
"dirs",
"fs2",
"getrandom 0.4.3",
"hex",
"nix 0.31.3",
@@ -3045,6 +3046,16 @@ dependencies = [
"percent-encoding",
]
[[package]]
name = "fs2"
version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213"
dependencies = [
"libc",
"winapi",
]
[[package]]
name = "fs_extra"
version = "1.3.0"