adk_server/rest/controllers/

ui.rs

use crate::ui_protocol::{
    TOOL_ENVELOPE_VERSION, UI_DEFAULT_PROTOCOL, UI_PROTOCOL_CAPABILITIES,
    UiProtocolDeprecationSpec, UiProtocolImplementationTier, UiProtocolSpecTrack,
};
use crate::ui_types::{
    McpAppsRenderOptions, McpUiBridgeSnapshot, McpUiHostCapabilities, McpUiHostInfo,
    McpUiPermissions, McpUiResourceCsp, default_mcp_ui_host_capabilities, default_mcp_ui_host_info,
    validate_mcp_apps_render_options,
};
use axum::Extension;
use axum::{Json, extract::Query, http::StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::collections::HashMap;
use std::hash::Hash;
use std::sync::{OnceLock, RwLock};
use tracing::{info, warn};

#[derive(Debug, Clone, Serialize)]
pub struct UiProtocolCapability {
    pub protocol: &'static str,
    pub versions: Vec<&'static str>,
    #[serde(rename = "implementationTier")]
    pub implementation_tier: UiProtocolImplementationTier,
    #[serde(rename = "specTrack")]
    pub spec_track: UiProtocolSpecTrack,
    pub summary: &'static str,
    pub features: Vec<&'static str>,
    pub limitations: Vec<&'static str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub deprecation: Option<UiProtocolDeprecation>,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UiProtocolDeprecation {
    pub stage: &'static str,
    pub announced_on: &'static str,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub sunset_target_on: Option<&'static str>,
    pub replacement_protocols: Vec<&'static str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub note: Option<&'static str>,
}

fn map_deprecation(
    spec: Option<&'static UiProtocolDeprecationSpec>,
) -> Option<UiProtocolDeprecation> {
    let spec = spec?;
    Some(UiProtocolDeprecation {
        stage: spec.stage,
        announced_on: spec.announced_on,
        sunset_target_on: spec.sunset_target_on,
        replacement_protocols: spec.replacement_protocols.to_vec(),
        note: spec.note,
    })
}

#[derive(Debug, Clone, Serialize)]
pub struct UiCapabilities {
    pub default_protocol: &'static str,
    pub protocols: Vec<UiProtocolCapability>,
    pub tool_envelope_version: &'static str,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UiResource {
    pub uri: String,
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    pub mime_type: String,
    #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
    pub meta: Option<Value>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UiResourceContent {
    pub uri: String,
    pub mime_type: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub text: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub blob: Option<String>,
    #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
    pub meta: Option<Value>,
}

#[derive(Debug, Clone, Serialize)]
pub struct UiResourceListResponse {
    pub resources: Vec<UiResource>,
}

#[derive(Debug, Clone, Serialize)]
pub struct UiResourceReadResponse {
    pub contents: Vec<UiResourceContent>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct RegisterUiResourceRequest {
    pub uri: String,
    pub name: String,
    #[serde(default)]
    pub description: Option<String>,
    pub mime_type: String,
    pub text: String,
    #[serde(rename = "_meta", default)]
    pub meta: Option<Value>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct ReadUiResourceQuery {
    pub uri: String,
}

#[derive(Debug, Clone)]
struct UiResourceEntry {
    resource: UiResource,
    content: UiResourceContent,
    /// The authenticated user that registered this resource, when the server has
    /// authentication configured. `None` means the resource was registered by an
    /// unauthenticated server and is visible to everyone, which preserves the
    /// behaviour of a server that has no auth to enforce.
    owner: Option<String>,
}

static UI_RESOURCE_REGISTRY: OnceLock<RwLock<HashMap<String, UiResourceEntry>>> = OnceLock::new();

const MCP_UI_DEFAULT_PROTOCOL_VERSION: &str = "2025-11-25";

fn resource_registry() -> &'static RwLock<HashMap<String, UiResourceEntry>> {
    UI_RESOURCE_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()))
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct McpUiBridgeSessionKey {
    app_name: String,
    user_id: String,
    session_id: String,
}

#[derive(Debug, Clone)]
struct McpUiBridgeSessionEntry {
    protocol_version: String,
    initialized: bool,
    app_info: Option<Value>,
    app_capabilities: Option<Value>,
    host_info: McpUiHostInfo,
    host_capabilities: McpUiHostCapabilities,
    host_context: Value,
    message_count: u64,
    last_message: Option<Value>,
    model_context: Vec<Value>,
    model_context_revision: u64,
    resource_list_revision: u64,
    tool_list_revision: u64,
    notification_count: u64,
    pending_notifications: Vec<McpUiBridgeNotification>,
}

static MCP_UI_BRIDGE_REGISTRY: OnceLock<
    RwLock<HashMap<McpUiBridgeSessionKey, McpUiBridgeSessionEntry>>,
> = OnceLock::new();

fn bridge_registry() -> &'static RwLock<HashMap<McpUiBridgeSessionKey, McpUiBridgeSessionEntry>> {
    MCP_UI_BRIDGE_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()))
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiBridgeNotification {
    pub notification_id: u64,
    pub method: String,
    pub revision: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<Value>,
}

