Merge remote-tracking branch 'origin/main' into kennylopez-agent-catalog-sharing

This commit is contained in:
kenny lopez
2026-07-23 13:34:47 -07:00
10 changed files with 445 additions and 46 deletions
+11
View File
@@ -1127,6 +1127,7 @@ async fn serve(
let (shutdown_tx, _) = tokio::sync::watch::channel(false);
let shutdown_flag = Arc::clone(&state.shutting_down);
let drain_conn_manager = Arc::clone(&state.conn_manager);
let tx = shutdown_tx.clone();
tokio::spawn(async move {
shutdown_signal().await;
@@ -1136,6 +1137,16 @@ async fn serve(
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
info!("Starting graceful drain (30s timeout)");
let _ = tx.send(true);
// Tell every connected client to reconnect NOW. Without this, upgraded
// WebSocket connections outlive the listener drain: clients ride the
// dying pod until the forced exit below and only learn about the
// restart from a TCP reset. The 1012 close frame turns a 35s silent
// death into an immediate, well-attributed reconnect.
let closed = drain_conn_manager.drain_all();
info!(
connections = closed,
"Sent restart close frame to all live WebSocket connections"
);
// Hard timeout: force exit if connections don't drain within 30s.
tokio::time::sleep(std::time::Duration::from_secs(30)).await;
tracing::error!("Drain timeout exceeded — forcing exit");
+14 -3
View File
@@ -299,9 +299,20 @@ async fn nip11_or_ws_handler(
let max_frame_bytes = state.config.max_frame_bytes;
match WebSocketUpgrade::from_request(req, &state).await {
Ok(ws) => limit_relay_websocket(ws, max_frame_bytes)
.on_upgrade(move |socket| handle_connection(socket, state, addr, tenant))
.into_response(),
Ok(ws) => {
// Shutting down: refuse new sockets instead of accepting a
// connection onto a dying pod. Readiness already returns 503, but
// that only stops K8s routing — direct and in-flight upgrades
// still reach here during the pre-drain grace window. Clients
// treat the refusal as a normal dial failure and retry, landing
// on a healthy pod.
if state.shutting_down.load(Ordering::Relaxed) {
return (StatusCode::SERVICE_UNAVAILABLE, "relay restarting").into_response();
}
limit_relay_websocket(ws, max_frame_bytes)
.on_upgrade(move |socket| handle_connection(socket, state, addr, tenant))
.into_response()
}
Err(_) => {
// Browser requesting HTML and Git web GUI is enabled → serve SPA.
if state.config.serve_git_web_gui {
+192
View File
@@ -181,6 +181,10 @@ where
/// Tracks active Nostr WebSocket connections and provides message routing by connection ID.
pub struct ConnectionManager {
connections: DashMap<Uuid, ConnEntry>,
/// Sticky drain flag set by [`Self::drain_all`]. Registrations that land
/// after the drain snapshot self-signal, so no upgrade-vs-shutdown
/// interleaving can produce a connection that misses the restart close.
draining: AtomicBool,
}
impl ConnectionManager {
@@ -188,6 +192,7 @@ impl ConnectionManager {
pub fn new() -> Self {
Self {
connections: DashMap::new(),
draining: AtomicBool::new(false),
}
}
@@ -208,6 +213,8 @@ impl ConnectionManager {
subscriptions: ConnectionSubscriptions,
grace_limit: u8,
) {
let drain_ctrl_tx = ctrl_tx.clone();
let drain_cancel = cancel.clone();
self.connections.insert(
conn_id,
ConnEntry {
@@ -221,6 +228,14 @@ impl ConnectionManager {
grace_limit,
},
);
// Insert-then-check pairs with drain_all's store-then-iterate: either
// the drain iteration sees this entry, or this check sees the flag.
// A registration that raced past the snapshot self-signals here, so
// no connection can outlive graceful shutdown unclosed.
if self.draining.load(Ordering::SeqCst) {
let _ = drain_ctrl_tx.try_send(Self::restart_close_frame());
drain_cancel.cancel();
}
}
/// Removes a connection from the registry.
@@ -318,6 +333,45 @@ impl ConnectionManager {
closed
}
/// Closes every live connection with a `1012 Service Restart` close frame.
///
/// Called when graceful shutdown starts draining. Without this, upgraded
/// WebSocket connections outlive the axum listener drain: clients ride the
/// dying pod until the forced exit and then learn about the restart from a
/// TCP reset (or, on an abrupt kill, from up to 60s of stall-watchdog
/// silence). The explicit close frame tells them to reconnect immediately
/// — and that the disconnect is a restart, not a policy action.
///
/// Uses the "queue frame on ctrl, then cancel" idiom (see
/// [`ConnectionManager::disconnect_pubkey`]): the send loop drains queued
/// control frames — including this close — before its cancel branch closes
/// the socket. Best-effort: a full control buffer still gets the close via
/// cancel, just without the restart code.
///
/// Returns the number of connections signalled.
pub fn drain_all(&self) -> usize {
// Store-then-iterate pairs with register's insert-then-check: a
// registration that misses this iteration observes the flag and
// self-signals instead. The flag is sticky — drain is one-way.
self.draining.store(true, Ordering::SeqCst);
let frame = Self::restart_close_frame();
let mut closed = 0usize;
for entry in self.connections.iter() {
let _ = entry.ctrl_tx.try_send(frame.clone());
entry.cancel.cancel();
closed += 1;
}
closed
}
/// The WS close frame announcing a graceful restart: 1012 Service Restart.
fn restart_close_frame() -> WsMessage {
WsMessage::Close(Some(axum::extract::ws::CloseFrame {
code: axum::extract::ws::close_code::RESTART,
reason: axum::extract::ws::Utf8Bytes::from_static("relay restarting"),
}))
}
/// Return the server-resolved community that the connection's host bound to.
pub fn community_for_conn(&self, conn_id: Uuid) -> Option<CommunityId> {
self.connections
@@ -1737,4 +1791,142 @@ mod tests {
"community-B session stays live — ban does not cross the tenant fence"
);
}
#[tokio::test]
async fn drain_all_sends_restart_close_and_cancels_every_conn() {
// Graceful shutdown must tell every live client to reconnect — across
// all communities — with a 1012 restart close frame queued ahead of
// the cancel-driven socket close.
let mgr = ConnectionManager::new();
let register = |community| {
let conn_id = Uuid::new_v4();
let (tx, _rx) = mpsc::channel(8);
let (ctrl_tx, ctrl_rx) = mpsc::channel(8);
let cancel = CancellationToken::new();
mgr.register(
conn_id,
tx,
ctrl_tx,
cancel.clone(),
community,
Arc::new(AtomicU8::new(0)),
Arc::new(Mutex::new(HashMap::new())),
3,
);
(ctrl_rx, cancel)
};
let (mut ctrl_a, cancel_a) = register(buzz_core::tenant::CommunityId::from_uuid(
Uuid::from_u128(0xa),
));
let (mut ctrl_b, cancel_b) = register(buzz_core::tenant::CommunityId::from_uuid(
Uuid::from_u128(0xb),
));
let closed = mgr.drain_all();
assert_eq!(closed, 2, "every connection is signalled, no tenant fence");
assert!(cancel_a.is_cancelled(), "community-A session is cancelled");
assert!(cancel_b.is_cancelled(), "community-B session is cancelled");
for ctrl_rx in [&mut ctrl_a, &mut ctrl_b] {
let frame = ctrl_rx.try_recv().expect("close frame delivered");
match frame {
WsMessage::Close(Some(close)) => {
assert_eq!(
close.code,
axum::extract::ws::close_code::RESTART,
"close code is 1012 Service Restart"
);
assert_eq!(close.reason.as_str(), "relay restarting");
}
other => panic!("expected a restart close frame, got {other:?}"),
}
}
}
#[tokio::test]
async fn drain_all_full_control_buffer_still_cancels() {
// Best-effort delivery: a wedged control channel must not block the
// drain — the cancel still closes the socket, just without the frame.
let mgr = ConnectionManager::new();
let conn_id = Uuid::new_v4();
let (tx, _rx) = mpsc::channel(8);
let (ctrl_tx, mut ctrl_rx) = mpsc::channel(1);
let cancel = CancellationToken::new();
mgr.register(
conn_id,
tx,
ctrl_tx.clone(),
cancel.clone(),
buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()),
Arc::new(AtomicU8::new(0)),
Arc::new(Mutex::new(HashMap::new())),
3,
);
// Wedge the 1-slot control channel.
ctrl_tx
.try_send(WsMessage::Text("wedge".into()))
.expect("fill control channel");
let closed = mgr.drain_all();
assert_eq!(closed, 1);
assert!(
cancel.is_cancelled(),
"cancel fires even when the close frame cannot be queued"
);
// Only the wedge frame is present — the close was dropped, not queued.
assert!(matches!(
ctrl_rx.try_recv().expect("wedge frame"),
WsMessage::Text(_)
));
assert!(ctrl_rx.try_recv().is_err(), "no second frame queued");
}
#[tokio::test]
async fn register_after_drain_self_signals_restart_close_and_cancel() {
// The shutdown-boundary race: an upgrade accepted before SIGTERM can
// finish its async admission check and register AFTER drain_all's
// one-shot snapshot. The sticky drain flag makes that interleaving
// deterministic — register itself queues the 1012 and cancels, so no
// late registration can ride out graceful shutdown unclosed.
let mgr = ConnectionManager::new();
// Drain with zero connections — sets the sticky flag.
assert_eq!(mgr.drain_all(), 0);
// Late registration lands after the snapshot.
let conn_id = Uuid::new_v4();
let (tx, _rx) = mpsc::channel(8);
let (ctrl_tx, mut ctrl_rx) = mpsc::channel(8);
let cancel = CancellationToken::new();
mgr.register(
conn_id,
tx,
ctrl_tx,
cancel.clone(),
buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()),
Arc::new(AtomicU8::new(0)),
Arc::new(Mutex::new(HashMap::new())),
3,
);
assert!(
cancel.is_cancelled(),
"late registration is cancelled by the sticky drain flag"
);
match ctrl_rx.try_recv().expect("close frame delivered") {
WsMessage::Close(Some(close)) => {
assert_eq!(
close.code,
axum::extract::ws::close_code::RESTART,
"late registration still gets the 1012 restart close"
);
assert_eq!(close.reason.as_str(), "relay restarting");
}
other => panic!("expected a restart close frame, got {other:?}"),
}
}
}
+1
View File
@@ -137,6 +137,7 @@ export default defineConfig({
"**/persona-sync.spec.ts",
"**/team-snapshot.spec.ts",
"**/agents-everywhere.live.spec.ts",
"**/relay-restart.live.spec.ts",
"**/parity-ancestor-island.spec.ts",
],
use: {
+53 -40
View File
@@ -534,8 +534,6 @@ export class RelayClient {
}
private async connect() {
// Clear any pending stability timer from a previous connection — a new
// connect attempt resets the clock and must re-arm the timer on success.
if (this.stabilityTimer !== null) {
window.clearTimeout(this.stabilityTimer);
this.stabilityTimer = null;
@@ -545,51 +543,66 @@ export class RelayClient {
this.hasConnectedOnce ? "reconnecting" : "connecting",
);
if (!this.relayUrl) {
this.relayUrl = await getRelayWsUrl();
}
const generation = ++this.connectionGeneration;
this.onMessageChannel = new Channel<unknown>((message) => {
void this.handleWsMessage(message, generation);
});
this.wsId = await invoke<number>("plugin:websocket|connect", {
url: this.relayUrl,
onMessage: this.onMessageChannel,
config: {},
});
await new Promise<void>((resolve, reject) => {
const timeout = window.setTimeout(() => {
this.authRequest = null;
void this.handleWsMessage(message, generation).catch((error) => {
if (generation !== this.connectionGeneration) return;
this.resetConnection(
new Error("Timed out while waiting for relay authentication."),
this.normalizeRelayError(error, "Relay connection errored."),
);
reject(new Error("Timed out while waiting for relay authentication."));
}, AUTH_TIMEOUT_MS);
this.authRequest = {
pendingEventId: "",
resolve,
reject,
timeout,
};
});
});
// Start a stability timer instead of resetting backoff immediately.
// The backoff resets to its base value only after BACKOFF_RESET_STABLE_MS
// of uninterrupted uptime, preventing fast reconnect loops from erasing
// the exponential backoff that throttles them.
this.stabilityTimer = window.setTimeout(() => {
this.stabilityTimer = null;
this.reconnectDelayMs = RECONNECT_BASE_DELAY_MS;
}, BACKOFF_RESET_STABLE_MS);
try {
if (!this.relayUrl) {
this.relayUrl = await getRelayWsUrl();
}
const wsId = await invoke<number>("plugin:websocket|connect", {
url: this.relayUrl,
onMessage: this.onMessageChannel,
config: {},
});
if (generation !== this.connectionGeneration) {
void closeWebSocket(wsId, "stale connection attempt");
throw new Error("Relay connection attempt was superseded.");
}
this.wsId = wsId;
await this.replayLiveSubscriptions();
this.connectionStateEmitter.set("connected");
this.stallWatchdog.start();
this.emitReconnectIfNeeded();
await new Promise<void>((resolve, reject) => {
const timeout = window.setTimeout(() => {
const error = new Error("Relay authentication timed out.");
this.authRequest = null;
this.resetConnection(error);
reject(error);
}, AUTH_TIMEOUT_MS);
this.authRequest = {
pendingEventId: "",
resolve,
reject,
timeout,
};
});
this.stabilityTimer = window.setTimeout(() => {
this.stabilityTimer = null;
this.reconnectDelayMs = RECONNECT_BASE_DELAY_MS;
}, BACKOFF_RESET_STABLE_MS);
await this.replayLiveSubscriptions();
this.connectionStateEmitter.set("connected");
this.stallWatchdog.start();
this.emitReconnectIfNeeded();
} catch (error) {
const connectionError = this.normalizeRelayError(
error,
"Failed to connect to relay.",
);
if (generation === this.connectionGeneration) {
this.resetConnection(connectionError);
}
throw connectionError;
}
}
private async subscribe(
+7
View File
@@ -250,6 +250,8 @@ type E2eConfig = {
openerError?: string;
/** Delay binding signatures so specs can exercise request supersession. */
nostrBindSignDelayMs?: number;
/** Reject successive mock WebSocket connect attempts, then resume. */
websocketConnectErrors?: string[];
stallWebsocketSends?: boolean;
userSearchDelayMs?: number;
// NIP-IA gate inputs — see tests/helpers/bridge.ts:MockBridgeOptions for
@@ -8496,6 +8498,11 @@ async function connectRealSocket(args: { url?: string; onMessage: unknown }) {
}
async function connectMockSocket(args: { onMessage: unknown }) {
const connectError = getConfig()?.mock?.websocketConnectErrors?.shift();
if (connectError) {
throw new Error(connectError);
}
if (mockWebsocketSendMutexWedged) {
return new Promise<number>(() => {});
}
+42 -3
View File
@@ -17,7 +17,7 @@ type OwnedProcess = { name: string; child: ChildProcess; logPath: string };
export class TwoRelayHarness {
readonly root: string;
readonly relays: readonly [RelaySpec, RelaySpec];
readonly relays: readonly RelaySpec[];
private readonly processes: OwnedProcess[] = [];
get ownedPids(): number[] {
@@ -26,12 +26,12 @@ export class TwoRelayHarness {
);
}
private constructor(root: string, relays: readonly [RelaySpec, RelaySpec]) {
private constructor(root: string, relays: readonly RelaySpec[]) {
this.root = root;
this.relays = relays;
}
static async create(relays: readonly [RelaySpec, RelaySpec]) {
static async create(relays: readonly RelaySpec[]) {
return new TwoRelayHarness(
await mkdtemp(join(tmpdir(), "buzz-ae-e2e-")),
relays,
@@ -46,6 +46,45 @@ export class TwoRelayHarness {
);
}
/**
* SIGTERM one relay by name and wait for it to exit on its own. Unlike
* `stop()`, this never escalates to SIGKILL: the relay's graceful drain
* (readiness 503 → 5s grace → 1012 close broadcast → listener drain) is
* exactly what restart tests exercise, so a forced kill would invalidate
* the run. Throws if the process does not exit within `timeoutMs`.
*/
async terminateRelayGracefully(name: string, timeoutMs = 45_000) {
const owned = [...this.processes]
.reverse()
.find(
(candidate) =>
candidate.name === name &&
candidate.child.exitCode === null &&
candidate.child.signalCode === null,
);
if (!owned) throw new Error(`no live relay process named ${name}`);
const exited = new Promise<void>((resolveExit) =>
owned.child.once("exit", () => resolveExit()),
);
this.signal(owned.child, "SIGTERM");
if (
!(await Promise.race([exited.then(() => true), delay(timeoutMs, false)]))
) {
throw new Error(
`${name} did not exit within ${timeoutMs}ms of SIGTERM — graceful drain is stuck`,
);
}
}
/** Start a fresh process for a relay spec on its original ports. */
async restartRelay(name: string, binary = process.env.BUZZ_E2E_RELAY_BIN) {
if (!binary)
throw new Error("BUZZ_E2E_RELAY_BIN is required for the live gate");
const relay = this.relays.find((candidate) => candidate.name === name);
if (!relay) throw new Error(`no relay spec named ${name}`);
await this.startRelay(binary, relay);
}
async startAcp(
name: string,
relayWsUrl: string,
+27
View File
@@ -80,6 +80,33 @@ test.beforeEach(async ({ page }) => {
await installMockBridge(page);
});
test("failed initial relay dial retries automatically", async ({ page }) => {
await installMockBridge(page, {
websocketConnectErrors: ["mock relay pod unavailable"],
});
await page.goto("/");
// App-shell preconnect owns a keep-alive request. The first native dial is
// rejected before a socket ID exists; the session must still enter its
// backoff loop and recover without a click, query, or reload.
await expect
.poll(
() =>
page.evaluate(() => {
const getState = (
window as Window & {
__BUZZ_E2E_GET_RELAY_CONNECTION_STATE__?: () => string;
}
).__BUZZ_E2E_GET_RELAY_CONNECTION_STATE__;
if (!getState) throw new Error("Relay state seam is not installed.");
return getState();
}),
{ timeout: 10_000 },
)
.toBe("connected");
await expect(page.getByTestId("channel-general")).toBeVisible();
});
test("passive relay watchdog does not write while the websocket is half-open", async ({
page,
}) => {
@@ -0,0 +1,96 @@
import { expect, test, type Page } from "@playwright/test";
import { installBridge } from "../helpers/bridge";
import { TwoRelayHarness, type RelaySpec } from "./helpers/twoRelayHarness";
// Live gate: boots a REAL buzz-relay process, points the app at it, SIGTERMs
// the relay mid-session, restarts it on the same port, and asserts the client
// converges back to "connected". This proves the full restart story end to
// end: the relay's graceful-drain 1012 close broadcast (server side) and the
// client's dial-failure retry + 1012 fast-reconnect (desktop side) — the two
// halves that synthetic mock-websocket specs cannot compose.
//
// Requires: BUZZ_E2E_RELAY_RESTART=1, BUZZ_E2E_RELAY_BIN, and
// BUZZ_E2E_DATABASE_URL (plus reachable Redis and media object store, same
// infra as the agents-everywhere live gate).
const enabled = process.env.BUZZ_E2E_RELAY_RESTART === "1";
function required(name: string, value: string | undefined): string {
if (!value) throw new Error(`${name} is required for the live gate`);
return value;
}
async function connectionState(page: Page): Promise<string> {
return page.evaluate(() => {
const win = window as Window & {
__BUZZ_E2E_GET_RELAY_CONNECTION_STATE__?: () => string;
};
return win.__BUZZ_E2E_GET_RELAY_CONNECTION_STATE__?.() ?? "uninstalled";
});
}
test.describe("relay restart live gate", () => {
test.skip(!enabled, "set BUZZ_E2E_RELAY_RESTART=1 to run live gate");
test("client reconnects after the relay is SIGTERMed and restarted", async ({
page,
}) => {
test.setTimeout(180_000);
const portBase = 26_000 + (process.pid % 3_000);
const spec: RelaySpec = {
name: "relay-restart",
ports: {
main: portBase,
health: portBase + 3_000,
metrics: portBase + 6_000,
},
databaseUrl: required(
"BUZZ_E2E_DATABASE_URL",
process.env.BUZZ_E2E_DATABASE_URL,
),
redisUrl:
process.env.BUZZ_E2E_REDIS_RESTART ?? "redis://127.0.0.1:6379/13",
};
const harness = await TwoRelayHarness.create([spec]);
try {
await harness.startRelays();
const relayHttpUrl = `http://127.0.0.1:${spec.ports.main}`;
await installBridge(page, {
mode: "relay",
user: "tyler",
relayHttpUrl,
relayWsUrl: `ws://127.0.0.1:${spec.ports.main}`,
});
await page.goto("/");
// Baseline: the app converges to a live authenticated session.
await expect
.poll(() => connectionState(page), { timeout: 60_000 })
.toBe("connected");
// Roll the pod. Graceful drain: readiness 503 → 5s grace → 1012 close
// broadcast → process exit. The client must observe the close (not a
// silent stall) and start retrying.
await harness.terminateRelayGracefully(spec.name);
await expect
.poll(() => connectionState(page), { timeout: 30_000 })
.not.toBe("connected");
// Bring the "new pod" up on the same address, exactly like a k8s
// restart behind a stable service endpoint.
await harness.restartRelay(spec.name);
// The client's retry loop must find the fresh relay and converge back
// to connected without any user interaction.
await expect
.poll(() => connectionState(page), { timeout: 60_000 })
.toBe("connected");
} catch (error) {
console.error(await harness.logs());
throw error;
} finally {
await harness.stop();
}
});
});
+2
View File
@@ -261,6 +261,8 @@ type MockBridgeOptions = {
openerError?: string;
/** Delay binding signatures so specs can exercise request supersession. */
nostrBindSignDelayMs?: number;
/** Reject successive mock WebSocket connect attempts, then resume. */
websocketConnectErrors?: string[];
stallWebsocketSends?: boolean;
userSearchDelayMs?: number;
/**