Rich UI: chunk queries for the manifest and stored templates (Occipital)

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."
+}

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<subc_protocol::ErrorBody> {
+    (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<crate::db::TrackedConnection>>,
+) -> 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;
     }

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<u8>)> {
+    #[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<Option<(String, Vec<u8>)>> = 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<Vec<u8>>) -> 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<u64>), 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<u64>,
+) -> Result<Value, ErrorBody> {
+    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<Value, ErrorBody> {
+    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<C>(body: &Value, db: Option<&Mutex<C>>) -> Result<Value, ErrorBody>
+where
+    C: Deref<Target = rusqlite::Connection>,
+{
+    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;