fn bridge_session_key(app_name: &str, user_id: &str, session_id: &str) -> McpUiBridgeSessionKey {
    McpUiBridgeSessionKey {
        app_name: app_name.to_string(),
        user_id: user_id.to_string(),
        session_id: session_id.to_string(),
    }
}

fn default_host_context(app_name: &str, user_id: &str, session_id: &str) -> Value {
    json!({
        "appName": app_name,
        "userId": user_id,
        "sessionId": session_id,
        "theme": "light",
        "locale": "en-US",
        "timeZone": "UTC",
        "platform": "adk-server",
        "displayMode": "inline",
        "availableDisplayModes": ["inline"]
    })
}

fn merge_host_context(
    target: &mut Value,
    patch: Option<Value>,
) -> Result<(), (StatusCode, String)> {
    let Some(patch) = patch else {
        return Ok(());
    };
    let patch_object = patch.as_object().ok_or_else(|| {
        (StatusCode::BAD_REQUEST, "hostContext must be a JSON object".to_string())
    })?;
    let target_object = target.as_object_mut().ok_or_else(|| {
        (StatusCode::INTERNAL_SERVER_ERROR, "host context store invalid".to_string())
    })?;
    for (key, value) in patch_object {
        target_object.insert(key.clone(), value.clone());
    }
    Ok(())
}

fn ensure_bridge_session<'a>(
    registry: &'a mut HashMap<McpUiBridgeSessionKey, McpUiBridgeSessionEntry>,
    app_name: &str,
    user_id: &str,
    session_id: &str,
) -> &'a mut McpUiBridgeSessionEntry {
    registry.entry(bridge_session_key(app_name, user_id, session_id)).or_insert_with(|| {
        McpUiBridgeSessionEntry {
            protocol_version: MCP_UI_DEFAULT_PROTOCOL_VERSION.to_string(),
            initialized: false,
            app_info: None,
            app_capabilities: None,
            host_info: default_mcp_ui_host_info(),
            host_capabilities: default_mcp_ui_host_capabilities(),
            host_context: default_host_context(app_name, user_id, session_id),
            message_count: 0,
            last_message: None,
            model_context: vec![],
            model_context_revision: 0,
            resource_list_revision: 0,
            tool_list_revision: 0,
            notification_count: 0,
            pending_notifications: vec![],
        }
    })
}

fn validate_ui_resource_uri(uri: &str) -> Result<(), (StatusCode, String)> {
    if !uri.starts_with("ui://") {
        return Err((
            StatusCode::BAD_REQUEST,
            "ui resource uri must start with 'ui://'".to_string(),
        ));
    }
    Ok(())
}

fn validate_ui_resource_mime(mime_type: &str) -> Result<(), (StatusCode, String)> {
    if mime_type != "text/html;profile=mcp-app" {
        return Err((
            StatusCode::BAD_REQUEST,
            "mimeType must be 'text/html;profile=mcp-app'".to_string(),
        ));
    }
    Ok(())
}

fn parse_ui_meta_options(
    meta: &Option<Value>,
) -> Result<McpAppsRenderOptions, (StatusCode, String)> {
    let Some(meta_value) = meta else {
        return Ok(McpAppsRenderOptions::default());
    };
    let meta_object = meta_value
        .as_object()
        .ok_or_else(|| (StatusCode::BAD_REQUEST, "_meta must be a JSON object".to_string()))?;
    let Some(ui_value) = meta_object.get("ui") else {
        return Ok(McpAppsRenderOptions::default());
    };
    let ui_object = ui_value
        .as_object()
        .ok_or_else(|| (StatusCode::BAD_REQUEST, "_meta.ui must be a JSON object".to_string()))?;

    let domain = ui_object
        .get("domain")
        .map(|domain_value| {
            domain_value.as_str().ok_or_else(|| {
                (StatusCode::BAD_REQUEST, "_meta.ui.domain must be a string".to_string())
            })
        })
        .transpose()?
        .map(ToString::to_string);

    let prefers_border = ui_object
        .get("prefersBorder")
        .map(|value| {
            value.as_bool().ok_or_else(|| {
                (StatusCode::BAD_REQUEST, "_meta.ui.prefersBorder must be a boolean".to_string())
            })
        })
        .transpose()?;

    let csp = ui_object
        .get("csp")
        .map(|value| {
            serde_json::from_value::<McpUiResourceCsp>(value.clone()).map_err(|error| {
                (
                    StatusCode::BAD_REQUEST,
                    format!("_meta.ui.csp must be an object with domain arrays: {}", error),
                )
            })
        })
        .transpose()?;

    let permissions = ui_object
        .get("permissions")
        .map(|value| {
            serde_json::from_value::<McpUiPermissions>(value.clone()).map_err(|error| {
                (
                    StatusCode::BAD_REQUEST,
                    format!("_meta.ui.permissions must be an object: {}", error),
                )
            })
        })
        .transpose()?;

    Ok(McpAppsRenderOptions { domain, prefers_border, csp, permissions })
}

