fix(keyless): address broker transport review

Signed-off-by: Joel Robotham <jrobotham@squareup.com>
This commit is contained in:
Joel Robotham
2026-08-31 13:24:01 -07:00
parent 70dfcf4eb7
commit 3d2ea5b89a
8 changed files with 325 additions and 51 deletions

1
Cargo.lock generated
View File

@@ -985,6 +985,7 @@ dependencies = [
"reqwest 0.13.4",
"serde_json",
"tokio",
"url",
]
[[package]]

View File

@@ -527,6 +527,28 @@ impl ObserverPublishQueue {
/// Pending coalesced chunks are flushed into the queue first, so a
/// publish slot never leaves merged chunk text stranded behind the tick.
fn next_frame(&mut self) -> Option<observer::ObserverEvent> {
self.next_frame_fitting(
|_| {},
|frame| serialized_len(frame) <= OBSERVER_MAX_PLAINTEXT_LEN,
)
}
/// Pack one frame against the broker action's encoded argument budget.
/// Individual oversized events may be elided, but a batch that only
/// overflows because of the broker's outer JSON escaping is split and its
/// overflow remains queued for the next publish slot.
fn next_broker_frame(&mut self) -> Option<observer::ObserverEvent> {
self.next_frame_fitting(
fit_broker_observer_event_to_budget,
broker_observer_event_fits,
)
}
fn next_frame_fitting(
&mut self,
mut fit_single: impl FnMut(&mut observer::ObserverEvent),
frame_fits: impl Fn(&observer::ObserverEvent) -> bool,
) -> Option<observer::ObserverEvent> {
for (source_events, ready) in self.coalescer.flush() {
self.enqueue(source_events, ready);
}
@@ -536,18 +558,21 @@ impl ObserverPublishQueue {
let mut kept: VecDeque<(usize, u64, observer::ObserverEvent)> =
VecDeque::with_capacity(self.events.len());
let mut gathering = true;
while let Some((bytes, source_events, event)) = self.events.pop_front() {
while let Some((old_bytes, source_events, mut event)) = self.events.pop_front() {
fit_single(&mut event);
let bytes = serialized_len(&event);
self.pending_bytes = self.pending_bytes - old_bytes + bytes;
if gathering && event.channel_id == channel {
picked.push(event);
if picked.len() > 1
&& serialized_len(&batch_envelope(&picked)) > OBSERVER_MAX_PLAINTEXT_LEN
{
let candidate = seal_batch(picked.clone());
if picked.len() > 1 && !frame_fits(&candidate) {
// Frame full: the overflow event stays queued and leads
// its channel's next slot.
let event = picked.pop().expect("len > 1");
kept.push_back((bytes, source_events, event));
gathering = false;
} else {
debug_assert!(frame_fits(&candidate));
self.pending_bytes -= bytes;
}
} else {
@@ -654,8 +679,7 @@ fn spawn_broker_observer_publisher(
Err(tokio::sync::broadcast::error::RecvError::Closed) => closed = true,
},
_ = publish_tick.tick() => {
if let Some(mut event) = queue.next_frame() {
fit_broker_observer_event_to_budget(&mut event);
if let Some(event) = queue.next_broker_frame() {
match serde_json::to_string(&event) {
Ok(payload) => {
let frame = buzz_sdk::broker::ObserverFrame {
@@ -6304,6 +6328,41 @@ mod observer_publish_queue_tests {
);
}
/// A broker frame has another JSON layer around the observer payload.
/// Many individually small, quote-heavy events can fit the relay
/// plaintext limit as one batch while exceeding the encoded broker action
/// limit. The overflow must remain queued, not collapse the whole batch to
/// an elision stub.
#[test]
fn broker_batches_split_before_outer_json_overflow_without_losing_events() {
let mut queue = queue_of(
(1..=31)
.map(|seq| {
let mut e = event(seq, "tool_call", Some("chan-a"));
e.payload = serde_json::json!({ "body": "\"".repeat(900) });
e
})
.collect(),
);
let mut frames = Vec::new();
while !queue.is_empty() {
frames.push(queue.next_broker_frame().expect("queue not empty"));
}
assert!(frames.len() > 1, "outer encoding must split this batch");
assert!(frames.iter().all(broker_observer_event_fits));
assert!(
frames
.iter()
.all(|frame| !serde_json::to_string(frame).unwrap().contains("[elided")),
"small source events must not be elided"
);
let published: Vec<u64> = frames.iter().flat_map(frame_seqs).collect();
assert_eq!(published, (1..=31).collect::<Vec<_>>());
assert_eq!(queue.dropped_events, 0);
}
/// The queue preserves the coalescer's ordering rule: a non-chunk event
/// force-flushes pending chunk text ahead of itself, so merged chunks can
/// never leapfrog a tool call that arrived after them.

View File

@@ -4,7 +4,8 @@
//! mode polls `channel.read` and routes keyless storage and live signals back
//! through the broker. Capabilities with no broker action remain relay-only.
use std::collections::{HashMap, HashSet, VecDeque};
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use buzz_broker_client::HttpBrokerClient;
@@ -103,13 +104,22 @@ pub struct BrokerRuntime {
client: HttpBrokerClient,
channel_ids: Vec<Uuid>,
filters: HashMap<Uuid, ChannelFilter>,
filter_tx: tokio::sync::watch::Sender<HashMap<Uuid, ChannelFilter>>,
event_rx: tokio::sync::mpsc::Receiver<BuzzEvent>,
poll_task: tokio::task::JoinHandle<()>,
placeholder_keys: Keys,
terminal_error: Arc<Mutex<Option<String>>>,
}
struct BrokerPoller {
client: HttpBrokerClient,
channel_ids: Vec<Uuid>,
filter_rx: tokio::sync::watch::Receiver<HashMap<Uuid, ChannelFilter>>,
cursors: HashMap<Uuid, String>,
pending: VecDeque<BuzzEvent>,
seen_current: HashSet<String>,
seen_previous: HashSet<String>,
poll: tokio::time::Interval,
placeholder_keys: Keys,
terminal_error: Option<String>,
terminal_error: Arc<Mutex<Option<String>>>,
}
/// Cloneable handle for broker-backed capabilities used outside the polling
@@ -251,23 +261,57 @@ impl BrokerRuntime {
let agent_pubkey = address.author_pubkey.as_str().to_string();
let mut poll = tokio::time::interval(poll_interval);
poll.set_missed_tick_behavior(MissedTickBehavior::Skip);
let filters = HashMap::new();
let (filter_tx, filter_rx) = tokio::sync::watch::channel(filters.clone());
let (event_tx, event_rx) = tokio::sync::mpsc::channel(256);
let terminal_error = Arc::new(Mutex::new(None));
let poller = BrokerPoller {
client: client.clone(),
channel_ids: channel_ids.clone(),
filter_rx,
cursors: HashMap::new(),
seen_current: HashSet::new(),
seen_previous: HashSet::new(),
poll,
terminal_error: Arc::clone(&terminal_error),
};
let poll_task = tokio::spawn(poller.run(event_tx));
Ok((
Self {
client,
channel_ids,
filters: HashMap::new(),
cursors: HashMap::new(),
pending: VecDeque::new(),
seen_current: HashSet::new(),
seen_previous: HashSet::new(),
poll,
filters,
filter_tx,
event_rx,
poll_task,
placeholder_keys,
terminal_error: None,
terminal_error,
},
agent_pubkey,
))
}
fn set_filter(&mut self, channel_id: Uuid, filter: Option<ChannelFilter>) {
if let Some(filter) = filter {
self.filters.insert(channel_id, filter);
} else {
self.filters.remove(&channel_id);
}
self.filter_tx.send_replace(self.filters.clone());
}
async fn next_event(&mut self) -> Option<BuzzEvent> {
self.event_rx.recv().await
}
}
impl Drop for BrokerRuntime {
fn drop(&mut self) {
self.poll_task.abort();
}
}
impl BrokerPoller {
fn remember(&mut self, event_id: String) -> bool {
if self.seen_current.contains(&event_id) || self.seen_previous.contains(&event_id) {
return false;
@@ -279,17 +323,15 @@ impl BrokerRuntime {
true
}
async fn next_event(&mut self) -> Option<BuzzEvent> {
async fn run(mut self, event_tx: tokio::sync::mpsc::Sender<BuzzEvent>) {
loop {
if let Some(event) = self.pending.pop_front() {
return Some(event);
}
self.poll.tick().await;
let filters = self.filter_rx.borrow().clone();
let subscriptions: Vec<(Uuid, bool)> = self
.channel_ids
.iter()
.filter_map(|channel_id| {
self.filters
filters
.get(channel_id)
.map(|filter| (*channel_id, filter.require_mention))
})
@@ -314,8 +356,9 @@ impl BrokerRuntime {
}
Err(error) if error.code == Some(BrokerErrorCode::Unauthenticated) => {
tracing::error!(%channel_id, "broker credential was rejected: {error}");
self.terminal_error = Some(error.to_string());
return None;
*self.terminal_error.lock().expect("terminal error lock") =
Some(error.to_string());
return;
}
Err(error) => {
tracing::warn!(%channel_id, "broker channel.read failed: {error}");
@@ -343,8 +386,13 @@ impl BrokerRuntime {
tracing::warn!(%channel_id, "broker returned an event for a different channel");
continue;
}
if self.remember(event.id.to_hex()) {
self.pending.push_back(BuzzEvent { channel_id, event });
if self.remember(event.id.to_hex())
&& event_tx
.send(BuzzEvent { channel_id, event })
.await
.is_err()
{
return;
}
}
}
@@ -466,7 +514,7 @@ impl RuntimeTransport {
match self {
Self::Local(relay) => relay.subscribe_channel(channel_id, filter).await,
Self::Broker(runtime) => {
runtime.filters.insert(channel_id, filter);
runtime.set_filter(channel_id, Some(filter));
Ok(())
}
}
@@ -485,7 +533,7 @@ impl RuntimeTransport {
.await
}
Self::Broker(runtime) => {
runtime.filters.insert(channel_id, filter);
runtime.set_filter(channel_id, Some(filter));
Ok(())
}
}
@@ -495,7 +543,7 @@ impl RuntimeTransport {
match self {
Self::Local(relay) => relay.unsubscribe_channel(channel_id).await,
Self::Broker(runtime) => {
runtime.filters.remove(&channel_id);
runtime.set_filter(channel_id, None);
Ok(())
}
}
@@ -511,10 +559,12 @@ impl RuntimeTransport {
pub async fn reconnect(&mut self) -> Result<(), RelayError> {
match self {
Self::Local(relay) => relay.reconnect().await,
Self::Broker(runtime) => match &runtime.terminal_error {
Some(error) => Err(RelayError::Http(error.clone())),
None => Ok(()),
},
Self::Broker(runtime) => {
match &*runtime.terminal_error.lock().expect("terminal error lock") {
Some(error) => Err(RelayError::Http(error.clone())),
None => Ok(()),
}
}
}
}
@@ -573,6 +623,7 @@ mod tests {
use axum::response::Response;
use axum::routing::post;
use axum::Router;
use nostr::{EventBuilder, Kind, Tag};
use std::sync::{Arc, Mutex};
use tokio::net::TcpListener;
@@ -674,4 +725,93 @@ mod tests {
]
);
}
#[tokio::test]
async fn polling_survives_main_loop_timer_cancellation_and_reaches_later_channels() {
let slow_channel = Uuid::new_v4();
let ready_channel = Uuid::new_v4();
let keys = Keys::generate();
let ready_event = EventBuilder::new(Kind::Custom(9), "ready")
.tags([Tag::parse(["h", &ready_channel.to_string()]).unwrap()])
.sign_with_keys(&keys)
.unwrap();
let pubkey = keys.public_key().to_hex();
let app = Router::new().route(
"/v1/action",
post(move |body: Bytes| {
let ready_event = ready_event.clone();
let pubkey = pubkey.clone();
async move {
let request: serde_json::Value = serde_json::from_slice(&body).unwrap();
let request_id = request["requestId"].as_str().unwrap();
let action = request["action"].as_str().unwrap();
let outcome = if action == "storage.address" {
serde_json::json!({
"authorPubkey": pubkey,
"kind": 30078,
"dTag": "a".repeat(64),
})
} else {
let channel = request["args"]["channelId"].as_str().unwrap();
if channel == slow_channel.to_string() {
tokio::time::sleep(Duration::from_millis(80)).await;
serde_json::json!({ "messages": [] })
} else {
serde_json::json!({ "messages": [ready_event] })
}
};
let response = serde_json::json!({
"type": "broker_result",
"protocolVersion": 1,
"requestId": request_id,
"status": "succeeded",
"action": action,
"outcome": outcome,
});
Response::builder()
.status(StatusCode::OK)
.body(Body::from(response.to_string()))
.unwrap()
}
}),
);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let (mut runtime, _) = BrokerRuntime::connect(
format!("http://{address}"),
"credential".into(),
vec![slow_channel, ready_channel],
Duration::from_millis(1),
Keys::generate(),
)
.await
.unwrap();
let filter = ChannelFilter {
kinds: Some(vec![9]),
require_mention: false,
};
runtime.set_filter(slow_channel, Some(filter.clone()));
runtime.set_filter(ready_channel, Some(filter));
let mut typing_timer = tokio::time::interval_at(
tokio::time::Instant::now() + Duration::from_millis(10),
Duration::from_millis(10),
);
let received = tokio::time::timeout(Duration::from_millis(300), async {
loop {
tokio::select! {
_ = typing_timer.tick() => {}
event = runtime.next_event() => break event,
}
}
})
.await
.expect("timer ticks must not restart the slow first channel")
.expect("poller stays open");
assert_eq!(received.channel_id, ready_channel);
assert_eq!(received.event.content, "ready");
}
}

View File

@@ -12,6 +12,7 @@ buzz-sdk = { workspace = true }
reqwest = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true }
url = { workspace = true }
[dev-dependencies]
axum = { workspace = true }

View File

@@ -5,6 +5,7 @@
//! back. Callers use [`buzz_sdk::broker::BrokerClientExt::execute`], which adds
//! correlation checks; this type never interprets a verdict.
use buzz_sdk::broker::actions::MAX_ENCODED_MESSAGE_BYTES;
use buzz_sdk::broker::{
BrokerClient, BrokerFuture, BrokerResponse, BrokerTransportError, Dispatch, PreparedRequest,
BROKER_ACTION_PATH, BROKER_CREDENTIAL_HEADER,
@@ -12,19 +13,13 @@ use buzz_sdk::broker::{
const DEFAULT_REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const MAX_RESPONSE_BYTES: usize = 2 * 1024 * 1024;
const MAX_CHANNEL_READ_MESSAGE_BYTES: usize = 128 * 1024;
// A channel.read page may contain 500 messages whose content alone is up to
// 64 KiB each. Keep the transport bounded while leaving room for the signed
// event envelopes, tags, and JSON encoding around that contract-valid page.
const MAX_CHANNEL_READ_RESPONSE_BYTES: usize = 64 * 1024 * 1024;
fn max_response_bytes(request: &PreparedRequest) -> usize {
request
.channel_read_limit()
.map_or(MAX_RESPONSE_BYTES, |limit| {
MAX_RESPONSE_BYTES
.saturating_add((limit as usize).saturating_mul(MAX_CHANNEL_READ_MESSAGE_BYTES))
.min(MAX_CHANNEL_READ_RESPONSE_BYTES)
.saturating_add((limit as usize).saturating_mul(MAX_ENCODED_MESSAGE_BYTES))
})
}
@@ -76,12 +71,12 @@ impl HttpBrokerClient {
let base_url = base_url.into();
let parsed = reqwest::Url::parse(&base_url)
.map_err(|error| BrokerClientConfigError(format!("invalid broker URL: {error}")))?;
let loopback = parsed.host_str().is_some_and(|host| {
host.eq_ignore_ascii_case("localhost")
|| host
.parse::<std::net::IpAddr>()
.is_ok_and(|address| address.is_loopback())
});
let loopback = match parsed.host() {
Some(url::Host::Domain(host)) => host.eq_ignore_ascii_case("localhost"),
Some(url::Host::Ipv4(address)) => address.is_loopback(),
Some(url::Host::Ipv6(address)) => address.is_loopback(),
None => false,
};
if parsed.scheme() != "https" && !(parsed.scheme() == "http" && loopback) {
return Err(BrokerClientConfigError(
"broker URL must use HTTPS (plain HTTP is allowed only for loopback)".into(),
@@ -207,7 +202,7 @@ mod tests {
use axum::response::Response;
use axum::routing::{any, post};
use axum::Router;
use buzz_sdk::broker::actions::MAX_CONTENT_BYTES;
use buzz_sdk::broker::actions::{MAX_CONTENT_BYTES, MAX_ENCODED_MESSAGE_BYTES};
use buzz_sdk::broker::{
ActionArgs, ActionOutcome, BrokerClientExt, BrokerMessage, BrokerRequest, BrokerResult,
ChannelReadArgs, MessagePage, MessagePostArgs,
@@ -383,6 +378,7 @@ mod tests {
assert!(HttpBrokerClient::new("http://broker.example", CRED).is_err());
assert!(HttpBrokerClient::new("https://broker.example", " ").is_err());
assert!(HttpBrokerClient::new("http://localhost:8787", CRED).is_ok());
assert!(HttpBrokerClient::new("http://[::1]:8787", CRED).is_ok());
}
#[tokio::test]
@@ -447,7 +443,7 @@ mod tests {
))
.unwrap();
assert!(body.len() > MAX_RESPONSE_BYTES);
assert!(body.len() < MAX_CHANNEL_READ_RESPONSE_BYTES);
assert!(body.len() < max_response_bytes(&request));
let (base, _) = spawn(StatusCode::OK, body).await;
let response = HttpBrokerClient::new(base, CRED)
@@ -466,9 +462,8 @@ mod tests {
assert_eq!(request.channel_read_limit(), Some(100));
assert_eq!(
response_cap,
MAX_RESPONSE_BYTES + 100 * MAX_CHANNEL_READ_MESSAGE_BYTES
MAX_RESPONSE_BYTES + 100 * MAX_ENCODED_MESSAGE_BYTES
);
assert!(response_cap < MAX_CHANNEL_READ_RESPONSE_BYTES);
let (base, _) = spawn(StatusCode::OK, "x".repeat(response_cap + 1)).await;
let error = HttpBrokerClient::new(base, CRED)
@@ -482,4 +477,34 @@ mod tests {
if detail.contains(&response_cap.to_string())
));
}
#[tokio::test]
async fn valid_maximally_escaped_channel_page_fits_the_request_cap() {
let request = channel_read_request();
let message = BrokerMessage(
EventBuilder::new(Kind::Custom(9), "\u{001b}".repeat(MAX_CONTENT_BYTES))
.tags([nostr::Tag::parse(["h", CHANNEL]).unwrap()])
.sign_with_keys(&Keys::generate())
.expect("fixture event signs"),
);
assert!(serde_json::to_vec(&message).unwrap().len() <= MAX_ENCODED_MESSAGE_BYTES);
let outcome = ActionOutcome::ChannelRead(MessagePage {
messages: vec![message; 100],
next_cursor: None,
});
let body = serde_json::to_string(&BrokerResponse::new(
request.request_id(),
BrokerResult::succeeded(outcome),
))
.unwrap();
assert!(body.len() > 32 * 1024 * 1024);
assert!(body.len() <= max_response_bytes(&request));
let (base, _) = spawn(StatusCode::OK, body).await;
HttpBrokerClient::new(base, CRED)
.unwrap()
.execute(&request)
.await
.expect("contract-valid escaped page");
}
}

View File

@@ -38,6 +38,15 @@ pub const MAX_ABOUT_CHARS: usize = 2_000;
/// Maximum bytes of message content, matching the SDK's channel-message cap.
pub const MAX_CONTENT_BYTES: usize = 64 * 1024;
/// Maximum JSON-encoded bytes of one signed event returned by `channel.read`.
///
/// Content can expand by up to 6x under JSON escaping (`U+0000` through
/// `U+001F`), and the returned event also carries its signed tags and fixed
/// Nostr envelope. 512 KiB therefore admits a maximally escaped 64 KiB
/// message with substantial tag headroom while keeping every accepted page
/// bounded independently of the host implementation.
pub const MAX_ENCODED_MESSAGE_BYTES: usize = 512 * 1024;
/// Maximum characters in a reaction payload (emoji or `:shortcode:`).
pub const MAX_EMOJI_CHARS: usize = 66;

View File

@@ -7,7 +7,7 @@ use serde::{Deserialize, Serialize};
use super::{
absent_or_valued, channel, channel_id, cursor, event_id, hex64_field, required, Action,
PubkeyHex, MAX_NAME_CHARS, MAX_PAGE_LIMIT,
PubkeyHex, MAX_CONTENT_BYTES, MAX_ENCODED_MESSAGE_BYTES, MAX_NAME_CHARS, MAX_PAGE_LIMIT,
};
use crate::SdkError;
use nostr::{Event, EventId, Kind, PublicKey, Tags, Timestamp};
@@ -326,6 +326,22 @@ impl ActionOutcome {
page.messages.len()
)));
}
for message in &page.messages {
if message.0.content.len() > MAX_CONTENT_BYTES {
return Err(SdkError::ContentTooLarge {
max: MAX_CONTENT_BYTES,
got: message.0.content.len(),
});
}
let encoded_len = serde_json::to_vec(message)
.map_err(|error| SdkError::InvalidInput(error.to_string()))?
.len();
if encoded_len > MAX_ENCODED_MESSAGE_BYTES {
return Err(SdkError::InvalidInput(format!(
"encoded channel message is {encoded_len} bytes, over the {MAX_ENCODED_MESSAGE_BYTES}-byte cap"
)));
}
}
page.next_cursor.as_deref().map(cursor).transpose()?;
}
Self::MessagePost(published)

View File

@@ -2459,6 +2459,29 @@ fn a_page_is_bounded_and_its_cursor_opaque() {
};
assert!(page(vec![], None).is_ok());
assert!(page(vec![signed_message(&signer)], Some("c1")).is_ok());
let escaped = BrokerMessage(
EventBuilder::new(
Kind::Custom(9),
"\u{001b}".repeat(actions::MAX_CONTENT_BYTES),
)
.tags([Tag::parse(["h", CHANNEL]).unwrap()])
.sign_with_keys(&signer)
.unwrap(),
);
assert!(page(vec![escaped], None).is_ok());
let oversized_content = BrokerMessage(
EventBuilder::new(Kind::Custom(9), "x".repeat(actions::MAX_CONTENT_BYTES + 1))
.sign_with_keys(&signer)
.unwrap(),
);
assert!(page(vec![oversized_content], None).is_err());
let oversized_tags = BrokerMessage(
EventBuilder::new(Kind::Custom(9), "hello")
.tags([Tag::parse(["x", &"\"".repeat(actions::MAX_ENCODED_MESSAGE_BYTES)]).unwrap()])
.sign_with_keys(&signer)
.unwrap(),
);
assert!(page(vec![oversized_tags], None).is_err());
assert!(page(vec![], Some("")).is_err());
assert!(page(vec![], Some("has space")).is_err());
assert!(page(