diff --git a/benchmarks/aft-search/slice-descriptors/land-rich-ui-s05-chunk-queries.json b/benchmarks/aft-search/slice-descriptors/land-rich-ui-s05-chunk-queries.json new file mode 100644 index 000000000..637757404 --- /dev/null +++ b/benchmarks/aft-search/slice-descriptors/land-rich-ui-s05-chunk-queries.json @@ -0,0 +1,12 @@ +{ + "slice_class": "non_ranking", + "targeted_mechanism": "none", + "kind": "harness", + "fixtures": [ + "harness-goldens" + ], + "slices": [ + "rich-ui-s05-chunk-queries" + ], + "note": "Rich UI chunk queries for Occipital (manifest and template chunks over tool-provider routes). No search ranking, routing or index source is touched; no changed path matches RANKING_FENCE_PREFIXES." +} diff --git a/crates/aft/src/subc/mod.rs b/crates/aft/src/subc/mod.rs index a7c3d1bf8..300c0a54a 100644 --- a/crates/aft/src/subc/mod.rs +++ b/crates/aft/src/subc/mod.rs @@ -243,6 +243,12 @@ mod remote_policy; pub(crate) mod result_body; mod retained_result; mod rich_ui; +// Drives the chunk queries and Occipital's route shapes through the real +// frame loop, so it lives beside the chunk module but runs as a child of this +// module to reach the loop and its test fixtures. +#[cfg(test)] +#[path = "rich_ui/chunks/route_tests.rs"] +mod rich_ui_chunk_route_tests; mod run_scope; mod scope_describe; #[cfg(test)] @@ -5764,6 +5770,22 @@ where let route = route_key(frame.header.channel, frame.header.epoch); let decoded = DecodedFrame { frame, phase_trace, parsed_body, confinement, launch_restore, codemode_policies }; if let (Some(identity), Ok(body)) = (routes.get(&route), decoded.body()) { + // Rich-UI chunk queries are provider reads that name no + // session or tool, so they are answered here, before the + // carrier gate and every tool-call admission step. + if let Some((operation, query)) = chunk_query(body) { + let db = shared_app.db(); + if let Err(error) = answer_chunk_query(&writer_tx, &dispatch_path_metrics, &decoded.frame, identity, operation, &query, db.as_deref()).await { + break Err(error); + } + continue; + } + if let Some(error) = occipital_session_refusal(identity, body) { + if let Err(error) = send_provider_error(&writer_tx, &dispatch_path_metrics, &decoded.frame, error).await { + break Err(error); + } + continue; + } if let Err(error) = carrier_read_only::admit(identity, body) { if let Err(error) = send_provider_error(&writer_tx, &dispatch_path_metrics, &decoded.frame, error).await { break Err(error); @@ -8175,6 +8197,21 @@ async fn handle_control_request( ) .await; } + // Occipital reaches AFT only through tool-provider v1 routes: their + // provider reads and admitted carrier calls. The legacy grammar has + // neither the chunk queries nor carrier admission, so refuse it. + if route_role == tool_provider::RouteRole::Legacy + && principal_id(&principal).as_deref() == Some(rich_ui::chunks::OCCIPITAL_PRINCIPAL) + { + return send_route_bind_error( + tx, + frame, + "route_refused", + "reserved:occipital binds only tool-provider v1 routes", + metrics, + ) + .await; + } let bind_project_root = resolve_route_project_root(&identity, scope.as_ref()); let mut bind_root_id = None; @@ -9440,6 +9477,121 @@ async fn send_provider_error( send_reliable_writer_frame(tx, metrics, response, "v1 refusal").await } +/// Whether the route was bound by the authenticated principal +/// `reserved:occipital`. This reads the route's own bind principal, not a +/// codemode run principal: routes in the route table never carry a run owner. +fn occipital_route(identity: &RouteIdentity) -> bool { + matches!(&identity.spawn_principal, AuthenticatedPrincipal::RouteBind { principal_id: Some(id), .. } + if id == rich_ui::chunks::OCCIPITAL_PRINCIPAL) +} + +/// Whether an Occipital route carries another principal's session: it is +/// scoped, and its stamp's owner is not Occipital itself. Only such a route +/// may take session calls. +fn occipital_carrier_route(identity: &RouteIdentity) -> bool { + identity.scope.as_ref().is_some_and(|scope| { + principal_id(&Some(scope.owner.clone())).as_deref() + != Some(rich_ui::chunks::OCCIPITAL_PRINCIPAL) + }) +} + +/// Whether a route request is a tool call: either `op: "tool.call"`, or the +/// bare form with no `op` whose `name` is not a provider operation. +fn route_request_is_tool_call(envelope: &Value) -> bool { + match envelope.get("op").and_then(Value::as_str) { + Some(operation) => operation == "tool.call", + None => !envelope + .get("name") + .and_then(Value::as_str) + .is_some_and(|name| { + tool_provider::recognized_operation(name) + || rich_ui::chunks::is_chunk_operation(name) + }), + } +} + +/// An Occipital route that carries no other session (unscoped, or scoped by a +/// stamp Occipital owns) admits provider reads only. A tool call there has no +/// session to run in, so it refuses `session_not_bound` before the carrier +/// gate or any other tool-call admission step runs. +fn occipital_session_refusal( + identity: &RouteIdentity, + envelope: &Value, +) -> Option { + (occipital_route(identity) + && !occipital_carrier_route(identity) + && route_request_is_tool_call(envelope)) + .then(|| rich_ui::chunks::refusal("session_not_bound")) +} + +/// The chunk query a route request names, either as `op` or, in the bare +/// form, as `name`, with the request body that carries its digest and index. +fn chunk_query(envelope: &Value) -> Option<(&str, Value)> { + match envelope.get("op").and_then(Value::as_str) { + Some(operation) if rich_ui::chunks::is_chunk_operation(operation) => { + let mut body = envelope.clone(); + if let Some(body) = body.as_object_mut() { + body.remove("op"); + } + Some((operation, body)) + } + Some(_) => None, + None => { + let name = envelope.get("name").and_then(Value::as_str)?; + rich_ui::chunks::is_chunk_operation(name).then(|| { + ( + name, + envelope + .get("arguments") + .cloned() + .unwrap_or_else(|| json!({})), + ) + }) + } + } +} + +/// Answer a rich-UI chunk query on the route it arrived on. Admission is by +/// the route's authenticated principal alone, scoped or unscoped; the query +/// names no session. A reply is a Response frame holding the chunk; every +/// refusal is a typed Error frame through [`send_provider_error`]. +async fn answer_chunk_query( + tx: &WriterSender, + metrics: &DispatchPathMetrics, + frame: &Frame, + identity: &RouteIdentity, + operation: &str, + body: &Value, + db: Option<&StdMutex>, +) -> Result<(), SubcError> { + let principal = match &identity.spawn_principal { + AuthenticatedPrincipal::RouteBind { principal_id, .. } => principal_id.as_deref(), + AuthenticatedPrincipal::FirstParty => None, + }; + let answer = rich_ui::chunks::admit_principal(principal).and_then(|()| { + if operation == rich_ui::chunks::MANIFEST_CHUNK { + rich_ui::chunks::manifest_chunk(body) + } else { + rich_ui::chunks::template_chunk(body, db) + } + }); + let answer = match answer { + Ok(answer) => answer, + Err(error) => return send_provider_error(tx, metrics, frame, error).await, + }; + let response = Frame::build_with_version( + frame.header.ver, + FrameType::Response, + frame.header.flags, + frame.header.channel, + frame.header.epoch, + frame.header.corr, + serde_json::to_vec(&answer).map_err(SubcError::Json)?, + ) + .map_err(SubcError::FrameBuild)?; + send_reliable_writer_frame(tx, metrics, response, "rich-ui chunk").await +} + #[allow(clippy::too_many_arguments)] async fn submit_provider_read( tx: &WriterSender, @@ -10111,6 +10263,9 @@ async fn handle_tool_call( phase_trace.received_at(), &envelope, ); + if let Some(error) = occipital_session_refusal(&identity, &envelope) { + return send_provider_error(tx, metrics, frame, error).await; + } if let Err(error) = carrier_read_only::admit(&identity, &envelope) { return send_provider_error(tx, metrics, frame, error).await; } diff --git a/crates/aft/src/subc/rich_ui/chunks.rs b/crates/aft/src/subc/rich_ui/chunks.rs new file mode 100644 index 000000000..da8b08698 --- /dev/null +++ b/crates/aft/src/subc/rich_ui/chunks.rs @@ -0,0 +1,189 @@ +//! The two chunk queries Occipital uses to fetch rich-UI assets by digest. +//! +//! `rich_ui.manifest_chunk {manifestDigest, index}` serves the running build's +//! canonical manifest bytes and `rich_ui.template_chunk {templateDigest, +//! index}` serves template bytes, either shipped in the binary or stored in +//! aft.db by a view record. Both are provider reads: they take only a digest +//! and a 0-based chunk index, never a session, and answer +//! `{digest, index, chunks, bytes, data}` where `data` is standard padded +//! base64 of at most [`CHUNK_BYTES`] decoded bytes, `chunks` is the total +//! chunk count and `bytes` is the whole asset's length. +//! +//! Only the authenticated principal `reserved:occipital` may call them. Every +//! refusal is an [`ErrorBody`] carrying one of the refusal codes listed in +//! `tests/fixtures/rich_ui/transport-contract.json` (`chunks.refusalCodes`), +//! with the fixed message from [`operations::fixed_message`], so no request +//! data ever appears in a refusal. + +use std::ops::Deref; +use std::sync::{Mutex, TryLockError}; + +use base64::Engine as _; +use serde_json::{json, Value}; +use subc_protocol::ErrorBody; + +use super::builder::CHUNK_BYTES; +use super::operations; + +pub(crate) const MANIFEST_CHUNK: &str = "rich_ui.manifest_chunk"; +pub(crate) const TEMPLATE_CHUNK: &str = "rich_ui.template_chunk"; + +/// The only principal admitted to the chunk queries (and the rich-UI carrier). +pub(crate) const OCCIPITAL_PRINCIPAL: &str = "reserved:occipital"; + +pub(crate) fn is_chunk_operation(operation: &str) -> bool { + matches!(operation, MANIFEST_CHUNK | TEMPLATE_CHUNK) +} + +/// A refusal with its code and the code's fixed, request-free message. +pub(crate) fn refusal(code: &str) -> ErrorBody { + ErrorBody::new(code, operations::fixed_message(code)) +} + +/// Refuse every principal other than Occipital, before any lookup runs. +pub(crate) fn admit_principal(principal_id: Option<&str>) -> Result<(), ErrorBody> { + if principal_id == Some(OCCIPITAL_PRINCIPAL) { + Ok(()) + } else { + Err(refusal("principal_not_allowed")) + } +} + +/// The manifest bytes the running build publishes, with their digest. The +/// production builder does not publish a manifest yet, so this is `None` +/// outside tests; tests inject the golden manifest through +/// [`set_test_manifest_for_test`]. +fn running_manifest() -> Option<(String, Vec)> { + #[cfg(test)] + { + if let Some(manifest) = TEST_MANIFEST + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + { + return Some(manifest); + } + } + super::builder::production_manifest().map(|manifest| (manifest.digest, manifest.bytes)) +} + +#[cfg(test)] +static TEST_MANIFEST: Mutex)>> = Mutex::new(None); + +/// Serializes tests that inject a manifest, since the injection is +/// process-global; the guard clears the injection when dropped. +#[cfg(test)] +static TEST_MANIFEST_SERIAL: Mutex<()> = Mutex::new(()); + +#[cfg(test)] +pub(crate) struct TestManifestGuard { + _serial: std::sync::MutexGuard<'static, ()>, +} + +#[cfg(test)] +impl Drop for TestManifestGuard { + fn drop(&mut self) { + *TEST_MANIFEST + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = None; + } +} + +/// Hold the returned guard while the injected manifest must stay in place. +/// `None` asserts the production state: no manifest at all. +#[cfg(test)] +pub(crate) fn set_test_manifest_for_test(bytes: Option>) -> TestManifestGuard { + let serial = TEST_MANIFEST_SERIAL + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *TEST_MANIFEST + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + bytes.map(|bytes| (super::builder::digest(&bytes), bytes)); + TestManifestGuard { _serial: serial } +} + +/// Template bytes embedded in the binary, by digest. The binary embeds no +/// template yet, so this lookup always returns `None` and every template is +/// served from aft.db. +fn shipped_template(_digest: &str) -> Option<&'static [u8]> { + None +} + +/// The request's digest and index. A digest that is missing or not a string +/// names no asset (`digest_unknown`); an index that is missing or not a +/// non-negative integer names no chunk (`chunk_index_out_of_range`, checked +/// once the asset is found). A malformed request therefore gets one of the +/// same refusal codes as a well-formed one, rather than a separate code. +fn request(body: &Value, digest_field: &str) -> Result<(String, Option), ErrorBody> { + let digest = body + .get(digest_field) + .and_then(Value::as_str) + .ok_or_else(|| refusal("digest_unknown"))?; + Ok((digest.to_owned(), body.get("index").and_then(Value::as_u64))) +} + +/// Build the reply for chunk `index` of `bytes`, whose digest is `digest`. +pub(crate) fn chunk_reply( + digest: &str, + bytes: &[u8], + index: Option, +) -> Result { + let chunks = bytes.len().div_ceil(CHUNK_BYTES); + let index = index + .and_then(|index| usize::try_from(index).ok()) + .filter(|index| *index < chunks) + .ok_or_else(|| refusal("chunk_index_out_of_range"))?; + let start = index * CHUNK_BYTES; + let end = (start + CHUNK_BYTES).min(bytes.len()); + Ok(json!({ + "digest": digest, + "index": index, + "chunks": chunks, + "bytes": bytes.len(), + "data": base64::engine::general_purpose::STANDARD.encode(&bytes[start..end]), + })) +} + +/// Answer `rich_ui.manifest_chunk`. It compares only against the running +/// build's manifest and never opens the database: with no manifest, or for +/// any other digest, it refuses `digest_unknown`. +pub(crate) fn manifest_chunk(body: &Value) -> Result { + let (digest, index) = request(body, "manifestDigest")?; + let (manifest_digest, bytes) = running_manifest() + .filter(|(manifest_digest, _)| *manifest_digest == digest) + .ok_or_else(|| refusal("digest_unknown"))?; + chunk_reply(&manifest_digest, &bytes, index) +} + +/// Answer `rich_ui.template_chunk`. The binary's shipped templates are +/// checked first and need no database. Otherwise the stored templates in the +/// process's one aft.db are read without waiting for its lock: no handle or a +/// poisoned mutex refuses `database_unavailable`, a held mutex refuses +/// `database_busy`, and only a lookup that ran and found nothing refuses +/// `digest_unknown`. A failed query is also `database_unavailable`, because +/// an unreadable database must never be reported as a missing digest. +pub(crate) fn template_chunk(body: &Value, db: Option<&Mutex>) -> Result +where + C: Deref, +{ + let (digest, index) = request(body, "templateDigest")?; + if let Some(bytes) = shipped_template(&digest) { + return chunk_reply(&digest, bytes, index); + } + let db = db.ok_or_else(|| refusal("database_unavailable"))?; + let conn = match db.try_lock() { + Ok(conn) => conn, + Err(TryLockError::WouldBlock) => return Err(refusal("database_busy")), + Err(TryLockError::Poisoned(_)) => return Err(refusal("database_unavailable")), + }; + let stored = crate::db::rich_ui::template_bytes(&conn, &digest) + .map_err(|_| refusal("database_unavailable"))?; + drop(conn); + let bytes = stored.ok_or_else(|| refusal("digest_unknown"))?; + chunk_reply(&digest, &bytes, index) +} + +#[cfg(test)] +#[path = "chunks/tests.rs"] +pub(crate) mod tests; diff --git a/crates/aft/src/subc/rich_ui/chunks/route_tests.rs b/crates/aft/src/subc/rich_ui/chunks/route_tests.rs new file mode 100644 index 000000000..e1771a639 --- /dev/null +++ b/crates/aft/src/subc/rich_ui/chunks/route_tests.rs @@ -0,0 +1,571 @@ +//! The rich-UI chunk queries and Occipital's route shapes, driven through the +//! real module frame loop: binds go through the channel-0 RouteBind arm and +//! requests through the routed-request arm, exactly as the daemon sends them. + +use std::collections::BTreeMap; + +use base64::Engine as _; +use sha2::{Digest, Sha256}; + +use super::rich_ui::chunks; +use super::*; + +const GOLDEN_MANIFEST: &[u8] = + include_bytes!("../../../../tests/fixtures/rich_ui/manifest.jcs.json"); + +fn sha256(bytes: &[u8]) -> String { + format!("sha256:{:x}", Sha256::digest(bytes)) +} + +/// Removing fixture actors releases process-wide root-cache read markers. +struct FixtureActors(Arc); + +impl Drop for FixtureActors { + fn drop(&mut self) { + for (root, ctx) in self.0.actor_entries() { + ctx.bash_background().shutdown(); + ctx.clear_db(); + self.0.remove_actor(&root); + } + } +} + +struct Wire { + daemon: tokio::io::DuplexStream, + module: Option>>, + ctx: Arc, + root: PathBuf, + _actors: FixtureActors, + _project: tempfile::TempDir, +} + +impl Wire { + async fn start() -> Self { + let (project, root_id) = test_support::test_root("rich-ui-chunk-route"); + let root = project.path().to_path_buf(); + let mut config = crate::config::Config::default(); + config.project_root = Some(root.clone()); + config.storage_dir = Some(root.join(".storage")); + let ctx = AppContext::new(Box::new(crate::parser::TreeSitterProvider::new()), config); + std::fs::create_dir_all(ctx.storage_dir()).unwrap(); + ctx.set_db(Arc::new(StdMutex::new( + crate::db::open(&ctx.storage_dir().join("aft.db")).unwrap(), + ))); + let ctx = Arc::new(ctx); + let executor = Arc::new(Executor::new()); + let actors = FixtureActors(executor.clone()); + executor.register_actor(root_id, ctx.clone()); + let app = ctx.app().clone(); + let module_root = root.clone(); + let storage = ctx.storage_dir().to_path_buf(); + let (daemon, module) = tokio::io::duplex(1024 * 1024); + let module = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (read, write) = tokio::io::split(module); + runtime.block_on(run_module_loop( + read, + write, + &module_root.join("absent-connection.json"), + app, + executor, + |request, ctx| { + if request.command == "configure" { + crate::commands::configure::handle_configure(&request, ctx) + } else { + Response::success(request.id, json!({"executed": true})) + } + }, + Some(test_support::bind_config_without_indexes(&module_root)), + false, + usize::MAX, + None, + &storage, + QuiescedLsp::ReapInBackground, + None, + None, + )) + }); + let mut wire = Self { + daemon, + module: Some(module), + ctx, + root, + _actors: actors, + _project: project, + }; + let hello = wire.receive(0, HELLO_CORR).await; + assert_eq!(hello.header.ty, FrameType::Hello); + wire.send(FrameType::HelloAck, 0, HELLO_CORR, json!({})) + .await; + wire + } + + async fn send(&mut self, ty: FrameType, channel: u16, corr: u64, body: Value) { + let frame = Frame::build( + ty, + control_flags(), + channel, + u32::from(channel != 0), + corr, + if matches!(ty, FrameType::Goodbye) { + Vec::new() + } else { + serde_json::to_vec(&body).unwrap() + }, + ) + .unwrap(); + write_frame(&mut self.daemon, &frame).await.unwrap(); + } + + async fn receive(&mut self, channel: u16, corr: u64) -> Frame { + tokio::time::timeout(Duration::from_secs(20), async { + loop { + let frame = read_frame(&mut self.daemon).await.unwrap().unwrap(); + if frame.header.channel == channel && frame.header.corr == corr { + return frame; + } + } + }) + .await + .expect("route reply timed out") + } + + /// Bind `channel` and return the bind reply frame. + async fn bind( + &mut self, + channel: u16, + target: RouteTarget, + principal: Principal, + scope: Option, + v1: bool, + ) -> Frame { + let bind = ModuleControlRequest::RouteBind { + route_channel: channel, + epoch: 1, + target, + identity: subc_protocol::BindIdentity::new( + &self.root, + "runner", + &format!("session-{channel}"), + ), + principal: Some(principal), + consumer_capabilities: None, + admission_facts: Default::default(), + scope, + role_versions: v1 + .then(|| BTreeMap::from([("tool-provider".to_string(), "v1".to_string())])), + call_mode: None, + }; + let corr = 1_000 + u64::from(channel); + self.send( + FrameType::Request, + 0, + corr, + serde_json::to_value(bind).unwrap(), + ) + .await; + self.receive(0, corr).await + } + + async fn bind_provider( + &mut self, + channel: u16, + principal: Principal, + scope: Option, + ) { + let ack = self + .bind( + channel, + RouteTarget::ToolProvider { + module_id: "aft".into(), + }, + principal, + scope, + true, + ) + .await; + assert!( + matches!( + serde_json::from_slice::(&ack.body), + Ok(ModuleControlResponse::RouteBindAck {}) + ), + "channel {channel} bind: {:?}", + String::from_utf8_lossy(&ack.body) + ); + } + + async fn request(&mut self, channel: u16, corr: u64, body: Value) -> (FrameType, Value) { + self.send(FrameType::Request, channel, corr, body).await; + let frame = self.receive(channel, corr).await; + ( + frame.header.ty, + serde_json::from_slice(&frame.body).unwrap(), + ) + } + + async fn stop(mut self) { + self.send(FrameType::Goodbye, 0, 9_999, json!({})).await; + self.module.take().unwrap().join().unwrap().unwrap(); + } +} + +fn occipital() -> Principal { + Principal::Reserved { + module_id: "occipital".into(), + } +} + +fn stamp(owner: Principal, scope_ref: &str) -> subc_protocol::scope::ScopeStamp { + serde_json::from_value(json!({ + "owner": owner, "ref": scope_ref, "scope_epoch": 1, + "kind": "head", "owner_authorized": true, + })) + .unwrap() +} + +/// Assert a typed protocol Error frame carrying `code` and its fixed message. +fn assert_refusal(reply: (FrameType, Value), code: &str) { + let (ty, body) = reply; + assert_eq!(ty, FrameType::Error, "{body}"); + assert_eq!(body["code"], code, "{body}"); + assert_eq!( + body["message"], + super::rich_ui::operations::fixed_message(code), + "{body}" + ); + assert!( + body.get("content").is_none() && body.get("structuredContent").is_none(), + "a chunk refusal is never a tool-result body: {body}" + ); +} + +const UNSCOPED: u16 = 1; +const SELF_OWNED: u16 = 2; +const CARRIED: u16 = 3; +const BROCA_UNSCOPED: u16 = 4; +const BROCA_SCOPED: u16 = 5; + +async fn occipital_wire() -> Wire { + let mut wire = Wire::start().await; + wire.bind_provider(UNSCOPED, occipital(), None).await; + wire.bind_provider( + SELF_OWNED, + occipital(), + Some(stamp(occipital(), "occipital-own")), + ) + .await; + let core = Principal::Reserved { + module_id: "prefrontal-core".into(), + }; + wire.bind_provider(CARRIED, occipital(), Some(stamp(core.clone(), "head-s"))) + .await; + let broca = Principal::Reserved { + module_id: "broca".into(), + }; + wire.bind_provider(BROCA_UNSCOPED, broca.clone(), None) + .await; + wire.bind_provider(BROCA_SCOPED, broca, Some(stamp(core, "worker-w"))) + .await; + wire +} + +#[tokio::test] +async fn chunk_queries_answer_occipital_on_every_route_shape_and_refuse_others() { + let _watcher_disabled = crate::test_env::without_project_watcher(); + let _manifest = chunks::set_test_manifest_for_test(Some(GOLDEN_MANIFEST.to_vec())); + let mut wire = occipital_wire().await; + let manifest_digest = sha256(GOLDEN_MANIFEST); + + // A 300,000-byte stored template whose first chunk boundary splits a + // two-byte UTF-8 character. + let asset = chunks::tests::split_utf8_asset(); + let template_digest = sha256(&asset); + let db = wire.ctx.app().db().expect("the shared App holds aft.db"); + db.lock() + .unwrap() + .execute( + "INSERT INTO rich_ui_templates (digest, bytes) VALUES (?1, ?2)", + rusqlite::params![template_digest, asset], + ) + .unwrap(); + + let mut corr = 100; + for channel in [UNSCOPED, SELF_OWNED, CARRIED] { + let mut replies = Vec::new(); + for index in 0..3u64 { + corr += 1; + let (ty, reply) = wire + .request( + channel, + corr, + json!({"op": "rich_ui.template_chunk", "templateDigest": template_digest, "index": index}), + ) + .await; + assert_eq!(ty, FrameType::Response, "{reply}"); + replies.push(reply); + } + let mut sizes = Vec::new(); + let joined = chunks::tests::reassemble(&template_digest, |index| { + let reply = replies[index as usize].clone(); + sizes.push( + base64::engine::general_purpose::STANDARD + .decode(reply["data"].as_str().unwrap()) + .unwrap() + .len(), + ); + reply + }); + assert_eq!(replies[0]["chunks"], 3); + assert_eq!(replies[0]["bytes"], 300_000); + assert_eq!(sizes, [131_072, 131_072, 37_856]); + assert_eq!(joined, asset); + assert_eq!(sha256(&joined), template_digest); + + corr += 1; + let (ty, reply) = wire + .request( + channel, + corr, + json!({"op": "rich_ui.manifest_chunk", "manifestDigest": manifest_digest, "index": 0}), + ) + .await; + assert_eq!(ty, FrameType::Response, "{reply}"); + let joined = chunks::tests::reassemble(&manifest_digest, |_| reply.clone()); + assert_eq!(joined, GOLDEN_MANIFEST); + + // The bare form names the query in `name` with an `arguments` body. + corr += 1; + let (ty, reply) = wire + .request( + channel, + corr, + json!({"name": "rich_ui.template_chunk", "arguments": {"templateDigest": template_digest, "index": 2}}), + ) + .await; + assert_eq!(ty, FrameType::Response, "{reply}"); + assert_eq!(reply["index"], 2); + + corr += 1; + assert_refusal( + wire.request( + channel, + corr, + json!({"op": "rich_ui.template_chunk", "templateDigest": template_digest, "index": 3}), + ) + .await, + "chunk_index_out_of_range", + ); + corr += 1; + assert_refusal( + wire.request( + channel, + corr, + json!({"op": "rich_ui.template_chunk", "templateDigest": sha256(b"missing"), "index": 0}), + ) + .await, + "digest_unknown", + ); + corr += 1; + assert_refusal( + wire.request( + channel, + corr, + json!({"op": "rich_ui.manifest_chunk", "manifestDigest": sha256(b"missing"), "index": 0}), + ) + .await, + "digest_unknown", + ); + } + + // Any other principal is refused before a lookup, scoped or unscoped. + for channel in [BROCA_UNSCOPED, BROCA_SCOPED] { + for body in [ + json!({"op": "rich_ui.manifest_chunk", "manifestDigest": manifest_digest, "index": 0}), + json!({"op": "rich_ui.template_chunk", "templateDigest": template_digest, "index": 0}), + json!({"name": "rich_ui.template_chunk", "arguments": {"templateDigest": template_digest, "index": 0}}), + ] { + corr += 1; + assert_refusal( + wire.request(channel, corr, body).await, + "principal_not_allowed", + ); + } + } + + // A held database lock is busy, never a miss; the lookup does not wait. + { + let _held = db.lock().unwrap(); + for channel in [UNSCOPED, CARRIED] { + corr += 1; + assert_refusal( + wire.request( + channel, + corr, + json!({"op": "rich_ui.template_chunk", "templateDigest": template_digest, "index": 0}), + ) + .await, + "database_busy", + ); + } + // The manifest query never touches the database, so it still answers. + corr += 1; + let (ty, reply) = wire + .request( + UNSCOPED, + corr, + json!({"op": "rich_ui.manifest_chunk", "manifestDigest": manifest_digest, "index": 0}), + ) + .await; + assert_eq!(ty, FrameType::Response, "{reply}"); + } + + // With no database handle the template lookup is unavailable, not a miss. + drop(db); + wire.ctx.clear_db(); + for channel in [UNSCOPED, CARRIED] { + corr += 1; + assert_refusal( + wire.request( + channel, + corr, + json!({"op": "rich_ui.template_chunk", "templateDigest": template_digest, "index": 0}), + ) + .await, + "database_unavailable", + ); + } + wire.stop().await; +} + +#[tokio::test] +async fn manifest_chunk_before_publication_refuses_digest_unknown_on_the_route() { + let _watcher_disabled = crate::test_env::without_project_watcher(); + let _manifest = chunks::set_test_manifest_for_test(None); + let mut wire = occipital_wire().await; + for (corr, channel) in [(10, UNSCOPED), (11, CARRIED)] { + assert_refusal( + wire.request( + channel, + corr, + json!({"op": "rich_ui.manifest_chunk", "manifestDigest": sha256(GOLDEN_MANIFEST), "index": 0}), + ) + .await, + "digest_unknown", + ); + } + wire.stop().await; +} + +#[tokio::test] +async fn non_carrier_occipital_routes_admit_provider_reads_only() { + let _watcher_disabled = crate::test_env::without_project_watcher(); + let mut wire = occipital_wire().await; + let mut corr = 200; + for channel in [UNSCOPED, SELF_OWNED] { + for body in [ + json!({"op": "role.describe"}), + json!({"op": "tool.catalog"}), + ] { + corr += 1; + let (ty, reply) = wire.request(channel, corr, body.clone()).await; + assert_eq!(ty, FrameType::Response, "{body}: {reply}"); + } + // Every tool-call form refuses `session_not_bound`, before the carrier + // gate would have judged the tool (so `write` gets the same code). + for body in [ + json!({"op": "tool.call", "name": "read", "arguments": {"path": "x"}}), + json!({"name": "read", "arguments": {"path": "x"}}), + json!({"op": "tool.call", "name": "write", "arguments": {"path": "x", "content": "x"}}), + json!({"op": "tool.call", "name": "bash_status", "arguments": {"taskId": "t"}}), + ] { + corr += 1; + let (ty, reply) = wire.request(channel, corr, body.clone()).await; + assert_eq!(ty, FrameType::Error, "{body}: {reply}"); + assert_eq!(reply["code"], "session_not_bound", "{body}: {reply}"); + } + } + // A carried route is not refused here: its tool calls continue to the + // carrier gate, which refuses an unlisted tool with its own code. + corr += 1; + let (ty, reply) = wire + .request( + CARRIED, + corr, + json!({"op": "tool.call", "name": "write", "arguments": {"path": "x", "content": "x"}}), + ) + .await; + assert_eq!(ty, FrameType::Error, "{reply}"); + assert_eq!(reply["code"], "carrier_read_only", "{reply}"); + wire.stop().await; +} + +#[tokio::test] +async fn occipital_legacy_and_management_binds_are_refused() { + let _watcher_disabled = crate::test_env::without_project_watcher(); + let mut wire = Wire::start().await; + let core = Principal::Reserved { + module_id: "prefrontal-core".into(), + }; + let mut channel = 10; + for scope in [ + None, + Some(stamp(occipital(), "occipital-own")), + Some(stamp(core, "head-s")), + ] { + for (target, v1) in [ + ( + RouteTarget::ToolProvider { + module_id: "aft".into(), + }, + false, + ), + ( + RouteTarget::ManagementSurface { + module_id: "aft".into(), + }, + false, + ), + ( + RouteTarget::ManagementSurface { + module_id: "aft".into(), + }, + true, + ), + ] { + channel += 1; + let reply = wire + .bind(channel, target, occipital(), scope.clone(), v1) + .await; + assert_eq!(reply.header.ty, FrameType::Error); + let body: Value = serde_json::from_slice(&reply.body).unwrap(); + assert_eq!(body["code"], "route_refused", "{body}"); + } + } + // The same legacy bind by another principal is still accepted. + channel += 1; + let ack = wire + .bind( + channel, + RouteTarget::ToolProvider { + module_id: "aft".into(), + }, + Principal::Reserved { + module_id: "broca".into(), + }, + None, + false, + ) + .await; + assert!( + matches!( + serde_json::from_slice::(&ack.body), + Ok(ModuleControlResponse::RouteBindAck {}) + ), + "{:?}", + String::from_utf8_lossy(&ack.body) + ); + wire.stop().await; +} diff --git a/crates/aft/src/subc/rich_ui/chunks/tests.rs b/crates/aft/src/subc/rich_ui/chunks/tests.rs new file mode 100644 index 000000000..515df7514 --- /dev/null +++ b/crates/aft/src/subc/rich_ui/chunks/tests.rs @@ -0,0 +1,250 @@ +use std::sync::{Arc, Mutex}; + +use rusqlite::Connection; +use serde_json::{json, Value}; +use sha2::{Digest, Sha256}; + +use super::*; + +const GOLDEN_MANIFEST: &[u8] = + include_bytes!("../../../../tests/fixtures/rich_ui/manifest.jcs.json"); + +fn sha256(bytes: &[u8]) -> String { + format!("sha256:{:x}", Sha256::digest(bytes)) +} + +fn migrated() -> Mutex> { + let mut conn = Connection::open_in_memory().unwrap(); + crate::db::run_migrations(&mut conn).unwrap(); + Mutex::new(Box::new(conn)) +} + +fn store_template(db: &Mutex>, bytes: &[u8]) -> String { + let digest = sha256(bytes); + db.lock() + .unwrap() + .execute( + "INSERT INTO rich_ui_templates (digest, bytes) VALUES (?1, ?2)", + rusqlite::params![digest, bytes], + ) + .unwrap(); + digest +} + +fn code(result: Result) -> String { + let error = result.expect_err("expected a refusal"); + assert_eq!( + error.message, + operations::fixed_message(&error.code), + "every chunk refusal carries its code's fixed message" + ); + error.code +} + +/// A 300,000-byte asset whose first chunk boundary (byte 131,072) falls +/// between the two bytes of a two-byte UTF-8 character. +pub(crate) fn split_utf8_asset() -> Vec { + let mut asset = b"a".to_vec(); + while asset.len() + 2 <= 300_000 { + asset.extend_from_slice("é".as_bytes()); + } + asset.resize(300_000, b'a'); + assert_eq!(asset.len(), 300_000); + assert_eq!(asset[CHUNK_BYTES - 1], 0xC3, "a chunk ends on a lead byte"); + assert_eq!( + asset[CHUNK_BYTES], 0xA9, + "the next chunk starts mid-character" + ); + asset +} + +/// Fetch every chunk of `digest` through `fetch`, checking each reply's +/// shape, and return the reassembled bytes. +pub(crate) fn reassemble(digest: &str, mut fetch: impl FnMut(u64) -> Value) -> Vec { + let first = fetch(0); + let chunks = first["chunks"].as_u64().unwrap(); + let total = first["bytes"].as_u64().unwrap(); + let mut joined = Vec::new(); + for index in 0..chunks { + let reply = if index == 0 { + first.clone() + } else { + fetch(index) + }; + let mut fields: Vec<&str> = reply + .as_object() + .unwrap() + .keys() + .map(String::as_str) + .collect(); + fields.sort_unstable(); + assert_eq!(fields, ["bytes", "chunks", "data", "digest", "index"]); + assert_eq!(reply["digest"], digest); + assert_eq!(reply["index"], index); + assert_eq!(reply["chunks"], chunks); + assert_eq!(reply["bytes"], total); + let data = base64::engine::general_purpose::STANDARD + .decode(reply["data"].as_str().unwrap()) + .unwrap(); + assert!(data.len() <= CHUNK_BYTES); + joined.extend(data); + } + assert_eq!(joined.len() as u64, total); + joined +} + +#[test] +fn chunk_queries_admit_only_the_occipital_principal() { + assert!(admit_principal(Some("reserved:occipital")).is_ok()); + for principal in [ + None, + Some("reserved:prefrontal-core"), + Some("reserved:basal"), + Some("direct"), + Some("unverified"), + Some("reserved:occipital-shadow"), + ] { + assert_eq!( + code(admit_principal(principal).map(|()| Value::Null)), + "principal_not_allowed", + "{principal:?}" + ); + } +} + +#[test] +fn chunk_operation_names_are_exact() { + assert!(is_chunk_operation("rich_ui.manifest_chunk")); + assert!(is_chunk_operation("rich_ui.template_chunk")); + assert!(!is_chunk_operation("rich_ui.chunk")); + assert!(!is_chunk_operation("tool.catalog")); +} + +#[test] +fn manifest_chunk_before_publication_refuses_every_digest() { + let _guard = set_test_manifest_for_test(None); + for digest in [ + json!(sha256(GOLDEN_MANIFEST)), + json!(sha256(b"")), + json!("not-a-digest"), + ] { + assert_eq!( + code(manifest_chunk( + &json!({"manifestDigest": digest, "index": 0}) + )), + "digest_unknown" + ); + } +} + +#[test] +fn manifest_chunk_serves_the_running_manifest_and_checks_the_index() { + let _guard = set_test_manifest_for_test(Some(GOLDEN_MANIFEST.to_vec())); + let digest = sha256(GOLDEN_MANIFEST); + assert_eq!( + digest, + include_str!("../../../../tests/fixtures/rich_ui/manifest.sha256").trim() + ); + let joined = reassemble(&digest, |index| { + manifest_chunk(&json!({"manifestDigest": digest, "index": index})).unwrap() + }); + assert_eq!(joined, GOLDEN_MANIFEST); + assert_eq!(sha256(&joined), digest); + let chunks = GOLDEN_MANIFEST.len().div_ceil(CHUNK_BYTES) as u64; + for index in [ + json!(chunks), + json!(-1), + json!(1.5), + json!("0"), + Value::Null, + ] { + assert_eq!( + code(manifest_chunk( + &json!({"manifestDigest": digest, "index": index}) + )), + "chunk_index_out_of_range", + "{index}" + ); + } + assert_eq!( + code(manifest_chunk( + &json!({"manifestDigest": sha256(b"other"), "index": 0}) + )), + "digest_unknown" + ); + assert_eq!(code(manifest_chunk(&json!({"index": 0}))), "digest_unknown"); +} + +#[test] +fn template_chunk_splits_a_300000_byte_asset_into_three_chunks() { + let db = migrated(); + let asset = split_utf8_asset(); + let digest = store_template(&db, &asset); + let mut sizes = Vec::new(); + let joined = reassemble(&digest, |index| { + let reply = template_chunk( + &json!({"templateDigest": digest, "index": index}), + Some(&db), + ) + .unwrap(); + sizes.push( + base64::engine::general_purpose::STANDARD + .decode(reply["data"].as_str().unwrap()) + .unwrap() + .len(), + ); + reply + }); + assert_eq!(sizes, [131_072, 131_072, 37_856]); + assert_eq!(joined, asset); + assert_eq!(sha256(&joined), digest); + assert_eq!( + code(template_chunk( + &json!({"templateDigest": digest, "index": 3}), + Some(&db) + )), + "chunk_index_out_of_range" + ); +} + +#[test] +fn template_chunk_reports_database_state_rather_than_a_miss() { + let request = json!({"templateDigest": sha256(b"never stored"), "index": 0}); + + // No database handle at all. + assert_eq!( + code(template_chunk::>(&request, None)), + "database_unavailable" + ); + + // The handle's lock is held: a busy database, never a miss. + let db = migrated(); + { + let _held = db.lock().unwrap(); + assert_eq!(code(template_chunk(&request, Some(&db))), "database_busy"); + } + + // A readable database that lacks the digest is the only real miss. + assert_eq!(code(template_chunk(&request, Some(&db))), "digest_unknown"); + + // A poisoned mutex cannot be trusted to hold a consistent connection. + let poisoned = Arc::new(migrated()); + let thread_db = poisoned.clone(); + let _ = std::thread::spawn(move || { + let _guard = thread_db.lock().unwrap(); + panic!("poison the database mutex"); + }) + .join(); + assert!(poisoned.is_poisoned()); + assert_eq!( + code(template_chunk(&request, Some(&*poisoned))), + "database_unavailable" + ); + + // A database whose query fails (here: no template table) is unreadable. + let unmigrated = Mutex::new(Box::new(Connection::open_in_memory().unwrap())); + assert_eq!( + code(template_chunk(&request, Some(&unmigrated))), + "database_unavailable" + ); +} diff --git a/crates/aft/src/subc/rich_ui/mod.rs b/crates/aft/src/subc/rich_ui/mod.rs index d6065b070..1c1ab01bc 100644 --- a/crates/aft/src/subc/rich_ui/mod.rs +++ b/crates/aft/src/subc/rich_ui/mod.rs @@ -14,14 +14,18 @@ //! - `operations`: versioned schemas, closed projections and safe messages. //! - `builder`: deterministic assembly and caps, with no production manifest //! until the publication prerequisites are reviewed. +//! - `chunks`: the `rich_ui.manifest_chunk` and `rich_ui.template_chunk` +//! provider reads that serve manifest and template bytes to Occipital. //! -//! Nothing here touches the daemon, the database or the wire. Until publication -//! the only users are tests, so the dead-code allowances are module-local. +//! Apart from `chunks`, nothing here touches the daemon, the database or the +//! wire. Until publication the other modules' only users are tests, so their +//! dead-code allowances are module-local. #[cfg_attr(not(test), allow(dead_code))] pub(crate) mod builder; #[cfg_attr(not(test), allow(dead_code))] pub(crate) mod canonical; +pub(crate) mod chunks; #[cfg_attr(not(test), allow(dead_code))] pub(crate) mod jcs_number; #[cfg_attr(not(test), allow(dead_code))]