fn validate_ui_meta(meta: &Option<Value>) -> Result<McpAppsRenderOptions, (StatusCode, String)> {
    let options = parse_ui_meta_options(meta)?;
    validate_mcp_apps_render_options(&options).map_err(|error| {
        (StatusCode::BAD_REQUEST, format!("Invalid _meta.ui options for mcp_apps: {}", error))
    })?;
    Ok(options)
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiInitializeParams {
    pub app_name: String,
    pub user_id: String,
    pub session_id: String,
    #[serde(default)]
    pub protocol_version: Option<String>,
    #[serde(default)]
    pub app_info: Option<Value>,
    #[serde(default)]
    pub app_capabilities: Option<Value>,
    #[serde(default)]
    pub host_context: Option<Value>,
    #[serde(default)]
    pub host_info: Option<McpUiHostInfo>,
    #[serde(default)]
    pub host_capabilities: Option<McpUiHostCapabilities>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiMessageParams {
    pub app_name: String,
    pub user_id: String,
    pub session_id: String,
    #[serde(default)]
    pub role: Option<String>,
    #[serde(default)]
    pub content: Vec<Value>,
    #[serde(default)]
    pub metadata: Option<Value>,
    #[serde(default)]
    pub host_context: Option<Value>,
}

#[derive(Debug, Clone, Copy, Default, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum McpUiModelContextUpdateMode {
    Append,
    #[default]
    Replace,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiUpdateModelContextParams {
    pub app_name: String,
    pub user_id: String,
    pub session_id: String,
    #[serde(default)]
    pub content: Vec<Value>,
    #[serde(default)]
    pub structured_content: Option<Value>,
    #[serde(default)]
    pub host_context: Option<Value>,
    #[serde(default)]
    pub mode: McpUiModelContextUpdateMode,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiInitializeResult {
    pub initialized: bool,
    pub protocol_version: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub app_info: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub app_capabilities: Option<Value>,
    pub host_info: McpUiHostInfo,
    pub host_capabilities: McpUiHostCapabilities,
    pub host_context: Value,
    pub message_count: u64,
    pub model_context: Vec<Value>,
    pub model_context_revision: u64,
    pub resource_list_revision: u64,
    pub tool_list_revision: u64,
    pub notifications: Vec<McpUiBridgeNotification>,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiMessageResult {
    pub accepted: bool,
    pub initialized: bool,
    pub protocol_version: String,
    pub message_count: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub last_message: Option<Value>,
    pub resource_list_revision: u64,
    pub tool_list_revision: u64,
    pub notifications: Vec<McpUiBridgeNotification>,
    pub host_info: McpUiHostInfo,
    pub host_capabilities: McpUiHostCapabilities,
    pub host_context: Value,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiUpdateModelContextResult {
    pub accepted: bool,
    pub initialized: bool,
    pub protocol_version: String,
    pub model_context: Vec<Value>,
    pub model_context_revision: u64,
    pub resource_list_revision: u64,
    pub tool_list_revision: u64,
    pub notifications: Vec<McpUiBridgeNotification>,
    pub host_info: McpUiHostInfo,
    pub host_capabilities: McpUiHostCapabilities,
    pub host_context: Value,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiPollNotificationsParams {
    pub app_name: String,
    pub user_id: String,
    pub session_id: String,
    #[serde(default = "default_true")]
    pub drain: bool,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiPollNotificationsResult {
    pub initialized: bool,
    pub protocol_version: String,
    pub resource_list_revision: u64,
    pub tool_list_revision: u64,
    pub notifications: Vec<McpUiBridgeNotification>,
    pub host_info: McpUiHostInfo,
    pub host_capabilities: McpUiHostCapabilities,
    pub host_context: Value,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiListChangedParams {
    pub app_name: String,
    pub user_id: String,
    pub session_id: String,
    #[serde(default)]
    pub params: Option<Value>,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct McpUiListChangedResult {
    pub accepted: bool,
    pub initialized: bool,
    pub protocol_version: String,
    pub method: String,
    pub revision: u64,
    pub pending_notification_count: usize,
    pub resource_list_revision: u64,
    pub tool_list_revision: u64,
    pub host_info: McpUiHostInfo,
    pub host_capabilities: McpUiHostCapabilities,
    pub host_context: Value,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub(crate) enum McpUiBridgeInput<T> {
    Direct(T),
    Rpc(McpUiBridgeRpcRequest<T>),
}

#[derive(Debug, Clone, Deserialize)]
pub(crate) struct McpUiBridgeRpcRequest<T> {
    #[serde(default)]
    id: Option<Value>,
    method: String,
    params: T,
}

enum McpUiBridgeResponseMode {
    Direct,
    Rpc { id: Option<Value> },
}

fn parse_bridge_input<T>(
    input: McpUiBridgeInput<T>,
    expected_method: &str,
) -> Result<(T, McpUiBridgeResponseMode), (StatusCode, String)> {
    match input {
        McpUiBridgeInput::Direct(params) => Ok((params, McpUiBridgeResponseMode::Direct)),
        McpUiBridgeInput::Rpc(request) => {
            if request.method != expected_method {
                return Err((
                    StatusCode::BAD_REQUEST,
                    format!(
                        "unexpected MCP Apps bridge method '{}', expected '{}'",
                        request.method, expected_method
                    ),
                ));
            }
            Ok((request.params, McpUiBridgeResponseMode::Rpc { id: request.id }))
        }
    }
}

fn bridge_result_json<T: Serialize>(
    mode: McpUiBridgeResponseMode,
    result: T,
) -> Result<Json<Value>, (StatusCode, String)> {
    let value = match mode {
        McpUiBridgeResponseMode::Direct => serde_json::to_value(result).map_err(|error| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("failed to serialize bridge response: {}", error),
            )
        })?,
        McpUiBridgeResponseMode::Rpc { id } => json!({
            "jsonrpc": "2.0",
            "id": id,
            "result": result
        }),
    };
    Ok(Json(value))
}

fn model_context_items(params: &McpUiUpdateModelContextParams) -> Vec<Value> {
    let mut items = params.content.clone();
    if let Some(structured_content) = params.structured_content.clone() {
        items.push(json!({
            "type": "structuredContent",
            "structuredContent": structured_content
        }));
    }
    items
}

fn default_true() -> bool {
    true
}

fn bridge_snapshot(session: &McpUiBridgeSessionEntry) -> McpUiBridgeSnapshot {
    McpUiBridgeSnapshot::new(
        session.protocol_version.clone(),
        session.initialized,
        session.host_info.clone(),
        session.host_capabilities.clone(),
        session.host_context.clone(),
    )
    .with_optional_app_metadata(session.app_info.clone(), session.app_capabilities.clone())
}

fn queued_notifications(session: &McpUiBridgeSessionEntry) -> Vec<McpUiBridgeNotification> {
    session.pending_notifications.clone()
}

fn queue_bridge_notification(
    session: &mut McpUiBridgeSessionEntry,
    method: &str,
    revision: u64,
    params: Option<Value>,
) {
    session.notification_count += 1;
    session.pending_notifications.push(McpUiBridgeNotification {
        notification_id: session.notification_count,
        method: method.to_string(),
        revision,
        params,
    });
}

pub(crate) fn initialize_mcp_ui_bridge(
    params: McpUiInitializeParams,
) -> Result<McpUiInitializeResult, (StatusCode, String)> {
    let app_name = params.app_name.clone();
    let user_id = params.user_id.clone();
    let session_id = params.session_id.clone();
    let mut registry = bridge_registry()
        .write()
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "bridge registry poisoned".to_string()))?;
    let session = ensure_bridge_session(&mut registry, &app_name, &user_id, &session_id);

    if let Some(protocol_version) = params.protocol_version.filter(|value| !value.trim().is_empty())
    {
        session.protocol_version = protocol_version;
    }
    if let Some(app_info) = params.app_info {
        session.app_info = Some(app_info);
    }
    if let Some(app_capabilities) = params.app_capabilities {
        session.app_capabilities = Some(app_capabilities);
    }
    if let Some(host_info) = params.host_info {
        session.host_info = host_info;
    }
    if let Some(host_capabilities) = params.host_capabilities {
        session.host_capabilities = host_capabilities;
    }
    merge_host_context(&mut session.host_context, params.host_context)?;
    session.initialized = true;

    info!(
        app_name = %app_name,
        user_id = %user_id,
        session_id = %session_id,
        protocol_version = %session.protocol_version,
        "mcp ui initialize handled"
    );

    let snapshot = bridge_snapshot(session);
    Ok(McpUiInitializeResult {
        initialized: snapshot.initialized,
        protocol_version: snapshot.protocol_version.clone(),
        app_info: snapshot.app_info.clone(),
        app_capabilities: snapshot.app_capabilities.clone(),
        host_info: snapshot.host_info.clone(),
        host_capabilities: snapshot.host_capabilities.clone(),
        host_context: snapshot.host_context.clone(),
        message_count: session.message_count,
        model_context: session.model_context.clone(),
        model_context_revision: session.model_context_revision,
        resource_list_revision: session.resource_list_revision,
        tool_list_revision: session.tool_list_revision,
        notifications: queued_notifications(session),
    })
}

pub(crate) fn message_mcp_ui_bridge(
    params: McpUiMessageParams,
) -> Result<McpUiMessageResult, (StatusCode, String)> {
    let app_name = params.app_name.clone();
    let user_id = params.user_id.clone();
    let session_id = params.session_id.clone();
    let host_context = params.host_context.clone();
    let mut registry = bridge_registry()
        .write()
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "bridge registry poisoned".to_string()))?;
    let session = ensure_bridge_session(&mut registry, &app_name, &user_id, &session_id);
    merge_host_context(&mut session.host_context, host_context)?;

    session.message_count += 1;
    let mut message = json!({
        "role": params.role.unwrap_or_else(|| "user".to_string()),
        "content": params.content
    });
    if let Some(metadata) = params.metadata {
        let object = message.as_object_mut().ok_or_else(|| {
            (StatusCode::INTERNAL_SERVER_ERROR, "message payload must be an object".to_string())
        })?;
        object.insert("metadata".to_string(), metadata);
    }
    session.last_message = Some(message);

    info!(
        app_name = %app_name,
        user_id = %user_id,
        session_id = %session_id,
        message_count = session.message_count,
        initialized = session.initialized,
        "mcp ui message handled"
    );

    let snapshot = bridge_snapshot(session);
    Ok(McpUiMessageResult {
        accepted: true,
        initialized: snapshot.initialized,
        protocol_version: snapshot.protocol_version,
        message_count: session.message_count,
        last_message: session.last_message.clone(),
        resource_list_revision: session.resource_list_revision,
        tool_list_revision: session.tool_list_revision,
        notifications: queued_notifications(session),
        host_info: snapshot.host_info,
        host_capabilities: snapshot.host_capabilities,
        host_context: snapshot.host_context,
    })
}

pub(crate) fn update_mcp_ui_bridge_model_context(
    params: McpUiUpdateModelContextParams,
) -> Result<McpUiUpdateModelContextResult, (StatusCode, String)> {
    let app_name = params.app_name.clone();
    let user_id = params.user_id.clone();
    let session_id = params.session_id.clone();
    let host_context = params.host_context.clone();
    let items = model_context_items(&params);
    let mut registry = bridge_registry()
        .write()
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "bridge registry poisoned".to_string()))?;
    let session = ensure_bridge_session(&mut registry, &app_name, &user_id, &session_id);
    merge_host_context(&mut session.host_context, host_context)?;

    match params.mode {
        McpUiModelContextUpdateMode::Replace => session.model_context = items,
        McpUiModelContextUpdateMode::Append => session.model_context.extend(items),
    }
    session.model_context_revision += 1;

    info!(
        app_name = %app_name,
        user_id = %user_id,
        session_id = %session_id,
        model_context_revision = session.model_context_revision,
        initialized = session.initialized,
        "mcp ui model context updated"
    );

    let snapshot = bridge_snapshot(session);
    Ok(McpUiUpdateModelContextResult {
        accepted: true,
        initialized: snapshot.initialized,
        protocol_version: snapshot.protocol_version,
        model_context: session.model_context.clone(),
        model_context_revision: session.model_context_revision,
        resource_list_revision: session.resource_list_revision,
        tool_list_revision: session.tool_list_revision,
        notifications: queued_notifications(session),
        host_info: snapshot.host_info,
        host_capabilities: snapshot.host_capabilities,
        host_context: snapshot.host_context,
    })
}

pub(crate) fn poll_mcp_ui_bridge_notifications(
    params: McpUiPollNotificationsParams,
) -> Result<McpUiPollNotificationsResult, (StatusCode, String)> {
    let mut registry = bridge_registry()
        .write()
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "bridge registry poisoned".to_string()))?;
    let session =
        ensure_bridge_session(&mut registry, &params.app_name, &params.user_id, &params.session_id);
    let notifications = queued_notifications(session);
    if params.drain {
        session.pending_notifications.clear();
    }
    let snapshot = bridge_snapshot(session);

    Ok(McpUiPollNotificationsResult {
        initialized: snapshot.initialized,
        protocol_version: snapshot.protocol_version,
        resource_list_revision: session.resource_list_revision,
        tool_list_revision: session.tool_list_revision,
        notifications,
        host_info: snapshot.host_info,
        host_capabilities: snapshot.host_capabilities,
        host_context: snapshot.host_context,
    })
}

fn notify_mcp_ui_bridge_list_changed(
    params: McpUiListChangedParams,
    method: &'static str,
    revision_selector: impl Fn(&mut McpUiBridgeSessionEntry) -> &mut u64,
) -> Result<McpUiListChangedResult, (StatusCode, String)> {
    let mut registry = bridge_registry()
        .write()
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "bridge registry poisoned".to_string()))?;
    let session =
        ensure_bridge_session(&mut registry, &params.app_name, &params.user_id, &params.session_id);
    let revision = {
        let target = revision_selector(session);
        *target += 1;
        *target
    };
    queue_bridge_notification(session, method, revision, params.params);
    let snapshot = bridge_snapshot(session);

    Ok(McpUiListChangedResult {
        accepted: true,
        initialized: snapshot.initialized,
        protocol_version: snapshot.protocol_version,
        method: method.to_string(),
        revision,
        pending_notification_count: session.pending_notifications.len(),
        resource_list_revision: session.resource_list_revision,
        tool_list_revision: session.tool_list_revision,
        host_info: snapshot.host_info,
        host_capabilities: snapshot.host_capabilities,
        host_context: snapshot.host_context,
    })
}

pub(crate) fn notify_mcp_ui_resource_list_changed(
    params: McpUiListChangedParams,
) -> Result<McpUiListChangedResult, (StatusCode, String)> {
    notify_mcp_ui_bridge_list_changed(
        params,
        "ui/notifications/resources/list_changed",
        |session| &mut session.resource_list_revision,
    )
}

pub(crate) fn notify_mcp_ui_tool_list_changed(
    params: McpUiListChangedParams,
) -> Result<McpUiListChangedResult, (StatusCode, String)> {
    notify_mcp_ui_bridge_list_changed(params, "ui/notifications/tools/list_changed", |session| {
        &mut session.tool_list_revision
    })
}

pub(crate) fn mark_mcp_ui_initialized(
    app_name: &str,
    user_id: &str,
    session_id: &str,
) -> Result<(), (StatusCode, String)> {
    let mut registry = bridge_registry()
        .write()
        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "bridge registry poisoned".to_string()))?;
    let session = ensure_bridge_session(&mut registry, app_name, user_id, session_id);
    session.initialized = true;
    Ok(())
}

/// Whether `owner` may be accessed by `caller`.
///
/// A resource registered without an owner stays readable by anyone, so a server
/// with no authentication behaves as before. Once a resource has an owner, only
/// that user may read or replace it: the registry is process-global and keyed by
/// URI alone, so without this check one authenticated caller could read or
/// overwrite another's registered HTML.
fn may_access_resource(owner: Option<&String>, caller: Option<&adk_core::RequestContext>) -> bool {
    match (owner, caller) {
        (None, _) => true,
        (Some(_), None) => true,
        (Some(owner), Some(caller)) => owner == &caller.user_id,
    }
}

/// Bridge parameters that name the user whose bridge state is addressed.
///
/// Implemented for every bridge params type so one rule binds them all to the
/// authenticated caller.
trait BridgeUser {
    fn user_id_mut(&mut self) -> &mut String;
}

/// Replaces the caller-supplied user with the authenticated one.
///
/// The bridge registry is keyed by `(app_name, user_id, session_id)` taken from the
/// request body, so a body-supplied user would let one authenticated caller read and
/// mutate another's bridge state. When authentication is configured, the body value
/// is not trusted. When it is not configured, there is no authenticated identity to
/// bind and the body value stands.
fn bind_authenticated_user<P: BridgeUser>(
    params: &mut P,
    request_context: Option<&adk_core::RequestContext>,
) {
    if let Some(context) = request_context {
        let supplied = params.user_id_mut();
        if supplied != &context.user_id {
            info!(
                supplied.user_id = %supplied,
                authenticated.user_id = %context.user_id,
                "ui bridge request named a different user, using the authenticated user"
            );
            *supplied = context.user_id.clone();
        }
    }
}

impl BridgeUser for McpUiInitializeParams {
    fn user_id_mut(&mut self) -> &mut String {
        &mut self.user_id
    }
}

impl BridgeUser for McpUiMessageParams {
    fn user_id_mut(&mut self) -> &mut String {
        &mut self.user_id
    }
}

impl BridgeUser for McpUiUpdateModelContextParams {
    fn user_id_mut(&mut self) -> &mut String {
        &mut self.user_id
    }
}

impl BridgeUser for McpUiPollNotificationsParams {
    fn user_id_mut(&mut self) -> &mut String {
        &mut self.user_id
    }
}

impl BridgeUser for McpUiListChangedParams {
    fn user_id_mut(&mut self) -> &mut String {
        &mut self.user_id
    }
}

/// GET /api/ui/capabilities
pub async fn ui_capabilities() -> Json<UiCapabilities> {
    Json(UiCapabilities {
        default_protocol: UI_DEFAULT_PROTOCOL,
        protocols: UI_PROTOCOL_CAPABILITIES
            .iter()
            .map(|spec| UiProtocolCapability {
                protocol: spec.protocol,
                versions: spec.versions.to_vec(),
                implementation_tier: spec.implementation_tier,
                spec_track: spec.spec_track,
                summary: spec.summary,
                features: spec.features.to_vec(),
                limitations: spec.limitations.to_vec(),
                deprecation: map_deprecation(spec.deprecation),
            })
            .collect(),
        tool_envelope_version: TOOL_ENVELOPE_VERSION,
    })
}

/// POST /api/ui/initialize
pub(crate) async fn ui_initialize(
    Extension(request_context): Extension<Option<adk_core::RequestContext>>,
    Json(input): Json<McpUiBridgeInput<McpUiInitializeParams>>,
) -> Result<Json<Value>, (StatusCode, String)> {
    let (mut params, response_mode) = parse_bridge_input(input, "ui/initialize")?;
    bind_authenticated_user(&mut params, request_context.as_ref());
    let result = initialize_mcp_ui_bridge(params)?;
    bridge_result_json(response_mode, result)
}

/// POST /api/ui/message
pub(crate) async fn ui_message(
    Extension(request_context): Extension<Option<adk_core::RequestContext>>,
    Json(input): Json<McpUiBridgeInput<McpUiMessageParams>>,
) -> Result<Json<Value>, (StatusCode, String)> {
    let (mut params, response_mode) = parse_bridge_input(input, "ui/message")?;
    bind_authenticated_user(&mut params, request_context.as_ref());
    let result = message_mcp_ui_bridge(params)?;
    bridge_result_json(response_mode, result)
}

/// POST /api/ui/update-model-context
pub(crate) async fn ui_update_model_context(
    Extension(request_context): Extension<Option<adk_core::RequestContext>>,
    Json(input): Json<McpUiBridgeInput<McpUiUpdateModelContextParams>>,
) -> Result<Json<Value>, (StatusCode, String)> {
    let (mut params, response_mode) = parse_bridge_input(input, "ui/update-model-context")?;
    bind_authenticated_user(&mut params, request_context.as_ref());
    let result = update_mcp_ui_bridge_model_context(params)?;
    bridge_result_json(response_mode, result)
}

/// POST /api/ui/notifications/poll
pub(crate) async fn ui_poll_notifications(
    Extension(request_context): Extension<Option<adk_core::RequestContext>>,
    Json(input): Json<McpUiBridgeInput<McpUiPollNotificationsParams>>,
) -> Result<Json<Value>, (StatusCode, String)> {
    let (mut params, response_mode) = parse_bridge_input(input, "ui/notifications/poll")?;
    bind_authenticated_user(&mut params, request_context.as_ref());
    let result = poll_mcp_ui_bridge_notifications(params)?;
    bridge_result_json(response_mode, result)
}

/// POST /api/ui/notifications/resources-list-changed
pub(crate) async fn ui_notify_resources_list_changed(
    Extension(request_context): Extension<Option<adk_core::RequestContext>>,
    Json(input): Json<McpUiBridgeInput<McpUiListChangedParams>>,
) -> Result<Json<Value>, (StatusCode, String)> {
    let (mut params, response_mode) =
        parse_bridge_input(input, "ui/notifications/resources/list_changed")?;
    bind_authenticated_user(&mut params, request_context.as_ref());
    let result = notify_mcp_ui_resource_list_changed(params)?;
    bridge_result_json(response_mode, result)
}

/// POST /api/ui/notifications/tools-list-changed
pub(crate) async fn ui_notify_tools_list_changed(
    Extension(request_context): Extension<Option<adk_core::RequestContext>>,
    Json(input): Json<McpUiBridgeInput<McpUiListChangedParams>>,
) -> Result<Json<Value>, (StatusCode, String)> {
    let (mut params, response_mode) =
        parse_bridge_input(input, "ui/notifications/tools/list_changed")?;
    bind_authenticated_user(&mut params, request_context.as_ref());
    let result = notify_mcp_ui_tool_list_changed(params)?;
    bridge_result_json(response_mode, result)
}

/// GET /api/ui/resources
pub async fn list_ui_resources(
    Extension(request_context): Extension<Option<adk_core::RequestContext>>,
) -> Json<UiResourceListResponse> {
    let resources: Vec<UiResource> = resource_registry()
        .read()
        .map(|registry| {
            registry
                .values()
                .filter(|entry| may_access_resource(entry.owner.as_ref(), request_context.as_ref()))
                .map(|entry| entry.resource.clone())
                .collect()
        })
        .unwrap_or_default();
    info!(resource_count = resources.len(), "ui resource list requested");
    Json(UiResourceListResponse { resources })
}

/// GET /api/ui/resources/read?uri=ui://...
pub async fn read_ui_resource(
    Extension(request_context): Extension<Option<adk_core::RequestContext>>,
    Query(query): Query<ReadUiResourceQuery>,
) -> Result<Json<UiResourceReadResponse>, (StatusCode, String)> {
    validate_ui_resource_uri(&query.uri)?;
    let guard = resource_registry().read().map_err(|_| {
        (StatusCode::INTERNAL_SERVER_ERROR, "resource registry poisoned".to_string())
    })?;
    let Some(entry) = guard.get(&query.uri) else {
        warn!(uri = %query.uri, "ui resource read failed: not found");
        return Err((StatusCode::NOT_FOUND, format!("resource not found: {}", query.uri)));
    };
    if !may_access_resource(entry.owner.as_ref(), request_context.as_ref()) {
        // Reported as not found so the URI's existence is not disclosed.
        warn!(uri = %query.uri, "ui resource read denied: owned by another user");
        return Err((StatusCode::NOT_FOUND, format!("resource not found: {}", query.uri)));
    }
    info!(uri = %query.uri, "ui resource read");
    Ok(Json(UiResourceReadResponse { contents: vec![entry.content.clone()] }))
}

/// POST /api/ui/resources/register
pub async fn register_ui_resource(
    Extension(request_context): Extension<Option<adk_core::RequestContext>>,
    Json(req): Json<RegisterUiResourceRequest>,
) -> Result<StatusCode, (StatusCode, String)> {
    validate_ui_resource_uri(&req.uri)?;
    validate_ui_resource_mime(&req.mime_type)?;
    let ui_meta_options = validate_ui_meta(&req.meta)?;

    let uri = req.uri.clone();
    let name = req.name.clone();
    let mime_type = req.mime_type.clone();
    let meta = req.meta.clone();
    let domain = ui_meta_options.domain.unwrap_or_else(|| "<none>".to_string());

    let entry = UiResourceEntry {
        resource: UiResource {
            uri: uri.clone(),
            name: name.clone(),
            description: req.description.clone(),
            mime_type: mime_type.clone(),
            meta: meta.clone(),
        },
        content: UiResourceContent {
            uri: uri.clone(),
            mime_type: mime_type.clone(),
            text: Some(req.text),
            blob: None,
            meta,
        },
        owner: request_context.as_ref().map(|context| context.user_id.clone()),
    };

    let mut registry = resource_registry().write().map_err(|_| {
        (StatusCode::INTERNAL_SERVER_ERROR, "resource registry poisoned".to_string())
    })?;
    // Registration replaces an existing URI, so an unauthorized overwrite would let
    // one caller substitute another's resource content.
    if let Some(existing) = registry.get(&uri)
        && !may_access_resource(existing.owner.as_ref(), request_context.as_ref())
    {
        warn!(uri = %uri, "ui resource registration denied: owned by another user");
        return Err((StatusCode::FORBIDDEN, format!("resource is owned by another user: {uri}")));
    }
    registry.insert(uri.clone(), entry);
    drop(registry);
    info!(
        uri = %uri,
        name = %name,
        mime_type = %mime_type,
        ui_domain = %domain,
        "ui resource registered"
    );

    Ok(StatusCode::CREATED)
}